Merge pull request #26615 from vespa-engine/olaa/move-instance-registry-to-athenzcredentialsmaintainer

Tenant certificate creation in AthenzCredentialsMaintainer

1 files changed, 73 insertions, 15 deletions

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 15be7accb7d..d22fd667202 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
@@ -19,6 +19,10 @@ import com.yahoo.vespa.athenz.identityprovider.client.CsrGenerator;
 import com.yahoo.vespa.athenz.identityprovider.client.DefaultIdentityDocumentClient;
 import com.yahoo.vespa.athenz.tls.AthenzIdentityVerifier;
 import com.yahoo.vespa.athenz.utils.SiaUtils;
+import com.yahoo.vespa.flags.BooleanFlag;
+import com.yahoo.vespa.flags.FetchVector;
+import com.yahoo.vespa.flags.FlagSource;
+import com.yahoo.vespa.flags.Flags;
 import com.yahoo.vespa.hosted.node.admin.component.ConfigServerInfo;
 import com.yahoo.vespa.hosted.node.admin.container.ContainerName;
 import com.yahoo.vespa.hosted.node.admin.nodeagent.NodeAgentContext;
@@ -47,6 +51,9 @@ import java.util.concurrent.ConcurrentHashMap;
 import java.util.logging.Level;
 import java.util.logging.Logger;
 
+import static com.yahoo.vespa.hosted.node.admin.maintenance.identity.AthenzCredentialsMaintainer.IdentityType.NODE;
+import static com.yahoo.vespa.hosted.node.admin.maintenance.identity.AthenzCredentialsMaintainer.IdentityType.TENANT;
+
 /**
  * A maintainer that is responsible for providing and refreshing Athenz credentials for a container.
  *
@@ -68,7 +75,7 @@ public class AthenzCredentialsMaintainer implements CredentialsMaintainer {
     private final String certificateDnsSuffix;
     private final ServiceIdentityProvider hostIdentityProvider;
     private final IdentityDocumentClient identityDocumentClient;
-    private final boolean useInternalZts;
+    private final BooleanFlag tenantServiceIdentityFlag;
 
     // Used as an optimization to ensure ZTS is not DDoS'ed on continuously failing refresh attempts
     private final Map<ContainerName, Instant> lastRefreshAttempt = new ConcurrentHashMap<>();
@@ -78,7 +85,7 @@ public class AthenzCredentialsMaintainer implements CredentialsMaintainer {
                                        ConfigServerInfo configServerInfo,
                                        String certificateDnsSuffix,
                                        ServiceIdentityProvider hostIdentityProvider,
-                                       boolean useInternalZts,
+                                       FlagSource flagSource,
                                        Clock clock) {
         this.ztsEndpoint = ztsEndpoint;
         this.ztsTrustStorePath = ztsTrustStorePath;
@@ -89,24 +96,33 @@ public class AthenzCredentialsMaintainer implements CredentialsMaintainer {
                 hostIdentityProvider,
                 new AthenzIdentityVerifier(Set.of(configServerInfo.getConfigServerIdentity())));
         this.clock = clock;
-        this.useInternalZts = useInternalZts;
+        this.tenantServiceIdentityFlag = Flags.NODE_ADMIN_TENANT_SERVICE_REGISTRY.bindTo(flagSource);
     }
 
     public boolean converge(NodeAgentContext context) {
+        var modified = false;
+        modified |= maintain(context, NODE);
+        if (shouldWriteTenantServiceIdentity(context))
+            modified |= maintain(context, TENANT);
+        return modified;
+    }
+
+    private boolean maintain(NodeAgentContext context, IdentityType identityType) {
         if (context.isDisabled(NodeAgentTask.CredentialsMaintainer)) return false;
 
         try {
             context.log(logger, Level.FINE, "Checking certificate");
-            ContainerPath containerSiaDirectory = context.paths().of(CONTAINER_SIA_DIRECTORY, context.users().vespa());
-            ContainerPath privateKeyFile = (ContainerPath) SiaUtils.getPrivateKeyFile(containerSiaDirectory, context.identity());
-            ContainerPath certificateFile = (ContainerPath) SiaUtils.getCertificateFile(containerSiaDirectory, context.identity());
-            ContainerPath identityDocumentFile = containerSiaDirectory.resolve("vespa-node-identity-document.json");
+            ContainerPath siaDirectory = context.paths().of(CONTAINER_SIA_DIRECTORY, context.users().vespa());
+            ContainerPath identityDocumentFile = siaDirectory.resolve(identityType.getIdentityDocument());
+            AthenzIdentity athenzIdentity = getAthenzIdentity(context, identityType, identityDocumentFile);
+            ContainerPath privateKeyFile = (ContainerPath) SiaUtils.getPrivateKeyFile(siaDirectory, athenzIdentity);
+            ContainerPath certificateFile = (ContainerPath) SiaUtils.getCertificateFile(siaDirectory, athenzIdentity);
             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, privateKeyFile, certificateFile, identityDocumentFile);
+                registerIdentity(context, privateKeyFile, certificateFile, identityDocumentFile, identityType);
                 return true;
             }
 
@@ -116,11 +132,11 @@ public class AthenzCredentialsMaintainer implements CredentialsMaintainer {
             var doc = EntityBindingsMapper.readSignedIdentityDocumentFromFile(identityDocumentFile);
             if (doc.outdated()) {
                 context.log(logger, "Identity document is outdated (version=%d)", doc.documentVersion());
-                registerIdentity(context, privateKeyFile, certificateFile, identityDocumentFile);
+                registerIdentity(context, privateKeyFile, certificateFile, identityDocumentFile, identityType);
                 return true;
             } else if (isCertificateExpired(expiry, now)) {
                 context.log(logger, "Certificate has expired (expiry=%s)", expiry.toString());
-                registerIdentity(context, privateKeyFile, certificateFile, identityDocumentFile);
+                registerIdentity(context, privateKeyFile, certificateFile, identityDocumentFile, identityType);
                 return true;
             }
 
@@ -134,7 +150,7 @@ public class AthenzCredentialsMaintainer implements CredentialsMaintainer {
                     return false;
                 } else {
                     lastRefreshAttempt.put(context.containerName(), now);
-                    refreshIdentity(context, privateKeyFile, certificateFile, identityDocumentFile, doc);
+                    refreshIdentity(context, privateKeyFile, certificateFile, identityDocumentFile, doc, identityType);
                     return true;
                 }
             }
@@ -182,9 +198,9 @@ public class AthenzCredentialsMaintainer implements CredentialsMaintainer {
                         now)) > 0;
     }
 
-    private void registerIdentity(NodeAgentContext context, ContainerPath privateKeyFile, ContainerPath certificateFile, ContainerPath identityDocumentFile) {
+    private void registerIdentity(NodeAgentContext context, ContainerPath privateKeyFile, ContainerPath certificateFile, ContainerPath identityDocumentFile, IdentityType identityType) {
         KeyPair keyPair = KeyUtils.generateKeypair(KeyAlgorithm.RSA);
-        SignedIdentityDocument doc = identityDocumentClient.getNodeIdentityDocument(context.hostname().value());
+        SignedIdentityDocument doc = signedIdentityDocument(context, identityType);
         CsrGenerator csrGenerator = new CsrGenerator(certificateDnsSuffix, doc.providerService().getFullName());
         Pkcs10Csr csr = csrGenerator.generateInstanceCsr(
                 context.identity(), doc.providerUniqueId(), doc.ipAddresses(), doc.clusterType(), keyPair);
@@ -214,7 +230,7 @@ public class AthenzCredentialsMaintainer implements CredentialsMaintainer {
                 .orElse(ztsEndpoint);
     }
     private void refreshIdentity(NodeAgentContext context, ContainerPath privateKeyFile, ContainerPath certificateFile,
-                                 ContainerPath identityDocumentFile, SignedIdentityDocument doc) {
+                                 ContainerPath identityDocumentFile, SignedIdentityDocument doc, IdentityType identityType) {
         KeyPair keyPair = KeyUtils.generateKeypair(KeyAlgorithm.RSA);
         CsrGenerator csrGenerator = new CsrGenerator(certificateDnsSuffix, doc.providerService().getFullName());
         Pkcs10Csr csr = csrGenerator.generateInstanceCsr(
@@ -239,7 +255,7 @@ public class AthenzCredentialsMaintainer implements CredentialsMaintainer {
             } catch (ZtsClientException e) {
                 if (e.getErrorCode() == 403 && e.getDescription().startsWith("Certificate revoked")) {
                     context.log(logger, Level.SEVERE, "Certificate cannot be refreshed as it is revoked by ZTS - re-registering the instance now", e);
-                    registerIdentity(context, privateKeyFile, certificateFile, identityDocumentFile);
+                    registerIdentity(context, privateKeyFile, certificateFile, identityDocumentFile, identityType);
                 } else {
                     throw e;
                 }
@@ -272,4 +288,46 @@ public class AthenzCredentialsMaintainer implements CredentialsMaintainer {
     private static boolean isCertificateExpired(Instant expiry, Instant now) {
         return now.isAfter(expiry.minus(EXPIRY_MARGIN));
     }
+
+    private SignedIdentityDocument signedIdentityDocument(NodeAgentContext context, IdentityType identityType) {
+        return switch (identityType) {
+            case NODE -> identityDocumentClient.getNodeIdentityDocument(context.hostname().value());
+            case TENANT -> identityDocumentClient.getTenantIdentityDocument(context.hostname().value());
+        };
+    }
+
+    private AthenzIdentity getAthenzIdentity(NodeAgentContext context, IdentityType identityType, ContainerPath identityDocumentFile) {
+        return switch (identityType) {
+            case NODE -> context.identity();
+            case TENANT -> getTenantIdentity(context, identityDocumentFile);
+        };
+    }
+
+    private AthenzIdentity getTenantIdentity(NodeAgentContext context, ContainerPath identityDocumentFile) {
+        if (Files.exists(identityDocumentFile)) {
+            return EntityBindingsMapper.readSignedIdentityDocumentFromFile(identityDocumentFile).serviceIdentity();
+        } else {
+            return identityDocumentClient.getTenantIdentityDocument(context.hostname().value()).serviceIdentity();
+        }
+    }
+
+    private boolean shouldWriteTenantServiceIdentity(NodeAgentContext context) {
+        return tenantServiceIdentityFlag
+                .with(FetchVector.Dimension.HOSTNAME, context.hostname().value())
+                .value();
+    }
+
+    enum IdentityType {
+        NODE("vespa-node-identity-document.json"),
+        TENANT("vespa-tenant-identity-document.json");
+
+        private String identityDocument;
+        IdentityType(String identityDocument) {
+            this.identityDocument = identityDocument;
+        }
+
+        public String getIdentityDocument() {
+            return identityDocument;
+        }
+    }
 }