Skip to content

Commit c4e1106

Browse files
nastranastra
authored
AWS: Don't fetch credential from endpoint if properties contain a valid credential (#12504) (#12515)
When the `VendedCredentialProvider` is created, the `properties` typically already contain a valid credential from the first time a table is loaded. This PR uses the credential from the properties and otherwise falls back to loading a valid credential from the refresh endpoint in case the credential in `properties` is incomplete or expired.
1 parent f057e87 commit c4e1106

File tree

2 files changed

+206
-1
lines changed

2 files changed

+206
-1
lines changed

‎aws/src/main/java/org/apache/iceberg/aws/s3/VendedCredentialsProvider.java‎

Lines changed: 36 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -22,8 +22,10 @@
2222
import java.time.temporal.ChronoUnit;
2323
import java.util.List;
2424
import java.util.Map;
25+
import java.util.Optional;
2526
import java.util.stream.Collectors;
2627
import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
28+
import org.apache.iceberg.relocated.com.google.common.base.Strings;
2729
import org.apache.iceberg.rest.ErrorHandlers;
2830
import org.apache.iceberg.rest.HTTPClient;
2931
import org.apache.iceberg.rest.RESTClient;
@@ -50,7 +52,7 @@ private VendedCredentialsProvider(Map<String, String> properties) {
5052
Preconditions.checkArgument(null != properties.get(URI), "Invalid URI: null");
5153
this.properties = properties;
5254
this.credentialCache =
53-
CachedSupplier.builder(this::refreshCredential)
55+
CachedSupplier.builder(() -> credentialFromProperties().orElseGet(this::refreshCredential))
5456
.cachedValueName(VendedCredentialsProvider.class.getName())
5557
.build();
5658
}
@@ -92,6 +94,39 @@ private LoadCredentialsResponse fetchCredentials() {
9294
ErrorHandlers.defaultErrorHandler());
9395
}
9496

97+
private Optional<RefreshResult<AwsCredentials>> credentialFromProperties() {
98+
String accessKeyId = properties.get(S3FileIOProperties.ACCESS_KEY_ID);
99+
String secretAccessKey = properties.get(S3FileIOProperties.SECRET_ACCESS_KEY);
100+
String sessionToken = properties.get(S3FileIOProperties.SESSION_TOKEN);
101+
String tokenExpiresAtMillis = properties.get(S3FileIOProperties.SESSION_TOKEN_EXPIRES_AT_MS);
102+
if (Strings.isNullOrEmpty(accessKeyId)
103+
|| Strings.isNullOrEmpty(secretAccessKey)
104+
|| Strings.isNullOrEmpty(sessionToken)
105+
|| Strings.isNullOrEmpty(tokenExpiresAtMillis)) {
106+
return Optional.empty();
107+
}
108+
109+
Instant expiresAt = Instant.ofEpochMilli(Long.parseLong(tokenExpiresAtMillis));
110+
Instant prefetchAt = expiresAt.minus(5, ChronoUnit.MINUTES);
111+
112+
if (Instant.now().isAfter(prefetchAt)) {
113+
return Optional.empty();
114+
}
115+
116+
return Optional.of(
117+
RefreshResult.builder(
118+
(AwsCredentials)
119+
AwsSessionCredentials.builder()
120+
.accessKeyId(accessKeyId)
121+
.secretAccessKey(secretAccessKey)
122+
.sessionToken(sessionToken)
123+
.expirationTime(expiresAt)
124+
.build())
125+
.staleTime(expiresAt)
126+
.prefetchTime(prefetchAt)
127+
.build());
128+
}
129+
95130
private RefreshResult<AwsCredentials> refreshCredential() {
96131
LoadCredentialsResponse response = fetchCredentials();
97132

‎aws/src/test/java/org/apache/iceberg/aws/s3/TestVendedCredentialsProvider.java‎

Lines changed: 170 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -302,6 +302,176 @@ public void multipleS3Credentials() {
302302
}
303303
}
304304

305+
@Test
306+
public void nonExpiredTokenInProperties() {
307+
HttpRequest mockRequest = request("/v1/credentials").withMethod(HttpMethod.GET.name());
308+
String expiresAt = Long.toString(Instant.now().plus(10, ChronoUnit.HOURS).toEpochMilli());
309+
Credential credentialFromProperties =
310+
ImmutableCredential.builder()
311+
.prefix("s3")
312+
.config(
313+
ImmutableMap.of(
314+
S3FileIOProperties.ACCESS_KEY_ID,
315+
"randomAccessKeyFromProperties",
316+
S3FileIOProperties.SECRET_ACCESS_KEY,
317+
"randomSecretAccessKeyFromProperties",
318+
S3FileIOProperties.SESSION_TOKEN,
319+
"sessionTokenFromProperties",
320+
S3FileIOProperties.SESSION_TOKEN_EXPIRES_AT_MS,
321+
expiresAt))
322+
.build();
323+
324+
Credential credential =
325+
ImmutableCredential.builder()
326+
.prefix("s3")
327+
.config(
328+
ImmutableMap.of(
329+
S3FileIOProperties.ACCESS_KEY_ID,
330+
"randomAccessKey",
331+
S3FileIOProperties.SECRET_ACCESS_KEY,
332+
"randomSecretAccessKey",
333+
S3FileIOProperties.SESSION_TOKEN,
334+
"sessionToken",
335+
S3FileIOProperties.SESSION_TOKEN_EXPIRES_AT_MS,
336+
Long.toString(Instant.now().plus(1, ChronoUnit.HOURS).toEpochMilli())))
337+
.build();
338+
LoadCredentialsResponse response =
339+
ImmutableLoadCredentialsResponse.builder().addCredentials(credential).build();
340+
341+
HttpResponse mockResponse =
342+
response(LoadCredentialsResponseParser.toJson(response)).withStatusCode(200);
343+
mockServer.when(mockRequest).respond(mockResponse);
344+
345+
try (VendedCredentialsProvider provider =
346+
VendedCredentialsProvider.create(
347+
ImmutableMap.of(
348+
VendedCredentialsProvider.URI,
349+
URI,
350+
S3FileIOProperties.ACCESS_KEY_ID,
351+
"randomAccessKeyFromProperties",
352+
S3FileIOProperties.SECRET_ACCESS_KEY,
353+
"randomSecretAccessKeyFromProperties",
354+
S3FileIOProperties.SESSION_TOKEN,
355+
"sessionTokenFromProperties",
356+
S3FileIOProperties.SESSION_TOKEN_EXPIRES_AT_MS,
357+
expiresAt))) {
358+
AwsCredentials awsCredentials = provider.resolveCredentials();
359+
360+
verifyCredentials(awsCredentials, credentialFromProperties);
361+
362+
for (int i = 0; i < 5; i++) {
363+
// resolving credentials multiple times should not hit the credentials endpoint again
364+
assertThat(provider.resolveCredentials()).isSameAs(awsCredentials);
365+
}
366+
}
367+
368+
// token endpoint isn't hit, because the credentials are extracted from the properties
369+
mockServer.verify(mockRequest, VerificationTimes.never());
370+
}
371+
372+
@Test
373+
public void expiredTokenInProperties() {
374+
HttpRequest mockRequest = request("/v1/credentials").withMethod(HttpMethod.GET.name());
375+
376+
Credential credential =
377+
ImmutableCredential.builder()
378+
.prefix("s3")
379+
.config(
380+
ImmutableMap.of(
381+
S3FileIOProperties.ACCESS_KEY_ID,
382+
"randomAccessKey",
383+
S3FileIOProperties.SECRET_ACCESS_KEY,
384+
"randomSecretAccessKey",
385+
S3FileIOProperties.SESSION_TOKEN,
386+
"sessionToken",
387+
S3FileIOProperties.SESSION_TOKEN_EXPIRES_AT_MS,
388+
Long.toString(Instant.now().plus(1, ChronoUnit.HOURS).toEpochMilli())))
389+
.build();
390+
LoadCredentialsResponse response =
391+
ImmutableLoadCredentialsResponse.builder().addCredentials(credential).build();
392+
393+
HttpResponse mockResponse =
394+
response(LoadCredentialsResponseParser.toJson(response)).withStatusCode(200);
395+
mockServer.when(mockRequest).respond(mockResponse);
396+
397+
try (VendedCredentialsProvider provider =
398+
VendedCredentialsProvider.create(
399+
ImmutableMap.of(
400+
VendedCredentialsProvider.URI,
401+
URI,
402+
S3FileIOProperties.ACCESS_KEY_ID,
403+
"randomAccessKeyFromProperties",
404+
S3FileIOProperties.SECRET_ACCESS_KEY,
405+
"randomSecretAccessKeyFromProperties",
406+
S3FileIOProperties.SESSION_TOKEN,
407+
"sessionTokenFromProperties",
408+
S3FileIOProperties.SESSION_TOKEN_EXPIRES_AT_MS,
409+
Long.toString(Instant.now().minus(1, ChronoUnit.HOURS).toEpochMilli())))) {
410+
AwsCredentials awsCredentials = provider.resolveCredentials();
411+
412+
verifyCredentials(awsCredentials, credential);
413+
414+
for (int i = 0; i < 5; i++) {
415+
// resolving credentials multiple times should not hit the credentials endpoint again
416+
assertThat(provider.resolveCredentials()).isSameAs(awsCredentials);
417+
}
418+
}
419+
420+
// token endpoint is hit once due to the properties containing an expired token
421+
mockServer.verify(mockRequest, VerificationTimes.once());
422+
}
423+
424+
@Test
425+
public void invalidTokenInProperties() {
426+
HttpRequest mockRequest = request("/v1/credentials").withMethod(HttpMethod.GET.name());
427+
428+
Credential credential =
429+
ImmutableCredential.builder()
430+
.prefix("s3")
431+
.config(
432+
ImmutableMap.of(
433+
S3FileIOProperties.ACCESS_KEY_ID,
434+
"randomAccessKey",
435+
S3FileIOProperties.SECRET_ACCESS_KEY,
436+
"randomSecretAccessKey",
437+
S3FileIOProperties.SESSION_TOKEN,
438+
"sessionToken",
439+
S3FileIOProperties.SESSION_TOKEN_EXPIRES_AT_MS,
440+
Long.toString(Instant.now().plus(1, ChronoUnit.HOURS).toEpochMilli())))
441+
.build();
442+
LoadCredentialsResponse response =
443+
ImmutableLoadCredentialsResponse.builder().addCredentials(credential).build();
444+
445+
HttpResponse mockResponse =
446+
response(LoadCredentialsResponseParser.toJson(response)).withStatusCode(200);
447+
mockServer.when(mockRequest).respond(mockResponse);
448+
449+
// token expiration is missing from the properties
450+
try (VendedCredentialsProvider provider =
451+
VendedCredentialsProvider.create(
452+
ImmutableMap.of(
453+
VendedCredentialsProvider.URI,
454+
URI,
455+
S3FileIOProperties.ACCESS_KEY_ID,
456+
"randomAccessKeyFromProperties",
457+
S3FileIOProperties.SECRET_ACCESS_KEY,
458+
"randomSecretAccessKeyFromProperties",
459+
S3FileIOProperties.SESSION_TOKEN,
460+
"sessionTokenFromProperties"))) {
461+
AwsCredentials awsCredentials = provider.resolveCredentials();
462+
463+
verifyCredentials(awsCredentials, credential);
464+
465+
for (int i = 0; i < 5; i++) {
466+
// resolving credentials multiple times should not hit the credentials endpoint again
467+
assertThat(provider.resolveCredentials()).isSameAs(awsCredentials);
468+
}
469+
}
470+
471+
// token endpoint is hit once due to the properties not containing the token's expiration
472+
mockServer.verify(mockRequest, VerificationTimes.once());
473+
}
474+
305475
private void verifyCredentials(AwsCredentials awsCredentials, Credential credential) {
306476
assertThat(awsCredentials).isInstanceOf(AwsSessionCredentials.class);
307477
AwsSessionCredentials creds = (AwsSessionCredentials) awsCredentials;

0 commit comments

Comments
 (0)