Make AwsCredentials Closeable, and close created clients in its close()

2 files changed, 17 insertions, 5 deletions

diff --git a/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/aws/AwsCredentials.java b/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/aws/AwsCredentials.java
index 8e9d00d7fa5..c9a5dbbcbfc 100644
--- a/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/aws/AwsCredentials.java
+++ b/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/aws/AwsCredentials.java
@@ -19,13 +19,14 @@ import java.util.Optional;
  *
  * @author tokle
  */
-public class AwsCredentials {
+public class AwsCredentials implements AutoCloseable {
 
     private final static Duration MIN_EXPIRY = Duration.ofMinutes(5);
     private final AthenzDomain athenzDomain;
     private final AwsRole awsRole;
     private final ZtsClient ztsClient;
     private final String externalId;
+    private final boolean close;
     private volatile AwsTemporaryCredentials credentials;
 
     public AwsCredentials(ZtsClient ztsClient, AthenzDomain athenzDomain, AwsRole awsRole) {
@@ -33,23 +34,28 @@ public class AwsCredentials {
     }
 
     public AwsCredentials(ZtsClient ztsClient, AthenzDomain athenzDomain, AwsRole awsRole, String externalId) {
+        this(ztsClient, athenzDomain, awsRole, externalId, false);
+    }
+
+    private AwsCredentials(ZtsClient ztsClient, AthenzDomain athenzDomain, AwsRole awsRole, String externalId, boolean close) {
         this.ztsClient = ztsClient;
         this.athenzDomain = athenzDomain;
         this.awsRole = awsRole;
         this.externalId = externalId;
+        this.close = close;
         this.credentials = get();
     }
 
     public AwsCredentials(URI ztsUrl, ServiceIdentityProvider identityProvider, AthenzDomain athenzDomain, AwsRole awsRole) {
-        this(new DefaultZtsClient.Builder(ztsUrl).withIdentityProvider(identityProvider).build(), athenzDomain, awsRole);
+        this(ztsUrl, identityProvider.getIdentitySslContext(), athenzDomain, awsRole);
     }
 
     public AwsCredentials(URI ztsUrl, SSLContext sslContext, AthenzDomain athenzDomain, AwsRole awsRole) {
-        this(new DefaultZtsClient.Builder(ztsUrl).withSslContext(sslContext).build(), athenzDomain, awsRole);
+        this(ztsUrl, sslContext, athenzDomain, awsRole, null);
     }
 
     public AwsCredentials(URI ztsUrl, SSLContext sslContext, AthenzDomain athenzDomain, AwsRole awsRole, String externalId) {
-        this(new DefaultZtsClient.Builder(ztsUrl).withSslContext(sslContext).build(), athenzDomain, awsRole, externalId);
+        this(new DefaultZtsClient.Builder(ztsUrl).withSslContext(sslContext).build(), athenzDomain, awsRole, externalId, true);
     }
 
     /**
@@ -70,4 +76,9 @@ public class AwsCredentials {
         return Duration.between(Instant.now(), expiration).toMinutes() < MIN_EXPIRY.toMinutes();
     }
 
+    @Override
+    public void close() {
+        if (close) ztsClient.close();
+    }
+
 }
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 b07f6da1a01..80a0ddff204 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
@@ -11,6 +11,7 @@ import com.yahoo.vespa.athenz.api.AthenzRoleInformation;
 import com.yahoo.vespa.athenz.api.AthenzService;
 import com.yahoo.vespa.athenz.api.OAuthCredentials;
 
+import java.io.Closeable;
 import java.time.Instant;
 import java.util.List;
 import java.util.Map;
@@ -20,7 +21,7 @@ import java.util.Set;
 /**
  * @author bjorncs
  */
-public interface ZmsClient extends AutoCloseable {
+public interface ZmsClient extends Closeable {
 
     void createTenancy(AthenzDomain tenantDomain, AthenzIdentity providerService, OAuthCredentials oAuthCredentials);