Merge pull request #25714 from vespa-engine/mortent/update-aws-account

Support updating athenz domain meta

5 files changed, 83 insertions, 2 deletions

diff --git a/controller-api/src/main/java/com/yahoo/vespa/hosted/controller/api/integration/athenz/AthenzDbMock.java b/controller-api/src/main/java/com/yahoo/vespa/hosted/controller/api/integration/athenz/AthenzDbMock.java
index 63dfff95c03..9a9c2af2d5d 100644
--- a/controller-api/src/main/java/com/yahoo/vespa/hosted/controller/api/integration/athenz/AthenzDbMock.java
+++ b/controller-api/src/main/java/com/yahoo/vespa/hosted/controller/api/integration/athenz/AthenzDbMock.java
@@ -28,7 +28,11 @@ public class AthenzDbMock {
     }
 
     public Domain getOrCreateDomain(AthenzDomain domain) {
-        return domains.computeIfAbsent(domain, Domain::new);
+        return this.getOrCreateDomain(domain, Map.of());
+    }
+
+    public Domain getOrCreateDomain(AthenzDomain domain, Map<String, Object> attributes) {
+        return domains.computeIfAbsent(domain, Domain::new).withAttributes(attributes);
     }
 
     public AthenzDbMock addHostedOperator(AthenzIdentity athenzIdentity) {
@@ -46,6 +50,7 @@ public class AthenzDbMock {
         public final List<Role> roles = new ArrayList<>();
         public final Map<String, Policy> policies = new HashMap<>();
         public boolean isVespaTenant = false;
+        public final Map<String, Object> attributes = new HashMap<>();
 
         public Domain(AthenzDomain name) {
             this.name = name;
@@ -72,6 +77,11 @@ public class AthenzDbMock {
             return this;
         }
 
+        public Domain withAttributes(Map<String, Object> attributes) {
+            this.attributes.putAll(attributes);
+            return this;
+        }
+
         /**
          * Simulates establishing Vespa tenancy in Athens.
          */
diff --git a/controller-api/src/main/java/com/yahoo/vespa/hosted/controller/api/integration/athenz/ZmsClientMock.java b/controller-api/src/main/java/com/yahoo/vespa/hosted/controller/api/integration/athenz/ZmsClientMock.java
index d5e815912c5..ed185f8af32 100644
--- a/controller-api/src/main/java/com/yahoo/vespa/hosted/controller/api/integration/athenz/ZmsClientMock.java
+++ b/controller-api/src/main/java/com/yahoo/vespa/hosted/controller/api/integration/athenz/ZmsClientMock.java
@@ -3,6 +3,7 @@ package com.yahoo.vespa.hosted.controller.api.integration.athenz;
 
 import com.yahoo.vespa.athenz.api.AthenzAssertion;
 import com.yahoo.vespa.athenz.api.AthenzDomain;
+import com.yahoo.vespa.athenz.api.AthenzDomainMeta;
 import com.yahoo.vespa.athenz.api.AthenzGroup;
 import com.yahoo.vespa.athenz.api.AthenzIdentity;
 import com.yahoo.vespa.athenz.api.AthenzPolicy;
@@ -139,6 +140,25 @@ public class ZmsClientMock implements ZmsClient {
     }
 
     @Override
+    public AthenzDomainMeta getDomainMeta(AthenzDomain domain) {
+        return Optional.ofNullable(athenz.domains.get(domain))
+                .map(d -> d.attributes)
+                .map(attrs -> {
+                    if (attrs.containsKey("account")) {
+                        return new AthenzDomainMeta((String)attrs.get("account"), domain.getName());
+                    }
+                    return null;
+                })
+                .orElse(null);
+    }
+
+    @Override
+    public void updateDomain(AthenzDomain domain, Map<String, Object> attributes) {
+        if (!athenz.domains.containsKey(domain)) throw new IllegalStateException("Domain does not exist: " + domain.getName());
+        athenz.domains.get(domain).withAttributes(attributes);
+    }
+
+    @Override
     public boolean hasAccess(AthenzResourceName resource, String action, AthenzIdentity identity) {
         log("hasAccess(resource=%s, action=%s, identity=%s)", resource, action, identity);
         if (resource.getDomain().equals(this.controllerIdentity.getDomain())) {
@@ -268,7 +288,11 @@ public class ZmsClientMock implements ZmsClient {
     }
 
     @Override
-    public void createSubdomain(AthenzDomain parent, String name, Map<String, Object> attributes) {}
+    public void createSubdomain(AthenzDomain parent, String name, Map<String, Object> attributes) {
+        AthenzDomain domain = new AthenzDomain(parent, name);
+        if (athenz.domains.containsKey(domain)) throw new IllegalStateException("Subdomain already exists: %s".formatted(domain.getName()));
+        athenz.getOrCreateDomain(domain, attributes);
+    }
 
     @Override
     public AthenzRoleInformation getFullRoleInformation(AthenzRole role) {
diff --git a/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/api/AthenzDomainMeta.java b/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/api/AthenzDomainMeta.java
new file mode 100644
index 00000000000..7808a689982
--- /dev/null
+++ b/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/api/AthenzDomainMeta.java
@@ -0,0 +1,12 @@
+// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+package com.yahoo.vespa.athenz.api;
+
+import com.fasterxml.jackson.annotation.JsonIgnoreProperties;
+import com.fasterxml.jackson.annotation.JsonProperty;
+
+@JsonIgnoreProperties(ignoreUnknown = true)
+public record AthenzDomainMeta(
+        @JsonProperty("account") String account,
+        @JsonProperty("name") String name
+) {
+}
diff --git a/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/client/zms/DefaultZmsClient.java b/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/client/zms/DefaultZmsClient.java
index 44ea5ef329f..d8d6d2c3486 100644
--- a/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/client/zms/DefaultZmsClient.java
+++ b/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/client/zms/DefaultZmsClient.java
@@ -5,6 +5,7 @@ import com.yahoo.athenz.auth.util.Crypto;
 import com.yahoo.security.KeyUtils;
 import com.yahoo.vespa.athenz.api.AthenzAssertion;
 import com.yahoo.vespa.athenz.api.AthenzDomain;
+import com.yahoo.vespa.athenz.api.AthenzDomainMeta;
 import com.yahoo.vespa.athenz.api.AthenzGroup;
 import com.yahoo.vespa.athenz.api.AthenzIdentity;
 import com.yahoo.vespa.athenz.api.AthenzPolicy;
@@ -33,6 +34,7 @@ import com.yahoo.vespa.athenz.utils.AthenzIdentities;
 import org.apache.http.Header;
 import org.apache.http.client.methods.HttpUriRequest;
 import org.apache.http.client.methods.RequestBuilder;
+import org.apache.http.entity.ContentType;
 import org.apache.http.entity.StringEntity;
 import org.apache.http.message.BasicHeader;
 
@@ -216,6 +218,34 @@ public class DefaultZmsClient extends ClientBase implements ZmsClient {
     }
 
     @Override
+    public AthenzDomainMeta getDomainMeta(AthenzDomain domain) {
+        HttpUriRequest request = RequestBuilder.get()
+                .setUri(zmsUrl.resolve("domain/%s".formatted(domain.getName())))
+                .build();
+        return execute(request, response -> readEntity(response, AthenzDomainMeta.class));
+    }
+
+    @Override
+    public void updateDomain(AthenzDomain domain, Map<String, Object> attributes) {
+        for (String attribute : attributes.keySet()) {
+            Object attrVal = attributes.get(attribute);
+
+            String val = attrVal instanceof String ? "\"" + attrVal.toString() + "\"" : attrVal.toString();
+            String domainMeta = """
+                    {
+                        "%s": %s
+                    }
+                    """
+                    .formatted(attribute, val);
+            HttpUriRequest request = RequestBuilder.put()
+                    .setUri(zmsUrl.resolve("domain/%s/meta/system/%s".formatted(domain.getName(), attribute)))
+                    .setEntity(new StringEntity(domainMeta, ContentType.APPLICATION_JSON))
+                    .build();
+            execute(request, response -> readEntity(response, Void.class));
+        }
+    }
+
+    @Override
     public boolean hasAccess(AthenzResourceName resource, String action, AthenzIdentity identity) {
         URI uri = zmsUrl.resolve(String.format("access/%s/%s?principal=%s",
                                                action, resource.toResourceNameString(), identity.getFullName()));
diff --git a/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/client/zms/ZmsClient.java b/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/client/zms/ZmsClient.java
index 4342b32e4c8..541aa4d9c60 100644
--- a/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/client/zms/ZmsClient.java
+++ b/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/client/zms/ZmsClient.java
@@ -2,6 +2,7 @@
 package com.yahoo.vespa.athenz.client.zms;
 
 import com.yahoo.vespa.athenz.api.AthenzDomain;
+import com.yahoo.vespa.athenz.api.AthenzDomainMeta;
 import com.yahoo.vespa.athenz.api.AthenzGroup;
 import com.yahoo.vespa.athenz.api.AthenzIdentity;
 import com.yahoo.vespa.athenz.api.AthenzPolicy;
@@ -52,6 +53,10 @@ public interface ZmsClient extends Closeable {
 
     List<AthenzDomain> getDomainListByAccount(String id);
 
+    AthenzDomainMeta getDomainMeta(AthenzDomain domain);
+
+    void updateDomain(AthenzDomain domain, Map<String, Object> attributes);
+
     boolean hasAccess(AthenzResourceName resource, String action, AthenzIdentity identity);
 
     void createPolicy(AthenzDomain athenzDomain, String athenzPolicy);