From d24b80a444aac9e6d4aefd43794a9857d6bf419c Mon Sep 17 00:00:00 2001 From: Chris Lu Date: Mon, 4 May 2026 16:15:41 -0700 Subject: [PATCH] Honor aws.endpoint for STS and S3 clients Two builders inside the server were always targeting real AWS: - AwsCredentialGenerator.StsAwsCredentialGenerator built its StsClient with .region(...).credentialsProvider(...) only. There is no endpointOverride and aws.endpoint isn't propagated, so every AssumeRole goes to sts..amazonaws.com. AWS_ENDPOINT_URL_STS / AWS_ENDPOINT_URL env vars and the matching system properties don't kick in for that builder shape -- I confirmed by pointing AWS_ENDPOINT_URL_STS at port 1 and getting back the same real-AWS InvalidClientTokenId 403, with zero traffic to the configured port (sniffer in front). - FileIOFactory.getS3Client built its S3Client the same way, with no endpointOverride and forcePathStyle(false), so any Iceberg read against a self-hosted S3-compatible backend would also miss the endpoint and reject path-style URLs. Add an aws.endpoint property + per-bucket s3.endpoint.N, plumb it through S3StorageConfig, and apply endpointOverride on both the StsClient and the S3Client when set. For the S3Client also flip to path-style when an endpoint is set, since SeaweedFS / MinIO / Ceph RGW / LocalStack typically can't issue wildcard certs for bucket.host. Default empty -> AWS path is unchanged. Two new unit tests cover "endpoint set -> endpointOverride called with that URI" and "endpoint unset -> endpointOverride never called". --- etc/conf/server.properties | 4 + .../aws/AwsCredentialGenerator.java | 7 +- .../credential/aws/S3StorageConfig.java | 1 + .../server/service/iceberg/FileIOFactory.java | 20 ++-- .../server/utils/ServerProperties.java | 4 + .../credential/CloudCredentialVendorTest.java | 92 +++++++++++++++++++ 6 files changed, 120 insertions(+), 8 deletions(-) diff --git a/etc/conf/server.properties b/etc/conf/server.properties index 88c9eb0784..bf394353ff 100644 --- a/etc/conf/server.properties +++ b/etc/conf/server.properties @@ -48,6 +48,8 @@ aws.accessKey= aws.secretKey= # The region of AWS. If running on-premises, pick the nearest region. aws.region= +# Optional STS/S3 endpoint override for self-hosted S3-compatible backends. +aws.endpoint= #### The following "s3.*" configs are legacy per-bucket hardcoded credential config @@ -60,6 +62,8 @@ s3.accessKey.0= s3.secretKey.0= # Test Only (If you provide a session token, it will just use those session creds, no downscoping) s3.sessionToken.0= +# Optional per-bucket endpoint override. +s3.endpoint.0= ## ADLS Storage Config (Multiple configs can be added by incrementing the index) adls.storageAccountName.0= diff --git a/server/src/main/java/io/unitycatalog/server/service/credential/aws/AwsCredentialGenerator.java b/server/src/main/java/io/unitycatalog/server/service/credential/aws/AwsCredentialGenerator.java index 98b8ce03bc..280b0751ef 100644 --- a/server/src/main/java/io/unitycatalog/server/service/credential/aws/AwsCredentialGenerator.java +++ b/server/src/main/java/io/unitycatalog/server/service/credential/aws/AwsCredentialGenerator.java @@ -3,6 +3,7 @@ import io.unitycatalog.server.model.AwsIamRoleResponse; import io.unitycatalog.server.persist.dao.CredentialDAO; import io.unitycatalog.server.service.credential.CredentialContext; +import java.net.URI; import java.time.Duration; import java.util.Optional; import java.util.UUID; @@ -85,7 +86,11 @@ public StsAwsCredentialGenerator(StsClientBuilder builder, S3StorageConfig confi credentialsProvider = DefaultCredentialsProvider.create(); } - this.stsClient = builder.region(region).credentialsProvider(credentialsProvider).build(); + StsClientBuilder configured = builder.region(region).credentialsProvider(credentialsProvider); + if (config.getEndpoint() != null && !config.getEndpoint().isEmpty()) { + configured = configured.endpointOverride(URI.create(config.getEndpoint())); + } + this.stsClient = configured.build(); this.staticAwsRoleArn = config.getAwsRoleArn(); } diff --git a/server/src/main/java/io/unitycatalog/server/service/credential/aws/S3StorageConfig.java b/server/src/main/java/io/unitycatalog/server/service/credential/aws/S3StorageConfig.java index ed9b688461..9b1ab7e80b 100644 --- a/server/src/main/java/io/unitycatalog/server/service/credential/aws/S3StorageConfig.java +++ b/server/src/main/java/io/unitycatalog/server/service/credential/aws/S3StorageConfig.java @@ -15,4 +15,5 @@ public class S3StorageConfig { private final String secretKey; private final String sessionToken; private final String credentialGenerator; + private final String endpoint; } diff --git a/server/src/main/java/io/unitycatalog/server/service/iceberg/FileIOFactory.java b/server/src/main/java/io/unitycatalog/server/service/iceberg/FileIOFactory.java index f4fdfd00cb..b7fa0ead7c 100644 --- a/server/src/main/java/io/unitycatalog/server/service/iceberg/FileIOFactory.java +++ b/server/src/main/java/io/unitycatalog/server/service/iceberg/FileIOFactory.java @@ -25,6 +25,7 @@ import software.amazon.awssdk.regions.Region; import software.amazon.awssdk.services.s3.S3Client; +import java.net.URI; import java.util.Map; import java.util.Set; @@ -86,20 +87,25 @@ protected S3FileIO getS3FileIO(NormalizedURL location) { S3StorageConfig s3StorageConfig = s3Configurations.get(location.getStorageBase()); S3FileIO s3FileIO = - new S3FileIO(() -> getS3Client(getAwsCredentialsProvider(location), - s3StorageConfig.getRegion())); + new S3FileIO(() -> getS3Client(getAwsCredentialsProvider(location), s3StorageConfig)); s3FileIO.initialize(Map.of()); return s3FileIO; } - protected S3Client getS3Client(AwsCredentialsProvider awsCredentialsProvider, String region) { - return S3Client.builder() - .region(Region.of(region)) + protected S3Client getS3Client( + AwsCredentialsProvider awsCredentialsProvider, S3StorageConfig s3StorageConfig) { + String endpoint = s3StorageConfig.getEndpoint(); + boolean hasEndpoint = endpoint != null && !endpoint.isEmpty(); + var builder = S3Client.builder() + .region(Region.of(s3StorageConfig.getRegion())) .credentialsProvider(awsCredentialsProvider) - .forcePathStyle(false) - .build(); + .forcePathStyle(hasEndpoint); + if (hasEndpoint) { + builder = builder.endpointOverride(URI.create(endpoint)); + } + return builder.build(); } private AwsCredentialsProvider getAwsCredentialsProvider(NormalizedURL location) { diff --git a/server/src/main/java/io/unitycatalog/server/utils/ServerProperties.java b/server/src/main/java/io/unitycatalog/server/utils/ServerProperties.java index 1ceebfa441..9503371fad 100644 --- a/server/src/main/java/io/unitycatalog/server/utils/ServerProperties.java +++ b/server/src/main/java/io/unitycatalog/server/utils/ServerProperties.java @@ -205,6 +205,7 @@ public enum Property { AWS_SECRET_KEY("aws.secretKey"), AWS_SESSION_TOKEN("aws.sessionToken"), AWS_REGION("aws.region"), + AWS_ENDPOINT("aws.endpoint"), INCLUDE_STACK_TRACE_IN_ERROR("server.include-stacktrace-in-error", "false", BOOLEAN_VALIDATOR); // The is not an exhaustive list. Some property keys like s3.bucketPath.0 with a numbering // suffix is not included. They are only accessed internally from functions like @@ -300,6 +301,7 @@ public S3StorageConfig getS3MasterRoleConfiguration() { .accessKey(get(Property.AWS_ACCESS_KEY)) .secretKey(get(Property.AWS_SECRET_KEY)) // Does not take AWS_SESSION_TOKEN as it's only part of a temporary credential. + .endpoint(get(Property.AWS_ENDPOINT)) .build(); } @@ -314,6 +316,7 @@ public Map getS3Configurations() { String secretKey = getProperty("s3.secretKey." + i); String sessionToken = getProperty("s3.sessionToken." + i); String credentialGenerator = getProperty("s3.credentialGenerator." + i); + String endpoint = getProperty("s3.endpoint." + i); if ((bucketPath == null || region == null || awsRoleArn == null) && (accessKey == null || secretKey == null || sessionToken == null)) { break; @@ -327,6 +330,7 @@ public Map getS3Configurations() { .secretKey(secretKey) .sessionToken(sessionToken) .credentialGenerator(credentialGenerator) + .endpoint(endpoint) .build(); s3BucketConfigMap.put(NormalizedURL.from(bucketPath), s3StorageConfig); i++; diff --git a/server/src/test/java/io/unitycatalog/server/service/credential/CloudCredentialVendorTest.java b/server/src/test/java/io/unitycatalog/server/service/credential/CloudCredentialVendorTest.java index 0d78413507..7dd3c9a2ec 100644 --- a/server/src/test/java/io/unitycatalog/server/service/credential/CloudCredentialVendorTest.java +++ b/server/src/test/java/io/unitycatalog/server/service/credential/CloudCredentialVendorTest.java @@ -26,6 +26,7 @@ import io.unitycatalog.server.service.credential.gcp.TestingCredentialGenerator; import io.unitycatalog.server.utils.NormalizedURL; import io.unitycatalog.server.utils.ServerProperties; +import java.net.URI; import java.util.Map; import java.util.Optional; import java.util.Set; @@ -287,4 +288,95 @@ public void testVendCredentialWithExternalLocationCredential() { assertThat(capturedRequest.externalId()).isEqualTo(expectedExternalId); } } + + @Test + public void testStsClientHonorsAwsEndpointOverride() { + final String S3_PATH = "s3://my-bucket/path/to/data"; + final String CUSTOM_ENDPOINT = "http://host.docker.internal:8333"; + + when(serverProperties.getS3Configurations()) + .thenReturn( + Map.of( + NormalizedURL.from("s3://my-bucket"), + S3StorageConfig.builder() + .bucketPath("s3://my-bucket") + .region("us-east-1") + .awsRoleArn("arn:aws:iam::000000000000:role/UCVendedRole") + .accessKey("admin") + .secretKey("admin") + .endpoint(CUSTOM_ENDPOINT) + .build())); + + StsClient mockStsClient = Mockito.mock(StsClient.class); + when(mockStsClient.assumeRole(any(AssumeRoleRequest.class))) + .thenReturn( + AssumeRoleResponse.builder() + .credentials( + Credentials.builder() + .accessKeyId("ak") + .secretAccessKey("sk") + .sessionToken("st") + .build()) + .build()); + + StsClientBuilder mockBuilder = Mockito.mock(StsClientBuilder.class); + when(mockBuilder.region(any())).thenReturn(mockBuilder); + when(mockBuilder.credentialsProvider(any())).thenReturn(mockBuilder); + when(mockBuilder.endpointOverride(any())).thenReturn(mockBuilder); + when(mockBuilder.build()).thenReturn(mockStsClient); + ArgumentCaptor endpointCaptor = ArgumentCaptor.forClass(URI.class); + + try (MockedStatic mockedStsClient = Mockito.mockStatic(StsClient.class)) { + mockedStsClient.when(StsClient::builder).thenReturn(mockBuilder); + AwsCredentialVendor awsCredentialVendor = new AwsCredentialVendor(serverProperties); + credentialsOperations = new CloudCredentialVendor(awsCredentialVendor, null, null); + vendCredential(S3_PATH, Set.of(CredentialContext.Privilege.SELECT)); + } + + verify(mockBuilder).endpointOverride(endpointCaptor.capture()); + assertThat(endpointCaptor.getValue()).isEqualTo(URI.create(CUSTOM_ENDPOINT)); + } + + @Test + public void testStsClientSkipsEndpointOverrideWhenUnset() { + final String S3_PATH = "s3://my-bucket/path/to/data"; + + when(serverProperties.getS3Configurations()) + .thenReturn( + Map.of( + NormalizedURL.from("s3://my-bucket"), + S3StorageConfig.builder() + .bucketPath("s3://my-bucket") + .region("us-east-1") + .awsRoleArn("arn:aws:iam::123:role/r") + .accessKey("ak") + .secretKey("sk") + .build())); + + StsClient mockStsClient = Mockito.mock(StsClient.class); + when(mockStsClient.assumeRole(any(AssumeRoleRequest.class))) + .thenReturn( + AssumeRoleResponse.builder() + .credentials( + Credentials.builder() + .accessKeyId("ak") + .secretAccessKey("sk") + .sessionToken("st") + .build()) + .build()); + + StsClientBuilder mockBuilder = Mockito.mock(StsClientBuilder.class); + when(mockBuilder.region(any())).thenReturn(mockBuilder); + when(mockBuilder.credentialsProvider(any())).thenReturn(mockBuilder); + when(mockBuilder.build()).thenReturn(mockStsClient); + + try (MockedStatic mockedStsClient = Mockito.mockStatic(StsClient.class)) { + mockedStsClient.when(StsClient::builder).thenReturn(mockBuilder); + AwsCredentialVendor awsCredentialVendor = new AwsCredentialVendor(serverProperties); + credentialsOperations = new CloudCredentialVendor(awsCredentialVendor, null, null); + vendCredential(S3_PATH, Set.of(CredentialContext.Privilege.SELECT)); + } + + verify(mockBuilder, Mockito.never()).endpointOverride(any()); + } }