From 200ec489f2781a4f1cbcfc6463b62af93fb2a165 Mon Sep 17 00:00:00 2001 From: Sebastian Herold Date: Wed, 14 Aug 2024 03:05:27 +0200 Subject: [PATCH 1/3] Add a plugin to allow to assume an AWS IAM role before using IAM-based auth. --- ...AssumeChainedRolesCredentialsProvider.java | 116 ++++++++++++++++++ 1 file changed, 116 insertions(+) create mode 100644 src/main/java/com/amazon/redshift/plugin/AssumeChainedRolesCredentialsProvider.java diff --git a/src/main/java/com/amazon/redshift/plugin/AssumeChainedRolesCredentialsProvider.java b/src/main/java/com/amazon/redshift/plugin/AssumeChainedRolesCredentialsProvider.java new file mode 100644 index 0000000..6dfbd89 --- /dev/null +++ b/src/main/java/com/amazon/redshift/plugin/AssumeChainedRolesCredentialsProvider.java @@ -0,0 +1,116 @@ +package com.amazon.redshift.plugin; + +import com.amazon.redshift.IPlugin; +import com.amazon.redshift.logger.LogLevel; +import com.amazon.redshift.logger.RedshiftLogger; +import com.amazonaws.auth.AWSCredentials; +import com.amazonaws.auth.AWSCredentialsProvider; +import com.amazonaws.auth.DefaultAWSCredentialsProviderChain; +import com.amazonaws.auth.STSAssumeRoleSessionCredentialsProvider; +import com.amazonaws.services.securitytoken.AWSSecurityTokenService; +import com.amazonaws.services.securitytoken.AWSSecurityTokenServiceClient; +import com.amazonaws.services.securitytoken.AWSSecurityTokenServiceClientBuilder; +import org.apache.commons.codec.binary.StringUtils; + +import static com.amazon.redshift.logger.LogLevel.INFO; + +/** + * Plugin to assume a role or a chain of roles to authenticate via IAM-based auth in Redshift. + * Set the role_arn parameter to the ARN of the role to assume. If you want to assume multiple roles in a chain, separate + * the ARNs with commas. The plugin will assume the roles in the order they are provided. + * Optionally, set the session_name parameter to a custom session name. If not provided, the default session name will be used. + */ +public class AssumeChainedRolesCredentialsProvider implements IPlugin { + + private static final String CACHE_KEY = AssumeChainedRolesCredentialsProvider.class.getName(); + private static final String DEFAULT_SESSION_NAME = CACHE_KEY; + private AWSCredentialsProvider credentialsProvider = null; + + private String roleArn = null; + private String sessionName = null; + private RedshiftLogger log = null; + + + @Override + public void addParameter(String key, String value) { + if (key.equals("role_arn")) { + roleArn = value; + } + + if (key.equals("session_name")) { + sessionName = value; + } + } + + @Override + public void setLogger(RedshiftLogger log) { + this.log = log; + } + + @Override + public String getPluginSpecificCacheKey() { + return CACHE_KEY; + } + + @Override + public void setGroupFederation(boolean groupFederation) { + } + + @Override + public String getIdpToken() { + return null; + } + + @Override + public String getCacheKey() { + return getPluginSpecificCacheKey(); + } + + @Override + public int getSubType() { + return 0; + } + + @Override + public AWSCredentials getCredentials() { + return getCredentialsProvider().getCredentials(); + } + + protected AWSCredentialsProvider getCredentialsProvider() { + if (credentialsProvider == null) { + log(INFO, "Creating new credentials provider"); + if (roleArn != null && !roleArn.isEmpty()) { + log(INFO, "Found roleArn: %s and sessionName: %s", roleArn, getNonEmptySessionName()); + String[] rolesArns = roleArn.split(","); + AWSSecurityTokenService stsClient = AWSSecurityTokenServiceClientBuilder.defaultClient(); + for (String roleArn : rolesArns) { + credentialsProvider = new STSAssumeRoleSessionCredentialsProvider.Builder(roleArn, getNonEmptySessionName()) + .withStsClient(stsClient) + .build(); + stsClient = AWSSecurityTokenServiceClientBuilder.standard() + .withCredentials(credentialsProvider) + .build(); + } + } else { + log(INFO, "No roleArn found, using DefaultAWSCredentialsProviderChain"); + credentialsProvider = new DefaultAWSCredentialsProviderChain(); + } + } + return credentialsProvider; + } + + private String getNonEmptySessionName() { + return sessionName != null ? sessionName : DEFAULT_SESSION_NAME; + } + + protected void log(LogLevel logLevel, String msg, Object... msgArgs) { + if (RedshiftLogger.isEnable()) { + log.log(logLevel, msg, msgArgs); + } + } + + @Override + public void refresh() { + getCredentialsProvider().refresh(); + } +} From 6a312d784d460139033aa0f9eb68ca732cb9b169 Mon Sep 17 00:00:00 2001 From: Sebastian Herold Date: Wed, 14 Aug 2024 03:06:06 +0200 Subject: [PATCH 2/3] Add a plugin to allow to assume an AWS IAM role before using IAM-based auth. --- .../redshift/plugin/AssumeChainedRolesCredentialsProvider.java | 2 -- 1 file changed, 2 deletions(-) diff --git a/src/main/java/com/amazon/redshift/plugin/AssumeChainedRolesCredentialsProvider.java b/src/main/java/com/amazon/redshift/plugin/AssumeChainedRolesCredentialsProvider.java index 6dfbd89..528c028 100644 --- a/src/main/java/com/amazon/redshift/plugin/AssumeChainedRolesCredentialsProvider.java +++ b/src/main/java/com/amazon/redshift/plugin/AssumeChainedRolesCredentialsProvider.java @@ -8,9 +8,7 @@ import com.amazonaws.auth.DefaultAWSCredentialsProviderChain; import com.amazonaws.auth.STSAssumeRoleSessionCredentialsProvider; import com.amazonaws.services.securitytoken.AWSSecurityTokenService; -import com.amazonaws.services.securitytoken.AWSSecurityTokenServiceClient; import com.amazonaws.services.securitytoken.AWSSecurityTokenServiceClientBuilder; -import org.apache.commons.codec.binary.StringUtils; import static com.amazon.redshift.logger.LogLevel.INFO; From 02349443eda3d76f7f18c350da2dc8a3aabaf90e Mon Sep 17 00:00:00 2001 From: Sebastian Herold Date: Wed, 14 Aug 2024 03:31:34 +0200 Subject: [PATCH 3/3] fix checkstyle warnings --- .../plugin/AssumeChainedRolesCredentialsProvider.java | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/main/java/com/amazon/redshift/plugin/AssumeChainedRolesCredentialsProvider.java b/src/main/java/com/amazon/redshift/plugin/AssumeChainedRolesCredentialsProvider.java index 528c028..ba7cf2d 100644 --- a/src/main/java/com/amazon/redshift/plugin/AssumeChainedRolesCredentialsProvider.java +++ b/src/main/java/com/amazon/redshift/plugin/AssumeChainedRolesCredentialsProvider.java @@ -31,11 +31,11 @@ public class AssumeChainedRolesCredentialsProvider implements IPlugin { @Override public void addParameter(String key, String value) { - if (key.equals("role_arn")) { + if ("role_arn".equals(key)) { roleArn = value; } - if (key.equals("session_name")) { + if ("session_name".equals(key)) { sessionName = value; } }