2 files changed, 40 insertions, 55 deletions

diff --git a/jdisc-cloud-aws/src/main/java/com/yahoo/jdisc/cloud/aws/AwsParameterStore.java b/jdisc-cloud-aws/src/main/java/com/yahoo/jdisc/cloud/aws/AwsParameterStore.java
index 1636c6aeb6d..e23dc54f4c6 100644
--- a/jdisc-cloud-aws/src/main/java/com/yahoo/jdisc/cloud/aws/AwsParameterStore.java
+++ b/jdisc-cloud-aws/src/main/java/com/yahoo/jdisc/cloud/aws/AwsParameterStore.java
@@ -9,55 +9,56 @@ import com.amazonaws.services.simplesystemsmanagement.AWSSimpleSystemsManagement
 import com.amazonaws.services.simplesystemsmanagement.AWSSimpleSystemsManagementClient;
 import com.amazonaws.services.simplesystemsmanagement.model.GetParametersRequest;
 import com.amazonaws.services.simplesystemsmanagement.model.GetParametersResult;
+import com.google.inject.Inject;
 import com.yahoo.cloud.config.SecretStoreConfig;
+import com.yahoo.component.AbstractComponent;
 import com.yahoo.container.jdisc.secretstore.SecretNotFoundException;
 import com.yahoo.container.jdisc.secretstore.SecretStore;
 
 /**
  * @author mortent
  */
-public class AwsParameterStore implements SecretStore {
+public class AwsParameterStore extends AbstractComponent implements SecretStore {
 
     private final VespaAwsCredentialsProvider credentialsProvider;
-    private final String roleToAssume;
-    private final String externalId;
-    private final String region;
+    private final SecretStoreConfig secretStoreConfig;
 
-    AwsParameterStore(VespaAwsCredentialsProvider credentialsProvider, String roleToAssume, String externalId, String region) {
-        this.credentialsProvider = credentialsProvider;
-        this.roleToAssume = roleToAssume;
-        this.externalId = externalId;
-        this.region = region;
+    @Inject
+    public AwsParameterStore(SecretStoreConfig secretStoreConfig) {
+        this.secretStoreConfig = secretStoreConfig;
+        this.credentialsProvider = new VespaAwsCredentialsProvider();
     }
 
     @Override
     public String getSecret(String key) {
-        AWSSecurityTokenService tokenService = AWSSecurityTokenServiceClientBuilder
-                .standard()
-                .withRegion(region)
-                .withCredentials(credentialsProvider)
-                .build();
+        for (var group : secretStoreConfig.groups()) {
+            AWSSecurityTokenService tokenService = AWSSecurityTokenServiceClientBuilder
+                    .standard()
+                    .withRegion(group.region())
+                    .withCredentials(credentialsProvider)
+                    .build();
 
-        STSAssumeRoleSessionCredentialsProvider assumeExtAccountRole = new STSAssumeRoleSessionCredentialsProvider
-                .Builder(roleToAssume, "vespa")
-                .withExternalId(externalId)
-                .withStsClient(tokenService)
-                .build();
+            STSAssumeRoleSessionCredentialsProvider assumeExtAccountRole = new STSAssumeRoleSessionCredentialsProvider
+                    .Builder(toRoleArn(group.awsId(), group.role()), "vespa")
+                    .withExternalId(group.externalId())
+                    .withStsClient(tokenService)
+                    .build();
 
-        AWSSimpleSystemsManagement client = AWSSimpleSystemsManagementClient.builder()
-                .withCredentials(assumeExtAccountRole)
-                .withRegion(region)
-                .build();
+            AWSSimpleSystemsManagement client = AWSSimpleSystemsManagementClient.builder()
+                    .withCredentials(assumeExtAccountRole)
+                    .withRegion(group.region())
+                    .build();
 
-        GetParametersRequest parametersRequest = new GetParametersRequest().withNames(key).withWithDecryption(true);
-        GetParametersResult parameters = client.getParameters(parametersRequest);
-        int count = parameters.getParameters().size();
-        if (count < 1) {
-            throw new SecretNotFoundException("Could not find secret " + key + " using role " + roleToAssume);
-        } else if (count > 1) {
-            throw new RuntimeException("Found too many parameters, expected 1, but found " + count);
+            GetParametersRequest parametersRequest = new GetParametersRequest().withNames(key).withWithDecryption(true);
+            GetParametersResult parameters = client.getParameters(parametersRequest);
+            int count = parameters.getParameters().size();
+            if (count == 1) {
+                return parameters.getParameters().get(0).getValue();
+            } else if (count > 1) {
+                throw new RuntimeException("Found too many parameters, expected 1, but found " + count);
+            }
         }
-        return parameters.getParameters().get(0).getValue();
+        throw new SecretNotFoundException("Could not find secret " + key + " in any configured secret store");
     }
 
     @Override
@@ -65,4 +66,8 @@ public class AwsParameterStore implements SecretStore {
         // TODO
         return getSecret(key);
     }
+
+    private String toRoleArn(String awsId, String role) {
+        return "arn:aws:iam::" + awsId + ":role/" + role;
+    }
 }
diff --git a/jdisc-cloud-aws/src/main/java/com/yahoo/jdisc/cloud/aws/AwsParameterStoreValidationHandler.java b/jdisc-cloud-aws/src/main/java/com/yahoo/jdisc/cloud/aws/AwsParameterStoreValidationHandler.java
index 5d5cad2f75d..8b6e3d52d37 100644
--- a/jdisc-cloud-aws/src/main/java/com/yahoo/jdisc/cloud/aws/AwsParameterStoreValidationHandler.java
+++ b/jdisc-cloud-aws/src/main/java/com/yahoo/jdisc/cloud/aws/AwsParameterStoreValidationHandler.java
@@ -28,22 +28,14 @@ import java.util.logging.Logger;
 public class AwsParameterStoreValidationHandler extends LoggingRequestHandler {
 
     private static final Logger log = Logger.getLogger(AwsParameterStoreValidationHandler.class.getName());
-    private final VespaAwsCredentialsProvider credentialsProvider;
-    private final SecretStoreConfig secretStoreConfig;
+    private final AwsParameterStore awsParameterStore;
 
     @Inject
-    public AwsParameterStoreValidationHandler(Context ctx, SecretStoreConfig secretStoreConfig) {
-        this(ctx, secretStoreConfig, new VespaAwsCredentialsProvider());
-    }
-
-
-    public AwsParameterStoreValidationHandler(Context ctx, SecretStoreConfig secretStoreConfig, VespaAwsCredentialsProvider credentialsProvider) {
+    public AwsParameterStoreValidationHandler(Context ctx, AwsParameterStore awsParameterStore) {
         super(ctx);
-        this.credentialsProvider = credentialsProvider;
-        this.secretStoreConfig = secretStoreConfig;
+        this.awsParameterStore = awsParameterStore;
     }
 
-
     @Override
     public HttpResponse handle(HttpRequest request) {
         try {
@@ -66,10 +58,7 @@ public class AwsParameterStoreValidationHandler extends LoggingRequestHandler {
         settings.toSlime(root.setObject("settings"));
 
         try {
-            var arn = "arn:aws:iam::" + settings.awsId + ":role/" + settings.role;
-            var region = getRegion(settings);
-            var store = new AwsParameterStore(this.credentialsProvider, arn, settings.externalId, region);
-            store.getSecret("vespa-secret");
+            awsParameterStore.getSecret("vespa-secret");
             root.setString("status", "ok");
         } catch (RuntimeException e) {
             root.setString("status", "error");
@@ -90,15 +79,6 @@ public class AwsParameterStoreValidationHandler extends LoggingRequestHandler {
         }
     }
 
-    private String getRegion(AwsSettings settings) {
-        return secretStoreConfig.groups()
-                .stream()
-                .filter(group -> group.name().equals(settings.name))
-                .map(SecretStoreConfig.Groups::region)
-                .findFirst()
-                .orElseThrow(() -> new RuntimeException("No secret store named '" + settings.name + "' configured in services.xml"));
-    }
-
     private static class AwsSettings {
         String name;
         String role;