Allow applications to re-use SSLContext returned by provider

- Use MutableX509KeyManager instead of re-creating SSLContext when credentials are updated.
- Make AthenzIdentityProviderImpl responsible for creating initial SSLContext.
- Use proper x509 certificates and keys in unit test instead of Mockito mocks.

4 files changed, 108 insertions, 55 deletions

diff --git a/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/identityprovider/client/AthenzCredentials.java b/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/identityprovider/client/AthenzCredentials.java
index a1d8a9ca258..d4494c1bd26 100644
--- a/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/identityprovider/client/AthenzCredentials.java
+++ b/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/identityprovider/client/AthenzCredentials.java
@@ -3,7 +3,6 @@ package com.yahoo.vespa.athenz.identityprovider.client;
 
 import com.yahoo.vespa.athenz.identityprovider.api.SignedIdentityDocument;
 
-import javax.net.ssl.SSLContext;
 import java.security.KeyPair;
 import java.security.cert.X509Certificate;
 
@@ -15,16 +14,13 @@ class AthenzCredentials {
     private final X509Certificate certificate;
     private final KeyPair keyPair;
     private final SignedIdentityDocument identityDocument;
-    private final SSLContext identitySslContext;
 
     AthenzCredentials(X509Certificate certificate,
                       KeyPair keyPair,
-                      SignedIdentityDocument identityDocument,
-                      SSLContext identitySslContext) {
+                      SignedIdentityDocument identityDocument) {
         this.certificate = certificate;
         this.keyPair = keyPair;
         this.identityDocument = identityDocument;
-        this.identitySslContext = identitySslContext;
     }
 
     X509Certificate getCertificate() {
@@ -39,7 +35,4 @@ class AthenzCredentials {
         return identityDocument;
     }
 
-    SSLContext getIdentitySslContext() {
-        return identitySslContext;
-    }
 }
diff --git a/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/identityprovider/client/AthenzCredentialsService.java b/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/identityprovider/client/AthenzCredentialsService.java
index 39d0db4affd..9e2d8bc548c 100644
--- a/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/identityprovider/client/AthenzCredentialsService.java
+++ b/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/identityprovider/client/AthenzCredentialsService.java
@@ -4,7 +4,7 @@ package com.yahoo.vespa.athenz.identityprovider.client;
 import com.yahoo.container.core.identity.IdentityConfig;
 import com.yahoo.security.KeyAlgorithm;
 import com.yahoo.security.KeyUtils;
-import com.yahoo.security.SslContextBuilder;
+import com.yahoo.security.Pkcs10Csr;
 import com.yahoo.vespa.athenz.api.AthenzService;
 import com.yahoo.vespa.athenz.client.zts.DefaultZtsClient;
 import com.yahoo.vespa.athenz.client.zts.InstanceIdentity;
@@ -14,12 +14,10 @@ import com.yahoo.vespa.athenz.identityprovider.api.EntityBindingsMapper;
 import com.yahoo.vespa.athenz.identityprovider.api.IdentityDocumentClient;
 import com.yahoo.vespa.athenz.identityprovider.api.SignedIdentityDocument;
 import com.yahoo.vespa.athenz.tls.AthenzIdentityVerifier;
-import com.yahoo.security.Pkcs10Csr;
 import com.yahoo.vespa.athenz.utils.SiaUtils;
 import com.yahoo.vespa.defaults.Defaults;
 
 import javax.net.ssl.SSLContext;
-import java.io.File;
 import java.net.URI;
 import java.nio.file.Files;
 import java.nio.file.Path;
@@ -31,7 +29,6 @@ import java.time.Clock;
 import java.time.Duration;
 import java.util.Optional;
 
-import static com.yahoo.security.KeyStoreType.JKS;
 import static java.util.Collections.singleton;
 
 /**
@@ -49,14 +46,12 @@ class AthenzCredentialsService {
     private final URI ztsEndpoint;
     private final AthenzService configserverIdentity;
     private final ServiceIdentityProvider nodeIdentityProvider;
-    private final File trustStoreJks;
     private final String hostname;
     private final CsrGenerator csrGenerator;
     private final Clock clock;
 
     AthenzCredentialsService(IdentityConfig identityConfig,
                              ServiceIdentityProvider nodeIdentityProvider,
-                             File trustStoreJks,
                              String hostname,
                              Clock clock) {
         this.tenantIdentity = new AthenzService(identityConfig.domain(), identityConfig.service());
@@ -64,7 +59,6 @@ class AthenzCredentialsService {
         this.ztsEndpoint = URI.create(identityConfig.ztsUrl());
         this.configserverIdentity = new AthenzService(identityConfig.configserverIdentityName());
         this.nodeIdentityProvider = nodeIdentityProvider;
-        this.trustStoreJks = trustStoreJks;
         this.hostname = hostname;
         this.csrGenerator = new CsrGenerator(identityConfig.athenzDnsSuffix(), identityConfig.configserverIdentityName());
         this.clock = clock;
@@ -94,9 +88,8 @@ class AthenzCredentialsService {
                             false,
                             csr);
             X509Certificate certificate = instanceIdentity.certificate();
-            SSLContext identitySslContext = createIdentitySslContext(keyPair.getPrivate(), certificate);
             writeCredentialsToDisk(keyPair.getPrivate(), certificate, document);
-            return new AthenzCredentials(certificate, keyPair, document, identitySslContext);
+            return new AthenzCredentials(certificate, keyPair, document);
         }
     }
 
@@ -117,9 +110,8 @@ class AthenzCredentialsService {
                             false,
                             csr);
             X509Certificate certificate = instanceIdentity.certificate();
-            SSLContext identitySslContext = createIdentitySslContext(newKeyPair.getPrivate(), certificate);
             writeCredentialsToDisk(newKeyPair.getPrivate(), certificate, document);
-            return new AthenzCredentials(certificate, newKeyPair, document, identitySslContext);
+            return new AthenzCredentials(certificate, newKeyPair, document);
         }
     }
 
@@ -134,8 +126,7 @@ class AthenzCredentialsService {
         if (Files.notExists(IDENTITY_DOCUMENT_FILE)) return Optional.empty();
         SignedIdentityDocument signedIdentityDocument = EntityBindingsMapper.readSignedIdentityDocumentFromFile(IDENTITY_DOCUMENT_FILE);
         KeyPair keyPair = new KeyPair(KeyUtils.extractPublicKey(privateKey.get()), privateKey.get());
-        SSLContext sslContext = createIdentitySslContext(privateKey.get(), certificate.get());
-        return Optional.of(new AthenzCredentials(certificate.get(), keyPair, signedIdentityDocument, sslContext));
+        return Optional.of(new AthenzCredentials(certificate.get(), keyPair, signedIdentityDocument));
     }
 
     private boolean isExpired(X509Certificate certificate) {
@@ -150,13 +141,6 @@ class AthenzCredentialsService {
         EntityBindingsMapper.writeSignedIdentityDocumentToFile(IDENTITY_DOCUMENT_FILE, identityDocument);
     }
 
-    private SSLContext createIdentitySslContext(PrivateKey privateKey, X509Certificate certificate) {
-        return new SslContextBuilder()
-                .withKeyStore(privateKey, certificate)
-                .withTrustStore(trustStoreJks.toPath(), JKS)
-                .build();
-    }
-
     private DefaultIdentityDocumentClient createIdentityDocumentClient() {
         return new DefaultIdentityDocumentClient(
                 configserverEndpoint,
diff --git a/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/identityprovider/client/AthenzIdentityProviderImpl.java b/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/identityprovider/client/AthenzIdentityProviderImpl.java
index ac255289883..e4633fb708b 100644
--- a/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/identityprovider/client/AthenzIdentityProviderImpl.java
+++ b/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/identityprovider/client/AthenzIdentityProviderImpl.java
@@ -12,8 +12,11 @@ import com.yahoo.container.jdisc.athenz.AthenzIdentityProvider;
 import com.yahoo.container.jdisc.athenz.AthenzIdentityProviderException;
 import com.yahoo.jdisc.Metric;
 import com.yahoo.log.LogLevel;
+import com.yahoo.security.KeyStoreBuilder;
 import com.yahoo.security.KeyStoreType;
+import com.yahoo.security.Pkcs10Csr;
 import com.yahoo.security.SslContextBuilder;
+import com.yahoo.security.tls.MutableX509KeyManager;
 import com.yahoo.vespa.athenz.api.AthenzDomain;
 import com.yahoo.vespa.athenz.api.AthenzRole;
 import com.yahoo.vespa.athenz.api.AthenzService;
@@ -22,13 +25,14 @@ import com.yahoo.vespa.athenz.client.zts.DefaultZtsClient;
 import com.yahoo.vespa.athenz.client.zts.ZtsClient;
 import com.yahoo.vespa.athenz.identity.ServiceIdentityProvider;
 import com.yahoo.vespa.athenz.identity.SiaIdentityProvider;
-import com.yahoo.security.Pkcs10Csr;
 import com.yahoo.vespa.athenz.utils.SiaUtils;
 import com.yahoo.vespa.defaults.Defaults;
 
 import javax.net.ssl.SSLContext;
-import java.io.File;
+import javax.net.ssl.X509ExtendedKeyManager;
 import java.net.URI;
+import java.nio.file.Path;
+import java.nio.file.Paths;
 import java.security.PrivateKey;
 import java.security.cert.X509Certificate;
 import java.time.Clock;
@@ -42,6 +46,9 @@ import java.util.concurrent.TimeUnit;
 import java.util.function.Function;
 import java.util.logging.Logger;
 
+import static com.yahoo.security.KeyStoreType.JKS;
+import static com.yahoo.security.KeyStoreType.PKCS12;
+
 /**
  * A {@link AthenzIdentityProvider} / {@link ServiceIdentityProvider} component that provides the tenant identity.
  *
@@ -59,10 +66,14 @@ public final class AthenzIdentityProviderImpl extends AbstractComponent implemen
     private final static Duration ROLE_SSL_CONTEXT_EXPIRY = Duration.ofHours(24);
     private final static Duration ROLE_TOKEN_EXPIRY = Duration.ofMinutes(30);
 
+    // TODO Make path to trust store config
+    private static final Path DEFAULT_TRUST_STORE = Paths.get(Defaults.getDefaults().underVespaHome("share/ssl/certs/yahoo_certificate_bundle.pem"));
+
     public static final String CERTIFICATE_EXPIRY_METRIC_NAME = "athenz-tenant-cert.expiry.seconds";
 
     private volatile AthenzCredentials credentials;
     private final Metric metric;
+    private final Path trustStore;
     private final AthenzCredentialsService athenzCredentialsService;
     private final ScheduledExecutorService scheduler;
     private final Clock clock;
@@ -70,6 +81,8 @@ public final class AthenzIdentityProviderImpl extends AbstractComponent implemen
     private final String dnsSuffix;
     private final URI ztsEndpoint;
 
+    private final MutableX509KeyManager identityKeyManager = new MutableX509KeyManager();
+    private final SSLContext identitySslContext;
     private final LoadingCache<AthenzRole, SSLContext> roleSslContextCache;
     private final LoadingCache<AthenzRole, ZToken> roleSpecificRoleTokenCache;
     private final LoadingCache<AthenzDomain, ZToken> domainSpecificRoleTokenCache;
@@ -79,9 +92,9 @@ public final class AthenzIdentityProviderImpl extends AbstractComponent implemen
     public AthenzIdentityProviderImpl(IdentityConfig config, Metric metric) {
         this(config,
              metric,
+             DEFAULT_TRUST_STORE,
              new AthenzCredentialsService(config,
-                                          createNodeIdentityProvider(config),
-                                          getDefaultTrustStoreLocation(),
+                                          createNodeIdentityProvider(config, DEFAULT_TRUST_STORE),
                                           Defaults.getDefaults().vespaHostname(),
                                           Clock.systemUTC()),
              new ScheduledThreadPoolExecutor(1),
@@ -92,10 +105,12 @@ public final class AthenzIdentityProviderImpl extends AbstractComponent implemen
 
     AthenzIdentityProviderImpl(IdentityConfig config,
                                Metric metric,
+                               Path trustStore,
                                AthenzCredentialsService athenzCredentialsService,
                                ScheduledExecutorService scheduler,
                                Clock clock) {
         this.metric = metric;
+        this.trustStore = trustStore;
         this.athenzCredentialsService = athenzCredentialsService;
         this.scheduler = scheduler;
         this.clock = clock;
@@ -106,6 +121,7 @@ public final class AthenzIdentityProviderImpl extends AbstractComponent implemen
         roleSpecificRoleTokenCache = createCache(ROLE_TOKEN_EXPIRY, this::createRoleToken);
         domainSpecificRoleTokenCache = createCache(ROLE_TOKEN_EXPIRY, this::createRoleToken);
         this.csrGenerator = new CsrGenerator(config.athenzDnsSuffix(), config.configserverIdentityName());
+        this.identitySslContext = createIdentitySslContext(identityKeyManager, trustStore);
         registerInstance();
     }
 
@@ -121,11 +137,18 @@ public final class AthenzIdentityProviderImpl extends AbstractComponent implemen
                 });
     }
 
+    private static SSLContext createIdentitySslContext(X509ExtendedKeyManager keyManager, Path trustStore) {
+        return new SslContextBuilder()
+                .withKeyManager(keyManager)
+                .withTrustStore(trustStore, JKS)
+                .build();
+    }
+
     private void registerInstance() {
         try {
-            credentials = athenzCredentialsService.registerInstance();
-            scheduler.scheduleAtFixedRate(this::refreshCertificate, UPDATE_PERIOD.toMinutes(), UPDATE_PERIOD.toMinutes(), TimeUnit.MINUTES);
-            scheduler.scheduleAtFixedRate(this::reportMetrics, 0, 5, TimeUnit.MINUTES);
+            updateIdentityCredentials(this.athenzCredentialsService.registerInstance());
+            this.scheduler.scheduleAtFixedRate(this::refreshCertificate, UPDATE_PERIOD.toMinutes(), UPDATE_PERIOD.toMinutes(), TimeUnit.MINUTES);
+            this.scheduler.scheduleAtFixedRate(this::reportMetrics, 0, 5, TimeUnit.MINUTES);
         } catch (Throwable t) {
             throw new AthenzIdentityProviderException("Could not retrieve Athenz credentials", t);
         }
@@ -148,7 +171,7 @@ public final class AthenzIdentityProviderImpl extends AbstractComponent implemen
 
     @Override
     public SSLContext getIdentitySslContext() {
-        return credentials.getIdentitySslContext();
+        return identitySslContext;
     }
 
     @Override
@@ -189,13 +212,22 @@ public final class AthenzIdentityProviderImpl extends AbstractComponent implemen
         return Collections.singletonList(credentials.getCertificate());
     }
 
+    private void updateIdentityCredentials(AthenzCredentials credentials) {
+        this.credentials = credentials;
+        this.identityKeyManager.updateKeystore(
+                KeyStoreBuilder.withType(PKCS12)
+                        .withKeyEntry("default", credentials.getKeyPair().getPrivate(), credentials.getCertificate())
+                        .build(),
+                new char[0]);
+    }
+
     private SSLContext createRoleSslContext(AthenzRole role) {
         Pkcs10Csr csr = csrGenerator.generateRoleCsr(identity, role, credentials.getIdentityDocument().providerUniqueId(), credentials.getKeyPair());
         try (ZtsClient client = createZtsClient()) {
             X509Certificate roleCertificate = client.getRoleCertificate(role, csr);
             return new SslContextBuilder()
                     .withKeyStore(credentials.getKeyPair().getPrivate(), roleCertificate)
-                    .withTrustStore(getDefaultTrustStoreLocation().toPath(), KeyStoreType.JKS)
+                    .withTrustStore(trustStore, KeyStoreType.JKS)
                     .build();
         }
     }
@@ -226,13 +258,9 @@ public final class AthenzIdentityProviderImpl extends AbstractComponent implemen
         }
     }
 
-    private static SiaIdentityProvider createNodeIdentityProvider(IdentityConfig config) {
+    private static SiaIdentityProvider createNodeIdentityProvider(IdentityConfig config, Path trustStore) {
         return new SiaIdentityProvider(
-                new AthenzService(config.nodeIdentityName()), SiaUtils.DEFAULT_SIA_DIRECTORY, getDefaultTrustStoreLocation());
-    }
-
-    private static File getDefaultTrustStoreLocation() {
-        return new File(Defaults.getDefaults().underVespaHome("share/ssl/certs/yahoo_certificate_bundle.jks"));
+                new AthenzService(config.nodeIdentityName()), SiaUtils.DEFAULT_SIA_DIRECTORY, trustStore.toFile());
     }
 
     private boolean isExpired(AthenzCredentials credentials) {
@@ -245,9 +273,9 @@ public final class AthenzIdentityProviderImpl extends AbstractComponent implemen
 
     void refreshCertificate() {
         try {
-            credentials = isExpired(credentials)
-                    ? athenzCredentialsService.registerInstance()
-                    : athenzCredentialsService.updateCredentials(credentials.getIdentityDocument(), credentials.getIdentitySslContext());
+            updateIdentityCredentials(isExpired(credentials)
+                                      ? athenzCredentialsService.registerInstance()
+                                      : athenzCredentialsService.updateCredentials(credentials.getIdentityDocument(), identitySslContext));
         } catch (Throwable t) {
             log.log(LogLevel.WARNING, "Failed to update credentials: " + t.getMessage(), t);
         }
diff --git a/vespa-athenz/src/test/java/com/yahoo/vespa/athenz/identityprovider/client/AthenzIdentityProviderImplTest.java b/vespa-athenz/src/test/java/com/yahoo/vespa/athenz/identityprovider/client/AthenzIdentityProviderImplTest.java
index 01dab2dada3..c584b803815 100644
--- a/vespa-athenz/src/test/java/com/yahoo/vespa/athenz/identityprovider/client/AthenzIdentityProviderImplTest.java
+++ b/vespa-athenz/src/test/java/com/yahoo/vespa/athenz/identityprovider/client/AthenzIdentityProviderImplTest.java
@@ -4,14 +4,30 @@ package com.yahoo.vespa.athenz.identityprovider.client;
 import com.yahoo.container.core.identity.IdentityConfig;
 import com.yahoo.container.jdisc.athenz.AthenzIdentityProviderException;
 import com.yahoo.jdisc.Metric;
+import com.yahoo.security.KeyAlgorithm;
+import com.yahoo.security.KeyStoreBuilder;
+import com.yahoo.security.KeyStoreType;
+import com.yahoo.security.KeyStoreUtils;
+import com.yahoo.security.KeyUtils;
+import com.yahoo.security.Pkcs10Csr;
+import com.yahoo.security.Pkcs10CsrBuilder;
+import com.yahoo.security.SignatureAlgorithm;
+import com.yahoo.security.X509CertificateBuilder;
 import com.yahoo.test.ManualClock;
+import org.junit.Before;
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.rules.TemporaryFolder;
 
+import javax.security.auth.x500.X500Principal;
+import java.io.IOException;
+import java.math.BigInteger;
+import java.nio.file.Path;
+import java.security.KeyPair;
 import java.security.cert.X509Certificate;
 import java.time.Duration;
 import java.time.Instant;
+import java.time.temporal.ChronoUnit;
 import java.util.Date;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.function.Supplier;
@@ -43,13 +59,36 @@ public class AthenzIdentityProviderImplTest {
                                        .ztsUrl("https:localhost:4443/zts/v1")
                                        .athenzDnsSuffix("dev-us-north-1.vespa.cloud"));
 
+    private final KeyPair caKeypair = KeyUtils.generateKeypair(KeyAlgorithm.EC);
+    private Path trustStoreFile;
+    private X509Certificate caCertificate;
+
+    @Before
+    public void createTrustStoreFile() throws IOException {
+        caCertificate = X509CertificateBuilder
+                .fromKeypair(
+                        caKeypair,
+                        new X500Principal("CN=mydummyca"),
+                        Instant.EPOCH,
+                        Instant.EPOCH.plus(10000, ChronoUnit.DAYS),
+                        SignatureAlgorithm.SHA256_WITH_ECDSA,
+                        BigInteger.ONE)
+                .build();
+        trustStoreFile = tempDir.newFile().toPath();
+        KeyStoreUtils.writeKeyStoreToFile(
+                KeyStoreBuilder.withType(KeyStoreType.JKS)
+                        .withKeyEntry("default", caKeypair.getPrivate(), caCertificate)
+                        .build(),
+                trustStoreFile);
+    }
+
     @Test(expected = AthenzIdentityProviderException.class)
     public void component_creation_fails_when_credentials_not_found() {
         AthenzCredentialsService credentialService = mock(AthenzCredentialsService.class);
         when(credentialService.registerInstance())
                 .thenThrow(new RuntimeException("athenz unavailable"));
 
-        new AthenzIdentityProviderImpl(IDENTITY_CONFIG, mock(Metric.class), credentialService, mock(ScheduledExecutorService.class), new ManualClock(Instant.EPOCH));
+        new AthenzIdentityProviderImpl(IDENTITY_CONFIG, mock(Metric.class), trustStoreFile ,credentialService, mock(ScheduledExecutorService.class), new ManualClock(Instant.EPOCH));
     }
 
     @Test
@@ -59,18 +98,19 @@ public class AthenzIdentityProviderImplTest {
 
         AthenzCredentialsService athenzCredentialsService = mock(AthenzCredentialsService.class);
 
-        X509Certificate certificate = getCertificate(getExpirationSupplier(clock));
+        KeyPair keyPair = KeyUtils.generateKeypair(KeyAlgorithm.EC);
+        X509Certificate certificate = getCertificate(keyPair, getExpirationSupplier(clock));
 
         when(athenzCredentialsService.registerInstance())
-                .thenReturn(new AthenzCredentials(certificate, null, null, null));
+                .thenReturn(new AthenzCredentials(certificate, keyPair, null));
 
         when(athenzCredentialsService.updateCredentials(any(), any()))
                 .thenThrow(new RuntimeException("#1"))
                 .thenThrow(new RuntimeException("#2"))
-                .thenReturn(new AthenzCredentials(certificate, null, null, null));
+                .thenReturn(new AthenzCredentials(certificate, keyPair, null));
 
         AthenzIdentityProviderImpl identityProvider =
-                new AthenzIdentityProviderImpl(IDENTITY_CONFIG, metric, athenzCredentialsService, mock(ScheduledExecutorService.class), clock);
+                new AthenzIdentityProviderImpl(IDENTITY_CONFIG, metric, trustStoreFile, athenzCredentialsService, mock(ScheduledExecutorService.class), clock);
 
         identityProvider.reportMetrics();
         verify(metric).set(eq(AthenzIdentityProviderImpl.CERTIFICATE_EXPIRY_METRIC_NAME), eq(certificateValidity.getSeconds()), any());
@@ -99,10 +139,18 @@ public class AthenzIdentityProviderImplTest {
         return () -> new Date(clock.instant().plus(certificateValidity).toEpochMilli());
     }
 
-    private X509Certificate getCertificate(Supplier<Date> expiry) {
-        X509Certificate x509Certificate = mock(X509Certificate.class);
-        when(x509Certificate.getNotAfter()).thenReturn(expiry.get());
-        return x509Certificate;
+    private X509Certificate getCertificate(KeyPair keyPair, Supplier<Date> expiry) {
+        Pkcs10Csr csr = Pkcs10CsrBuilder.fromKeypair(new X500Principal("CN=dummy"), keyPair, SignatureAlgorithm.SHA256_WITH_ECDSA)
+                .build();
+        return X509CertificateBuilder
+                .fromCsr(csr,
+                         caCertificate.getSubjectX500Principal(),
+                         Instant.EPOCH,
+                         expiry.get().toInstant(),
+                         caKeypair.getPrivate(),
+                         SignatureAlgorithm.SHA256_WITH_ECDSA,
+                         BigInteger.ONE)
+                .build();
     }
 
 }