Skip to content

Commit

Permalink
Allow using configured S3 credentials with IAM role
Browse files Browse the repository at this point in the history 
Previously, IAM roles could only be used with the default
provider chain (typically, EC2 instance credentials).
  • Loading branch information
@electrum
electrum committed Apr 7, 2020
1 parent ffd54c9 commit 33a61a6
Show file tree
Hide file tree
Showing 2 changed files with 83 additions and 26 deletions.
53 changes: 32 additions & 21 deletions presto-hive/src/main/java/io/prestosql/plugin/hive/s3/PrestoS3FileSystem.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -59,6 +59,7 @@
import com.amazonaws.services.s3.transfer.TransferManagerBuilder;
import com.amazonaws.services.s3.transfer.Upload;
import com.google.common.annotations.VisibleForTesting;
import com.google.common.base.Splitter;
import com.google.common.collect.AbstractSequentialIterator;
import com.google.common.collect.Iterators;
import com.google.common.io.Closer;
Expand Down Expand Up @@ -794,23 +795,31 @@ private static Optional<EncryptionMaterialsProvider> createEncryptionMaterialsPr

private AWSCredentialsProvider createAwsCredentialsProvider(URI uri, Configuration conf)
{
Optional<AWSCredentials> credentials = getAwsCredentials(uri, conf);
// credentials embedded in the URI take precedence and are used alone
Optional<AWSCredentials> credentials = getEmbeddedAwsCredentials(uri);
if (credentials.isPresent()) {
return new AWSStaticCredentialsProvider(credentials.get());
}

if (iamRole != null) {
return new STSAssumeRoleSessionCredentialsProvider.Builder(this.iamRole, "presto-session")
.withExternalId(this.externalId)
.build();
}

// a custom credential provider is also used alone
String providerClass = conf.get(S3_CREDENTIALS_PROVIDER);
if (!isNullOrEmpty(providerClass)) {
return getCustomAWSCredentialsProvider(uri, conf, providerClass);
}

return DefaultAWSCredentialsProviderChain.getInstance();
// use configured credentials or default chain with optional role
AWSCredentialsProvider provider = getAwsCredentials(conf)
.map(value -> (AWSCredentialsProvider) new AWSStaticCredentialsProvider(value))
.orElseGet(DefaultAWSCredentialsProviderChain::getInstance);

if (iamRole != null) {
provider = new STSAssumeRoleSessionCredentialsProvider.Builder(iamRole, "presto-session")
.withExternalId(externalId)
.withLongLivedCredentialsProvider(provider)
.build();
}

return provider;
}

private static AWSCredentialsProvider getCustomAWSCredentialsProvider(URI uri, Configuration conf, String providerClass)
Expand All @@ -827,22 +836,24 @@ private static AWSCredentialsProvider getCustomAWSCredentialsProvider(URI uri, C
}
}

private static Optional<AWSCredentials> getAwsCredentials(URI uri, Configuration conf)
private static Optional<AWSCredentials> getEmbeddedAwsCredentials(URI uri)
{
String accessKey = conf.get(S3_ACCESS_KEY);
String secretKey = conf.get(S3_SECRET_KEY);

String userInfo = uri.getUserInfo();
if (userInfo != null) {
int index = userInfo.indexOf(':');
if (index < 0) {
accessKey = userInfo;
}
else {
accessKey = userInfo.substring(0, index);
secretKey = userInfo.substring(index + 1);
String userInfo = nullToEmpty(uri.getUserInfo());
List<String> parts = Splitter.on(':').limit(2).splitToList(userInfo);
if (parts.size() == 2) {
String accessKey = parts.get(0);
String secretKey = parts.get(1);
if (!accessKey.isEmpty() && !secretKey.isEmpty()) {
return Optional.of(new BasicAWSCredentials(accessKey, secretKey));
}
}
return Optional.empty();
}

private static Optional<AWSCredentials> getAwsCredentials(Configuration conf)
{
String accessKey = conf.get(S3_ACCESS_KEY);
String secretKey = conf.get(S3_SECRET_KEY);

if (isNullOrEmpty(accessKey) || isNullOrEmpty(secretKey)) {
return Optional.empty();
Expand Down
56 changes: 51 additions & 5 deletions presto-hive/src/test/java/io/prestosql/plugin/hive/s3/TestPrestoS3FileSystem.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,8 @@
import com.amazonaws.services.s3.model.EncryptionMaterialsProvider;
import com.amazonaws.services.s3.model.GetObjectMetadataRequest;
import com.amazonaws.services.s3.model.ObjectMetadata;
import com.amazonaws.services.securitytoken.AWSSecurityTokenService;
import com.amazonaws.services.securitytoken.AWSSecurityTokenServiceClient;
import com.google.common.base.VerifyException;
import io.prestosql.plugin.hive.s3.PrestoS3FileSystem.UnrecoverableS3OperationException;
import org.apache.hadoop.conf.Configuration;
Expand Down Expand Up @@ -88,6 +90,19 @@ public class TestPrestoS3FileSystem
{
private static final int HTTP_RANGE_NOT_SATISFIABLE = 416;

@Test
public void testEmbeddedCredentials()
throws Exception
{
Configuration config = new Configuration(false);
try (PrestoS3FileSystem fs = new PrestoS3FileSystem()) {
AWSCredentials credentials = getStaticCredentials(config, fs, "s3n://testAccess:testSecret@test-bucket/");
assertEquals(credentials.getAWSAccessKeyId(), "testAccess");
assertEquals(credentials.getAWSSecretKey(), "testSecret");
assertThat(credentials).isNotInstanceOf(AWSSessionCredentials.class);
}
}

@Test
public void testStaticCredentials()
throws Exception
Expand Down Expand Up @@ -135,20 +150,51 @@ public void testEndpointWithPinToCurrentRegionConfiguration()
}

@Test
public void testAssumeRoleCredentials()
public void testAssumeRoleDefaultCredentials()
throws Exception
{
Configuration config = new Configuration(false);
config.set(S3_IAM_ROLE, "role");
config.set(S3_IAM_ROLE, "test_role");

try (PrestoS3FileSystem fs = new PrestoS3FileSystem()) {
fs.initialize(new URI("s3n://test-bucket/"), config);
AWSCredentialsProvider awsCredentialsProvider = getAwsCredentialsProvider(fs);
assertInstanceOf(awsCredentialsProvider, STSAssumeRoleSessionCredentialsProvider.class);
assertEquals(getFieldValue(awsCredentialsProvider, "roleArn", String.class), "role");
AWSCredentialsProvider tokenService = getStsCredentialsProvider(fs, "test_role");
assertInstanceOf(tokenService, DefaultAWSCredentialsProviderChain.class);
}
}

@Test
public void testAssumeRoleStaticCredentials()
throws Exception
{
Configuration config = new Configuration(false);
config.set(S3_ACCESS_KEY, "test_access_key");
config.set(S3_SECRET_KEY, "test_secret_key");
config.set(S3_IAM_ROLE, "test_role");

try (PrestoS3FileSystem fs = new PrestoS3FileSystem()) {
fs.initialize(new URI("s3n://test-bucket/"), config);
AWSCredentialsProvider tokenService = getStsCredentialsProvider(fs, "test_role");
assertInstanceOf(tokenService, AWSStaticCredentialsProvider.class);

AWSCredentials credentials = tokenService.getCredentials();
assertEquals(credentials.getAWSAccessKeyId(), "test_access_key");
assertEquals(credentials.getAWSSecretKey(), "test_secret_key");
}
}

private static AWSCredentialsProvider getStsCredentialsProvider(PrestoS3FileSystem fs, String expectedRole)
{
AWSCredentialsProvider awsCredentialsProvider = getAwsCredentialsProvider(fs);
assertInstanceOf(awsCredentialsProvider, STSAssumeRoleSessionCredentialsProvider.class);

assertEquals(getFieldValue(awsCredentialsProvider, "roleArn", String.class), expectedRole);

AWSSecurityTokenService tokenService = getFieldValue(awsCredentialsProvider, "securityTokenService", AWSSecurityTokenService.class);
assertInstanceOf(tokenService, AWSSecurityTokenServiceClient.class);
return getFieldValue(tokenService, "awsCredentialsProvider", AWSCredentialsProvider.class);
}

@Test
public void testAssumeRoleCredentialsWithExternalId()
throws Exception
Expand Down

0 comments on commit 33a61a6

Please sign in to comment.