Skip to content

Commit 5dc1050

Browse files
rohangolirohangoli
committed
Handle Non-Standard STS XML Response
1 parent 69e402e commit 5dc1050

File tree

10 files changed

+492
-15
lines changed

10 files changed

+492
-15
lines changed

‎polaris-core/src/main/java/org/apache/polaris/core/storage/aws/AwsCredentialsStorageIntegration.java‎

Lines changed: 35 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -109,19 +109,41 @@ public AccessConfig getSubscopedCreds(
109109
storageConfig.getIgnoreSSLVerification()));
110110

111111
AssumeRoleResponse response = stsClient.assumeRole(request.build());
112-
accessConfig.put(StorageAccessProperty.AWS_KEY_ID, response.credentials().accessKeyId());
113-
accessConfig.put(
114-
StorageAccessProperty.AWS_SECRET_KEY, response.credentials().secretAccessKey());
115-
accessConfig.put(StorageAccessProperty.AWS_TOKEN, response.credentials().sessionToken());
116-
Optional.ofNullable(response.credentials().expiration())
117-
.ifPresent(
118-
i -> {
119-
accessConfig.put(
120-
StorageAccessProperty.EXPIRATION_TIME, String.valueOf(i.toEpochMilli()));
121-
accessConfig.put(
122-
StorageAccessProperty.AWS_SESSION_TOKEN_EXPIRES_AT_MS,
123-
String.valueOf(i.toEpochMilli()));
124-
});
112+
if (response != null && response.credentials() != null) {
113+
accessConfig.put(StorageAccessProperty.AWS_KEY_ID, response.credentials().accessKeyId());
114+
accessConfig.put(
115+
StorageAccessProperty.AWS_SECRET_KEY, response.credentials().secretAccessKey());
116+
accessConfig.put(StorageAccessProperty.AWS_TOKEN, response.credentials().sessionToken());
117+
Optional.ofNullable(response.credentials().expiration())
118+
.ifPresent(
119+
i -> {
120+
accessConfig.put(
121+
StorageAccessProperty.EXPIRATION_TIME, String.valueOf(i.toEpochMilli()));
122+
accessConfig.put(
123+
StorageAccessProperty.AWS_SESSION_TOKEN_EXPIRES_AT_MS,
124+
String.valueOf(i.toEpochMilli()));
125+
});
126+
} else {
127+
// Try to recover by reading raw STS body captured by interceptor
128+
try {
129+
String raw = org.apache.polaris.core.storage.aws.StsResponseCapture.getLastBody();
130+
if (raw != null && !raw.isBlank()) {
131+
try {
132+
var parsed =
133+
org.apache.polaris.core.storage.aws.StsXmlParser.parseToAccessConfig(raw);
134+
// merge parsed credentials into accessConfig builder
135+
parsed.credentials().forEach((k, v) -> accessConfig.putCredential(k, v));
136+
parsed.internalProperties().forEach((k, v) -> accessConfig.putInternalProperty(k, v));
137+
parsed.extraProperties().forEach((k, v) -> accessConfig.putExtraProperty(k, v));
138+
parsed.expiresAt().ifPresent(accessConfig::expiresAt);
139+
} catch (Exception ignore) {
140+
// parsing failed - ignore and fallthrough
141+
}
142+
}
143+
} finally {
144+
org.apache.polaris.core.storage.aws.StsResponseCapture.clear();
145+
}
146+
}
125147
}
126148

127149
if (region != null) {

‎polaris-core/src/main/java/org/apache/polaris/core/storage/aws/StsResponseCapture.java‎

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,42 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one
3+
* or more contributor license agreements. See the NOTICE file
4+
* distributed with this work for additional information
5+
* regarding copyright ownership. The ASF licenses this file
6+
* to you under the Apache License, Version 2.0 (the
7+
* "License"); you may not use this file except in compliance
8+
* with the License. You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing,
13+
* software distributed under the License is distributed on an
14+
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15+
* KIND, either express or implied. See the License for the
16+
* specific language governing permissions and limitations
17+
* under the License.
18+
*/
19+
package org.apache.polaris.core.storage.aws;
20+
21+
/**
22+
* Simple thread-local holder for the last raw STS HTTP response body captured by an
23+
* ExecutionInterceptor. This is intended as a pragmatic bridge so synchronous SDK calls can consult
24+
* the raw response body when unmarshalling produced an unexpected null result.
25+
*/
26+
public final class StsResponseCapture {
27+
private static final ThreadLocal<String> LAST_BODY = new ThreadLocal<>();
28+
29+
private StsResponseCapture() {}
30+
31+
public static void setLastBody(String body) {
32+
LAST_BODY.set(body);
33+
}
34+
35+
public static String getLastBody() {
36+
return LAST_BODY.get();
37+
}
38+
39+
public static void clear() {
40+
LAST_BODY.remove();
41+
}
42+
}

‎polaris-core/src/main/java/org/apache/polaris/core/storage/aws/StsResponseCaptureInterceptor.java‎

Lines changed: 95 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,95 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one
3+
* or more contributor license agreements. See the NOTICE file
4+
* distributed with this work for additional information
5+
* regarding copyright ownership. The ASF licenses this file
6+
* to you under the Apache License, Version 2.0 (the
7+
* "License"); you may not use this file except in compliance
8+
* with the License. You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing,
13+
* software distributed under the License is distributed on an
14+
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15+
* KIND, either express or implied. See the License for the
16+
* specific language governing permissions and limitations
17+
* under the License.
18+
*/
19+
package org.apache.polaris.core.storage.aws;
20+
21+
import java.io.ByteArrayOutputStream;
22+
import java.io.InputStream;
23+
import java.lang.reflect.Method;
24+
import java.nio.charset.StandardCharsets;
25+
import java.util.Optional;
26+
import software.amazon.awssdk.core.interceptor.Context.AfterTransmission;
27+
import software.amazon.awssdk.core.interceptor.ExecutionAttributes;
28+
import software.amazon.awssdk.core.interceptor.ExecutionInterceptor;
29+
30+
/**
31+
* ExecutionInterceptor that captures the raw HTTP response body for STS calls and saves it into
32+
* StsResponseCapture (thread-local). This allows calling code to inspect the raw response when the
33+
* SDK's unmarshalling yields null credentials.
34+
*/
35+
public class StsResponseCaptureInterceptor implements ExecutionInterceptor {
36+
37+
@Override
38+
public void afterTransmission(
39+
AfterTransmission context, ExecutionAttributes executionAttributes) {
40+
try {
41+
// Use reflection to call context.httpResponse() because SDK versions expose different
42+
// types/APIs
43+
Method httpRespMethod = context.getClass().getMethod("httpResponse");
44+
Object httpResp = httpRespMethod.invoke(context);
45+
if (httpResp != null) {
46+
try {
47+
Optional<InputStream> content = OptionalUtils.safeGetContent(httpResp);
48+
if (content.isPresent()) {
49+
try (InputStream in = content.get();
50+
ByteArrayOutputStream out = new ByteArrayOutputStream()) {
51+
byte[] buf = new byte[8192];
52+
int r;
53+
while ((r = in.read(buf)) != -1) {
54+
out.write(buf, 0, r);
55+
}
56+
String resp = new String(out.toByteArray(), StandardCharsets.UTF_8);
57+
StsResponseCapture.setLastBody(resp);
58+
}
59+
}
60+
} catch (Exception e) {
61+
// best-effort; don't fail the call because of capture problems
62+
}
63+
}
64+
} catch (Throwable t) {
65+
// swallow - capture is non-fatal
66+
}
67+
}
68+
}
69+
70+
// Small utility to safely read content from SdkHttpResponse across SDK versions.
71+
final class OptionalUtils {
72+
static Optional<InputStream> safeGetContent(Object httpResp) {
73+
try {
74+
Method contentMethod = httpResp.getClass().getMethod("content");
75+
Object val = contentMethod.invoke(httpResp);
76+
if (val == null) return Optional.empty();
77+
if (val instanceof Optional) {
78+
Optional<?> anyOpt = (Optional<?>) val;
79+
if (anyOpt.isPresent() && anyOpt.get() instanceof InputStream) {
80+
return Optional.of((InputStream) anyOpt.get());
81+
}
82+
return Optional.empty();
83+
}
84+
// Some SDKs may return an InputStream directly
85+
if (val instanceof InputStream) {
86+
return Optional.of((InputStream) val);
87+
}
88+
} catch (NoSuchMethodException nsme) {
89+
// ignore
90+
} catch (Exception e) {
91+
// ignore other reflection errors
92+
}
93+
return Optional.empty();
94+
}
95+
}

‎polaris-core/src/main/java/org/apache/polaris/core/storage/aws/StsXmlParser.java‎

Lines changed: 104 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,104 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one
3+
* or more contributor license agreements. See the NOTICE file
4+
* distributed with this work for additional information
5+
* regarding copyright ownership. The ASF licenses this file
6+
* to you under the Apache License, Version 2.0 (the
7+
* "License"); you may not use this file except in compliance
8+
* with the License. You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing,
13+
* software distributed under the License is distributed on an
14+
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15+
* KIND, either express or implied. See the License for the
16+
* specific language governing permissions and limitations
17+
* under the License.
18+
*/
19+
package org.apache.polaris.core.storage.aws;
20+
21+
import java.time.Instant;
22+
import java.util.Objects;
23+
import org.apache.polaris.core.storage.AccessConfig;
24+
25+
/** Utility to parse STS AssumeRoleResponse XML (namespaced or not) into an AccessConfig. */
26+
public final class StsXmlParser {
27+
private StsXmlParser() {}
28+
29+
public static AccessConfig parseToAccessConfig(String xml) throws Exception {
30+
Objects.requireNonNull(xml);
31+
var dbf = javax.xml.parsers.DocumentBuilderFactory.newInstance();
32+
dbf.setNamespaceAware(true);
33+
var db = dbf.newDocumentBuilder();
34+
try (java.io.ByteArrayInputStream in =
35+
new java.io.ByteArrayInputStream(xml.getBytes(java.nio.charset.StandardCharsets.UTF_8))) {
36+
org.w3c.dom.Document doc = db.parse(in);
37+
javax.xml.xpath.XPath xPath = javax.xml.xpath.XPathFactory.newInstance().newXPath();
38+
39+
String accessKeyId =
40+
(String)
41+
xPath.evaluate(
42+
"//*[local-name() = 'Credentials']/*[local-name() = 'AccessKeyId']/text()",
43+
doc,
44+
javax.xml.xpath.XPathConstants.STRING);
45+
String secretAccessKey =
46+
(String)
47+
xPath.evaluate(
48+
"//*[local-name() = 'Credentials']/*[local-name() = 'SecretAccessKey']/text()",
49+
doc,
50+
javax.xml.xpath.XPathConstants.STRING);
51+
String sessionToken =
52+
(String)
53+
xPath.evaluate(
54+
"//*[local-name() = 'Credentials']/*[local-name() = 'SessionToken']/text()",
55+
doc,
56+
javax.xml.xpath.XPathConstants.STRING);
57+
String expiration =
58+
(String)
59+
xPath.evaluate(
60+
"//*[local-name() = 'Credentials']/*[local-name() = 'Expiration']/text()",
61+
doc,
62+
javax.xml.xpath.XPathConstants.STRING);
63+
64+
if (accessKeyId == null || accessKeyId.isBlank()) {
65+
throw new IllegalArgumentException("No AccessKeyId found in STS response XML");
66+
}
67+
68+
AccessConfig.Builder builder = AccessConfig.builder();
69+
builder.putCredential(
70+
org.apache.polaris.core.storage.StorageAccessProperty.AWS_KEY_ID.getPropertyName(),
71+
accessKeyId.trim());
72+
if (secretAccessKey != null && !secretAccessKey.isBlank()) {
73+
builder.putCredential(
74+
org.apache.polaris.core.storage.StorageAccessProperty.AWS_SECRET_KEY.getPropertyName(),
75+
secretAccessKey.trim());
76+
}
77+
if (sessionToken != null && !sessionToken.isBlank()) {
78+
builder.putCredential(
79+
org.apache.polaris.core.storage.StorageAccessProperty.AWS_TOKEN.getPropertyName(),
80+
sessionToken.trim());
81+
}
82+
if (expiration != null && !expiration.isBlank()) {
83+
try {
84+
Instant i = Instant.parse(expiration.trim());
85+
builder.putCredential(
86+
org.apache.polaris.core.storage.StorageAccessProperty.EXPIRATION_TIME
87+
.getPropertyName(),
88+
String.valueOf(i.toEpochMilli()));
89+
builder.putCredential(
90+
org.apache.polaris.core.storage.StorageAccessProperty.AWS_SESSION_TOKEN_EXPIRES_AT_MS
91+
.getPropertyName(),
92+
String.valueOf(i.toEpochMilli()));
93+
builder.expiresAt(i);
94+
} catch (Exception e) {
95+
builder.putExtraProperty(
96+
org.apache.polaris.core.storage.StorageAccessProperty.EXPIRATION_TIME
97+
.getPropertyName(),
98+
expiration.trim());
99+
}
100+
}
101+
return builder.build();
102+
}
103+
}
104+
}

‎polaris-core/src/test/java/org/apache/polaris/core/storage/aws/StsResponseCaptureInterceptorTest.java‎

Lines changed: 96 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,96 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one
3+
* or more contributor license agreements. See the NOTICE file
4+
* distributed with this work for additional information
5+
* regarding copyright ownership. The ASF licenses this file
6+
* to you under the Apache License, Version 2.0 (the
7+
* "License"); you may not use this file except in compliance
8+
* with the License. You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing,
13+
* software distributed under the License is distributed on an
14+
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15+
* KIND, either express or implied. See the License for the
16+
* specific language governing permissions and limitations
17+
* under the License.
18+
*/
19+
package org.apache.polaris.core.storage.aws;
20+
21+
import static org.assertj.core.api.Assertions.assertThat;
22+
23+
import java.io.ByteArrayInputStream;
24+
import java.util.Optional;
25+
import org.junit.jupiter.api.BeforeEach;
26+
import org.junit.jupiter.api.Test;
27+
import software.amazon.awssdk.core.interceptor.ExecutionAttributes;
28+
29+
public class StsResponseCaptureInterceptorTest {
30+
31+
@BeforeEach
32+
public void before() {
33+
StsResponseCapture.clear();
34+
}
35+
36+
// Local interface used by the dynamic proxy to provide content()
37+
@SuppressWarnings("unused")
38+
private interface ContentHolder {
39+
java.util.Optional<java.io.InputStream> content();
40+
}
41+
42+
@Test
43+
public void testAfterTransmissionCapturesBody() {
44+
StsResponseCaptureInterceptor interceptor = new StsResponseCaptureInterceptor();
45+
try {
46+
Class<?> afterCls = software.amazon.awssdk.core.interceptor.Context.AfterTransmission.class;
47+
Class<?> httpRespType = afterCls.getMethod("httpResponse").getReturnType();
48+
49+
// Build a response proxy that implements both the SDK response return type and ContentHolder
50+
Object respProxy =
51+
java.lang.reflect.Proxy.newProxyInstance(
52+
httpRespType.getClassLoader(),
53+
new Class[] {httpRespType, ContentHolder.class},
54+
(proxy, method, args) -> {
55+
if ("content".equals(method.getName())) {
56+
return Optional.of(
57+
new ByteArrayInputStream(
58+
"raw-body-xyz".getBytes(java.nio.charset.StandardCharsets.UTF_8)));
59+
}
60+
return null;
61+
});
62+
63+
// Now build an AfterTransmission proxy that returns the respProxy from httpResponse()
64+
Object ctxProxy =
65+
java.lang.reflect.Proxy.newProxyInstance(
66+
afterCls.getClassLoader(),
67+
new Class[] {afterCls},
68+
(proxy, method, args) -> {
69+
if ("httpResponse".equals(method.getName())) {
70+
return respProxy;
71+
}
72+
return null;
73+
});
74+
75+
interceptor.afterTransmission(
76+
(software.amazon.awssdk.core.interceptor.Context.AfterTransmission) ctxProxy,
77+
new ExecutionAttributes());
78+
} catch (Throwable t) {
79+
throw new AssertionError(t);
80+
}
81+
// Because we cast via reflection, ensure the captured body is set
82+
assertThat(StsResponseCapture.getLastBody()).isEqualTo("raw-body-xyz");
83+
}
84+
85+
@Test
86+
public void testAfterTransmissionSilentlyIgnoresUnknownContext() {
87+
StsResponseCaptureInterceptor interceptor = new StsResponseCaptureInterceptor();
88+
// Create a mock AfterTransmission that returns null for httpResponse()
89+
software.amazon.awssdk.core.interceptor.Context.AfterTransmission nullRespContext =
90+
org.mockito.Mockito.mock(
91+
software.amazon.awssdk.core.interceptor.Context.AfterTransmission.class);
92+
org.mockito.Mockito.when(nullRespContext.httpResponse()).thenReturn(null);
93+
interceptor.afterTransmission(nullRespContext, new ExecutionAttributes());
94+
assertThat(StsResponseCapture.getLastBody()).isNull();
95+
}
96+
}

0 commit comments

Comments
 (0)