Make AthenzCredentialMaintainer shared

5 files changed, 70 insertions, 80 deletions

diff --git a/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/component/Environment.java b/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/component/Environment.java
index b860863c08d..857f796fb5c 100644
--- a/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/component/Environment.java
+++ b/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/component/Environment.java
@@ -41,7 +41,7 @@ public class Environment {
     private final Path trustStorePath;
     private final DockerNetworking dockerNetworking;
 
-    private Environment(ConfigServerConfig configServerConfig,
+    private Environment(ConfigServerInfo configServerInfo,
                         Path trustStorePath,
                         String environment,
                         String region,
@@ -58,9 +58,7 @@ public class Environment {
                         AthenzService nodeAthenzIdentity,
                         boolean nodeAgentCertEnabled,
                         DockerNetworking dockerNetworking) {
-        Objects.requireNonNull(configServerConfig, "configServerConfig cannot be null");
-
-        this.configServerInfo = new ConfigServerInfo(configServerConfig);
+        this.configServerInfo = Objects.requireNonNull(configServerInfo, "configServerConfig cannot be null");
         this.environment = Objects.requireNonNull(environment, "environment cannot be null");;
         this.region = Objects.requireNonNull(region, "region cannot be null");;
         this.system = Objects.requireNonNull(system, "system cannot be null");;
@@ -189,7 +187,7 @@ public class Environment {
     }
 
     public static class Builder {
-        private ConfigServerConfig configServerConfig;
+        private ConfigServerInfo configServerInfo;
         private String environment;
         private String region;
         private String system;
@@ -208,7 +206,12 @@ public class Environment {
         private DockerNetworking dockerNetworking;
 
         public Builder configServerConfig(ConfigServerConfig configServerConfig) {
-            this.configServerConfig = configServerConfig;
+            this.configServerInfo = new ConfigServerInfo(configServerConfig);
+            return this;
+        }
+
+        public Builder configServerInfo(ConfigServerInfo configServerInfo) {
+            this.configServerInfo = configServerInfo;
             return this;
         }
 
@@ -293,7 +296,7 @@ public class Environment {
         }
 
         public Environment build() {
-            return new Environment(configServerConfig,
+            return new Environment(configServerInfo,
                                    trustStorePath,
                                    environment,
                                    region,
diff --git a/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/maintenance/identity/AthenzCredentialsMaintainer.java b/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/maintenance/identity/AthenzCredentialsMaintainer.java
index 216878b92f4..9dab204df06 100644
--- a/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/maintenance/identity/AthenzCredentialsMaintainer.java
+++ b/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/maintenance/identity/AthenzCredentialsMaintainer.java
@@ -21,9 +21,9 @@ import com.yahoo.vespa.athenz.identityprovider.client.CsrGenerator;
 import com.yahoo.vespa.athenz.tls.AthenzIdentityVerifier;
 import com.yahoo.vespa.athenz.utils.SiaUtils;
 import com.yahoo.vespa.hosted.dockerapi.ContainerName;
-import com.yahoo.vespa.hosted.node.admin.component.Environment;
+import com.yahoo.vespa.hosted.node.admin.component.ConfigServerInfo;
 import com.yahoo.vespa.hosted.node.admin.nodeagent.NodeAgentContext;
-import com.yahoo.vespa.hosted.node.admin.util.PrefixLogger;
+import com.yahoo.vespa.hosted.node.admin.task.util.file.FileFinder;
 
 import javax.net.ssl.SSLContext;
 import java.io.IOException;
@@ -39,6 +39,8 @@ import java.security.cert.X509Certificate;
 import java.time.Clock;
 import java.time.Duration;
 import java.time.Instant;
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
 import java.util.logging.Logger;
 
 import static java.util.Collections.singleton;
@@ -58,57 +60,47 @@ public class AthenzCredentialsMaintainer {
 
     private static final Path CONTAINER_SIA_DIRECTORY = Paths.get("/var/lib/sia");
 
-    private final boolean enabled;
-    private final PrefixLogger log;
-    private final String hostname;
-    private final Path trustStorePath;
-    private final Path privateKeyFile;
-    private final Path certificateFile;
-    private final Path identityDocumentFile;
-    private final AthenzService containerIdentity;
     private final URI ztsEndpoint;
+    private final Path trustStorePath;
+    private final AthenzService configserverIdentity;
     private final Clock clock;
     private final ServiceIdentityProvider hostIdentityProvider;
     private final IdentityDocumentClient identityDocumentClient;
     private final CsrGenerator csrGenerator;
-    private final AthenzService configserverIdentity;
 
-    private Instant lastRefreshAttempt = Instant.EPOCH; // Used as an optimization to ensure ZTS is not DDoS'ed on continuously failing refresh attempts
+    // Used as an optimization to ensure ZTS is not DDoS'ed on continuously failing refresh attempts
+    private Map<ContainerName, Instant> lastRefreshAttempt = new ConcurrentHashMap<>();
 
-    public AthenzCredentialsMaintainer(String hostname,
-                                       Environment environment,
+    public AthenzCredentialsMaintainer(URI ztsEndpoint,
+                                       Path trustStorePath,
+                                       ConfigServerInfo configServerInfo,
+                                       String certificateDnsSuffix,
                                        ServiceIdentityProvider hostIdentityProvider) {
-        ContainerName containerName = ContainerName.fromHostname(hostname);
-        Path containerSiaDirectory = environment.pathInNodeAdminFromPathInNode(containerName, CONTAINER_SIA_DIRECTORY);
-        this.enabled = environment.isNodeAgentCertEnabled();
-        this.log = PrefixLogger.getNodeAgentLogger(AthenzCredentialsMaintainer.class, containerName);
-        this.hostname = hostname;
-        this.containerIdentity = environment.getNodeAthenzIdentity();
-        this.ztsEndpoint = environment.getZtsUri();
-        this.configserverIdentity = environment.getConfigserverAthenzIdentity();
-        this.csrGenerator = new CsrGenerator(environment.getCertificateDnsSuffix(), configserverIdentity.getFullName());
-        this.trustStorePath = environment.getTrustStorePath();
-        this.privateKeyFile = SiaUtils.getPrivateKeyFile(containerSiaDirectory, containerIdentity);
-        this.certificateFile = SiaUtils.getCertificateFile(containerSiaDirectory, containerIdentity);
-        this.identityDocumentFile = containerSiaDirectory.resolve("vespa-node-identity-document.json");
+        this.ztsEndpoint = ztsEndpoint;
+        this.trustStorePath = trustStorePath;
+        this.configserverIdentity = configServerInfo.getConfigServerIdentity();
+        this.csrGenerator = new CsrGenerator(certificateDnsSuffix, configserverIdentity.getFullName());
         this.hostIdentityProvider = hostIdentityProvider;
-        this.identityDocumentClient =
-                new DefaultIdentityDocumentClient(
-                        environment.getConfigserverLoadBalancerEndpoint(),
-                        hostIdentityProvider,
-                        new AthenzIdentityVerifier(singleton(configserverIdentity)));
+        this.identityDocumentClient = new DefaultIdentityDocumentClient(
+                configServerInfo.getLoadBalancerEndpoint(),
+                hostIdentityProvider,
+                new AthenzIdentityVerifier(singleton(configserverIdentity)));
         this.clock = Clock.systemUTC();
     }
 
     public void converge(NodeAgentContext context) {
         try {
             context.log(logger, LogLevel.DEBUG, "Checking certificate");
+            Path containerSiaDirectory = context.pathOnHostFromPathInNode(CONTAINER_SIA_DIRECTORY);
+            Path privateKeyFile = SiaUtils.getPrivateKeyFile(containerSiaDirectory, context.identity());
+            Path certificateFile = SiaUtils.getCertificateFile(containerSiaDirectory, context.identity());
+            Path identityDocumentFile = containerSiaDirectory.resolve("vespa-node-identity-document.json");
             if (!Files.exists(privateKeyFile) || !Files.exists(certificateFile) || !Files.exists(identityDocumentFile)) {
                 context.log(logger, "Certificate/private key/identity document file does not exist");
                 Files.createDirectories(privateKeyFile.getParent());
                 Files.createDirectories(certificateFile.getParent());
                 Files.createDirectories(identityDocumentFile.getParent());
-                registerIdentity(context);
+                registerIdentity(context, privateKeyFile, certificateFile, identityDocumentFile);
                 return;
             }
 
@@ -117,21 +109,21 @@ public class AthenzCredentialsMaintainer {
             Instant expiry = certificate.getNotAfter().toInstant();
             if (isCertificateExpired(expiry, now)) {
                 context.log(logger, "Certificate has expired (expiry=%s)", expiry.toString());
-                registerIdentity(context);
+                registerIdentity(context, privateKeyFile, certificateFile, identityDocumentFile);
                 return;
             }
 
             Duration age = Duration.between(certificate.getNotBefore().toInstant(), now);
             if (shouldRefreshCredentials(age)) {
                 context.log(logger, "Certificate is ready to be refreshed (age=%s)", age.toString());
-                if (shouldThrottleRefreshAttempts(now)) {
+                if (shouldThrottleRefreshAttempts(context.containerName(), now)) {
                     context.log(logger, LogLevel.WARNING, String.format(
                             "Skipping refresh attempt as last refresh was on %s (less than %s ago)",
-                            lastRefreshAttempt.toString(), REFRESH_BACKOFF.toString()));
+                            lastRefreshAttempt.get(context.containerName()).toString(), REFRESH_BACKOFF.toString()));
                     return;
                 } else {
-                    lastRefreshAttempt = now;
-                    refreshIdentity(context);
+                    lastRefreshAttempt.put(context.containerName(), now);
+                    refreshIdentity(context, privateKeyFile, certificateFile, identityDocumentFile);
                     return;
                 }
             }
@@ -141,45 +133,40 @@ public class AthenzCredentialsMaintainer {
         }
     }
 
-    public void clearCredentials() {
-        if (!enabled) return;
-        try {
-            if (Files.deleteIfExists(privateKeyFile))
-                log.info(String.format("Deleted private key file (path=%s)", privateKeyFile));
-            if (Files.deleteIfExists(certificateFile))
-                log.info(String.format("Deleted certificate file (path=%s)", certificateFile));
-            if (Files.deleteIfExists(identityDocumentFile))
-                log.info(String.format("Deleted identity document file (path=%s)", certificateFile));
-        } catch (IOException e) {
-            throw new UncheckedIOException(e);
-        }
+    public void clearCredentials(NodeAgentContext context) {
+        FileFinder.files(context.pathOnHostFromPathInNode(CONTAINER_SIA_DIRECTORY))
+                .deleteRecursively();
+        lastRefreshAttempt.remove(context.containerName());
     }
 
     private boolean shouldRefreshCredentials(Duration age) {
         return age.compareTo(REFRESH_PERIOD) >= 0;
     }
 
-    private boolean shouldThrottleRefreshAttempts(Instant now) {
-        return REFRESH_BACKOFF.compareTo(Duration.between(lastRefreshAttempt, now)) > 0;
+    private boolean shouldThrottleRefreshAttempts(ContainerName containerName, Instant now) {
+        return REFRESH_BACKOFF.compareTo(
+                Duration.between(
+                        lastRefreshAttempt.getOrDefault(containerName, Instant.EPOCH),
+                        now)) > 0;
     }
 
     @SuppressWarnings("deprecation")
-    private void registerIdentity(NodeAgentContext context) {
+    private void registerIdentity(NodeAgentContext context, Path privateKeyFile, Path certificateFile, Path identityDocumentFile) {
         KeyPair keyPair = KeyUtils.generateKeypair(KeyAlgorithm.RSA);
-        SignedIdentityDocument signedIdentityDocument = identityDocumentClient.getNodeIdentityDocument(hostname);
+        SignedIdentityDocument signedIdentityDocument = identityDocumentClient.getNodeIdentityDocument(context.hostname().value());
         com.yahoo.vespa.athenz.tls.Pkcs10Csr csr = csrGenerator.generateInstanceCsr(
-                containerIdentity, signedIdentityDocument.providerUniqueId(), signedIdentityDocument.ipAddresses(), keyPair);
+                context.identity(), signedIdentityDocument.providerUniqueId(), signedIdentityDocument.ipAddresses(), keyPair);
         try (ZtsClient ztsClient = new DefaultZtsClient(ztsEndpoint, hostIdentityProvider)) {
             InstanceIdentity instanceIdentity =
                     ztsClient.registerInstance(
                             configserverIdentity,
-                            containerIdentity,
+                            context.identity(),
                             signedIdentityDocument.providerUniqueId().asDottedString(),
                             EntityBindingsMapper.toAttestationData(signedIdentityDocument),
                             false,
                             csr);
             EntityBindingsMapper.writeSignedIdentityDocumentToFile(identityDocumentFile, signedIdentityDocument);
-            writePrivateKeyAndCertificate(keyPair.getPrivate(), instanceIdentity.certificate());
+            writePrivateKeyAndCertificate(privateKeyFile, keyPair.getPrivate(), certificateFile, instanceIdentity.certificate());
             context.log(logger, "Instance successfully registered and credentials written to file");
         } catch (IOException e) {
             throw new UncheckedIOException(e);
@@ -187,30 +174,31 @@ public class AthenzCredentialsMaintainer {
     }
 
     @SuppressWarnings("deprecation")
-    private void refreshIdentity(NodeAgentContext context) {
+    private void refreshIdentity(NodeAgentContext context, Path privateKeyFile, Path certificateFile, Path identityDocumentFile) {
         SignedIdentityDocument identityDocument = EntityBindingsMapper.readSignedIdentityDocumentFromFile(identityDocumentFile);
         KeyPair keyPair = KeyUtils.generateKeypair(KeyAlgorithm.RSA);
-        com.yahoo.vespa.athenz.tls.Pkcs10Csr csr = csrGenerator.generateInstanceCsr(containerIdentity, identityDocument.providerUniqueId(), identityDocument.ipAddresses(), keyPair);
+        com.yahoo.vespa.athenz.tls.Pkcs10Csr csr = csrGenerator.generateInstanceCsr(
+                context.identity(), identityDocument.providerUniqueId(), identityDocument.ipAddresses(), keyPair);
         SSLContext containerIdentitySslContext =
                 new SslContextBuilder()
                         .withKeyStore(privateKeyFile, certificateFile)
                         .withTrustStore(trustStorePath, KeyStoreType.JKS)
                         .build();
         try {
-            try (ZtsClient ztsClient = new DefaultZtsClient(ztsEndpoint, containerIdentity, containerIdentitySslContext)) {
+            try (ZtsClient ztsClient = new DefaultZtsClient(ztsEndpoint, context.identity(), containerIdentitySslContext)) {
                 InstanceIdentity instanceIdentity =
                         ztsClient.refreshInstance(
                                 configserverIdentity,
-                                containerIdentity,
+                                context.identity(),
                                 identityDocument.providerUniqueId().asDottedString(),
                                 false,
                                 csr);
-                writePrivateKeyAndCertificate(keyPair.getPrivate(), instanceIdentity.certificate());
+                writePrivateKeyAndCertificate(privateKeyFile, keyPair.getPrivate(), certificateFile, instanceIdentity.certificate());
                 context.log(logger, "Instance successfully refreshed and credentials written to file");
             } catch (ZtsClientException e) {
                 if (e.getErrorCode() == 403 && e.getDescription().startsWith("Certificate revoked")) {
                     context.log(logger, LogLevel.ERROR, "Certificate cannot be refreshed as it is revoked by ZTS - re-registering the instance now", e);
-                    registerIdentity(context);
+                    registerIdentity(context, privateKeyFile, certificateFile, identityDocumentFile);
                 } else {
                     throw e;
                 }
@@ -220,7 +208,9 @@ public class AthenzCredentialsMaintainer {
         }
     }
 
-    private void writePrivateKeyAndCertificate(PrivateKey privateKey, X509Certificate certificate) throws IOException {
+
+    private static void writePrivateKeyAndCertificate(
+            Path privateKeyFile, PrivateKey privateKey, Path certificateFile, X509Certificate certificate) throws IOException {
         Path tempPrivateKeyFile = toTempPath(privateKeyFile);
         Files.write(tempPrivateKeyFile, KeyUtils.toPem(privateKey).getBytes());
         Path tempCertificateFile = toTempPath(certificateFile);
diff --git a/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/nodeagent/NodeAgentImpl.java b/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/nodeagent/NodeAgentImpl.java
index a7ed5392e14..a76d0c8b91b 100644
--- a/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/nodeagent/NodeAgentImpl.java
+++ b/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/nodeagent/NodeAgentImpl.java
@@ -76,7 +76,7 @@ public class NodeAgentImpl implements NodeAgent {
     private final Environment environment;
     private final Clock clock;
     private final Duration timeBetweenEachConverge;
-    private final AthenzCredentialsMaintainer athenzCredentialsMaintainer;
+    private final Optional<AthenzCredentialsMaintainer> athenzCredentialsMaintainer;
 
     private int numberOfUnhandledException = 0;
     private Instant lastConverge;
@@ -120,7 +120,7 @@ public class NodeAgentImpl implements NodeAgent {
             final Environment environment,
             final Clock clock,
             final Duration timeBetweenEachConverge,
-            final AthenzCredentialsMaintainer athenzCredentialsMaintainer) {
+            final Optional<AthenzCredentialsMaintainer> athenzCredentialsMaintainer) {
         this.context = context;
         this.nodeRepository = nodeRepository;
         this.orchestrator = orchestrator;
@@ -516,7 +516,7 @@ public class NodeAgentImpl implements NodeAgent {
                 startServicesIfNeeded();
                 resumeNodeIfNeeded(node);
 
-                athenzCredentialsMaintainer.converge(context);
+                athenzCredentialsMaintainer.ifPresent(maintainer -> maintainer.converge(context));
 
                 doBeforeConverge(node);
 
@@ -544,7 +544,7 @@ public class NodeAgentImpl implements NodeAgent {
             case dirty:
                 removeContainerIfNeededUpdateContainerState(node, container);
                 context.log(logger, "State is " + node.getState() + ", will delete application storage and mark node as ready");
-                athenzCredentialsMaintainer.clearCredentials();
+                athenzCredentialsMaintainer.ifPresent(maintainer -> maintainer.clearCredentials(context));
                 storageMaintainer.archiveNodeStorage(context);
                 updateNodeRepoWithCurrentAttributes(node);
                 nodeRepository.setNodeState(context.hostname().value(), Node.State.ready);
diff --git a/node-admin/src/test/java/com/yahoo/vespa/hosted/node/admin/integrationTests/DockerTester.java b/node-admin/src/test/java/com/yahoo/vespa/hosted/node/admin/integrationTests/DockerTester.java
index ef369a21c37..bca9dbc44fa 100644
--- a/node-admin/src/test/java/com/yahoo/vespa/hosted/node/admin/integrationTests/DockerTester.java
+++ b/node-admin/src/test/java/com/yahoo/vespa/hosted/node/admin/integrationTests/DockerTester.java
@@ -14,7 +14,6 @@ import com.yahoo.vespa.hosted.node.admin.docker.DockerNetworking;
 import com.yahoo.vespa.hosted.node.admin.docker.DockerOperations;
 import com.yahoo.vespa.hosted.node.admin.docker.DockerOperationsImpl;
 import com.yahoo.vespa.hosted.node.admin.maintenance.acl.AclMaintainer;
-import com.yahoo.vespa.hosted.node.admin.maintenance.identity.AthenzCredentialsMaintainer;
 import com.yahoo.vespa.hosted.node.admin.nodeadmin.NodeAdminImpl;
 import com.yahoo.vespa.hosted.node.admin.nodeadmin.NodeAdminStateUpdaterImpl;
 import com.yahoo.vespa.hosted.node.admin.nodeagent.NodeAgent;
@@ -93,13 +92,12 @@ public class DockerTester implements AutoCloseable {
         DockerOperations dockerOperations = new DockerOperationsImpl(dockerMock, environment, processExecuter);
         StorageMaintainerMock storageMaintainer = new StorageMaintainerMock(dockerOperations, null, environment, callOrderVerifier);
         AclMaintainer aclMaintainer = mock(AclMaintainer.class);
-        AthenzCredentialsMaintainer athenzCredentialsMaintainer = mock(AthenzCredentialsMaintainer.class);
 
 
         MetricReceiverWrapper mr = new MetricReceiverWrapper(MetricReceiver.nullImplementation);
         Function<String, NodeAgent> nodeAgentFactory = (hostName) -> new NodeAgentImpl(
                 NodeAgentContextImplTest.nodeAgentFromHostname(fileSystem, hostName), nodeRepositoryMock,
-                orchestratorMock, dockerOperations, storageMaintainer, aclMaintainer, environment, clock, NODE_AGENT_SCAN_INTERVAL, athenzCredentialsMaintainer);
+                orchestratorMock, dockerOperations, storageMaintainer, aclMaintainer, environment, clock, NODE_AGENT_SCAN_INTERVAL, Optional.empty());
         nodeAdmin = new NodeAdminImpl(nodeAgentFactory, aclMaintainer, mr, Clock.systemUTC());
         nodeAdminStateUpdater = new NodeAdminStateUpdaterImpl(nodeRepositoryMock, orchestratorMock, storageMaintainer,
                 nodeAdmin, DOCKER_HOST_HOSTNAME, clock, NODE_ADMIN_CONVERGE_STATE_INTERVAL,
diff --git a/node-admin/src/test/java/com/yahoo/vespa/hosted/node/admin/nodeagent/NodeAgentImplTest.java b/node-admin/src/test/java/com/yahoo/vespa/hosted/node/admin/nodeagent/NodeAgentImplTest.java
index b2facd3dd6d..5cc881fe62d 100644
--- a/node-admin/src/test/java/com/yahoo/vespa/hosted/node/admin/nodeagent/NodeAgentImplTest.java
+++ b/node-admin/src/test/java/com/yahoo/vespa/hosted/node/admin/nodeagent/NodeAgentImplTest.java
@@ -6,7 +6,6 @@ import com.yahoo.config.provision.NodeType;
 import com.yahoo.metrics.simple.MetricReceiver;
 import com.yahoo.test.ManualClock;
 import com.yahoo.vespa.hosted.dockerapi.Container;
-import com.yahoo.vespa.hosted.dockerapi.ContainerName;
 import com.yahoo.vespa.hosted.dockerapi.ContainerResources;
 import com.yahoo.vespa.hosted.dockerapi.ContainerStats;
 import com.yahoo.vespa.hosted.dockerapi.exception.DockerException;
@@ -761,7 +760,7 @@ public class NodeAgentImplTest {
         doNothing().when(storageMaintainer).writeMetricsConfig(any(), any());
 
         return new NodeAgentImpl(context, nodeRepository, orchestrator, dockerOperations,
-                storageMaintainer, aclMaintainer, environment, clock, NODE_AGENT_SCAN_INTERVAL, athenzCredentialsMaintainer);
+                storageMaintainer, aclMaintainer, environment, clock, NODE_AGENT_SCAN_INTERVAL, Optional.of(athenzCredentialsMaintainer));
     }
 
     private void mockGetContainer(DockerImage dockerImage, boolean isRunning) {