port multi-account support to aws2 helper (#392)

Ports the changes in #345 to `iep-module-aws2`.

iep-module-aws2/src/main/java/com/netflix/iep/aws2/AwsClientFactory.java

@@ -27,21 +27,25 @@
 import software.amazon.awssdk.services.sts.StsClient;
 import software.amazon.awssdk.services.sts.auth.StsAssumeRoleCredentialsProvider;
 import software.amazon.awssdk.services.sts.model.AssumeRoleRequest;
+import software.amazon.awssdk.utils.SdkAutoCloseable;
 
 import javax.inject.Inject;
 import javax.inject.Singleton;
 import java.lang.reflect.Method;
 import java.time.Duration;
+import java.util.concurrent.ConcurrentHashMap;
 import java.util.function.Consumer;
 
 /**
  * Factory for creating instances of AWS clients.
  */
 @Singleton
-public class AwsClientFactory {
+public class AwsClientFactory implements AutoCloseable {
 
   private static final Logger LOGGER = LoggerFactory.getLogger(AwsClientFactory.class);
 
+  private final ConcurrentHashMap<String, SdkAutoCloseable> clients = new ConcurrentHashMap<>();
+
   private final Config config;
   private final String region;
 
@@ -109,8 +113,24 @@ private Config getConfig(String name, String suffix) {
         : config.getConfig(cfgPrefix + ".default." + suffix);
   }
 
-  private AwsCredentialsProvider createAssumeRoleProvider(Config cfg, AwsCredentialsProvider p) {
-    final String arn = cfg.getString("role-arn");
+  private String createRoleArn(String arnPattern, String accountId) {
+    final boolean needsSubstitution = arnPattern.contains("{account}");
+    if (accountId == null) {
+      if (needsSubstitution) {
+        throw new IllegalArgumentException("missing account id for ARN pattern: " + arnPattern);
+      }
+      return arnPattern;
+    } else if (needsSubstitution) {
+      return arnPattern.replace("{account}", accountId);
+    } else {
+      LOGGER.warn("requested account, {}, is not used by ARN pattern: {}", accountId, arnPattern);
+      return arnPattern;
+    }
+  }
+
+  private AwsCredentialsProvider createAssumeRoleProvider(
+      Config cfg, String accountId, AwsCredentialsProvider p) {
+    final String arn = createRoleArn(cfg.getString("role-arn"), accountId);
     final String name = cfg.getString("role-session-name");
     final StsClient stsClient = StsClient.builder()
         .credentialsProvider(p)
@@ -126,14 +146,17 @@ private AwsCredentialsProvider createAssumeRoleProvider(Config cfg, AwsCredentia
         .build();
   }
 
-  AwsCredentialsProvider createCredentialsProvider(String name) {
+  AwsCredentialsProvider createCredentialsProvider(String name, String accountId) {
     final AwsCredentialsProvider dflt = DefaultCredentialsProvider.builder()
         .asyncCredentialUpdateEnabled(true)
         .build();
     final Config cfg = getConfig(name, "credentials");
     if (cfg.hasPath("role-arn")) {
-      return createAssumeRoleProvider(cfg, dflt);
+      return createAssumeRoleProvider(cfg, accountId, dflt);
     } else {
+      if (accountId != null) {
+        LOGGER.warn("requested account, {}, ignored, no role ARN configured", accountId);
+      }
       return dflt;
     }
   }
@@ -151,21 +174,161 @@ private Region chooseRegion(String name, Class<?> cls) {
     return Region.of(endpointRegion);
   }
 
+  /**
+   * Create a new instance of an AWS client of the specified type. The name of the config
+   * block will be based on the package for the class name. For example, if requesting an
+   * instance of {@code software.amazon.awssdk.services.ec2.Ec2Client} the config name used
+   * will be {@code ec2}.
+   *
+   * @param cls
+   *     Class for the AWS client type to create, e.g. {@code Ec2Client.class}.
+   * @return
+   *     AWS client instance.
+   */
   public <T> T newInstance(Class<T> cls) {
     return newInstance(getDefaultName(cls), cls);
   }
 
-  @SuppressWarnings("unchecked")
+  /**
+   * Create a new instance of an AWS client of the specified type.
+   *
+   * @param name
+   *     Name of the client. This is used to load config settings specific to the name.
+   * @param cls
+   *     Class for the AWS client type to create, e.g. {@code Ec2Client.class}.
+   * @return
+   *     AWS client instance.
+   */
   public <T> T newInstance(String name, Class<T> cls) {
+    return newInstance(name, cls, null);
+  }
+
+  /**
+   * Create a new instance of an AWS client. The name of the config
+   * block will be based on the package for the class name. For example, if requesting an
+   * instance of {@code software.amazon.awssdk.services.ec2.Ec2Client} the config name used
+   * will be {@code ec2}.
+   *
+   * @param cls
+   *     Class for the AWS client type to create, e.g. {@code Ec2Client.class}.
+   * @param accountId
+   *     The AWS account id to use when assuming to a role. If null, then the account
+   *     id should be specified directly in the role-arn setting or leave out the setting
+   *     to use the default credentials provider.
+   * @return
+   *     AWS client instance.
+   */
+  public <T> T newInstance(Class<T> cls, String accountId) {
+    return newInstance(getDefaultName(cls), cls, accountId);
+  }
+
+  /**
+   * Create a new instance of an AWS client. This method will always create a new instance.
+   * If you want to create or reuse an existing instance, then see
+   * {@link #getInstance(String, Class, String)}.
+   *
+   * @param name
+   *     Name of the client. This is used to load config settings specific to the name.
+   * @param cls
+   *     Class for the AWS client type to create, e.g. {@code Ec2Client.class}.
+   * @param accountId
+   *     The AWS account id to use when assuming to a role. If null, then the account
+   *     id should be specified directly in the role-arn setting or leave out the setting
+   *     to use the default credentials provider.
+   * @return
+   *     AWS client instance.
+   */
+  @SuppressWarnings("unchecked")
+  public <T> T newInstance(String name, Class<T> cls, String accountId) {
     try {
       Method builderMethod = cls.getMethod("builder");
       return (T) ((AwsClientBuilder) builderMethod.invoke(null))
-          .credentialsProvider(createCredentialsProvider(name))
+          .credentialsProvider(createCredentialsProvider(name, accountId))
           .region(chooseRegion(name, cls))
           .overrideConfiguration(createClientConfig(name))
           .build();
     } catch (Exception e) {
       throw new RuntimeException("failed to create instance of " + cls.getName(), e);
     }
   }
+
+  /**
+   * Get a shared instance of an AWS client of the specified type. The name of the config
+   * block will be based on the package for the class name. For example, if requesting an
+   * instance of {@code software.amazon.awssdk.services.ec2.Ec2Client} the config name used
+   * will be {@code ec2}.
+   *
+   * @param cls
+   *     Class for the AWS client type to create, e.g. {@code Ec2Client.class}.
+   * @return
+   *     AWS client instance.
+   */
+  public <T> T getInstance(Class<T> cls) {
+    return getInstance(getDefaultName(cls), cls);
+  }
+
+  /**
+   * Get a shared instance of an AWS client.
+   *
+   * @param name
+   *     Name of the client. This is used to load config settings specific to the name.
+   * @param cls
+   *     Class for the AWS client type to create, e.g. {@code Ec2Client.class}.
+   * @return
+   *     AWS client instance.
+   */
+  public <T> T getInstance(String name, Class<T> cls) {
+    return getInstance(name, cls, null);
+  }
+
+  /**
+   * Get a shared instance of an AWS client. The name of the config
+   * block will be based on the package for the class name. For example, if requesting an
+   * instance of {@code software.amazon.awssdk.services.ec2.Ec2Client} the config name used
+   * will be {@code ec2}.
+   *
+   * @param cls
+   *     Class for the AWS client type to create, e.g. {@code Ec2Client.class}.
+   * @param accountId
+   *     The AWS account id to use when assuming to a role. If null, then the account
+   *     id should be specified directly in the role-arn setting or leave out the setting
+   *     to use the default credentials provider.
+   * @return
+   *     AWS client instance.
+   */
+  public <T> T getInstance(Class<T> cls, String accountId) {
+    return getInstance(getDefaultName(cls), cls, accountId);
+  }
+
+  /**
+   * Get a shared instance of an AWS client.
+   *
+   * @param name
+   *     Name of the client. This is used to load config settings specific to the name.
+   * @param cls
+   *     Class for the AWS client type to create, e.g. {@code Ec2Client.class}.
+   * @param accountId
+   *     The AWS account id to use when assuming to a role. If null, then the account
+   *     id should be specified directly in the role-arn setting or leave out the setting
+   *     to use the default credentials provider.
+   * @return
+   *     AWS client instance.
+   */
+  @SuppressWarnings("unchecked")
+  public <T> T getInstance(String name, Class<T> cls, String accountId) {
+    try {
+      final String key = name + ":" + cls.getName() + ":" + accountId;
+      return (T) clients.computeIfAbsent(key,
+          k -> (SdkAutoCloseable) newInstance(name, cls, accountId));
+    } catch (Exception e) {
+      throw new RuntimeException("failed to get instance of " + cls.getName(), e);
+    }
+  }
+
+  /**
+   * Cleanup resources used by shared clients.
+   */
+  @Override public void close() throws Exception {
+    clients.values().forEach(SdkAutoCloseable::close);
+  }
 }

iep-module-aws2/src/test/java/com/netflix/iep/aws2/AwsClientFactoryTest.java

@@ -25,6 +25,7 @@
 import software.amazon.awssdk.auth.credentials.DefaultCredentialsProvider;
 import software.amazon.awssdk.core.client.config.ClientOverrideConfiguration;
 import software.amazon.awssdk.core.client.config.SdkAdvancedClientOption;
+import software.amazon.awssdk.regions.Region;
 import software.amazon.awssdk.services.ec2.Ec2AsyncClient;
 import software.amazon.awssdk.services.ec2.Ec2Client;
 import software.amazon.awssdk.services.ec2.model.DescribeAddressesRequest;
@@ -88,7 +89,7 @@ public void createClientConfigOverrideWithDefaults() {
   @Test
   public void createCredentialsProvider() {
     AwsClientFactory factory = new AwsClientFactory(config);
-    AwsCredentialsProvider creds = factory.createCredentialsProvider(null);
+    AwsCredentialsProvider creds = factory.createCredentialsProvider(null, null);
     Assert.assertTrue(creds instanceof DefaultCredentialsProvider);
   }
 
@@ -105,7 +106,31 @@ private AssumeRoleRequest getRequest(AwsCredentialsProvider creds) throws Except
   @Test
   public void createCredentialsProviderOverride() throws Exception {
     AwsClientFactory factory = new AwsClientFactory(config);
-    AwsCredentialsProvider creds = factory.createCredentialsProvider("ec2-test");
+    AwsCredentialsProvider creds = factory.createCredentialsProvider("ec2-test", null);
+    Assert.assertTrue(creds instanceof StsAssumeRoleCredentialsProvider);
+    Assert.assertEquals("arn:aws:iam::1234567890:role/IepTest", getRequest(creds).roleArn());
+    Assert.assertEquals("iep", getRequest(creds).roleSessionName());
+  }
+
+  @Test
+  public void createCredentialsProviderForAccount() throws Exception {
+    AwsClientFactory factory = new AwsClientFactory(config);
+    AwsCredentialsProvider creds = factory.createCredentialsProvider("ec2-account", "123");
+    Assert.assertTrue(creds instanceof StsAssumeRoleCredentialsProvider);
+    Assert.assertEquals("arn:aws:iam::123:role/IepTest", getRequest(creds).roleArn());
+    Assert.assertEquals("ieptest", getRequest(creds).roleSessionName());
+  }
+
+  @Test(expected = IllegalArgumentException.class)
+  public void createCredentialsProviderForAccountNull() throws Exception {
+    AwsClientFactory factory = new AwsClientFactory(config);
+    factory.createCredentialsProvider("ec2-account", null);
+  }
+
+  @Test
+  public void createCredentialsProviderForAccountIgnored() throws Exception {
+    AwsClientFactory factory = new AwsClientFactory(config);
+    AwsCredentialsProvider creds = factory.createCredentialsProvider("ec2-test", "123");
     Assert.assertTrue(creds instanceof StsAssumeRoleCredentialsProvider);
     Assert.assertEquals("arn:aws:iam::1234567890:role/IepTest", getRequest(creds).roleArn());
     Assert.assertEquals("iep", getRequest(creds).roleSessionName());
@@ -118,6 +143,12 @@ public void newInstanceInterface() throws Exception {
     Assert.assertNotNull(ec2);
   }
 
+  @Test(expected = RuntimeException.class)
+  public void newInstanceBadClass() throws Exception {
+    AwsClientFactory factory = new AwsClientFactory(config);
+    factory.newInstance(AwsClientFactory.class);
+  }
+
   @Test
   public void newInstanceName() throws Exception {
     AwsClientFactory factory = new AwsClientFactory(config);
@@ -139,6 +170,46 @@ public void newInstanceNameAsync() throws Exception {
     Assert.assertNotNull(ec2);
   }
 
+  @Test
+  public void getInstanceInterface() throws Exception {
+    AwsClientFactory factory = new AwsClientFactory(config);
+    Ec2Client ec2 = factory.getInstance(Ec2Client.class);
+    Assert.assertNotNull(ec2);
+    Assert.assertSame(ec2, factory.getInstance(Ec2Client.class));
+    Assert.assertNotSame(ec2, factory.newInstance(Ec2Client.class));
+  }
+
+  @Test
+  public void getInstanceClient() throws Exception {
+    AwsClientFactory factory = new AwsClientFactory(config);
+    Ec2Client ec2 = factory.getInstance(Ec2Client.class);
+    Assert.assertNotNull(ec2);
+    Assert.assertSame(ec2, factory.getInstance(Ec2Client.class));
+  }
+
+  @Test(expected = RuntimeException.class)
+  public void getInstanceBadClass() throws Exception {
+    AwsClientFactory factory = new AwsClientFactory(config);
+    factory.getInstance(AwsClientFactory.class);
+  }
+
+  @Test
+  public void getInstanceName() throws Exception {
+    AwsClientFactory factory = new AwsClientFactory(config);
+    Ec2Client ec2 = factory.getInstance("ec2-test", Ec2Client.class);
+    Assert.assertNotNull(ec2);
+    Assert.assertSame(ec2, factory.getInstance("ec2-test", Ec2Client.class));
+    Assert.assertNotSame(ec2, factory.getInstance(Ec2Client.class));
+  }
+
+  @Test
+  public void closeClients() throws Exception {
+    // Verifies the close completes without throwing
+    AwsClientFactory factory = new AwsClientFactory(config);
+    factory.getInstance(Ec2Client.class);
+    factory.close();
+  }
+
   @Test
   public void settingsUserAgentPrefix() {
     AwsClientFactory factory = new AwsClientFactory(config);

iep-module-aws2/src/test/resources/aws-client-factory.conf

@@ -27,6 +27,14 @@ netflix.iep.aws {
     }
   }
 
+  // Test account substitution
+  ec2-account {
+    credentials {
+      role-arn = "arn:aws:iam::{account}:role/IepTest"
+      role-session-name = "ieptest"
+    }
+  }
+
   timeouts {
     client {
       api-call-attempt-timeout = 42s