Skip to content

Refactor: Use per-request STS credentials #1629

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 4 commits into from
May 23, 2025
Merged

Refactor: Use per-request STS credentials #1629

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
775750e
Refactor: Use per-request STS credentials
dimas-b May 20, 2025
744a40d
review: fix imports
dimas-b May 22, 2025
649c173
review: rename vars to credentialsProvider
dimas-b May 22, 2025
27c7d84
review: fix more imports
dimas-b May 23, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
39 changes: 24 additions & 15 deletions ...e/src/main/java/org/apache/polaris/core/storage/aws/AwsCredentialsStorageIntegration.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -33,6 +33,7 @@
import org.apache.polaris.core.storage.InMemoryStorageIntegration;
import org.apache.polaris.core.storage.StorageAccessProperty;
import org.apache.polaris.core.storage.StorageUtil;
import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider;
import software.amazon.awssdk.policybuilder.iam.IamConditionOperator;
import software.amazon.awssdk.policybuilder.iam.IamEffect;
import software.amazon.awssdk.policybuilder.iam.IamPolicy;
Expand All @@ -46,10 +47,17 @@
public class AwsCredentialsStorageIntegration
extends InMemoryStorageIntegration<AwsStorageConfigurationInfo> {
private final StsClient stsClient;
private final Optional<AwsCredentialsProvider> credentialsProvider;

public AwsCredentialsStorageIntegration(StsClient stsClient) {
this(stsClient, Optional.empty());
}

public AwsCredentialsStorageIntegration(
StsClient stsClient, Optional<AwsCredentialsProvider> credentialsProvider) {
super(AwsCredentialsStorageIntegration.class.getName());
this.stsClient = stsClient;
this.credentialsProvider = credentialsProvider;
}

/** {@inheritDoc} */
Expand All @@ -60,21 +68,22 @@ public EnumMap<StorageAccessProperty, String> getSubscopedCreds(
boolean allowListOperation,
@Nonnull Set<String> allowedReadLocations,
@Nonnull Set<String> allowedWriteLocations) {
AssumeRoleResponse response =
stsClient.assumeRole(
AssumeRoleRequest.builder()
.externalId(storageConfig.getExternalId())
.roleArn(storageConfig.getRoleARN())
.roleSessionName("PolarisAwsCredentialsStorageIntegration")
.policy(
policyString(
storageConfig.getRoleARN(),
allowListOperation,
allowedReadLocations,
allowedWriteLocations)
.toJson())
.durationSeconds(loadConfig(STORAGE_CREDENTIAL_DURATION_SECONDS))
.build());
AssumeRoleRequest.Builder request =
AssumeRoleRequest.builder()
.externalId(storageConfig.getExternalId())
.roleArn(storageConfig.getRoleARN())
.roleSessionName("PolarisAwsCredentialsStorageIntegration")
.policy(
policyString(
storageConfig.getRoleARN(),
allowListOperation,
allowedReadLocations,
allowedWriteLocations)
.toJson())
.durationSeconds(loadConfig(STORAGE_CREDENTIAL_DURATION_SECONDS));
credentialsProvider.ifPresent(
cp -> request.overrideConfiguration(b -> b.credentialsProvider(cp)));
AssumeRoleResponse response = stsClient.assumeRole(request.build());
EnumMap<StorageAccessProperty, String> credentialMap =
new EnumMap<>(StorageAccessProperty.class);
credentialMap.put(StorageAccessProperty.AWS_KEY_ID, response.credentials().accessKeyId());
Expand Down
5 changes: 4 additions & 1 deletion .../service/src/test/java/org/apache/polaris/service/quarkus/admin/PolarisAuthzTestBase.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,7 @@
import java.util.Date;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.Set;
import java.util.function.Function;
import java.util.stream.Collectors;
Expand Down Expand Up @@ -219,7 +220,9 @@ public Map<String, String> getConfigOverrides() {
public static void setUpMocks() {
PolarisStorageIntegrationProviderImpl mock =
new PolarisStorageIntegrationProviderImpl(
Mockito::mock, () -> GoogleCredentials.create(new AccessToken("abc", new Date())));
Mockito::mock,
Optional.empty(),
() -> GoogleCredentials.create(new AccessToken("abc", new Date())));
QuarkusMock.installMockForType(mock, PolarisStorageIntegrationProviderImpl.class);
}

Expand Down
15 changes: 12 additions & 3 deletions ...c/main/java/org/apache/polaris/service/storage/PolarisStorageIntegrationProviderImpl.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,7 @@
import jakarta.inject.Inject;
import java.util.EnumMap;
import java.util.Map;
import java.util.Optional;
import java.util.Set;
import java.util.function.Supplier;
import org.apache.polaris.core.PolarisDiagnostics;
Expand All @@ -39,22 +40,30 @@
import org.apache.polaris.core.storage.aws.AwsCredentialsStorageIntegration;
import org.apache.polaris.core.storage.azure.AzureCredentialsStorageIntegration;
import org.apache.polaris.core.storage.gcp.GcpCredentialsStorageIntegration;
import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider;
import software.amazon.awssdk.services.sts.StsClient;

@ApplicationScoped
public class PolarisStorageIntegrationProviderImpl implements PolarisStorageIntegrationProvider {

private final Supplier<StsClient> stsClientSupplier;
private final Optional<AwsCredentialsProvider> stsCredentials;
private final Supplier<GoogleCredentials> gcpCredsProvider;

@Inject
public PolarisStorageIntegrationProviderImpl(StorageConfiguration storageConfiguration) {
this(storageConfiguration.stsClientSupplier(), storageConfiguration.gcpCredentialsSupplier());
this(
storageConfiguration.stsClientSupplier(false),
Optional.ofNullable(storageConfiguration.stsCredentials()),
storageConfiguration.gcpCredentialsSupplier());
}

public PolarisStorageIntegrationProviderImpl(
Supplier<StsClient> stsClientSupplier, Supplier<GoogleCredentials> gcpCredsProvider) {
Supplier<StsClient> stsClientSupplier,
Optional<AwsCredentialsProvider> stsCredentials,
Supplier<GoogleCredentials> gcpCredsProvider) {
this.stsClientSupplier = stsClientSupplier;
this.stsCredentials = stsCredentials;
this.gcpCredsProvider = gcpCredsProvider;
}

Expand All @@ -71,7 +80,7 @@ public PolarisStorageIntegrationProviderImpl(
case S3:
storageIntegration =
(PolarisStorageIntegration<T>)
new AwsCredentialsStorageIntegration(stsClientSupplier.get());
new AwsCredentialsStorageIntegration(stsClientSupplier.get(), stsCredentials);
break;
case GCS:
storageIntegration =
Expand Down
26 changes: 19 additions & 7 deletions service/common/src/main/java/org/apache/polaris/service/storage/StorageConfiguration.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,8 @@
import java.util.function.Supplier;
import org.slf4j.LoggerFactory;
import software.amazon.awssdk.auth.credentials.AwsBasicCredentials;
import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider;
import software.amazon.awssdk.auth.credentials.DefaultCredentialsProvider;
import software.amazon.awssdk.auth.credentials.StaticCredentialsProvider;
import software.amazon.awssdk.services.sts.StsClient;
import software.amazon.awssdk.services.sts.StsClientBuilder;
Expand Down Expand Up @@ -61,21 +63,31 @@ public interface StorageConfiguration {
Optional<Duration> gcpAccessTokenLifespan();

default Supplier<StsClient> stsClientSupplier() {
return stsClientSupplier(true);
}

default Supplier<StsClient> stsClientSupplier(boolean withCredentials) {
return Suppliers.memoize(
() -> {
StsClientBuilder stsClientBuilder = StsClient.builder();
if (awsAccessKey().isPresent() && awsSecretKey().isPresent()) {
LoggerFactory.getLogger(StorageConfiguration.class)
.warn("Using hard-coded AWS credentials - this is not recommended for production");
StaticCredentialsProvider awsCredentialsProvider =
StaticCredentialsProvider.create(
AwsBasicCredentials.create(awsAccessKey().get(), awsSecretKey().get()));
stsClientBuilder.credentialsProvider(awsCredentialsProvider);
if (withCredentials) {
stsClientBuilder.credentialsProvider(stsCredentials());
}
return stsClientBuilder.build();
});
}

default AwsCredentialsProvider stsCredentials() {
if (awsAccessKey().isPresent() && awsSecretKey().isPresent()) {
LoggerFactory.getLogger(StorageConfiguration.class)
Copy link
Contributor

@singhpk234 singhpk234 May 21, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

any reason to keep Logger getter in this method rather outside ?

Copy link
Contributor Author

@dimas-b dimas-b May 22, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

oops, it was just a copy-paste.... will fix.

Copy link
Contributor Author

@dimas-b dimas-b May 22, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actually, refactoring this in current PR is awkward - I would not like to add a logger field to this interface. Adding a class looks like an overkill.

Would you mind if I moved this to a production readiness check in a follow-up PR?

Copy link
Contributor

@singhpk234 singhpk234 May 22, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

sounds fair !

.warn("Using hard-coded AWS credentials - this is not recommended for production");
return StaticCredentialsProvider.create(
AwsBasicCredentials.create(awsAccessKey().get(), awsSecretKey().get()));
} else {
return DefaultCredentialsProvider.create();
Copy link
Contributor

@singhpk234 singhpk234 May 21, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[orthogonal] Thoughts on supportingdynamically loading credential provider implementation ?

Copy link
Contributor Author

@dimas-b dimas-b May 22, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The idea (for now) is for custom builds to manage cred. provided via a custom impl. of PolarisStorageIntegrationProvider.

I suppose this will evolve as we move forward with @XJDKC 's proposal: https://lists.apache.org/thread/ph1tvn3lzvn8kh8fnhc6k585qmw2m12r

Copy link
Member

@XJDKC XJDKC May 23, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes, I have a branch that's based on the proposal, and it's working. I'm trying to clean it and add more tests, then I will open a PR. Welcome to review the PR and leave your comments!
https://github.com/XJDKC/polaris/tree/rxing-catalog-federation-sigv4-poc

}
}

default Supplier<GoogleCredentials> gcpCredentialsSupplier() {
return Suppliers.memoize(
() -> {
Expand Down
13 changes: 12 additions & 1 deletion ...ice/common/src/test/java/org/apache/polaris/service/storage/StorageConfigurationTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,7 @@
import org.mockito.ArgumentCaptor;
import org.mockito.MockedStatic;
import org.mockito.Mockito;
import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider;
import software.amazon.awssdk.auth.credentials.StaticCredentialsProvider;
import software.amazon.awssdk.services.sts.StsClient;
import software.amazon.awssdk.services.sts.StsClientBuilder;
Expand Down Expand Up @@ -105,7 +106,7 @@ public void testSingletonStsClientWithStaticCredentials() {
staticMock.when(StsClient::builder).thenReturn(mockBuilder);

StorageConfiguration config = configWithAwsCredentialsAndGcpToken();
Supplier<StsClient> supplier = config.stsClientSupplier();
Supplier<StsClient> supplier = config.stsClientSupplier(true);
StsClient client1 = supplier.get();
StsClient client2 = supplier.get();

Expand All @@ -119,6 +120,16 @@ public void testSingletonStsClientWithStaticCredentials() {
}
}

@Test
public void testStaticStsCredentials() {
StorageConfiguration config = configWithAwsCredentialsAndGcpToken();
AwsCredentialsProvider credentialsProvider = config.stsCredentials();
assertThat(credentialsProvider).isInstanceOf(StaticCredentialsProvider.class);
assertThat(credentialsProvider.resolveCredentials().accessKeyId()).isEqualTo(TEST_ACCESS_KEY);
assertThat(credentialsProvider.resolveCredentials().secretAccessKey())
.isEqualTo(TEST_SECRET_KEY);
}

@Test
public void testCreateGcpCredentialsFromStaticToken() {
Supplier<GoogleCredentials> supplier =
Expand Down
2 changes: 2 additions & 0 deletions service/common/src/testFixtures/java/org/apache/polaris/service/TestServices.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,7 @@
import java.util.Date;
import java.util.HashMap;
import java.util.Map;
import java.util.Optional;
import java.util.Set;
import org.apache.polaris.core.PolarisCallContext;
import org.apache.polaris.core.PolarisDiagnostics;
Expand Down Expand Up @@ -150,6 +151,7 @@ public TestServices build() {
PolarisStorageIntegrationProviderImpl storageIntegrationProvider =
new PolarisStorageIntegrationProviderImpl(
() -> stsClient,
Optional.empty(),
() -> GoogleCredentials.create(new AccessToken(GCP_ACCESS_TOKEN, new Date())));
InMemoryPolarisMetaStoreManagerFactory metaStoreManagerFactory =
new InMemoryPolarisMetaStoreManagerFactory(
Expand Down
Loading