Revert "Revert "Merge pull request #26859 from vespa-engine/revert-26805-olaa/versioned-feature-flag""

This reverts commit e10612188741a2d5ce9c69b80c35015b38ea1627.

5 files changed, 122 insertions, 635 deletions

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 77aaf17419d..b5579cbfcd8 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
@@ -11,9 +11,7 @@ import com.yahoo.container.jdisc.athenz.AthenzIdentityProvider;
 import com.yahoo.container.jdisc.athenz.AthenzIdentityProviderException;
 import com.yahoo.jdisc.Metric;
 import com.yahoo.metrics.ContainerMetrics;
-import com.yahoo.security.AutoReloadingX509KeyManager;
 import com.yahoo.security.KeyStoreBuilder;
-import com.yahoo.security.KeyUtils;
 import com.yahoo.security.MutableX509KeyManager;
 import com.yahoo.security.Pkcs10Csr;
 import com.yahoo.security.SslContextBuilder;
@@ -26,9 +24,9 @@ import com.yahoo.vespa.athenz.api.ZToken;
 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.identityprovider.api.VespaUniqueInstanceId;
-import com.yahoo.vespa.athenz.tls.AthenzX509CertificateUtils;
+import com.yahoo.vespa.athenz.identity.SiaIdentityProvider;
 import com.yahoo.vespa.athenz.utils.SiaUtils;
+import com.yahoo.vespa.defaults.Defaults;
 
 import javax.net.ssl.SSLContext;
 import javax.net.ssl.X509ExtendedKeyManager;
@@ -40,6 +38,7 @@ import java.security.cert.X509Certificate;
 import java.time.Clock;
 import java.time.Duration;
 import java.time.Instant;
+import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
@@ -66,6 +65,7 @@ public final class AthenzIdentityProviderImpl extends AbstractComponent implemen
 
     // TODO Make some of these values configurable through config. Match requested expiration of register/update requests.
     // TODO These should match the requested expiration
+    static final Duration UPDATE_PERIOD = Duration.ofDays(1);
     static final Duration AWAIT_TERMINTATION_TIMEOUT = Duration.ofSeconds(90);
     private final static Duration ROLE_SSL_CONTEXT_EXPIRY = Duration.ofHours(2);
     // TODO CMS expects 10min or less token ttl. Use 10min default until we have configurable expiry
@@ -73,17 +73,20 @@ public final class AthenzIdentityProviderImpl extends AbstractComponent implemen
 
     // TODO Make path to trust store paths config
     private static final Path CLIENT_TRUST_STORE = Paths.get("/opt/yahoo/share/ssl/certs/yahoo_certificate_bundle.pem");
+    private static final Path ATHENZ_TRUST_STORE = Paths.get("/opt/yahoo/share/ssl/certs/athenz_certificate_bundle.pem");
 
     public static final String CERTIFICATE_EXPIRY_METRIC_NAME = ContainerMetrics.ATHENZ_TENANT_CERT_EXPIRY_SECONDS.baseName();
 
+    private volatile AthenzCredentials credentials;
     private final Metric metric;
     private final Path trustStore;
+    private final AthenzCredentialsService athenzCredentialsService;
     private final ScheduledExecutorService scheduler;
     private final Clock clock;
     private final AthenzService identity;
     private final URI ztsEndpoint;
 
-    private final AutoReloadingX509KeyManager autoReloadingX509KeyManager;
+    private final MutableX509KeyManager identityKeyManager = new MutableX509KeyManager();
     private final SSLContext identitySslContext;
     private final LoadingCache<AthenzRole, X509Certificate> roleSslCertCache;
     private final Map<AthenzRole, MutableX509KeyManager> roleKeyManagerCache;
@@ -95,32 +98,40 @@ public final class AthenzIdentityProviderImpl extends AbstractComponent implemen
 
     @Inject
     public AthenzIdentityProviderImpl(IdentityConfig config, Metric metric) {
-        this(config, metric, CLIENT_TRUST_STORE, new ScheduledThreadPoolExecutor(1), Clock.systemUTC(), createAutoReloadingX509KeyManager(config));
+        this(config,
+             metric,
+                CLIENT_TRUST_STORE,
+             new AthenzCredentialsService(config,
+                                          createNodeIdentityProvider(config),
+                                          Defaults.getDefaults().vespaHostname(),
+                                          Clock.systemUTC()),
+             new ScheduledThreadPoolExecutor(1),
+             Clock.systemUTC());
     }
 
     // Test only
     AthenzIdentityProviderImpl(IdentityConfig config,
                                Metric metric,
                                Path trustStore,
+                               AthenzCredentialsService athenzCredentialsService,
                                ScheduledExecutorService scheduler,
-                               Clock clock,
-                               AutoReloadingX509KeyManager autoReloadingX509KeyManager) {
+                               Clock clock) {
         this.metric = metric;
         this.trustStore = trustStore;
+        this.athenzCredentialsService = athenzCredentialsService;
         this.scheduler = scheduler;
         this.clock = clock;
         this.identity = new AthenzService(config.domain(), config.service());
         this.ztsEndpoint = URI.create(config.ztsUrl());
-        this.roleSslCertCache = crateAutoReloadableCache(ROLE_SSL_CONTEXT_EXPIRY, this::requestRoleCertificate, this.scheduler);
-        this.roleKeyManagerCache = new HashMap<>();
-        this.roleSpecificRoleTokenCache = createCache(ROLE_TOKEN_EXPIRY, this::createRoleToken);
-        this.domainSpecificRoleTokenCache = createCache(ROLE_TOKEN_EXPIRY, this::createRoleToken);
-        this.domainSpecificAccessTokenCache = createCache(ROLE_TOKEN_EXPIRY, this::createAccessToken);
-        this.roleSpecificAccessTokenCache = createCache(ROLE_TOKEN_EXPIRY, this::createAccessToken);
+        roleSslCertCache = crateAutoReloadableCache(ROLE_SSL_CONTEXT_EXPIRY, this::requestRoleCertificate, this.scheduler);
+        roleKeyManagerCache = new HashMap<>();
+        roleSpecificRoleTokenCache = createCache(ROLE_TOKEN_EXPIRY, this::createRoleToken);
+        domainSpecificRoleTokenCache = createCache(ROLE_TOKEN_EXPIRY, this::createRoleToken);
+        domainSpecificAccessTokenCache = createCache(ROLE_TOKEN_EXPIRY, this::createAccessToken);
+        roleSpecificAccessTokenCache = createCache(ROLE_TOKEN_EXPIRY, this::createAccessToken);
         this.csrGenerator = new CsrGenerator(config.athenzDnsSuffix(), config.configserverIdentityName());
-        this.autoReloadingX509KeyManager = autoReloadingX509KeyManager;
-        this.identitySslContext = createIdentitySslContext(autoReloadingX509KeyManager, trustStore);
-        this.scheduler.scheduleAtFixedRate(this::reportMetrics, 0, 5, TimeUnit.MINUTES);
+        this.identitySslContext = createIdentitySslContext(identityKeyManager, trustStore);
+        registerInstance();
     }
 
     private static <KEY, VALUE> LoadingCache<KEY, VALUE> createCache(Duration expiry, Function<KEY, VALUE> cacheLoader) {
@@ -154,6 +165,16 @@ public final class AthenzIdentityProviderImpl extends AbstractComponent implemen
                 .build();
     }
 
+    private void registerInstance() {
+        try {
+            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);
+        }
+    }
+
     @Override
     public AthenzService identity() {
         return identity;
@@ -176,13 +197,13 @@ public final class AthenzIdentityProviderImpl extends AbstractComponent implemen
 
     @Override
     public X509CertificateWithKey getIdentityCertificateWithKey() {
-        var copy = this.autoReloadingX509KeyManager.getCurrentCertificateWithKey();
-        return new X509CertificateWithKey(copy.certificate(), copy.privateKey());
+        AthenzCredentials copy = this.credentials;
+        return new X509CertificateWithKey(copy.getCertificate(), copy.getKeyPair().getPrivate());
     }
 
-    @Override public Path certificatePath() { return SiaUtils.getCertificateFile(identity); }
+    @Override public Path certificatePath() { return athenzCredentialsService.certificatePath(); }
 
-    @Override public Path privateKeyPath() { return SiaUtils.getPrivateKeyFile(identity); }
+    @Override public Path privateKeyPath() { return athenzCredentialsService.privateKeyPath(); }
 
     @Override
     public SSLContext getRoleSslContext(String domain, String role) {
@@ -241,7 +262,7 @@ public final class AthenzIdentityProviderImpl extends AbstractComponent implemen
 
     @Override
     public PrivateKey getPrivateKey() {
-        return autoReloadingX509KeyManager.getPrivateKey(AutoReloadingX509KeyManager.CERTIFICATE_ALIAS);
+        return credentials.getKeyPair().getPrivate();
     }
 
     @Override
@@ -251,7 +272,7 @@ public final class AthenzIdentityProviderImpl extends AbstractComponent implemen
 
     @Override
     public List<X509Certificate> getIdentityCertificate() {
-        return List.of(autoReloadingX509KeyManager.getCertificateChain(AutoReloadingX509KeyManager.CERTIFICATE_ALIAS));
+        return Collections.singletonList(credentials.getCertificate());
     }
 
     @Override
@@ -267,15 +288,19 @@ public final class AthenzIdentityProviderImpl extends AbstractComponent implemen
         }
     }
 
+    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 X509Certificate requestRoleCertificate(AthenzRole role) {
-        var credentials = autoReloadingX509KeyManager.getCurrentCertificateWithKey();
-        var athenzUniqueInstanceId = VespaUniqueInstanceId.fromDottedString(
-                AthenzX509CertificateUtils.getInstanceId(credentials.certificate())
-                        .orElseThrow()
-        );
-        var keyPair = KeyUtils.toKeyPair(credentials.privateKey());
+        var doc = credentials.getIdentityDocument().identityDocument();
         Pkcs10Csr csr = csrGenerator.generateRoleCsr(
-                identity, role, athenzUniqueInstanceId, null, keyPair);
+                identity, role, doc.providerUniqueId(), doc.clusterType(), credentials.getKeyPair());
         try (ZtsClient client = createZtsClient()) {
             X509Certificate roleCertificate = client.getRoleCertificate(role, csr);
             updateRoleKeyManager(role, roleCertificate);
@@ -288,7 +313,7 @@ public final class AthenzIdentityProviderImpl extends AbstractComponent implemen
         MutableX509KeyManager keyManager = roleKeyManagerCache.computeIfAbsent(role, r -> new MutableX509KeyManager());
         keyManager.updateKeystore(
                 KeyStoreBuilder.withType(PKCS12)
-                        .withKeyEntry("default", autoReloadingX509KeyManager.getCurrentCertificateWithKey().privateKey(), certificate)
+                        .withKeyEntry("default", credentials.getKeyPair().getPrivate(), certificate)
                         .build(),
                 new char[0]);
     }
@@ -321,11 +346,6 @@ public final class AthenzIdentityProviderImpl extends AbstractComponent implemen
         return new DefaultZtsClient.Builder(ztsEndpoint).withSslContext(getIdentitySslContext()).build();
     }
 
-    private static AutoReloadingX509KeyManager createAutoReloadingX509KeyManager(IdentityConfig config) {
-        var tenantIdentity = new AthenzService(config.domain(), config.service());
-        return AutoReloadingX509KeyManager.fromPemFiles(SiaUtils.getPrivateKeyFile(tenantIdentity), SiaUtils.getCertificateFile(tenantIdentity));
-    }
-
     @Override
     public void deconstruct() {
         try {
@@ -336,13 +356,32 @@ public final class AthenzIdentityProviderImpl extends AbstractComponent implemen
         }
     }
 
-    private static Instant getExpirationTime(X509Certificate certificate) {
-        return certificate.getNotAfter().toInstant();
+    private static SiaIdentityProvider createNodeIdentityProvider(IdentityConfig config) {
+        return new SiaIdentityProvider(
+                new AthenzService(config.nodeIdentityName()), SiaUtils.DEFAULT_SIA_DIRECTORY, CLIENT_TRUST_STORE);
+    }
+
+    private boolean isExpired(AthenzCredentials credentials) {
+        return clock.instant().isAfter(getExpirationTime(credentials));
+    }
+
+    private static Instant getExpirationTime(AthenzCredentials credentials) {
+        return credentials.getCertificate().getNotAfter().toInstant();
+    }
+
+    void refreshCertificate() {
+        try {
+            updateIdentityCredentials(isExpired(credentials)
+                                      ? athenzCredentialsService.registerInstance()
+                                      : athenzCredentialsService.updateCredentials(credentials.getIdentityDocument(), identitySslContext));
+        } catch (Throwable t) {
+            log.log(Level.WARNING, "Failed to update credentials: " + t.getMessage(), t);
+        }
     }
 
     void reportMetrics() {
         try {
-            Instant expirationTime = getExpirationTime(autoReloadingX509KeyManager.getCurrentCertificateWithKey().certificate());
+            Instant expirationTime = getExpirationTime(credentials);
             Duration remainingLifetime = Duration.between(clock.instant(), expirationTime);
             metric.set(CERTIFICATE_EXPIRY_METRIC_NAME, remainingLifetime.getSeconds(), null);
         } catch (Throwable t) {
diff --git a/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/identityprovider/client/AthenzIdentityProviderProvider.java b/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/identityprovider/client/AthenzIdentityProviderProvider.java
deleted file mode 100644
index 66dad931815..00000000000
--- a/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/identityprovider/client/AthenzIdentityProviderProvider.java
+++ /dev/null
@@ -1,38 +0,0 @@
-package com.yahoo.vespa.athenz.identityprovider.client;
-
-import com.yahoo.container.core.identity.IdentityConfig;
-import com.yahoo.container.di.componentgraph.Provider;
-import com.yahoo.container.jdisc.athenz.AthenzIdentityProvider;
-import com.yahoo.jdisc.Metric;
-
-import javax.inject.Inject;
-import java.nio.file.Files;
-import java.nio.file.Path;
-import java.nio.file.Paths;
-
-/**
- * @author olaa
- */
-public class AthenzIdentityProviderProvider implements Provider<AthenzIdentityProvider> {
-
-    private final Path NODE_ADMIN_MANAGED_IDENTITY_DOCUMENT = Paths.get("/var/lib/sia/vespa-tenant-identity-document.json");
-    private final AthenzIdentityProvider athenzIdentityProvider;
-
-    @Inject
-    public AthenzIdentityProviderProvider(IdentityConfig config, Metric metric) {
-        if (Files.exists(NODE_ADMIN_MANAGED_IDENTITY_DOCUMENT))
-            athenzIdentityProvider = new AthenzIdentityProviderImpl(config, metric);
-        else
-            athenzIdentityProvider = new LegacyAthenzIdentityProviderImpl(config, metric);
-    }
-
-    @Override
-    public void deconstruct() {
-        athenzIdentityProvider.deconstruct();
-    }
-
-    @Override
-    public AthenzIdentityProvider get() {
-        return athenzIdentityProvider;
-    }
-}
diff --git a/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/identityprovider/client/LegacyAthenzIdentityProviderImpl.java b/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/identityprovider/client/LegacyAthenzIdentityProviderImpl.java
deleted file mode 100644
index d699564a4ee..00000000000
--- a/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/identityprovider/client/LegacyAthenzIdentityProviderImpl.java
+++ /dev/null
@@ -1,392 +0,0 @@
-// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
-package com.yahoo.vespa.athenz.identityprovider.client;
-
-import com.google.common.cache.CacheBuilder;
-import com.google.common.cache.CacheLoader;
-import com.google.common.cache.LoadingCache;
-import com.yahoo.component.AbstractComponent;
-import com.yahoo.component.annotation.Inject;
-import com.yahoo.container.core.identity.IdentityConfig;
-import com.yahoo.container.jdisc.athenz.AthenzIdentityProvider;
-import com.yahoo.container.jdisc.athenz.AthenzIdentityProviderException;
-import com.yahoo.jdisc.Metric;
-import com.yahoo.metrics.ContainerMetrics;
-import com.yahoo.security.KeyStoreBuilder;
-import com.yahoo.security.MutableX509KeyManager;
-import com.yahoo.security.Pkcs10Csr;
-import com.yahoo.security.SslContextBuilder;
-import com.yahoo.security.X509CertificateWithKey;
-import com.yahoo.vespa.athenz.api.AthenzAccessToken;
-import com.yahoo.vespa.athenz.api.AthenzDomain;
-import com.yahoo.vespa.athenz.api.AthenzRole;
-import com.yahoo.vespa.athenz.api.AthenzService;
-import com.yahoo.vespa.athenz.api.ZToken;
-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.vespa.athenz.utils.SiaUtils;
-import com.yahoo.vespa.defaults.Defaults;
-
-import javax.net.ssl.SSLContext;
-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;
-import java.time.Duration;
-import java.time.Instant;
-import java.util.Collections;
-import java.util.HashMap;
-import java.util.List;
-import java.util.Map;
-import java.util.concurrent.ScheduledExecutorService;
-import java.util.concurrent.ScheduledThreadPoolExecutor;
-import java.util.concurrent.TimeUnit;
-import java.util.function.Function;
-import java.util.logging.Level;
-import java.util.logging.Logger;
-
-import static com.yahoo.security.KeyStoreType.PKCS12;
-
-/**
- * A {@link AthenzIdentityProvider} / {@link ServiceIdentityProvider} component that provides the tenant identity.
- *
- * @author mortent
- * @author bjorncs
- */
-// This class should probably not implement ServiceIdentityProvider,
-// as that interface is intended for providing the node's identity, not the tenant's application identity.
-public final class LegacyAthenzIdentityProviderImpl extends AbstractComponent implements AthenzIdentityProvider, ServiceIdentityProvider {
-
-    private static final Logger log = Logger.getLogger(LegacyAthenzIdentityProviderImpl.class.getName());
-
-    // TODO Make some of these values configurable through config. Match requested expiration of register/update requests.
-    // TODO These should match the requested expiration
-    static final Duration UPDATE_PERIOD = Duration.ofDays(1);
-    static final Duration AWAIT_TERMINTATION_TIMEOUT = Duration.ofSeconds(90);
-    private final static Duration ROLE_SSL_CONTEXT_EXPIRY = Duration.ofHours(2);
-    // TODO CMS expects 10min or less token ttl. Use 10min default until we have configurable expiry
-    private final static Duration ROLE_TOKEN_EXPIRY = Duration.ofMinutes(10);
-
-    // TODO Make path to trust store paths config
-    private static final Path CLIENT_TRUST_STORE = Paths.get("/opt/yahoo/share/ssl/certs/yahoo_certificate_bundle.pem");
-    private static final Path ATHENZ_TRUST_STORE = Paths.get("/opt/yahoo/share/ssl/certs/athenz_certificate_bundle.pem");
-
-    public static final String CERTIFICATE_EXPIRY_METRIC_NAME = ContainerMetrics.ATHENZ_TENANT_CERT_EXPIRY_SECONDS.baseName();
-
-    private volatile AthenzCredentials credentials;
-    private final Metric metric;
-    private final Path trustStore;
-    private final AthenzCredentialsService athenzCredentialsService;
-    private final ScheduledExecutorService scheduler;
-    private final Clock clock;
-    private final AthenzService identity;
-    private final URI ztsEndpoint;
-
-    private final MutableX509KeyManager identityKeyManager = new MutableX509KeyManager();
-    private final SSLContext identitySslContext;
-    private final LoadingCache<AthenzRole, X509Certificate> roleSslCertCache;
-    private final Map<AthenzRole, MutableX509KeyManager> roleKeyManagerCache;
-    private final LoadingCache<AthenzRole, ZToken> roleSpecificRoleTokenCache;
-    private final LoadingCache<AthenzDomain, ZToken> domainSpecificRoleTokenCache;
-    private final LoadingCache<AthenzDomain, AthenzAccessToken> domainSpecificAccessTokenCache;
-    private final LoadingCache<List<AthenzRole>, AthenzAccessToken> roleSpecificAccessTokenCache;
-    private final CsrGenerator csrGenerator;
-
-    @Inject
-    public LegacyAthenzIdentityProviderImpl(IdentityConfig config, Metric metric) {
-        this(config,
-             metric,
-                CLIENT_TRUST_STORE,
-             new AthenzCredentialsService(config,
-                                          createNodeIdentityProvider(config),
-                                          Defaults.getDefaults().vespaHostname(),
-                                          Clock.systemUTC()),
-             new ScheduledThreadPoolExecutor(1),
-             Clock.systemUTC());
-    }
-
-    // Test only
-    LegacyAthenzIdentityProviderImpl(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;
-        this.identity = new AthenzService(config.domain(), config.service());
-        this.ztsEndpoint = URI.create(config.ztsUrl());
-        roleSslCertCache = crateAutoReloadableCache(ROLE_SSL_CONTEXT_EXPIRY, this::requestRoleCertificate, this.scheduler);
-        roleKeyManagerCache = new HashMap<>();
-        roleSpecificRoleTokenCache = createCache(ROLE_TOKEN_EXPIRY, this::createRoleToken);
-        domainSpecificRoleTokenCache = createCache(ROLE_TOKEN_EXPIRY, this::createRoleToken);
-        domainSpecificAccessTokenCache = createCache(ROLE_TOKEN_EXPIRY, this::createAccessToken);
-        roleSpecificAccessTokenCache = createCache(ROLE_TOKEN_EXPIRY, this::createAccessToken);
-        this.csrGenerator = new CsrGenerator(config.athenzDnsSuffix(), config.configserverIdentityName());
-        this.identitySslContext = createIdentitySslContext(identityKeyManager, trustStore);
-        registerInstance();
-    }
-
-    private static <KEY, VALUE> LoadingCache<KEY, VALUE> createCache(Duration expiry, Function<KEY, VALUE> cacheLoader) {
-        return CacheBuilder.newBuilder()
-                .refreshAfterWrite(expiry.dividedBy(2).toMinutes(), TimeUnit.MINUTES)
-                .expireAfterWrite(expiry.toMinutes(), TimeUnit.MINUTES)
-                .build(new CacheLoader<KEY, VALUE>() {
-                    @Override
-                    public VALUE load(KEY key) {
-                        return cacheLoader.apply(key);
-                    }
-                });
-    }
-
-    private static <KEY, VALUE> LoadingCache<KEY, VALUE> crateAutoReloadableCache(Duration expiry, Function<KEY, VALUE> cacheLoader, ScheduledExecutorService scheduler) {
-        LoadingCache<KEY, VALUE> cache = createCache(expiry, cacheLoader);
-
-        // The cache above will reload it's contents if and only if a request for the key is made. Scheduling
-        // a cache reloader to reload all keys in this cache.
-        scheduler.scheduleAtFixedRate(() -> { cache.asMap().keySet().forEach(cache::getUnchecked);},
-                                      expiry.dividedBy(4).toMinutes(),
-                                      expiry.dividedBy(4).toMinutes(),
-                                      TimeUnit.MINUTES);
-        return cache;
-    }
-
-    private static SSLContext createIdentitySslContext(X509ExtendedKeyManager keyManager, Path trustStore) {
-        return new SslContextBuilder()
-                .withKeyManager(keyManager)
-                .withTrustStore(trustStore)
-                .build();
-    }
-
-    private void registerInstance() {
-        try {
-            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);
-        }
-    }
-
-    @Override
-    public AthenzService identity() {
-        return identity;
-    }
-
-    @Override
-    public String domain() {
-        return identity.getDomain().getName();
-    }
-
-    @Override
-    public String service() {
-        return identity.getName();
-    }
-
-    @Override
-    public SSLContext getIdentitySslContext() {
-        return identitySslContext;
-    }
-
-    @Override
-    public X509CertificateWithKey getIdentityCertificateWithKey() {
-        AthenzCredentials copy = this.credentials;
-        return new X509CertificateWithKey(copy.getCertificate(), copy.getKeyPair().getPrivate());
-    }
-
-    @Override public Path certificatePath() { return athenzCredentialsService.certificatePath(); }
-
-    @Override public Path privateKeyPath() { return athenzCredentialsService.privateKeyPath(); }
-
-    @Override
-    public SSLContext getRoleSslContext(String domain, String role) {
-        try {
-            AthenzRole athenzRole = new AthenzRole(new AthenzDomain(domain), role);
-            // Make sure to request a certificate which triggers creating a new key manager for this role
-            X509Certificate x509Certificate = getRoleCertificate(athenzRole);
-            MutableX509KeyManager keyManager = roleKeyManagerCache.get(athenzRole);
-            return new SslContextBuilder()
-                    .withKeyManager(keyManager)
-                    .withTrustStore(trustStore)
-                    .build();
-        } catch (Exception e) {
-            throw new AthenzIdentityProviderException("Could not retrieve role certificate: " + e.getMessage(), e);
-        }
-    }
-
-    @Override
-    public String getRoleToken(String domain) {
-        try {
-            return domainSpecificRoleTokenCache.get(new AthenzDomain(domain)).getRawToken();
-        } catch (Exception e) {
-            throw new AthenzIdentityProviderException("Could not retrieve role token: " + e.getMessage(), e);
-        }
-    }
-
-    @Override
-    public String getRoleToken(String domain, String role) {
-        try {
-            return roleSpecificRoleTokenCache.get(new AthenzRole(domain, role)).getRawToken();
-        } catch (Exception e) {
-            throw new AthenzIdentityProviderException("Could not retrieve role token: " + e.getMessage(), e);
-        }
-    }
-
-    @Override
-    public String getAccessToken(String domain) {
-        try {
-            return domainSpecificAccessTokenCache.get(new AthenzDomain(domain)).value();
-        } catch (Exception e) {
-            throw new AthenzIdentityProviderException("Could not retrieve access token: " + e.getMessage(), e);
-        }
-    }
-
-    @Override
-    public String getAccessToken(String domain, List<String> roles) {
-        try {
-            List<AthenzRole> roleList = roles.stream()
-                    .map(roleName -> new AthenzRole(domain, roleName))
-                    .toList();
-            return roleSpecificAccessTokenCache.get(roleList).value();
-        } catch (Exception e) {
-            throw new AthenzIdentityProviderException("Could not retrieve access token: " + e.getMessage(), e);
-        }
-    }
-
-    @Override
-    public PrivateKey getPrivateKey() {
-        return credentials.getKeyPair().getPrivate();
-    }
-
-    @Override
-    public Path trustStorePath() {
-        return trustStore;
-    }
-
-    @Override
-    public List<X509Certificate> getIdentityCertificate() {
-        return Collections.singletonList(credentials.getCertificate());
-    }
-
-    @Override
-    public X509Certificate getRoleCertificate(String domain, String role) {
-        return getRoleCertificate(new AthenzRole(new AthenzDomain(domain), role));
-    }
-
-    private X509Certificate getRoleCertificate(AthenzRole athenzRole) {
-        try {
-            return roleSslCertCache.get(athenzRole);
-        } catch (Exception e) {
-            throw new AthenzIdentityProviderException("Could not retrieve role certificate: " + e.getMessage(), e);
-        }
-    }
-
-    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 X509Certificate requestRoleCertificate(AthenzRole role) {
-        var doc = credentials.getIdentityDocument().identityDocument();
-        Pkcs10Csr csr = csrGenerator.generateRoleCsr(
-                identity, role, doc.providerUniqueId(), doc.clusterType(), credentials.getKeyPair());
-        try (ZtsClient client = createZtsClient()) {
-            X509Certificate roleCertificate = client.getRoleCertificate(role, csr);
-            updateRoleKeyManager(role, roleCertificate);
-            log.info(String.format("Requester role certificate for role %s, expires: %s", role.toResourceNameString(), roleCertificate.getNotAfter().toInstant().toString()));
-            return roleCertificate;
-        }
-    }
-
-    private void updateRoleKeyManager(AthenzRole role, X509Certificate certificate) {
-        MutableX509KeyManager keyManager = roleKeyManagerCache.computeIfAbsent(role, r -> new MutableX509KeyManager());
-        keyManager.updateKeystore(
-                KeyStoreBuilder.withType(PKCS12)
-                        .withKeyEntry("default", credentials.getKeyPair().getPrivate(), certificate)
-                        .build(),
-                new char[0]);
-    }
-
-    private ZToken createRoleToken(AthenzRole athenzRole) {
-        try (ZtsClient client = createZtsClient()) {
-            return client.getRoleToken(athenzRole, ROLE_TOKEN_EXPIRY);
-        }
-    }
-
-    private ZToken createRoleToken(AthenzDomain domain) {
-        try (ZtsClient client = createZtsClient()) {
-            return client.getRoleToken(domain, ROLE_TOKEN_EXPIRY);
-        }
-    }
-
-    private AthenzAccessToken createAccessToken(AthenzDomain domain) {
-        try (ZtsClient client = createZtsClient()) {
-            return client.getAccessToken(domain);
-        }
-    }
-
-    private AthenzAccessToken createAccessToken(List<AthenzRole> roles) {
-        try (ZtsClient client = createZtsClient()) {
-            return client.getAccessToken(roles);
-        }
-    }
-
-    private DefaultZtsClient createZtsClient() {
-        return new DefaultZtsClient.Builder(ztsEndpoint).withSslContext(getIdentitySslContext()).build();
-    }
-
-    @Override
-    public void deconstruct() {
-        try {
-            scheduler.shutdownNow();
-            scheduler.awaitTermination(AWAIT_TERMINTATION_TIMEOUT.getSeconds(), TimeUnit.SECONDS);
-        } catch (InterruptedException e) {
-            throw new RuntimeException(e);
-        }
-    }
-
-    private static SiaIdentityProvider createNodeIdentityProvider(IdentityConfig config) {
-        return new SiaIdentityProvider(
-                new AthenzService(config.nodeIdentityName()), SiaUtils.DEFAULT_SIA_DIRECTORY, CLIENT_TRUST_STORE);
-    }
-
-    private boolean isExpired(AthenzCredentials credentials) {
-        return clock.instant().isAfter(getExpirationTime(credentials));
-    }
-
-    private static Instant getExpirationTime(AthenzCredentials credentials) {
-        return credentials.getCertificate().getNotAfter().toInstant();
-    }
-
-    void refreshCertificate() {
-        try {
-            updateIdentityCredentials(isExpired(credentials)
-                                      ? athenzCredentialsService.registerInstance()
-                                      : athenzCredentialsService.updateCredentials(credentials.getIdentityDocument(), identitySslContext));
-        } catch (Throwable t) {
-            log.log(Level.WARNING, "Failed to update credentials: " + t.getMessage(), t);
-        }
-    }
-
-    void reportMetrics() {
-        try {
-            Instant expirationTime = getExpirationTime(credentials);
-            Duration remainingLifetime = Duration.between(clock.instant(), expirationTime);
-            metric.set(CERTIFICATE_EXPIRY_METRIC_NAME, remainingLifetime.getSeconds(), null);
-        } catch (Throwable t) {
-            log.log(Level.WARNING, "Failed to update metrics: " + 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 108da9e0136..c9d2ea581bb 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
@@ -2,8 +2,8 @@
 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.AutoReloadingX509KeyManager;
 import com.yahoo.security.KeyAlgorithm;
 import com.yahoo.security.KeyStoreBuilder;
 import com.yahoo.security.KeyStoreType;
@@ -13,13 +13,13 @@ import com.yahoo.security.Pkcs10Csr;
 import com.yahoo.security.Pkcs10CsrBuilder;
 import com.yahoo.security.SignatureAlgorithm;
 import com.yahoo.security.X509CertificateBuilder;
-import com.yahoo.security.X509CertificateWithKey;
 import com.yahoo.test.ManualClock;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.io.TempDir;
 
 import javax.security.auth.x500.X500Principal;
+
 import java.io.File;
 import java.io.IOException;
 import java.math.BigInteger;
@@ -33,12 +33,17 @@ import java.util.Date;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.function.Supplier;
 
+import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
+/**
+ * @author mortent
+ * @author bjorncs
+ */
 public class AthenzIdentityProviderImplTest {
 
     @TempDir
@@ -80,25 +85,58 @@ public class AthenzIdentityProviderImplTest {
     }
 
     @Test
-    void certificate_expiry_metric_is_reported() {
+    void component_creation_fails_when_credentials_not_found() {
+        assertThrows(AthenzIdentityProviderException.class, () -> {
+            AthenzCredentialsService credentialService = mock(AthenzCredentialsService.class);
+            when(credentialService.registerInstance())
+                    .thenThrow(new RuntimeException("athenz unavailable"));
+
+            new AthenzIdentityProviderImpl(IDENTITY_CONFIG, mock(Metric.class), trustStoreFile, credentialService, mock(ScheduledExecutorService.class), new ManualClock(Instant.EPOCH));
+        });
+    }
+
+    @Test
+    void metrics_updated_on_refresh() {
         ManualClock clock = new ManualClock(Instant.EPOCH);
         Metric metric = mock(Metric.class);
-        AutoReloadingX509KeyManager keyManager = mock(AutoReloadingX509KeyManager.class);
+
+        AthenzCredentialsService athenzCredentialsService = mock(AthenzCredentialsService.class);
+
         KeyPair keyPair = KeyUtils.generateKeypair(KeyAlgorithm.EC);
         X509Certificate certificate = getCertificate(keyPair, getExpirationSupplier(clock));
-        when(keyManager.getCurrentCertificateWithKey()).thenReturn(new X509CertificateWithKey(certificate, keyPair.getPrivate()));
 
-        AthenzIdentityProviderImpl identityProvider = new AthenzIdentityProviderImpl(IDENTITY_CONFIG, metric, trustStoreFile, mock(ScheduledExecutorService.class), clock, keyManager);
+        when(athenzCredentialsService.registerInstance())
+                .thenReturn(new AthenzCredentials(certificate, keyPair, null));
+
+        when(athenzCredentialsService.updateCredentials(any(), any()))
+                .thenThrow(new RuntimeException("#1"))
+                .thenThrow(new RuntimeException("#2"))
+                .thenReturn(new AthenzCredentials(certificate, keyPair, null));
+
+        AthenzIdentityProviderImpl identityProvider =
+                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());
 
+        // Advance 1 day, refresh fails, cert is 1 day old
         clock.advance(Duration.ofDays(1));
+        identityProvider.refreshCertificate();
         identityProvider.reportMetrics();
         verify(metric).set(eq(AthenzIdentityProviderImpl.CERTIFICATE_EXPIRY_METRIC_NAME), eq(certificateValidity.minus(Duration.ofDays(1)).getSeconds()), any());
 
+        // Advance 1 more day, refresh fails, cert is 2 days old
         clock.advance(Duration.ofDays(1));
+        identityProvider.refreshCertificate();
         identityProvider.reportMetrics();
         verify(metric).set(eq(AthenzIdentityProviderImpl.CERTIFICATE_EXPIRY_METRIC_NAME), eq(certificateValidity.minus(Duration.ofDays(2)).getSeconds()), any());
+
+        // Advance 1 more day, refresh succeds, cert is new
+        clock.advance(Duration.ofDays(1));
+        identityProvider.refreshCertificate();
+        identityProvider.reportMetrics();
+        verify(metric).set(eq(AthenzIdentityProviderImpl.CERTIFICATE_EXPIRY_METRIC_NAME), eq(certificateValidity.getSeconds()), any());
+
     }
 
     private Supplier<Date> getExpirationSupplier(ManualClock clock) {
diff --git a/vespa-athenz/src/test/java/com/yahoo/vespa/athenz/identityprovider/client/LegacyAthenzIdentityProviderImplTest.java b/vespa-athenz/src/test/java/com/yahoo/vespa/athenz/identityprovider/client/LegacyAthenzIdentityProviderImplTest.java
deleted file mode 100644
index 75dc42cd4a6..00000000000
--- a/vespa-athenz/src/test/java/com/yahoo/vespa/athenz/identityprovider/client/LegacyAthenzIdentityProviderImplTest.java
+++ /dev/null
@@ -1,160 +0,0 @@
-// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
-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.jupiter.api.BeforeEach;
-import org.junit.jupiter.api.Test;
-import org.junit.jupiter.api.io.TempDir;
-
-import javax.security.auth.x500.X500Principal;
-
-import java.io.File;
-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;
-
-import static org.junit.jupiter.api.Assertions.assertThrows;
-import static org.mockito.ArgumentMatchers.any;
-import static org.mockito.ArgumentMatchers.eq;
-import static org.mockito.Mockito.mock;
-import static org.mockito.Mockito.verify;
-import static org.mockito.Mockito.when;
-
-/**
- * @author mortent
- * @author bjorncs
- */
-public class LegacyAthenzIdentityProviderImplTest {
-
-    @TempDir
-    public File tempDir;
-
-    public static final Duration certificateValidity = Duration.ofDays(30);
-
-    private static final IdentityConfig IDENTITY_CONFIG =
-            new IdentityConfig(new IdentityConfig.Builder()
-                                       .service("tenantService")
-                                       .domain("tenantDomain")
-                                       .nodeIdentityName("vespa.tenant")
-                                       .configserverIdentityName("vespa.configserver")
-                                       .loadBalancerAddress("cfg")
-                                       .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;
-
-    @BeforeEach
-    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 = File.createTempFile("junit", null, tempDir).toPath();
-        KeyStoreUtils.writeKeyStoreToFile(
-                KeyStoreBuilder.withType(KeyStoreType.JKS)
-                        .withKeyEntry("default", caKeypair.getPrivate(), caCertificate)
-                        .build(),
-                trustStoreFile);
-    }
-
-    @Test
-    void component_creation_fails_when_credentials_not_found() {
-        assertThrows(AthenzIdentityProviderException.class, () -> {
-            AthenzCredentialsService credentialService = mock(AthenzCredentialsService.class);
-            when(credentialService.registerInstance())
-                    .thenThrow(new RuntimeException("athenz unavailable"));
-
-            new LegacyAthenzIdentityProviderImpl(IDENTITY_CONFIG, mock(Metric.class), trustStoreFile, credentialService, mock(ScheduledExecutorService.class), new ManualClock(Instant.EPOCH));
-        });
-    }
-
-    @Test
-    void metrics_updated_on_refresh() {
-        ManualClock clock = new ManualClock(Instant.EPOCH);
-        Metric metric = mock(Metric.class);
-
-        AthenzCredentialsService athenzCredentialsService = mock(AthenzCredentialsService.class);
-
-        KeyPair keyPair = KeyUtils.generateKeypair(KeyAlgorithm.EC);
-        X509Certificate certificate = getCertificate(keyPair, getExpirationSupplier(clock));
-
-        when(athenzCredentialsService.registerInstance())
-                .thenReturn(new AthenzCredentials(certificate, keyPair, null));
-
-        when(athenzCredentialsService.updateCredentials(any(), any()))
-                .thenThrow(new RuntimeException("#1"))
-                .thenThrow(new RuntimeException("#2"))
-                .thenReturn(new AthenzCredentials(certificate, keyPair, null));
-
-        LegacyAthenzIdentityProviderImpl identityProvider =
-                new LegacyAthenzIdentityProviderImpl(IDENTITY_CONFIG, metric, trustStoreFile, athenzCredentialsService, mock(ScheduledExecutorService.class), clock);
-
-        identityProvider.reportMetrics();
-        verify(metric).set(eq(LegacyAthenzIdentityProviderImpl.CERTIFICATE_EXPIRY_METRIC_NAME), eq(certificateValidity.getSeconds()), any());
-
-        // Advance 1 day, refresh fails, cert is 1 day old
-        clock.advance(Duration.ofDays(1));
-        identityProvider.refreshCertificate();
-        identityProvider.reportMetrics();
-        verify(metric).set(eq(LegacyAthenzIdentityProviderImpl.CERTIFICATE_EXPIRY_METRIC_NAME), eq(certificateValidity.minus(Duration.ofDays(1)).getSeconds()), any());
-
-        // Advance 1 more day, refresh fails, cert is 2 days old
-        clock.advance(Duration.ofDays(1));
-        identityProvider.refreshCertificate();
-        identityProvider.reportMetrics();
-        verify(metric).set(eq(LegacyAthenzIdentityProviderImpl.CERTIFICATE_EXPIRY_METRIC_NAME), eq(certificateValidity.minus(Duration.ofDays(2)).getSeconds()), any());
-
-        // Advance 1 more day, refresh succeds, cert is new
-        clock.advance(Duration.ofDays(1));
-        identityProvider.refreshCertificate();
-        identityProvider.reportMetrics();
-        verify(metric).set(eq(LegacyAthenzIdentityProviderImpl.CERTIFICATE_EXPIRY_METRIC_NAME), eq(certificateValidity.getSeconds()), any());
-
-    }
-
-    private Supplier<Date> getExpirationSupplier(ManualClock clock) {
-        return () -> new Date(clock.instant().plus(certificateValidity).toEpochMilli());
-    }
-
-    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();
-    }
-
-}