From 95e723b7ea1ad74fa799bec42300e3aaa072ca0c Mon Sep 17 00:00:00 2001
From: Morten Tokle <mortent@verizonmedia.com>
Date: Mon, 21 Jun 2021 10:16:34 +0200
Subject: Add optional reason to membership request

---
 .../yahoo/vespa/athenz/client/zms/DefaultZmsClient.java  | 16 +++++++++++-----
 .../com/yahoo/vespa/athenz/client/zms/ZmsClient.java     |  3 ++-
 2 files changed, 13 insertions(+), 6 deletions(-)

(limited to 'vespa-athenz')

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 5817eb0c8d2..7503b5a39ed 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
@@ -32,6 +32,8 @@ import java.net.URI;
 import java.time.Instant;
 import java.util.Collections;
 import java.util.List;
+import java.util.Objects;
+import java.util.Optional;
 import java.util.OptionalInt;
 import java.util.Set;
 import java.util.function.Function;
@@ -111,13 +113,17 @@ public class DefaultZmsClient extends ClientBase implements ZmsClient {
     }
 
     @Override
-    public void addRoleMember(AthenzRole role, AthenzIdentity member) {
+    public void addRoleMember(AthenzRole role, AthenzIdentity member, Optional<String> reason) {
         URI uri = zmsUrl.resolve(String.format("domain/%s/role/%s/member/%s", role.domain().getName(), role.roleName(), member.getFullName()));
         MembershipEntity membership = new MembershipEntity.RoleMembershipEntity(member.getFullName(), true, role.roleName(), null);
-        HttpUriRequest request = RequestBuilder.put(uri)
-                .setEntity(toJsonStringEntity(membership))
-                .build();
-        execute(request, response -> readEntity(response, Void.class));
+
+
+        RequestBuilder requestBuilder = RequestBuilder.put(uri)
+                .setEntity(toJsonStringEntity(membership));
+        if (reason.filter(s -> !s.isBlank()).isPresent()) {
+            requestBuilder.addHeader("Y-Audit-Ref", reason.get());
+        }
+        execute(requestBuilder.build(), response -> readEntity(response, Void.class));
     }
 
     @Override
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 245078e3679..03afc9278cc 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
@@ -12,6 +12,7 @@ import com.yahoo.vespa.athenz.api.OktaIdentityToken;
 
 import java.time.Instant;
 import java.util.List;
+import java.util.Optional;
 import java.util.Set;
 
 /**
@@ -31,7 +32,7 @@ public interface ZmsClient extends AutoCloseable {
     void deleteProviderResourceGroup(AthenzDomain tenantDomain, AthenzIdentity providerService, String resourceGroup,
                                      OktaIdentityToken identityToken, OktaAccessToken accessToken);
 
-    void addRoleMember(AthenzRole role, AthenzIdentity member);
+    void addRoleMember(AthenzRole role, AthenzIdentity member, Optional<String> reason);
 
     void deleteRoleMember(AthenzRole role, AthenzIdentity member);
 
-- 
cgit v1.2.3