Control rollout of new identity document layout with feature flag

1 files changed, 22 insertions, 4 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 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");