Merge pull request #26908 from vespa-engine/olaa/athenzcredsmaintainer-fetch-roles

AthenzCredentialsMaintainer maintains role certificates

3 files changed, 115 insertions, 7 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 d5b48bc2609..1c16340641d 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
@@ -9,6 +9,7 @@ import com.yahoo.security.Pkcs10Csr;
 import com.yahoo.security.SslContextBuilder;
 import com.yahoo.security.X509CertificateUtils;
 import com.yahoo.vespa.athenz.api.AthenzIdentity;
+import com.yahoo.vespa.athenz.api.AthenzRole;
 import com.yahoo.vespa.athenz.client.zts.DefaultZtsClient;
 import com.yahoo.vespa.athenz.client.zts.InstanceIdentity;
 import com.yahoo.vespa.athenz.client.zts.ZtsClient;
@@ -47,6 +48,7 @@ import java.security.cert.X509Certificate;
 import java.time.Clock;
 import java.time.Duration;
 import java.time.Instant;
+import java.util.List;
 import java.util.Map;
 import java.util.Optional;
 import java.util.Set;
@@ -122,6 +124,7 @@ public class AthenzCredentialsMaintainer implements CredentialsMaintainer {
         if (context.isDisabled(NodeAgentTask.CredentialsMaintainer)) return false;
 
         try {
+            var modified = false;
             context.log(logger, Level.FINE, "Checking certificate");
             ContainerPath siaDirectory = context.paths().of(CONTAINER_SIA_DIRECTORY, context.users().vespa());
             ContainerPath identityDocumentFile = siaDirectory.resolve(identityType.getIdentityDocument());
@@ -137,7 +140,7 @@ public class AthenzCredentialsMaintainer implements CredentialsMaintainer {
                 Files.createDirectories(certificateFile.getParent());
                 Files.createDirectories(identityDocumentFile.getParent());
                 registerIdentity(context, privateKeyFile, certificateFile, identityDocumentFile, identityType, athenzIdentity);
-                return true;
+                modified = true;
             }
 
             X509Certificate certificate = readCertificateFromFile(certificateFile);
@@ -147,11 +150,11 @@ public class AthenzCredentialsMaintainer implements CredentialsMaintainer {
             if (refreshIdentityDocument(doc, context)) {
                 context.log(logger, "Identity document is outdated (version=%d)", doc.documentVersion());
                 registerIdentity(context, privateKeyFile, certificateFile, identityDocumentFile, identityType, athenzIdentity);
-                return true;
+                modified = true;
             } else if (isCertificateExpired(expiry, now)) {
                 context.log(logger, "Certificate has expired (expiry=%s)", expiry.toString());
                 registerIdentity(context, privateKeyFile, certificateFile, identityDocumentFile, identityType, athenzIdentity);
-                return true;
+                modified = true;
             }
 
             Duration age = Duration.between(certificate.getNotBefore().toInstant(), now);
@@ -161,20 +164,79 @@ public class AthenzCredentialsMaintainer implements CredentialsMaintainer {
                     context.log(logger, Level.WARNING, String.format(
                             "Skipping refresh attempt as last refresh was on %s (less than %s ago)",
                             lastRefreshAttempt.get(context.containerName()).toString(), REFRESH_BACKOFF.toString()));
-                    return false;
                 } else {
                     lastRefreshAttempt.put(context.containerName(), now);
                     refreshIdentity(context, privateKeyFile, certificateFile, identityDocumentFile, doc.identityDocument(), identityType, athenzIdentity);
-                    return true;
+                    modified = true;
                 }
             }
-            context.log(logger, Level.FINE, "Certificate is still valid");
-            return false;
+
+            modified |= maintainRoleCertificates(context, siaDirectory, privateKeyFile, certificateFile, athenzIdentity, doc.identityDocument());
+            return modified;
         } catch (IOException e) {
             throw new UncheckedIOException(e);
         }
     }
 
+    private boolean maintainRoleCertificates(NodeAgentContext context,
+                                             ContainerPath siaDirectory,
+                                             ContainerPath privateKeyFile,
+                                             ContainerPath certificateFile,
+                                             AthenzIdentity identity,
+                                             IdentityDocument identityDocument) {
+        var modified = false;
+
+        for (var role : getRoleList(context)) {
+                try {
+                    var roleCertificatePath = siaDirectory.resolve("certs")
+                            .resolve(String.format("%s.cert.pem", role));
+                    if (!Files.exists(roleCertificatePath)) {
+                        writeRoleCertificate(context, privateKeyFile, certificateFile, roleCertificatePath, identity, identityDocument, role);
+                        modified = true;
+                    } else if (shouldRefreshCertificate(context, roleCertificatePath)) {
+                        writeRoleCertificate(context, privateKeyFile, certificateFile, roleCertificatePath, identity, identityDocument, role);
+                        modified = true;
+                    }
+                } catch (IOException e) {
+                    context.log(logger, Level.WARNING, "Failed to maintain role certificate " + role, e);
+                }
+        }
+        return modified;
+    }
+
+    private boolean shouldRefreshCertificate(NodeAgentContext context, ContainerPath certificatePath) throws IOException {
+        var certificate = readCertificateFromFile(certificatePath);
+        var now = clock.instant();
+        var shouldRefresh = now.isAfter(certificate.getNotAfter().toInstant()) ||
+                now.isBefore(certificate.getNotBefore().toInstant().plus(REFRESH_PERIOD));
+        return !shouldThrottleRefreshAttempts(context.containerName(), now) &&
+                shouldRefresh;
+    }
+
+    private void writeRoleCertificate(NodeAgentContext context,
+                                      ContainerPath privateKeyFile,
+                                      ContainerPath certificateFile,
+                                      ContainerPath roleCertificatePath,
+                                      AthenzIdentity identity,
+                                      IdentityDocument identityDocument,
+                                      String role) throws IOException {
+        HostnameVerifier ztsHostNameVerifier = (hostname, sslSession) -> true;
+        var athenzRole = AthenzRole.fromResourceNameString(role);
+        var privateKey = KeyUtils.fromPemEncodedPrivateKey(new String(Files.readAllBytes(privateKeyFile)));
+
+        var containerIdentitySslContext = new SslContextBuilder().withKeyStore(privateKeyFile, certificateFile)
+                .withTrustStore(ztsTrustStorePath)
+                .build();
+        try (ZtsClient ztsClient = new DefaultZtsClient.Builder(ztsEndpoint(identityDocument)).withSslContext(containerIdentitySslContext).withHostnameVerifier(ztsHostNameVerifier).build()) {
+            var csrGenerator = new CsrGenerator(certificateDnsSuffix, identityDocument.providerService().getFullName());
+            var csr = csrGenerator.generateRoleCsr(
+                    identity, athenzRole, identityDocument.providerUniqueId(), identityDocument.clusterType(), KeyUtils.toKeyPair(privateKey));
+            var roleCertificate = ztsClient.getRoleCertificate(athenzRole, csr);
+            writeFile(roleCertificatePath, X509CertificateUtils.toPem(roleCertificate));
+            context.log(logger, "Role certificate successfully retrieved written to file " + roleCertificatePath.pathInContainer());
+        }
+    }
+
     private boolean refreshIdentityDocument(SignedIdentityDocument signedIdentityDocument, NodeAgentContext context) {
         int expectedVersion = documentVersion(context);
         return signedIdentityDocument.outdated() || signedIdentityDocument.documentVersion() != expectedVersion;
@@ -375,6 +437,15 @@ public class AthenzCredentialsMaintainer implements CredentialsMaintainer {
                 : SignedIdentityDocument.LEGACY_DEFAULT_DOCUMENT_VERSION;
     }
 
+    private List<String> getRoleList(NodeAgentContext context) {
+        try {
+            return identityDocumentClient.getNodeRoles(context.hostname().value());
+        } catch (Exception e) {
+            context.log(logger, Level.WARNING, "Failed to retrieve role list", e);
+            return List.of();
+        }
+    }
+
     enum IdentityType {
         NODE("vespa-node-identity-document.json"),
         TENANT("vespa-tenant-identity-document.json");
diff --git a/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/identityprovider/api/IdentityDocumentClient.java b/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/identityprovider/api/IdentityDocumentClient.java
index a3c2f0264d3..522f40bc37d 100644
--- a/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/identityprovider/api/IdentityDocumentClient.java
+++ b/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/identityprovider/api/IdentityDocumentClient.java
@@ -1,6 +1,7 @@
 // Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 package com.yahoo.vespa.athenz.identityprovider.api;
 
+import java.util.List;
 import java.util.Optional;
 import java.util.OptionalInt;
 
@@ -12,4 +13,5 @@ import java.util.OptionalInt;
 public interface IdentityDocumentClient {
     SignedIdentityDocument getNodeIdentityDocument(String host, int documentVersion);
     Optional<SignedIdentityDocument> getTenantIdentityDocument(String host, int documentVersion);
+    List<String> getNodeRoles(String hostname);
 }
diff --git a/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/identityprovider/client/DefaultIdentityDocumentClient.java b/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/identityprovider/client/DefaultIdentityDocumentClient.java
index f95a3335c24..81aa6e5bd2a 100644
--- a/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/identityprovider/client/DefaultIdentityDocumentClient.java
+++ b/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/identityprovider/client/DefaultIdentityDocumentClient.java
@@ -7,6 +7,7 @@ import com.yahoo.vespa.athenz.identity.ServiceIdentityProvider;
 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.identityprovider.api.bindings.RoleListEntity;
 import com.yahoo.vespa.athenz.identityprovider.api.bindings.SignedIdentityDocumentEntity;
 import org.apache.http.client.config.RequestConfig;
 import org.apache.http.client.methods.CloseableHttpResponse;
@@ -23,6 +24,7 @@ import java.io.IOException;
 import java.io.UncheckedIOException;
 import java.net.URI;
 import java.time.Duration;
+import java.util.List;
 import java.util.Optional;
 import java.util.function.Supplier;
 
@@ -66,6 +68,39 @@ public class DefaultIdentityDocumentClient implements IdentityDocumentClient {
         return getIdentityDocument(host, "tenant", documentVersion);
     }
 
+    @Override
+    public List<String> getNodeRoles(String hostname) {
+        try (var client = createHttpClient(sslContextSupplier.get(), hostnameVerifier)) {
+            var uri = configserverUri
+                    .resolve(IDENTITY_DOCUMENT_API)
+                    .resolve("roles/")
+                    .resolve(hostname);
+
+            var request = RequestBuilder.get()
+                    .setUri(uri)
+                    .addHeader("Connection", "close")
+                    .addHeader("Accept", "application/json")
+                    .build();
+            try (var response = client.execute(request)) {
+                String responseContent = EntityUtils.toString(response.getEntity());
+                int statusCode = response.getStatusLine().getStatusCode();
+                if (statusCode >= 200 && statusCode <= 299) {
+                    var rolesEntity = objectMapper.readValue(responseContent, RoleListEntity.class);
+                    return rolesEntity.roles();
+                } else {
+                    throw new RuntimeException(
+                            String.format(
+                                    "Failed to retrieve roles for host %s: %d - %s",
+                                    hostname,
+                                    statusCode,
+                                    responseContent));
+                }
+            }
+        } catch (IOException e) {
+            throw new UncheckedIOException(e);
+        }
+    }
+
     private Optional<SignedIdentityDocument> getIdentityDocument(String host, String type, int documentVersion) {
 
         try (CloseableHttpClient client = createHttpClient(sslContextSupplier.get(), hostnameVerifier)) {