obelix74 commented on code in PR #4707:
URL: https://github.com/apache/polaris/pull/4707#discussion_r3455139236
##########
polaris-core/src/test/java/org/apache/polaris/core/storage/gcp/GcpCredentialsStorageIntegrationTest.java:
##########
@@ -502,4 +515,147 @@ private static String expressionAt(JsonNode parsedRules,
int ruleIndex) {
.path("expression")
.asText();
}
+
+ private static RealmConfig configWith(Map<String, String> values) {
+ RealmConfigurationSource source = (rc, name) -> values.get(name);
+ return new RealmConfigImpl(source, () -> "test-realm");
+ }
+
+ @Test
+ void resolveAttributionParamsDisabledByDefault() {
+
assertThat(GcpCredentialsStorageIntegration.resolveAttributionParams(EMPTY_REALM_CONFIG))
+ .isEmpty();
+ }
+
+ @Test
+ void resolveAttributionParamsReturnsParamsWhenFullyConfigured() {
+ RealmConfig config =
+ configWith(
+ Map.of(
+ "GCS_PRINCIPAL_ATTRIBUTION_ENABLED", "true",
+ "GCS_PRINCIPAL_ATTRIBUTION_WIF_AUDIENCE",
+
"//iam.googleapis.com/projects/123/locations/global/workloadIdentityPools/p/providers/pr",
+ "GCS_PRINCIPAL_ATTRIBUTION_TOKEN_ISSUER",
"https://issuer.example.com";,
+ "GCS_PRINCIPAL_ATTRIBUTION_SIGNING_KEY_FILE", "/tmp/key.pem"));
+ Optional<GcpAttributionParams> result =
+ GcpCredentialsStorageIntegration.resolveAttributionParams(config);
+ assertThat(result).isPresent();
+ assertThat(result.get().wifAudience())
+ .isEqualTo(
+
"//iam.googleapis.com/projects/123/locations/global/workloadIdentityPools/p/providers/pr");
+
assertThat(result.get().tokenIssuer()).isEqualTo("https://issuer.example.com";);
+ assertThat(result.get().signingKeyFile()).isEqualTo("/tmp/key.pem");
+ }
+
+ @Test
+ void resolveAttributionParamsThrowsWhenEnabledButAudienceMissing() {
+ RealmConfig config =
+ configWith(
+ Map.of(
+ "GCS_PRINCIPAL_ATTRIBUTION_ENABLED", "true",
+ "GCS_PRINCIPAL_ATTRIBUTION_TOKEN_ISSUER",
"https://issuer.example.com";,
+ "GCS_PRINCIPAL_ATTRIBUTION_SIGNING_KEY_FILE", "/tmp/key.pem"));
+ assertThatThrownBy(() ->
GcpCredentialsStorageIntegration.resolveAttributionParams(config))
+ .isInstanceOf(IllegalStateException.class)
+ .hasMessageContaining("GCS_PRINCIPAL_ATTRIBUTION_WIF_AUDIENCE");
+ }
+
+ @Test
+ void resolveAttributionParamsThrowsWhenMultipleRequiredParamsMissing() {
+ RealmConfig config =
configWith(Map.of("GCS_PRINCIPAL_ATTRIBUTION_ENABLED", "true"));
+ assertThatThrownBy(() ->
GcpCredentialsStorageIntegration.resolveAttributionParams(config))
+ .isInstanceOf(IllegalStateException.class)
+ .hasMessageContaining("GCS_PRINCIPAL_ATTRIBUTION_WIF_AUDIENCE")
+ .hasMessageContaining("GCS_PRINCIPAL_ATTRIBUTION_TOKEN_ISSUER")
+ .hasMessageContaining("GCS_PRINCIPAL_ATTRIBUTION_SIGNING_KEY_FILE");
+ }
+
+ @Test
+ void testWifAttributionFlowUsesFederatedCredentials(@TempDir Path tempDir)
throws Exception {
+ String serviceAccount = "[email protected]";
+ GcpStorageConfigurationInfo config =
+ GcpStorageConfigurationInfo.builder()
+ .addAllAllowedLocations(List.of("gs://bucket/path"))
+ .gcpServiceAccount(serviceAccount)
+ .build();
+
+ AtomicReference<GoogleCredentials> capturedSource = new
AtomicReference<>();
+
+ IamCredentialsClient mockIamClient =
Mockito.mock(IamCredentialsClient.class);
+ GenerateAccessTokenResponse mockResponse =
+ GenerateAccessTokenResponse.newBuilder()
+ .setAccessToken("wif-impersonated-token")
+ .setExpireTime(
+ Timestamp.newBuilder().setSeconds(System.currentTimeMillis() /
1000 + 3600).build())
+ .build();
+
Mockito.when(mockIamClient.generateAccessToken(Mockito.any(GenerateAccessTokenRequest.class)))
+ .thenReturn(mockResponse);
+
+ GoogleCredentials mockCreds = Mockito.mock(GoogleCredentials.class);
+
Mockito.when(mockCreds.createScoped(Mockito.any(String.class))).thenReturn(mockCreds);
+
+ GcpCredentialOps testOps =
+ new GcpCredentialOps() {
+ @Override
+ public IamCredentialsClient
createIamCredentialsClient(GoogleCredentials credentials) {
+ capturedSource.set(credentials);
+ return mockIamClient;
+ }
+
+ @Override
+ public AccessToken refreshAccessToken(DownscopedCredentials
credentials) {
+ return new AccessToken("downscoped-token", new Date());
+ }
+ };
+
+ KeyPairGenerator gen = KeyPairGenerator.getInstance("RSA");
+ gen.initialize(2048);
+ KeyPair keyPair = gen.generateKeyPair();
+ String pem =
+ "-----BEGIN PRIVATE KEY-----\n"
+ + Base64.getMimeEncoder(64, "\n".getBytes(UTF_8))
+ .encodeToString(keyPair.getPrivate().getEncoded())
+ + "\n-----END PRIVATE KEY-----\n";
+ Path keyFile = tempDir.resolve("attribution-key.pem");
+ Files.writeString(keyFile, pem);
+
+ GcpAttributionParams params =
+ GcpAttributionParams.of(
+ "https://issuer.example.com";,
+
"//iam.googleapis.com/projects/123/locations/global/workloadIdentityPools/p/providers/pr",
+ keyFile.toString(),
+ "key-1");
+
+ GcpCredentialsStorageIntegration integration =
+ new GcpCredentialsStorageIntegration(
+ mockCreds,
+ ServiceOptions.getFromServiceLoader(HttpTransportFactory.class,
NetHttpTransport::new),
+ null,
+ config,
+ EMPTY_REALM_CONFIG,
+ testOps,
+ Optional.of(params));
+
+ integration.getStorageAccessConfig(
+ toGrants(
+ Set.of("gs://bucket/path"), Set.of("gs://bucket/path"),
Set.of("gs://bucket/path")),
+ Optional.empty(),
+ CredentialVendingContext.builder()
+ .realm(Optional.of("realm1"))
+ .principalName(Optional.of("principal1"))
+ .build());
+
+ Mockito.verify(mockIamClient)
+ .generateAccessToken(
+ Mockito.argThat(
+ request ->
+ request
+ .getName()
+ .equals(
+
GcpCredentialsStorageIntegration.SERVICE_ACCOUNT_PREFIX
Review Comment:
@adnanhemani Can you please approve the PR if you are ok with this?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
