Revert "Revert "Validate register and refresh""

10 files changed, 343 insertions, 32 deletions

diff --git a/athenz-identity-provider-service/src/main/java/com/yahoo/vespa/hosted/athenz/instanceproviderservice/instanceconfirmation/InstanceValidator.java b/athenz-identity-provider-service/src/main/java/com/yahoo/vespa/hosted/athenz/instanceproviderservice/instanceconfirmation/InstanceValidator.java
index f1a93e58526..2eae26a814d 100644
--- a/athenz-identity-provider-service/src/main/java/com/yahoo/vespa/hosted/athenz/instanceproviderservice/instanceconfirmation/InstanceValidator.java
+++ b/athenz-identity-provider-service/src/main/java/com/yahoo/vespa/hosted/athenz/instanceproviderservice/instanceconfirmation/InstanceValidator.java
@@ -40,6 +40,9 @@ public class InstanceValidator {
     static final String SERVICE_PROPERTIES_SERVICE_KEY = "identity.service";
     static final String INSTANCE_ID_DELIMITER = ".instanceid.athenz.";
 
+    public static final String SAN_IPS_ATTRNAME = "sanIP";
+    public static final String SAN_DNS_ATTRNAME = "sanDNS";
+
     private final IdentityDocumentSigner signer;
     private final KeyProvider keyProvider;
     private final SuperModelProvider superModelProvider;
@@ -96,7 +99,7 @@ public class InstanceValidator {
         log.log(LogLevel.INFO, () -> String.format("Accepting refresh for instance with identity '%s', provider '%s', instanceId '%s'.",
                                                    new AthenzService(confirmation.domain, confirmation.service).getFullName(),
                                                    confirmation.provider,
-                                                   confirmation.attributes.get("sanDNS")));
+                                                   confirmation.attributes.get(SAN_DNS_ATTRNAME)));
         try {
             return validateAttributes(confirmation, getVespaUniqueInstanceId(confirmation));
         } catch (Exception e) {
@@ -107,7 +110,7 @@ public class InstanceValidator {
 
     private VespaUniqueInstanceId getVespaUniqueInstanceId(InstanceConfirmation instanceConfirmation) {
         // Find a list of SAN DNS
-        List<String> sanDNS = Optional.ofNullable(instanceConfirmation.attributes.get("sanDNS"))
+        List<String> sanDNS = Optional.ofNullable(instanceConfirmation.attributes.get(SAN_DNS_ATTRNAME))
                 .map(s -> s.split(","))
                 .map(Arrays::asList)
                 .map(List::stream)
@@ -124,7 +127,7 @@ public class InstanceValidator {
 
     private boolean validateAttributes(InstanceConfirmation confirmation, VespaUniqueInstanceId vespaUniqueInstanceId) {
         if(vespaUniqueInstanceId == null) {
-            log.log(LogLevel.WARNING, "Unabe to find unique instance ID in refresh request: " + confirmation.toString());
+            log.log(LogLevel.WARNING, "Unable to find unique instance ID in refresh request: " + confirmation.toString());
             return false;
         }
 
@@ -140,7 +143,7 @@ public class InstanceValidator {
         }
 
         // Find list of ipaddresses
-        List<InetAddress> ips = Optional.ofNullable(confirmation.attributes.get("sanIP"))
+        List<InetAddress> ips = Optional.ofNullable(confirmation.attributes.get(SAN_IPS_ATTRNAME))
                 .map(s -> s.split(","))
                 .map(Arrays::asList)
                 .map(List::stream)
diff --git a/athenz-identity-provider-service/src/main/java/com/yahoo/vespa/hosted/ca/Certificates.java b/athenz-identity-provider-service/src/main/java/com/yahoo/vespa/hosted/ca/Certificates.java
index 308127e29c7..b5a2405167a 100644
--- a/athenz-identity-provider-service/src/main/java/com/yahoo/vespa/hosted/ca/Certificates.java
+++ b/athenz-identity-provider-service/src/main/java/com/yahoo/vespa/hosted/ca/Certificates.java
@@ -4,14 +4,17 @@ package com.yahoo.vespa.hosted.ca;
 import com.yahoo.security.Pkcs10Csr;
 import com.yahoo.security.SubjectAlternativeName;
 import com.yahoo.security.X509CertificateBuilder;
+import com.yahoo.security.X509CertificateUtils;
 import com.yahoo.vespa.athenz.identityprovider.api.VespaUniqueInstanceId;
 
 import java.security.PrivateKey;
 import java.security.cert.X509Certificate;
 import java.time.Clock;
 import java.time.Duration;
+import java.util.List;
 import java.util.Objects;
 import java.util.Optional;
+import java.util.stream.Collectors;
 
 import static com.yahoo.security.SignatureAlgorithm.SHA256_WITH_ECDSA;
 import static com.yahoo.security.SubjectAlternativeName.Type.DNS_NAME;
@@ -53,14 +56,22 @@ public class Certificates {
 
     /** Returns instance ID parsed from the Subject Alternative Names in given csr */
     public static String instanceIdFrom(Pkcs10Csr csr) {
-        return csr.getSubjectAlternativeNames().stream()
-                  .filter(san -> san.getType() == DNS_NAME)
-                  .map(SubjectAlternativeName::getValue)
-                  .map(Certificates::parseInstanceId)
-                  .flatMap(Optional::stream)
-                  .map(VespaUniqueInstanceId::asDottedString)
-                  .findFirst()
-                  .orElseThrow(() -> new IllegalArgumentException("No instance ID found in CSR"));
+        return getInstanceIdFromSAN(csr.getSubjectAlternativeNames())
+                .orElseThrow(() -> new IllegalArgumentException("No instance ID found in CSR"));
+    }
+
+    public static Optional<String> instanceIdFrom(X509Certificate certificate) {
+        return getInstanceIdFromSAN(X509CertificateUtils.getSubjectAlternativeNames(certificate));
+    }
+
+    private static Optional<String> getInstanceIdFromSAN(List<SubjectAlternativeName> subjectAlternativeNames) {
+        return subjectAlternativeNames.stream()
+                .filter(san -> san.getType() == DNS_NAME)
+                .map(SubjectAlternativeName::getValue)
+                .map(Certificates::parseInstanceId)
+                .flatMap(Optional::stream)
+                .map(VespaUniqueInstanceId::asDottedString)
+                .findFirst();
     }
 
     private static Optional<VespaUniqueInstanceId> parseInstanceId(String dnsName) {
@@ -74,4 +85,11 @@ public class Certificates {
         }
     }
 
+    public static String getSubjectAlternativeNames(Pkcs10Csr csr, SubjectAlternativeName.Type sanType) {
+        return csr.getSubjectAlternativeNames().stream()
+                .map(SubjectAlternativeName::decode)
+                .filter(san -> san.getType() == sanType)
+                .map(SubjectAlternativeName::getValue)
+                .collect(Collectors.joining(","));
+    }
 }
diff --git a/athenz-identity-provider-service/src/main/java/com/yahoo/vespa/hosted/ca/instance/InstanceRegistration.java b/athenz-identity-provider-service/src/main/java/com/yahoo/vespa/hosted/ca/instance/InstanceRegistration.java
index 2a2b702d21b..8d9f44e4564 100644
--- a/athenz-identity-provider-service/src/main/java/com/yahoo/vespa/hosted/ca/instance/InstanceRegistration.java
+++ b/athenz-identity-provider-service/src/main/java/com/yahoo/vespa/hosted/ca/instance/InstanceRegistration.java
@@ -2,6 +2,7 @@
 package com.yahoo.vespa.hosted.ca.instance;
 
 import com.yahoo.security.Pkcs10Csr;
+import com.yahoo.vespa.athenz.identityprovider.api.SignedIdentityDocument;
 
 import java.util.Objects;
 
@@ -16,10 +17,10 @@ public class InstanceRegistration {
     private final String provider;
     private final String domain;
     private final String service;
-    private final String attestationData;
+    private final SignedIdentityDocument attestationData;
     private final Pkcs10Csr csr;
 
-    public InstanceRegistration(String provider, String domain, String service, String attestationData, Pkcs10Csr csr) {
+    public InstanceRegistration(String provider, String domain, String service, SignedIdentityDocument attestationData, Pkcs10Csr csr) {
         this.provider = Objects.requireNonNull(provider, "provider must be non-null");
         this.domain = Objects.requireNonNull(domain, "domain must be non-null");
         this.service = Objects.requireNonNull(service, "service must be non-null");
@@ -43,7 +44,7 @@ public class InstanceRegistration {
     }
 
     /** Host document describing this instance (received from config server) */
-    public String attestationData() {
+    public SignedIdentityDocument attestationData() {
         return attestationData;
     }
 
diff --git a/athenz-identity-provider-service/src/main/java/com/yahoo/vespa/hosted/ca/restapi/CertificateAuthorityApiHandler.java b/athenz-identity-provider-service/src/main/java/com/yahoo/vespa/hosted/ca/restapi/CertificateAuthorityApiHandler.java
index 1b9bdcdb987..4c01b0943e4 100644
--- a/athenz-identity-provider-service/src/main/java/com/yahoo/vespa/hosted/ca/restapi/CertificateAuthorityApiHandler.java
+++ b/athenz-identity-provider-service/src/main/java/com/yahoo/vespa/hosted/ca/restapi/CertificateAuthorityApiHandler.java
@@ -6,16 +6,24 @@ import com.yahoo.container.jdisc.HttpRequest;
 import com.yahoo.container.jdisc.HttpResponse;
 import com.yahoo.container.jdisc.LoggingRequestHandler;
 import com.yahoo.container.jdisc.secretstore.SecretStore;
+import com.yahoo.jdisc.http.servlet.ServletRequest;
+import com.yahoo.log.LogLevel;
 import com.yahoo.restapi.ErrorResponse;
 import com.yahoo.restapi.Path;
 import com.yahoo.restapi.SlimeJsonResponse;
 import com.yahoo.security.KeyUtils;
+import com.yahoo.security.SubjectAlternativeName;
 import com.yahoo.security.X509CertificateUtils;
 import com.yahoo.slime.Slime;
+import com.yahoo.vespa.athenz.api.AthenzService;
+import com.yahoo.vespa.athenz.identityprovider.api.EntityBindingsMapper;
 import com.yahoo.vespa.config.SlimeUtils;
 import com.yahoo.vespa.hosted.athenz.instanceproviderservice.config.AthenzProviderServiceConfig;
+import com.yahoo.vespa.hosted.athenz.instanceproviderservice.instanceconfirmation.InstanceConfirmation;
+import com.yahoo.vespa.hosted.athenz.instanceproviderservice.instanceconfirmation.InstanceValidator;
 import com.yahoo.vespa.hosted.ca.Certificates;
 import com.yahoo.vespa.hosted.ca.instance.InstanceIdentity;
+import com.yahoo.vespa.hosted.ca.instance.InstanceRefresh;
 import com.yahoo.yolean.Exceptions;
 
 import java.io.IOException;
@@ -23,6 +31,10 @@ import java.io.UncheckedIOException;
 import java.security.PrivateKey;
 import java.security.cert.X509Certificate;
 import java.time.Clock;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
+import java.util.Objects;
 import java.util.Optional;
 import java.util.function.Function;
 import java.util.logging.Level;
@@ -43,18 +55,20 @@ public class CertificateAuthorityApiHandler extends LoggingRequestHandler {
     private final Certificates certificates;
     private final String caPrivateKeySecretName;
     private final String caCertificateSecretName;
+    private final InstanceValidator instanceValidator;
 
     @Inject
-    public CertificateAuthorityApiHandler(Context ctx, SecretStore secretStore, AthenzProviderServiceConfig athenzProviderServiceConfig) {
-        this(ctx, secretStore, new Certificates(Clock.systemUTC()), athenzProviderServiceConfig);
+    public CertificateAuthorityApiHandler(Context ctx, SecretStore secretStore, AthenzProviderServiceConfig athenzProviderServiceConfig, InstanceValidator instanceValidator) {
+        this(ctx, secretStore, new Certificates(Clock.systemUTC()), athenzProviderServiceConfig, instanceValidator);
     }
 
-    CertificateAuthorityApiHandler(Context ctx, SecretStore secretStore, Certificates certificates, AthenzProviderServiceConfig athenzProviderServiceConfig) {
+    CertificateAuthorityApiHandler(Context ctx, SecretStore secretStore, Certificates certificates, AthenzProviderServiceConfig athenzProviderServiceConfig, InstanceValidator instanceValidator) {
         super(ctx);
         this.secretStore = secretStore;
         this.certificates = certificates;
         this.caPrivateKeySecretName = athenzProviderServiceConfig.secretName();
         this.caCertificateSecretName = athenzProviderServiceConfig.domain() + ".ca.cert";
+        this.instanceValidator = instanceValidator;
     }
 
     @Override
@@ -81,6 +95,14 @@ public class CertificateAuthorityApiHandler extends LoggingRequestHandler {
 
     private HttpResponse registerInstance(HttpRequest request) {
         var instanceRegistration = deserializeRequest(request, InstanceSerializer::registrationFromSlime);
+
+        InstanceConfirmation confirmation = new InstanceConfirmation(instanceRegistration.provider(), instanceRegistration.domain(), instanceRegistration.service(), EntityBindingsMapper.toSignedIdentityDocumentEntity(instanceRegistration.attestationData()));
+        confirmation.set(InstanceValidator.SAN_IPS_ATTRNAME, Certificates.getSubjectAlternativeNames(instanceRegistration.csr(), SubjectAlternativeName.Type.IP_ADDRESS));
+        confirmation.set(InstanceValidator.SAN_DNS_ATTRNAME, Certificates.getSubjectAlternativeNames(instanceRegistration.csr(), SubjectAlternativeName.Type.DNS_NAME));
+        if (!instanceValidator.isValidInstance(confirmation)) {
+            log.log(LogLevel.INFO, "Invalid instance registration for " + instanceRegistration.toString());
+            return ErrorResponse.forbidden("Unable to launch service: " +instanceRegistration.service());
+        }
         var certificate = certificates.create(instanceRegistration.csr(), caCertificate(), caPrivateKey());
         var instanceId = Certificates.instanceIdFrom(instanceRegistration.csr());
         var identity = new InstanceIdentity(instanceRegistration.provider(), instanceRegistration.service(), instanceId,
@@ -91,21 +113,65 @@ public class CertificateAuthorityApiHandler extends LoggingRequestHandler {
     private HttpResponse refreshInstance(HttpRequest request, String provider, String service, String instanceId) {
         var instanceRefresh = deserializeRequest(request, InstanceSerializer::refreshFromSlime);
         var instanceIdFromCsr = Certificates.instanceIdFrom(instanceRefresh.csr());
+        var athenzService = new AthenzService(request.getJDiscRequest().getUserPrincipal().getName());
         if (!instanceIdFromCsr.equals(instanceId)) {
             throw new IllegalArgumentException("Mismatch between instance ID in URL path and instance ID in CSR " +
                                                "[instanceId=" + instanceId + ",instanceIdFromCsr=" + instanceIdFromCsr +
                                                "]");
         }
+
+        // Verify that the csr instance id matches one of the certificates in the chain
+        refreshesSameInstanceId(instanceIdFromCsr, request);
+
+
+        // Validate that there is no privilege escalation (can only refresh same service)
+        refreshesSameService(instanceRefresh, athenzService);
+
+        InstanceConfirmation instanceConfirmation = new InstanceConfirmation(provider, athenzService.getDomain().getName(), athenzService.getName(), null);
+        instanceConfirmation.set(InstanceValidator.SAN_IPS_ATTRNAME, Certificates.getSubjectAlternativeNames(instanceRefresh.csr(), SubjectAlternativeName.Type.IP_ADDRESS));
+        instanceConfirmation.set(InstanceValidator.SAN_DNS_ATTRNAME, Certificates.getSubjectAlternativeNames(instanceRefresh.csr(), SubjectAlternativeName.Type.DNS_NAME));
+        if(!instanceValidator.isValidRefresh(instanceConfirmation)) {
+            return ErrorResponse.forbidden("Unable to refresh cert: " + instanceRefresh.csr().getSubject().toString());
+        }
+
         var certificate = certificates.create(instanceRefresh.csr(), caCertificate(), caPrivateKey());
         var identity = new InstanceIdentity(provider, service, instanceIdFromCsr, Optional.of(certificate));
         return new SlimeJsonResponse(InstanceSerializer.identityToSlime(identity));
     }
 
+    public void refreshesSameInstanceId(String csrInstanceId, HttpRequest request) {
+        String certificateInstanceId = getRequestCertificateChain(request).stream()
+                .map(Certificates::instanceIdFrom)
+                .filter(Optional::isPresent)
+                .map(Optional::get)
+                .findAny().orElseThrow(() -> new IllegalArgumentException("No client certificate with instance id in request."));
+
+        if(! Objects.equals(certificateInstanceId, csrInstanceId)) {
+            throw new IllegalArgumentException("Mismatch between instance ID in client certificate and instance ID in CSR " +
+                                               "[instanceId=" + certificateInstanceId + ",instanceIdFromCsr=" + csrInstanceId +
+                                               "]");
+        }
+    }
+
+    private void refreshesSameService(InstanceRefresh instanceRefresh, AthenzService athenzService) {
+        List<String> commonNames = X509CertificateUtils.getCommonNames(instanceRefresh.csr().getSubject());
+        if(commonNames.size() != 1 && !Objects.equals(commonNames.get(0), athenzService.getFullName())) {
+            throw new IllegalArgumentException(String.format("Invalid request, trying to refresh service %s using service %s.", instanceRefresh.csr().getSubject().getName(), athenzService.getFullName()));
+        }
+    }
+
     /** Returns CA certificate from secret store */
     private X509Certificate caCertificate() {
         return X509CertificateUtils.fromPem(secretStore.getSecret(caCertificateSecretName));
     }
 
+    private List<X509Certificate> getRequestCertificateChain(HttpRequest request) {
+        return Optional.ofNullable(request.getJDiscRequest().context().get(ServletRequest.JDISC_REQUEST_X509CERT))
+                .map(X509Certificate[].class::cast)
+                .map(Arrays::asList)
+                .orElse(Collections.emptyList());
+    }
+
     /** Returns CA private key from secret store */
     private PrivateKey caPrivateKey() {
         return KeyUtils.fromPemEncodedPrivateKey(secretStore.getSecret(caPrivateKeySecretName));
diff --git a/athenz-identity-provider-service/src/main/java/com/yahoo/vespa/hosted/ca/restapi/InstanceSerializer.java b/athenz-identity-provider-service/src/main/java/com/yahoo/vespa/hosted/ca/restapi/InstanceSerializer.java
index a2537cd68f1..f08e5e37417 100644
--- a/athenz-identity-provider-service/src/main/java/com/yahoo/vespa/hosted/ca/restapi/InstanceSerializer.java
+++ b/athenz-identity-provider-service/src/main/java/com/yahoo/vespa/hosted/ca/restapi/InstanceSerializer.java
@@ -3,12 +3,23 @@ package com.yahoo.vespa.hosted.ca.restapi;
 
 import com.yahoo.security.Pkcs10CsrUtils;
 import com.yahoo.security.X509CertificateUtils;
+import com.yahoo.slime.ArrayTraverser;
 import com.yahoo.slime.Cursor;
 import com.yahoo.slime.Slime;
+import com.yahoo.text.StringUtilities;
+import com.yahoo.vespa.athenz.api.AthenzService;
+import com.yahoo.vespa.athenz.identityprovider.api.IdentityType;
+import com.yahoo.vespa.athenz.identityprovider.api.SignedIdentityDocument;
+import com.yahoo.vespa.athenz.identityprovider.api.VespaUniqueInstanceId;
+import com.yahoo.vespa.config.SlimeUtils;
 import com.yahoo.vespa.hosted.ca.instance.InstanceIdentity;
 import com.yahoo.vespa.hosted.ca.instance.InstanceRefresh;
 import com.yahoo.vespa.hosted.ca.instance.InstanceRegistration;
 
+import java.time.Instant;
+import java.util.HashSet;
+import java.util.Set;
+
 /**
  * @author mpolden
  */
@@ -23,6 +34,17 @@ public class InstanceSerializer {
     private static final String INSTANCE_ID_FIELD = "instanceId";
     private static final String X509_CERTIFICATE_FIELD = "x509Certificate";
 
+    private static final String IDD_SIGNATURE_FIELD = "signature";
+    private static final String IDD_SIGNING_KEY_VERSION_FIELD = "signing-key-version";
+    private static final String IDD_PROVIDER_UNIQUE_ID_FIELD = "provider-unique-id";
+    private static final String IDD_PROVIDER_SERVICE_FIELD = "provider-service";
+    private static final String IDD_DOCUMENT_VERSION_FIELD = "document-version";
+    private static final String IDD_CONFIGSERVER_HOSTNAME_FIELD = "configserver-hostname";
+    private static final String IDD_INSTANCE_HOSTNAME_FIELD = "instance-hostname";
+    private static final String IDD_CREATED_AT_FIELD = "created-at";
+    private static final String IDD_IPADDRESSES_FIELD = "ip-addresses";
+    private static final String IDD_IDENTITY_TYPE_FIELD = "identity-type";
+
     private InstanceSerializer() {}
 
     public static InstanceRegistration registrationFromSlime(Slime slime) {
@@ -30,7 +52,7 @@ public class InstanceSerializer {
         return new InstanceRegistration(requireField(PROVIDER_FIELD, root).asString(),
                                         requireField(DOMAIN_FIELD, root).asString(),
                                         requireField(SERVICE_FIELD, root).asString(),
-                                        requireField(ATTESTATION_DATA_FIELD, root).asString(),
+                                        attestationDataToIdentityDocument(StringUtilities.unescape(requireField(ATTESTATION_DATA_FIELD, root).asString())),
                                         Pkcs10CsrUtils.fromPem(requireField(CSR_FIELD, root).asString()));
     }
 
@@ -51,6 +73,26 @@ public class InstanceSerializer {
         return slime;
     }
 
+    private static SignedIdentityDocument attestationDataToIdentityDocument(String attestationData) {
+        Slime slime = SlimeUtils.jsonToSlime(attestationData);
+        Cursor root = slime.get();
+        String signature = requireField(IDD_SIGNATURE_FIELD, root).asString();
+        long signingKeyVersion = requireField(IDD_SIGNING_KEY_VERSION_FIELD, root).asLong();
+        VespaUniqueInstanceId providerUniqueId = VespaUniqueInstanceId.fromDottedString(requireField(IDD_PROVIDER_UNIQUE_ID_FIELD, root).asString());
+        AthenzService athenzService = new AthenzService(requireField(IDD_PROVIDER_SERVICE_FIELD, root).asString());
+        long documentVersion = requireField(IDD_DOCUMENT_VERSION_FIELD, root).asLong();
+        String configserverHostname = requireField(IDD_CONFIGSERVER_HOSTNAME_FIELD, root).asString();
+        String instanceHostname = requireField(IDD_INSTANCE_HOSTNAME_FIELD, root).asString();
+        double createdAtTimestamp = requireField(IDD_CREATED_AT_FIELD, root).asDouble();
+        Instant createdAt = Instant.ofEpochSecond(Double.valueOf(createdAtTimestamp).longValue());
+        Set<String> ips = new HashSet<>();
+        requireField(IDD_IPADDRESSES_FIELD, root).traverse((ArrayTraverser)  (__, entry) -> ips.add(entry.asString()));
+        IdentityType identityType = IdentityType.fromId(requireField(IDD_IDENTITY_TYPE_FIELD, root).asString());
+
+        return new SignedIdentityDocument(signature, (int)signingKeyVersion, providerUniqueId, athenzService, (int)documentVersion,
+                                          configserverHostname, instanceHostname, createdAt, ips, identityType);
+    }
+
     private static Cursor requireField(String fieldName, Cursor root) {
         var field = root.field(fieldName);
         if (!field.valid()) throw new IllegalArgumentException("Missing required field '" + fieldName + "'");
diff --git a/athenz-identity-provider-service/src/test/java/com/yahoo/vespa/hosted/ca/restapi/CertificateAuthorityApiTest.java b/athenz-identity-provider-service/src/test/java/com/yahoo/vespa/hosted/ca/restapi/CertificateAuthorityApiTest.java
index e377009b18c..88ee154dee8 100644
--- a/athenz-identity-provider-service/src/test/java/com/yahoo/vespa/hosted/ca/restapi/CertificateAuthorityApiTest.java
+++ b/athenz-identity-provider-service/src/test/java/com/yahoo/vespa/hosted/ca/restapi/CertificateAuthorityApiTest.java
@@ -2,25 +2,34 @@
 package com.yahoo.vespa.hosted.ca.restapi;
 
 import com.yahoo.application.container.handler.Request;
+import com.yahoo.jdisc.http.servlet.ServletRequest;
 import com.yahoo.security.KeyAlgorithm;
 import com.yahoo.security.KeyUtils;
 import com.yahoo.security.Pkcs10Csr;
 import com.yahoo.security.Pkcs10CsrUtils;
 import com.yahoo.security.X509CertificateUtils;
+import com.yahoo.text.StringUtilities;
+import com.yahoo.vespa.athenz.api.AthenzPrincipal;
 import com.yahoo.vespa.athenz.api.AthenzService;
 import com.yahoo.vespa.athenz.client.zts.DefaultZtsClient;
 import com.yahoo.vespa.config.SlimeUtils;
 import com.yahoo.vespa.hosted.ca.CertificateTester;
+import org.apache.http.client.ResponseHandler;
+import org.apache.http.client.methods.HttpUriRequest;
 import org.junit.Before;
 import org.junit.Test;
 
 import javax.net.ssl.SSLContext;
 import java.net.URI;
 import java.nio.charset.StandardCharsets;
+import java.security.Principal;
+import java.security.cert.X509Certificate;
 import java.util.List;
+import java.util.Optional;
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertTrue;
+import static org.junit.Assert.fail;
 
 /**
  * @author mpolden
@@ -29,6 +38,8 @@ public class CertificateAuthorityApiTest extends ContainerTester {
 
     private static final String INSTANCE_ID = "1.cluster1.default.app1.tenant1.us-north-1.prod.node";
     private static final String INSTANCE_ID_WITH_SUFFIX = INSTANCE_ID + ".instanceid.athenz.dev-us-north-1.vespa.aws.oath.cloud";
+    private static final String INVALID_INSTANCE_ID = "1.cluster1.default.otherapp.othertenant.us-north-1.prod.node";
+    private static final String INVALID_INSTANCE_ID_WITH_SUFFIX = INVALID_INSTANCE_ID + ".instanceid.athenz.dev-us-north-1.vespa.aws.oath.cloud";
 
     @Before
     public void before() {
@@ -44,24 +55,47 @@ public class CertificateAuthorityApiTest extends ContainerTester {
                                            Request.Method.POST));
 
         // POST instance registration with ZTS client
-        var ztsClient = new DefaultZtsClient(URI.create("http://localhost:12345/ca/v1/"), SSLContext.getDefault());
+        var ztsClient = new TestZtsClient(new AthenzPrincipal(new AthenzService("vespa.external.tenant-host")), null, URI.create("http://localhost:12345/ca/v1/"), SSLContext.getDefault());
         var instanceIdentity = ztsClient.registerInstance(new AthenzService("vespa.external", "provider_prod_us-north-1"),
                                                           new AthenzService("vespa.external", "tenant"),
-                                                          "identity document generated by config server",
+                                                          getAttestationData(),
                                                           csr);
         assertEquals("CN=Vespa CA", instanceIdentity.certificate().getIssuerX500Principal().getName());
     }
 
+    private X509Certificate registerInstance() throws Exception {
+        // POST instance registration
+        var csr = CertificateTester.createCsr(List.of("node1.example.com", INSTANCE_ID_WITH_SUFFIX));
+        assertIdentityResponse(new Request("http://localhost:12345/ca/v1/instance/",
+                                           instanceRegistrationJson(csr),
+                                           Request.Method.POST));
+
+        // POST instance registration with ZTS client
+        var ztsClient = new TestZtsClient(new AthenzPrincipal(new AthenzService("vespa.external.tenant-host")), null, URI.create("http://localhost:12345/ca/v1/"), SSLContext.getDefault());
+        var instanceIdentity = ztsClient.registerInstance(new AthenzService("vespa.external", "provider_prod_us-north-1"),
+                                                          new AthenzService("vespa.external", "tenant"),
+                                                          getAttestationData(),
+                                                          csr);
+        return instanceIdentity.certificate();
+    }
+
     @Test
     public void refresh_instance() throws Exception {
+        // Register instance to get cert
+        var certificate = registerInstance();
+
         // POST instance refresh
         var csr = CertificateTester.createCsr(List.of("node1.example.com", INSTANCE_ID_WITH_SUFFIX));
-        assertIdentityResponse(new Request("http://localhost:12345/ca/v1/instance/vespa.external.provider_prod_us-north-1/vespa.external/tenant/" + INSTANCE_ID,
+        var principal = new AthenzPrincipal(new AthenzService("vespa.external.tenant"));
+        var request = new Request("http://localhost:12345/ca/v1/instance/vespa.external.provider_prod_us-north-1/vespa.external/tenant/" + INSTANCE_ID,
                                            instanceRefreshJson(csr),
-                                           Request.Method.POST));
+                                           Request.Method.POST,
+                                           principal);
+        request.getAttributes().put(ServletRequest.JDISC_REQUEST_X509CERT, new X509Certificate[]{certificate});
+        assertIdentityResponse(request);
 
         // POST instance refresh with ZTS client
-        var ztsClient = new DefaultZtsClient(URI.create("http://localhost:12345/ca/v1/"), SSLContext.getDefault());
+        var ztsClient = new TestZtsClient(principal, certificate, URI.create("http://localhost:12345/ca/v1/"), SSLContext.getDefault());
         var instanceIdentity = ztsClient.refreshInstance(new AthenzService("vespa.external", "provider_prod_us-north-1"),
                                                           new AthenzService("vespa.external", "tenant"),
                                                           INSTANCE_ID,
@@ -70,7 +104,7 @@ public class CertificateAuthorityApiTest extends ContainerTester {
     }
 
     @Test
-    public void invalid_requests() {
+    public void invalid_requests() throws Exception {
         // POST instance registration with missing fields
         assertResponse(400, "{\"error-code\":\"BAD_REQUEST\",\"message\":\"POST http://localhost:12345/ca/v1/instance/ failed: Missing required field 'provider'\"}",
                        new Request("http://localhost:12345/ca/v1/instance/",
@@ -91,11 +125,28 @@ public class CertificateAuthorityApiTest extends ContainerTester {
                                    Request.Method.POST));
 
         // POST instance refresh where instanceId does not match CSR dnsName
+        var principal = new AthenzPrincipal(new AthenzService("vespa.external.tenant"));
         csr = CertificateTester.createCsr(List.of("node1.example.com", INSTANCE_ID_WITH_SUFFIX));
         assertResponse(400, "{\"error-code\":\"BAD_REQUEST\",\"message\":\"POST http://localhost:12345/ca/v1/instance/vespa.external.provider_prod_us-north-1/vespa.external/tenant/foobar failed: Mismatch between instance ID in URL path and instance ID in CSR [instanceId=foobar,instanceIdFromCsr=1.cluster1.default.app1.tenant1.us-north-1.prod.node]\"}",
                        new Request("http://localhost:12345/ca/v1/instance/vespa.external.provider_prod_us-north-1/vespa.external/tenant/foobar",
                                    instanceRefreshJson(csr),
-                                   Request.Method.POST));
+                                   Request.Method.POST,
+                                   principal));
+
+        // POST instance refresh using zts client where client cert does not contain instanceid
+        var certificate = registerInstance();
+        var ztsClient = new TestZtsClient(principal, certificate, URI.create("http://localhost:12345/ca/v1/"), SSLContext.getDefault());
+        try {
+            var invalidCsr = CertificateTester.createCsr(List.of("node1.example.com", INVALID_INSTANCE_ID_WITH_SUFFIX));
+            var instanceIdentity = ztsClient.refreshInstance(new AthenzService("vespa.external", "provider_prod_us-north-1"),
+                                                             new AthenzService("vespa.external", "tenant"),
+                                                             INSTANCE_ID,
+                                                             invalidCsr);
+            fail("Refresh instance should have failed");
+        } catch (Exception e) {
+            String expectedMessage = "Received error from ZTS: code=0, message=\"POST http://localhost:12345/ca/v1/instance/vespa.external.provider_prod_us-north-1/vespa.external/tenant/1.cluster1.default.app1.tenant1.us-north-1.prod.node failed: Mismatch between instance ID in URL path and instance ID in CSR [instanceId=1.cluster1.default.app1.tenant1.us-north-1.prod.node,instanceIdFromCsr=1.cluster1.default.otherapp.othertenant.us-north-1.prod.node]\"";
+            assertEquals(expectedMessage, e.getMessage());
+        }
     }
 
     private void setCaCertificateAndKey() {
@@ -132,10 +183,52 @@ public class CertificateAuthorityApiTest extends ContainerTester {
                "  \"provider\": \"vespa.external.provider_prod_us-north-1\",\n" +
                "  \"domain\": \"vespa.external\",\n" +
                "  \"service\": \"tenant\",\n" +
-               "  \"attestationData\": \"identity document generated by config server\",\n" +
+               "  \"attestationData\": \""+getAttestationData()+"\",\n" +
                "  \"csr\": \"" + csrPem + "\"\n" +
                "}";
         return json.getBytes(StandardCharsets.UTF_8);
     }
 
+    private static String getAttestationData () {
+        var json = "{\n" +
+                   "  \"signature\": \"SIGNATURE\",\n" +
+                   "  \"signing-key-version\": 0,\n" +
+                   "  \"provider-unique-id\": \"0.default.default.application.tenant.us-north-1.dev.tenant\",\n" +
+                   "  \"provider-service\": \"domain.service\",\n" +
+                   "  \"document-version\": 1,\n" +
+                   "  \"configserver-hostname\": \"localhost\",\n" +
+                   "  \"instance-hostname\": \"docker-container\",\n" +
+                   "  \"created-at\": 1572000079.00000,\n" +
+                   "  \"ip-addresses\": [\n" +
+                   "    \"::1\"\n" +
+                   "  ],\n" +
+                   "  \"identity-type\": \"tenant\"\n" +
+                   "}";
+        return StringUtilities.escape(json);
+    }
+
+    /*
+    Zts client that adds principal as header (since setting up ssl in test is cumbersome)
+     */
+    private static class TestZtsClient extends DefaultZtsClient {
+
+        private final Principal principal;
+        private final X509Certificate certificate;
+
+        public TestZtsClient(Principal principal, X509Certificate certificate, URI ztsUrl, SSLContext sslContext) {
+            super(ztsUrl, sslContext);
+            this.principal = principal;
+            this.certificate = certificate;
+        }
+
+        @Override
+        protected <T> T execute(HttpUriRequest request, ResponseHandler<T> responseHandler) {
+            request.addHeader("PRINCIPAL", principal.getName());
+            Optional.ofNullable(certificate).ifPresent(cert -> {
+                var pem = X509CertificateUtils.toPem(certificate);
+                request.addHeader("CERTIFICATE", StringUtilities.escape(pem));
+            });
+            return super.execute(request, responseHandler);
+        }
+    }
 }
diff --git a/athenz-identity-provider-service/src/test/java/com/yahoo/vespa/hosted/ca/restapi/ContainerTester.java b/athenz-identity-provider-service/src/test/java/com/yahoo/vespa/hosted/ca/restapi/ContainerTester.java
index 139314b0f86..0eda6bd946b 100644
--- a/athenz-identity-provider-service/src/test/java/com/yahoo/vespa/hosted/ca/restapi/ContainerTester.java
+++ b/athenz-identity-provider-service/src/test/java/com/yahoo/vespa/hosted/ca/restapi/ContainerTester.java
@@ -67,13 +67,20 @@ public class ContainerTester {
                "    <ztsUrl>https://localhost:123/</ztsUrl>\n" +
                "  </config>\n" +
                "  <component id='com.yahoo.vespa.hosted.ca.restapi.mock.SecretStoreMock'/>\n" +
+               "  <component id='com.yahoo.vespa.hosted.ca.restapi.mock.InstanceValidatorMock'/>\n" +
                "  <handler id='com.yahoo.vespa.hosted.ca.restapi.CertificateAuthorityApiHandler'>\n" +
                "    <binding>http://*/ca/v1/*</binding>\n" +
                "  </handler>\n" +
                "  <http>\n" +
                "    <server id='default' port='12345'/>\n" +
+               "    <filtering>\n" +
+               "      <request-chain id=\"my-default-chain\">\n" +
+               "        <filter id='com.yahoo.vespa.hosted.ca.restapi.mock.PrincipalFromHeaderFilter' />\n" +
+               "        <binding>http://*/*</binding>\n" +
+               "      </request-chain>\n" +
+               "    </filtering>\n" +
                "  </http>\n" +
                "</container>";
     }
 
-}
+}
\ No newline at end of file
diff --git a/athenz-identity-provider-service/src/test/java/com/yahoo/vespa/hosted/ca/restapi/InstanceSerializerTest.java b/athenz-identity-provider-service/src/test/java/com/yahoo/vespa/hosted/ca/restapi/InstanceSerializerTest.java
index 83ea9249ad0..545d7d9eab7 100644
--- a/athenz-identity-provider-service/src/test/java/com/yahoo/vespa/hosted/ca/restapi/InstanceSerializerTest.java
+++ b/athenz-identity-provider-service/src/test/java/com/yahoo/vespa/hosted/ca/restapi/InstanceSerializerTest.java
@@ -4,6 +4,12 @@ package com.yahoo.vespa.hosted.ca.restapi;
 import com.yahoo.security.Pkcs10CsrUtils;
 import com.yahoo.security.X509CertificateUtils;
 import com.yahoo.slime.Slime;
+import com.yahoo.text.StringUtilities;
+import com.yahoo.vespa.athenz.api.AthenzService;
+import com.yahoo.vespa.athenz.identityprovider.api.EntityBindingsMapper;
+import com.yahoo.vespa.athenz.identityprovider.api.IdentityType;
+import com.yahoo.vespa.athenz.identityprovider.api.SignedIdentityDocument;
+import com.yahoo.vespa.athenz.identityprovider.api.VespaUniqueInstanceId;
 import com.yahoo.vespa.config.SlimeUtils;
 import com.yahoo.vespa.hosted.ca.CertificateTester;
 import com.yahoo.vespa.hosted.ca.instance.InstanceIdentity;
@@ -14,6 +20,8 @@ import org.junit.Test;
 import java.io.IOException;
 import java.io.UncheckedIOException;
 import java.nio.charset.StandardCharsets;
+import java.time.Instant;
+import java.util.Collections;
 import java.util.Optional;
 
 import static org.junit.Assert.assertEquals;
@@ -27,15 +35,27 @@ public class InstanceSerializerTest {
     public void deserialize_instance_registration() {
         var csr = CertificateTester.createCsr();
         var csrPem = Pkcs10CsrUtils.toPem(csr);
-        var json = "{\n" +
+        SignedIdentityDocument signedIdentityDocument = new SignedIdentityDocument(
+                "signature",
+                0,
+                new VespaUniqueInstanceId(0, "cluster", "instance", "application", "tenant", "region", "prod", IdentityType.NODE),
+                new AthenzService("domain", "service"),
+                0,
+                "configserverhostname",
+                "instancehostname",
+                Instant.ofEpochSecond(1572000079),
+                Collections.emptySet(),
+                IdentityType.NODE);
+
+        var json = String.format("{\n" +
                    "  \"provider\": \"provider_prod_us-north-1\",\n" +
                    "  \"domain\": \"vespa.external\",\n" +
                    "  \"service\": \"tenant\",\n" +
-                   "  \"attestationData\": \"identity document from configserevr\",\n" +
+                   "  \"attestationData\":\"%s\",\n" +
                    "  \"csr\": \"" + csrPem + "\"\n" +
-                   "}";
+                   "}", StringUtilities.escape(EntityBindingsMapper.toAttestationData(signedIdentityDocument)));
         var instanceRegistration = new InstanceRegistration("provider_prod_us-north-1", "vespa.external",
-                                                            "tenant", "identity document from configserevr",
+                                                            "tenant", signedIdentityDocument,
                                                             csr);
         var deserialized = InstanceSerializer.registrationFromSlime(SlimeUtils.jsonToSlime(json));
         assertEquals(instanceRegistration, deserialized);
diff --git a/athenz-identity-provider-service/src/test/java/com/yahoo/vespa/hosted/ca/restapi/mock/InstanceValidatorMock.java b/athenz-identity-provider-service/src/test/java/com/yahoo/vespa/hosted/ca/restapi/mock/InstanceValidatorMock.java
new file mode 100644
index 00000000000..16f489c818f
--- /dev/null
+++ b/athenz-identity-provider-service/src/test/java/com/yahoo/vespa/hosted/ca/restapi/mock/InstanceValidatorMock.java
@@ -0,0 +1,27 @@
+// Copyright 2019 Oath Inc. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+package com.yahoo.vespa.hosted.ca.restapi.mock;
+
+import com.yahoo.vespa.hosted.athenz.instanceproviderservice.instanceconfirmation.InstanceConfirmation;
+import com.yahoo.vespa.hosted.athenz.instanceproviderservice.instanceconfirmation.InstanceValidator;
+
+/**
+ * @author mortent
+ */
+public class InstanceValidatorMock extends InstanceValidator {
+
+    public InstanceValidatorMock() {
+        super(null, null, null, null);
+    }
+
+    @Override
+    public boolean isValidInstance(InstanceConfirmation instanceConfirmation) {
+        return instanceConfirmation.attributes.get(SAN_DNS_ATTRNAME) != null &&
+               instanceConfirmation.attributes.get(SAN_IPS_ATTRNAME) != null;
+    }
+
+    @Override
+    public boolean isValidRefresh(InstanceConfirmation confirmation) {
+        return confirmation.attributes.get(SAN_DNS_ATTRNAME) != null &&
+               confirmation.attributes.get(SAN_IPS_ATTRNAME) != null;
+    }
+}
diff --git a/athenz-identity-provider-service/src/test/java/com/yahoo/vespa/hosted/ca/restapi/mock/PrincipalFromHeaderFilter.java b/athenz-identity-provider-service/src/test/java/com/yahoo/vespa/hosted/ca/restapi/mock/PrincipalFromHeaderFilter.java
new file mode 100644
index 00000000000..d9ee4c8bb9b
--- /dev/null
+++ b/athenz-identity-provider-service/src/test/java/com/yahoo/vespa/hosted/ca/restapi/mock/PrincipalFromHeaderFilter.java
@@ -0,0 +1,34 @@
+// Copyright 2019 Oath Inc. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+package com.yahoo.vespa.hosted.ca.restapi.mock;
+
+import com.yahoo.jdisc.handler.ResponseHandler;
+import com.yahoo.jdisc.http.filter.DiscFilterRequest;
+import com.yahoo.jdisc.http.filter.SecurityRequestFilter;
+import com.yahoo.jdisc.http.servlet.ServletRequest;
+import com.yahoo.security.X509CertificateUtils;
+import com.yahoo.text.StringUtilities;
+import com.yahoo.vespa.athenz.api.AthenzPrincipal;
+import com.yahoo.vespa.athenz.api.AthenzService;
+
+import java.security.cert.X509Certificate;
+import java.util.Optional;
+
+/**
+ * Read principal from http header
+ *
+ * @author mortent
+ */
+public class PrincipalFromHeaderFilter implements SecurityRequestFilter {
+
+    @Override
+    public void filter(DiscFilterRequest request, ResponseHandler handler) {
+        String principal = request.getHeader("PRINCIPAL");
+        request.setUserPrincipal(new AthenzPrincipal(new AthenzService(principal)));
+
+        Optional<String> certificate = Optional.ofNullable(request.getHeader("CERTIFICATE"));
+        certificate.ifPresent(cert -> {
+            var x509cert = X509CertificateUtils.fromPem(StringUtilities.unescape(cert));
+            request.setAttribute(ServletRequest.JDISC_REQUEST_X509CERT, new X509Certificate[]{x509cert});
+        });
+    }
+}