1 files changed, 36 insertions, 31 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;
+    }
 }