Control rollout of new identity document layout with feature flag

5 files changed, 40 insertions, 13 deletions

diff --git a/flags/src/main/java/com/yahoo/vespa/flags/Flags.java b/flags/src/main/java/com/yahoo/vespa/flags/Flags.java
index 0a5cd1aacac..e89724dfb16 100644
--- a/flags/src/main/java/com/yahoo/vespa/flags/Flags.java
+++ b/flags/src/main/java/com/yahoo/vespa/flags/Flags.java
@@ -410,6 +410,12 @@ public class Flags {
             "Takes effect at redeployment",
             ZONE_ID);
 
+    public static final UnboundBooleanFlag NEW_IDDOC_LAYOUT = defineFeatureFlag(
+            "new_iddoc_layout", false, List.of("tokle", "bjorncs", "olaa"), "2023-04-24", "2023-05-31",
+            "Whether to use new identity document lauoyt",
+            "Takes effect on node reboot",
+            HOSTNAME);
+
     /** WARNING: public for testing: All flags should be defined in {@link Flags}. */
     public static UnboundBooleanFlag defineFeatureFlag(String flagId, boolean defaultValue, List<String> owners,
                                                        String createdAt, String expiresAt, String description,
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 2cb2fb12668..c9c76e1edd3 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
@@ -77,6 +77,7 @@ public class AthenzCredentialsMaintainer implements CredentialsMaintainer {
     private final ServiceIdentityProvider hostIdentityProvider;
     private final IdentityDocumentClient identityDocumentClient;
     private final BooleanFlag tenantServiceIdentityFlag;
+    private final BooleanFlag useNewIdentityDocumentLayout;
 
     // Used as an optimization to ensure ZTS is not DDoS'ed on continuously failing refresh attempts
     private final Map<ContainerName, Instant> lastRefreshAttempt = new ConcurrentHashMap<>();
@@ -98,6 +99,7 @@ public class AthenzCredentialsMaintainer implements CredentialsMaintainer {
                 new AthenzIdentityVerifier(Set.of(configServerInfo.getConfigServerIdentity())));
         this.clock = clock;
         this.tenantServiceIdentityFlag = Flags.NODE_ADMIN_TENANT_SERVICE_REGISTRY.bindTo(flagSource);
+        this.useNewIdentityDocumentLayout = Flags.NEW_IDDOC_LAYOUT.bindTo(flagSource);
     }
 
     public boolean converge(NodeAgentContext context) {
@@ -131,7 +133,7 @@ public class AthenzCredentialsMaintainer implements CredentialsMaintainer {
             Instant now = clock.instant();
             Instant expiry = certificate.getNotAfter().toInstant();
             var doc = EntityBindingsMapper.readSignedIdentityDocumentFromFile(identityDocumentFile);
-            if (doc.outdated()) {
+            if (refreshIdentityDocument(doc, context)) {
                 context.log(logger, "Identity document is outdated (version=%d)", doc.documentVersion());
                 registerIdentity(context, privateKeyFile, certificateFile, identityDocumentFile, identityType, athenzIdentity);
                 return true;
@@ -162,6 +164,11 @@ public class AthenzCredentialsMaintainer implements CredentialsMaintainer {
         }
     }
 
+    private boolean refreshIdentityDocument(SignedIdentityDocument signedIdentityDocument, NodeAgentContext context) {
+        int expectedVersion = documentVersion(context);
+        return signedIdentityDocument.outdated() || signedIdentityDocument.documentVersion() != expectedVersion;
+    }
+
     public void clearCredentials(NodeAgentContext context) {
         FileFinder.files(context.paths().of(CONTAINER_SIA_DIRECTORY))
                 .deleteRecursively(context);
@@ -293,8 +300,8 @@ public class AthenzCredentialsMaintainer implements CredentialsMaintainer {
 
     private SignedIdentityDocument signedIdentityDocument(NodeAgentContext context, IdentityType identityType) {
         return switch (identityType) {
-            case NODE -> identityDocumentClient.getNodeIdentityDocument(context.hostname().value());
-            case TENANT -> identityDocumentClient.getTenantIdentityDocument(context.hostname().value());
+            case NODE -> identityDocumentClient.getNodeIdentityDocument(context.hostname().value(), documentVersion(context));
+            case TENANT -> identityDocumentClient.getTenantIdentityDocument(context.hostname().value(), documentVersion(context));
         };
     }
 
@@ -309,7 +316,7 @@ public class AthenzCredentialsMaintainer implements CredentialsMaintainer {
         if (Files.exists(identityDocumentFile)) {
             return EntityBindingsMapper.readSignedIdentityDocumentFromFile(identityDocumentFile).identityDocument().serviceIdentity();
         } else {
-            return identityDocumentClient.getTenantIdentityDocument(context.hostname().value()).identityDocument().serviceIdentity();
+            return identityDocumentClient.getTenantIdentityDocument(context.hostname().value(), documentVersion(context)).identityDocument().serviceIdentity();
         }
     }
 
@@ -319,6 +326,17 @@ public class AthenzCredentialsMaintainer implements CredentialsMaintainer {
                 .value();
     }
 
+    /*
+    Get the document version to ask for
+     */
+    private int documentVersion(NodeAgentContext context) {
+        return useNewIdentityDocumentLayout
+                .with(FetchVector.Dimension.HOSTNAME, context.hostname().value())
+                .value()
+                ? SignedIdentityDocument.DEFAULT_DOCUMENT_VERSION
+                : SignedIdentityDocument.LEGACY_DEFAULT_DOCUMENT_VERSION;
+    }
+
     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 5a0f77ec765..0e13cba8de9 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,12 +1,14 @@
 // 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.OptionalInt;
+
 /**
  * A client that communicates that fetches an identity document.
  *
  * @author bjorncs
  */
 public interface IdentityDocumentClient {
-    SignedIdentityDocument getNodeIdentityDocument(String host);
-    SignedIdentityDocument getTenantIdentityDocument(String host);
+    SignedIdentityDocument getNodeIdentityDocument(String host, int documentVersion);
+    SignedIdentityDocument getTenantIdentityDocument(String host, int documentVersion);
 }
diff --git a/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/identityprovider/client/AthenzCredentialsService.java b/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/identityprovider/client/AthenzCredentialsService.java
index 6b167dcde21..1858653c9b4 100644
--- a/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/identityprovider/client/AthenzCredentialsService.java
+++ b/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/identityprovider/client/AthenzCredentialsService.java
@@ -75,7 +75,8 @@ class AthenzCredentialsService {
         }
         KeyPair keyPair = KeyUtils.generateKeypair(KeyAlgorithm.RSA);
         IdentityDocumentClient identityDocumentClient = createIdentityDocumentClient();
-        SignedIdentityDocument signedDocument = identityDocumentClient.getTenantIdentityDocument(hostname);
+        // Use legacy version for now.
+        SignedIdentityDocument signedDocument = identityDocumentClient.getTenantIdentityDocument(hostname, SignedIdentityDocument.LEGACY_DEFAULT_DOCUMENT_VERSION);
         IdentityDocument document = signedDocument.identityDocument();
         Pkcs10Csr csr = csrGenerator.generateInstanceCsr(
                 tenantIdentity,
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 36ae978c853..48fc021dced 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
@@ -56,16 +56,16 @@ public class DefaultIdentityDocumentClient implements IdentityDocumentClient {
     }
 
     @Override
-    public SignedIdentityDocument getNodeIdentityDocument(String host) {
-        return getIdentityDocument(host, "node");
+    public SignedIdentityDocument getNodeIdentityDocument(String host, int documentVersion) {
+        return getIdentityDocument(host, "node", documentVersion);
     }
 
     @Override
-    public SignedIdentityDocument getTenantIdentityDocument(String host) {
-        return getIdentityDocument(host, "tenant");
+    public SignedIdentityDocument getTenantIdentityDocument(String host, int documentVersion) {
+        return getIdentityDocument(host, "tenant", documentVersion);
     }
 
-    private SignedIdentityDocument getIdentityDocument(String host, String type) {
+    private SignedIdentityDocument getIdentityDocument(String host, String type, int documentVersion) {
 
         try (CloseableHttpClient client = createHttpClient(sslContextSupplier.get(), hostnameVerifier)) {
             URI uri = configserverUri
@@ -76,7 +76,7 @@ public class DefaultIdentityDocumentClient implements IdentityDocumentClient {
                     .setUri(uri)
                     .addHeader("Connection", "close")
                     .addHeader("Accept", "application/json")
-                    .addParameter("documentVersion", Integer.toString(SignedIdentityDocument.DEFAULT_DOCUMENT_VERSION))
+                    .addParameter("documentVersion", Integer.toString(documentVersion))
                     .build();
             try (CloseableHttpResponse response = client.execute(request)) {
                 String responseContent = EntityUtils.toString(response.getEntity());