Add InstanceValidator tests

3 files changed, 196 insertions, 21 deletions

diff --git a/athenz-identity-provider-service/src/main/java/com/yahoo/vespa/hosted/athenz/instanceproviderservice/impl/InstanceValidator.java b/athenz-identity-provider-service/src/main/java/com/yahoo/vespa/hosted/athenz/instanceproviderservice/impl/InstanceValidator.java
index d738b34041a..01a88cbcfcf 100644
--- a/athenz-identity-provider-service/src/main/java/com/yahoo/vespa/hosted/athenz/instanceproviderservice/impl/InstanceValidator.java
+++ b/athenz-identity-provider-service/src/main/java/com/yahoo/vespa/hosted/athenz/instanceproviderservice/impl/InstanceValidator.java
@@ -27,6 +27,8 @@ import java.util.logging.Logger;
 public class InstanceValidator {
 
     private static final Logger log = Logger.getLogger(InstanceValidator.class.getName());
+    static final String SERVICE_PROPERTIES_DOMAIN_KEY = "identity.domain";
+    static final String SERVICE_PROPERTIES_SERVICE_KEY = "identity.service";
 
     private final KeyProvider keyProvider;
     private final SuperModelProvider superModelProvider;
@@ -47,8 +49,7 @@ public class InstanceValidator {
         }
 
         log.log(LogLevel.INFO, () -> String.format("Validating instance %s.", providerUniqueId));
-        PublicKey publicKey = keyProvider.getPublicKey(signedIdentityDocument.signingKeyVersion);
-        if (isSignatureValid(publicKey, signedIdentityDocument.rawIdentityDocument, signedIdentityDocument.signature)) {
+        if (isInstanceSignatureValid(instanceConfirmation)) {
             log.log(LogLevel.INFO, () -> String.format("Instance %s is valid.", providerUniqueId));
             return true;
         }
@@ -56,7 +57,14 @@ public class InstanceValidator {
         return false;
     }
 
-    public static boolean isSignatureValid(PublicKey publicKey, String rawIdentityDocument, String signature) {
+    boolean isInstanceSignatureValid(InstanceConfirmation instanceConfirmation) {
+        SignedIdentityDocument signedIdentityDocument = instanceConfirmation.signedIdentityDocument;
+
+        PublicKey publicKey = keyProvider.getPublicKey(signedIdentityDocument.signingKeyVersion);
+        return isSignatureValid(publicKey, signedIdentityDocument.rawIdentityDocument, signedIdentityDocument.signature);
+    }
+
+    static boolean isSignatureValid(PublicKey publicKey, String rawIdentityDocument, String signature) {
         try {
             Signature signatureVerifier = Signature.getInstance("SHA512withRSA");
             signatureVerifier.initVerify(publicKey);
@@ -68,7 +76,7 @@ public class InstanceValidator {
     }
 
     // If/when we dont care about logging exactly whats wrong, this can be simplified
-    private boolean isSameIdentityAsInServicesXml(ApplicationId applicationId, String domain, String service) {
+    boolean isSameIdentityAsInServicesXml(ApplicationId applicationId, String domain, String service) {
         Optional<ApplicationInfo> applicationInfo = superModelProvider.getSuperModel().getApplicationInfo(applicationId);
 
         if (!applicationInfo.isPresent()) {
@@ -81,8 +89,8 @@ public class InstanceValidator {
                 .getHosts()
                 .stream()
                 .flatMap(hostInfo -> hostInfo.getServices().stream())
-                .filter(serviceInfo -> serviceInfo.getProperty("identity.domain").isPresent())
-                .filter(serviceInfo -> serviceInfo.getProperty("identity.service").isPresent())
+                .filter(serviceInfo -> serviceInfo.getProperty(SERVICE_PROPERTIES_DOMAIN_KEY).isPresent())
+                .filter(serviceInfo -> serviceInfo.getProperty(SERVICE_PROPERTIES_SERVICE_KEY).isPresent())
                 .findFirst();
 
         if (!matchingServiceInfo.isPresent()) {
@@ -90,8 +98,8 @@ public class InstanceValidator {
             return false;
         }
 
-        String domainInConfig = matchingServiceInfo.get().getProperty("identity.domain").get();
-        String serviceInConfig = matchingServiceInfo.get().getProperty("identity.service").get();
+        String domainInConfig = matchingServiceInfo.get().getProperty(SERVICE_PROPERTIES_DOMAIN_KEY).get();
+        String serviceInConfig = matchingServiceInfo.get().getProperty(SERVICE_PROPERTIES_SERVICE_KEY).get();
         if (domainInConfig.equals(domain) && serviceInConfig.equals(service)) {
             return true;
         } else {
diff --git a/athenz-identity-provider-service/src/test/java/com/yahoo/vespa/hosted/athenz/instanceproviderservice/AthenzInstanceProviderServiceTest.java b/athenz-identity-provider-service/src/test/java/com/yahoo/vespa/hosted/athenz/instanceproviderservice/AthenzInstanceProviderServiceTest.java
index 6fac9b549b5..a3cc82cd917 100644
--- a/athenz-identity-provider-service/src/test/java/com/yahoo/vespa/hosted/athenz/instanceproviderservice/AthenzInstanceProviderServiceTest.java
+++ b/athenz-identity-provider-service/src/test/java/com/yahoo/vespa/hosted/athenz/instanceproviderservice/AthenzInstanceProviderServiceTest.java
@@ -1,7 +1,6 @@
 // Copyright 2017 Yahoo Holdings. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 package com.yahoo.vespa.hosted.athenz.instanceproviderservice;
 
-import com.fasterxml.jackson.core.JsonProcessingException;
 import com.fasterxml.jackson.databind.ObjectMapper;
 import com.yahoo.config.provision.Environment;
 import com.yahoo.config.provision.RegionName;
@@ -45,9 +44,7 @@ import org.junit.Test;
 
 import javax.net.ssl.SSLContext;
 import java.io.IOException;
-import java.io.UnsupportedEncodingException;
 import java.math.BigInteger;
-import java.security.InvalidKeyException;
 import java.security.KeyManagementException;
 import java.security.KeyPair;
 import java.security.KeyPairGenerator;
@@ -56,7 +53,6 @@ import java.security.NoSuchAlgorithmException;
 import java.security.PrivateKey;
 import java.security.PublicKey;
 import java.security.Signature;
-import java.security.SignatureException;
 import java.security.cert.CertificateException;
 import java.security.cert.X509Certificate;
 import java.time.Instant;
@@ -192,11 +188,7 @@ public class AthenzInstanceProviderServiceTest {
                                                "localhost/zts",
                                                1));
             return new StringEntity(mapper.writeValueAsString(instanceConfirmation));
-        } catch (JsonProcessingException
-                | NoSuchAlgorithmException
-                | UnsupportedEncodingException
-                | SignatureException
-                | InvalidKeyException e) {
+        } catch (Exception e) {
             throw new RuntimeException(e);
         }
     }
@@ -205,10 +197,14 @@ public class AthenzInstanceProviderServiceTest {
 
         private final KeyPair keyPair;
 
-        public AutoGeneratedKeyProvider() throws IOException, NoSuchAlgorithmException {
-            KeyPairGenerator rsa = KeyPairGenerator.getInstance("RSA");
-            rsa.initialize(2048);
-            keyPair = rsa.genKeyPair();
+        public AutoGeneratedKeyProvider() {
+            try {
+                KeyPairGenerator rsa = KeyPairGenerator.getInstance("RSA");
+                rsa.initialize(2048);
+                keyPair = rsa.genKeyPair();
+            } catch (NoSuchAlgorithmException e) {
+                throw new RuntimeException(e);
+            }
         }
 
         @Override
diff --git a/athenz-identity-provider-service/src/test/java/com/yahoo/vespa/hosted/athenz/instanceproviderservice/impl/InstanceValidatorTest.java b/athenz-identity-provider-service/src/test/java/com/yahoo/vespa/hosted/athenz/instanceproviderservice/impl/InstanceValidatorTest.java
new file mode 100644
index 00000000000..c1fab319ebf
--- /dev/null
+++ b/athenz-identity-provider-service/src/test/java/com/yahoo/vespa/hosted/athenz/instanceproviderservice/impl/InstanceValidatorTest.java
@@ -0,0 +1,171 @@
+package com.yahoo.vespa.hosted.athenz.instanceproviderservice.impl;
+
+import com.fasterxml.jackson.databind.ObjectMapper;
+import com.yahoo.config.model.api.ApplicationInfo;
+import com.yahoo.config.model.api.HostInfo;
+import com.yahoo.config.model.api.Model;
+import com.yahoo.config.model.api.ServiceInfo;
+import com.yahoo.config.model.api.SuperModel;
+import com.yahoo.config.model.api.SuperModelProvider;
+import com.yahoo.config.provision.ApplicationId;
+import com.yahoo.vespa.hosted.athenz.instanceproviderservice.AthenzInstanceProviderServiceTest.AutoGeneratedKeyProvider;
+import com.yahoo.vespa.hosted.athenz.instanceproviderservice.impl.model.IdentityDocument;
+import com.yahoo.vespa.hosted.athenz.instanceproviderservice.impl.model.InstanceConfirmation;
+import com.yahoo.vespa.hosted.athenz.instanceproviderservice.impl.model.ProviderUniqueId;
+import com.yahoo.vespa.hosted.athenz.instanceproviderservice.impl.model.SignedIdentityDocument;
+import org.junit.Test;
+
+import java.security.PrivateKey;
+import java.security.Signature;
+import java.time.Instant;
+import java.util.Base64;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.function.Function;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
+import java.util.stream.Stream;
+
+import static com.yahoo.vespa.hosted.athenz.instanceproviderservice.impl.InstanceValidator.SERVICE_PROPERTIES_DOMAIN_KEY;
+import static com.yahoo.vespa.hosted.athenz.instanceproviderservice.impl.InstanceValidator.SERVICE_PROPERTIES_SERVICE_KEY;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+/**
+ * @author valerijf
+ */
+public class InstanceValidatorTest {
+
+    private final ApplicationId applicationId = ApplicationId.from("tenant", "application", "instance");
+    private final String domain = "domain";
+    private final String service = "service";
+
+    @Test
+    public void valid_signature() throws Exception {
+        KeyProvider keyProvider = new AutoGeneratedKeyProvider();
+        InstanceValidator instanceValidator = new InstanceValidator(keyProvider, null);
+        InstanceConfirmation instanceConfirmation = createInstanceConfirmation(
+                keyProvider.getPrivateKey(0), applicationId, domain, service);
+
+        assertTrue(instanceValidator.isInstanceSignatureValid(instanceConfirmation));
+    }
+
+    @Test
+    public void invalid_signature() throws Exception {
+        KeyProvider keyProvider = new AutoGeneratedKeyProvider();
+        InstanceValidator instanceValidator = new InstanceValidator(keyProvider, null);
+
+        KeyProvider fakeKeyProvider = new AutoGeneratedKeyProvider();
+        InstanceConfirmation instanceConfirmation = createInstanceConfirmation(
+                fakeKeyProvider.getPrivateKey(0), applicationId, domain, service);
+
+        assertFalse(instanceValidator.isInstanceSignatureValid(instanceConfirmation));
+    }
+
+    @Test
+    public void application_does_not_exist() {
+        SuperModelProvider superModelProvider = mockSuperModelProvider();
+        InstanceValidator instanceValidator = new InstanceValidator(null, superModelProvider);
+
+        assertFalse(instanceValidator.isSameIdentityAsInServicesXml(applicationId, domain, service));
+    }
+
+    @Test
+    public void application_does_not_have_domain_set() {
+        SuperModelProvider superModelProvider = mockSuperModelProvider(
+                mockApplicationInfo(applicationId, 5, Collections.emptyList()));
+        InstanceValidator instanceValidator = new InstanceValidator(null, superModelProvider);
+
+        assertFalse(instanceValidator.isSameIdentityAsInServicesXml(applicationId, domain, service));
+    }
+
+    @Test
+    public void application_has_wrong_domain() {
+        ServiceInfo serviceInfo = new ServiceInfo("serviceName", "type", Collections.emptyList(),
+                Collections.singletonMap(SERVICE_PROPERTIES_DOMAIN_KEY, "not-domain"), "confId", "hostName");
+
+        SuperModelProvider superModelProvider = mockSuperModelProvider(
+                mockApplicationInfo(applicationId, 5, Collections.singletonList(serviceInfo)));
+        InstanceValidator instanceValidator = new InstanceValidator(null, superModelProvider);
+
+        assertFalse(instanceValidator.isSameIdentityAsInServicesXml(applicationId, domain, service));
+    }
+
+    @Test
+    public void application_has_same_domain_and_service() {
+        Map<String, String> properties = new HashMap<>();
+        properties.put(SERVICE_PROPERTIES_DOMAIN_KEY, domain);
+        properties.put(SERVICE_PROPERTIES_SERVICE_KEY, service);
+
+        ServiceInfo serviceInfo = new ServiceInfo("serviceName", "type", Collections.emptyList(),
+                properties, "confId", "hostName");
+
+        SuperModelProvider superModelProvider = mockSuperModelProvider(
+                mockApplicationInfo(applicationId, 5, Collections.singletonList(serviceInfo)));
+        InstanceValidator instanceValidator = new InstanceValidator(null, superModelProvider);
+
+        assertTrue(instanceValidator.isSameIdentityAsInServicesXml(applicationId, domain, service));
+    }
+
+    private static InstanceConfirmation createInstanceConfirmation(PrivateKey privateKey, ApplicationId applicationId,
+                                                                   String domain, String service) {
+        IdentityDocument identityDocument = new IdentityDocument(
+                new ProviderUniqueId(applicationId.tenant().value(), applicationId.application().value(),
+                        "environment", "region", applicationId.instance().value(), "cluster-id", 0),
+                "hostname",
+                "instance-hostname",
+                Instant.now());
+
+        try {
+            ObjectMapper mapper = Utils.getMapper();
+            String encodedIdentityDocument =
+                    Base64.getEncoder().encodeToString(mapper.writeValueAsString(identityDocument).getBytes());
+            Signature sigGenerator = Signature.getInstance("SHA512withRSA");
+            sigGenerator.initSign(privateKey);
+            sigGenerator.update(encodedIdentityDocument.getBytes());
+
+            return new InstanceConfirmation(
+                    "provider", domain, service,
+                    new SignedIdentityDocument(encodedIdentityDocument,
+                            Base64.getEncoder().encodeToString(sigGenerator.sign()),
+                            0,
+                            identityDocument.providerUniqueId.asString(),
+                            "dnssuffix",
+                            "service",
+                            "localhost/zts",
+                            1));
+        } catch (Exception e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    private SuperModelProvider mockSuperModelProvider(ApplicationInfo... appInfos) {
+        SuperModel superModel = new SuperModel(Stream.of(appInfos)
+                .collect(Collectors.groupingBy(
+                        appInfo -> appInfo.getApplicationId().tenant(),
+                        Collectors.toMap(
+                                ApplicationInfo::getApplicationId,
+                                Function.identity()
+                        )
+                )));
+
+        SuperModelProvider superModelProvider = mock(SuperModelProvider.class);
+        when(superModelProvider.getSuperModel()).thenReturn(superModel);
+        return superModelProvider;
+    }
+
+    private ApplicationInfo mockApplicationInfo(ApplicationId appId, int numHosts, List<ServiceInfo> serviceInfo) {
+        List<HostInfo> hosts = IntStream.range(0, numHosts)
+                .mapToObj(i -> new HostInfo("host-" + i + "." + appId.toShortString() + ".yahoo.com", serviceInfo))
+                .collect(Collectors.toList());
+
+        Model model = mock(Model.class);
+        when(model.getHosts()).thenReturn(hosts);
+
+        return new ApplicationInfo(appId, 0, model);
+    }
+}
\ No newline at end of file