Skip to content

Commit

Permalink
Add stale time config to InstanceProfileCredentialsProvider (#5758)
Browse files Browse the repository at this point in the history
* Add stale time config to InstanceProfileCredentialsProvider to allow for refreshing credentials earlier due to stale value. This will help prevent returning invalid credentials when an error is encountered during asynchronous refresh.

* fix spotbug

* add configurable retry policy

* javadoc

* remove InstanceProfileCredentialsRetryPolicy

* default retryPolicy set in StaticResourcesEndpointProvider ctor
  • Loading branch information
L-Applin authored Jan 24, 2025
1 parent 943db51 commit 3c74396
Show file tree
Hide file tree
Showing 5 changed files with 122 additions and 4 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -88,6 +88,8 @@ public final class InstanceProfileCredentialsProvider

private final String profileName;

private final Duration staleTime;

/**
* @see #builder()
*/
Expand All @@ -108,6 +110,8 @@ private InstanceProfileCredentialsProvider(BuilderImpl builder) {
.profileName(profileName)
.build();

this.staleTime = Validate.getOrDefault(builder.staleTime, () -> Duration.ofSeconds(1));

if (Boolean.TRUE.equals(builder.asyncCredentialUpdateEnabled)) {
Validate.paramNotBlank(builder.asyncThreadName, "asyncThreadName");
this.credentialsCache = CachedSupplier.builder(this::refreshCredentials)
Expand Down Expand Up @@ -174,7 +178,7 @@ private Instant staleTime(Instant expiration) {
return null;
}

return expiration.minusSeconds(1);
return expiration.minus(staleTime);
}

private Instant prefetchTime(Instant expiration) {
Expand Down Expand Up @@ -340,6 +344,18 @@ public interface Builder extends HttpCredentialsProvider.Builder<InstanceProfile
*/
Builder profileName(String profileName);

/**
* Configure the amount of time before the moment of expiration of credentials for which to consider the credentials to
* be stale. A higher value can lead to a higher rate of request being made to the Amazon EC2 Instance Metadata Service.
* The default is 1 sec.
* <p>Increasing this value to a higher value (10s or more) may help with situations where a higher load on the instance
* metadata service causes it to return 503s error, for which the SDK may not be able to recover fast enough and
* returns expired credentials.
*
* @param duration the amount of time before expiration for when to consider the credentials to be stale and need refresh
*/
Builder staleTime(Duration duration);

/**
* Build a {@link InstanceProfileCredentialsProvider} from the provided configuration.
*/
Expand All @@ -355,6 +371,7 @@ static final class BuilderImpl implements Builder {
private String asyncThreadName;
private Supplier<ProfileFile> profileFile;
private String profileName;
private Duration staleTime;

private BuilderImpl() {
asyncThreadName("instance-profile-credentials-provider");
Expand All @@ -367,6 +384,7 @@ private BuilderImpl(InstanceProfileCredentialsProvider provider) {
this.asyncThreadName = provider.asyncThreadName;
this.profileFile = provider.profileFile;
this.profileName = provider.profileName;
this.staleTime = provider.staleTime;
}

Builder clock(Clock clock) {
Expand Down Expand Up @@ -435,6 +453,16 @@ public void setProfileName(String profileName) {
profileName(profileName);
}

@Override
public Builder staleTime(Duration duration) {
this.staleTime = duration;
return this;
}

public void setStaleTime(Duration duration) {
staleTime(duration);
}

@Override
public InstanceProfileCredentialsProvider build() {
return new InstanceProfileCredentialsProvider(this);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -17,10 +17,15 @@

import java.io.IOException;
import software.amazon.awssdk.annotations.SdkInternalApi;
import software.amazon.awssdk.auth.credentials.ContainerCredentialsProvider;
import software.amazon.awssdk.auth.credentials.InstanceProfileCredentialsProvider;
import software.amazon.awssdk.http.HttpStatusFamily;
import software.amazon.awssdk.regions.util.ResourcesEndpointRetryParameters;
import software.amazon.awssdk.regions.util.ResourcesEndpointRetryPolicy;

/**
* Retry policy shared by {@link InstanceProfileCredentialsProvider} and {@link ContainerCredentialsProvider#}.
*/
@SdkInternalApi
public final class ContainerCredentialsRetryPolicy implements ResourcesEndpointRetryPolicy {

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -24,23 +24,27 @@
import java.util.Optional;
import software.amazon.awssdk.annotations.SdkInternalApi;
import software.amazon.awssdk.regions.util.ResourcesEndpointProvider;
import software.amazon.awssdk.regions.util.ResourcesEndpointRetryPolicy;
import software.amazon.awssdk.utils.Validate;

@SdkInternalApi
public final class StaticResourcesEndpointProvider implements ResourcesEndpointProvider {
private final URI endpoint;
private final Map<String, String> headers;
private final Duration connectionTimeout;
private final ResourcesEndpointRetryPolicy retryPolicy;

private StaticResourcesEndpointProvider(URI endpoint,
Map<String, String> additionalHeaders,
Duration customTimeout) {
Map<String, String> additionalHeaders,
Duration customTimeout,
ResourcesEndpointRetryPolicy retryPolicy) {
this.endpoint = Validate.paramNotNull(endpoint, "endpoint");
this.headers = ResourcesEndpointProvider.super.headers();
if (additionalHeaders != null) {
this.headers.putAll(additionalHeaders);
}
this.connectionTimeout = customTimeout;
this.retryPolicy = Validate.getOrDefault(retryPolicy, () -> ResourcesEndpointRetryPolicy.NO_RETRY);
}

@Override
Expand All @@ -58,10 +62,16 @@ public Map<String, String> headers() {
return Collections.unmodifiableMap(headers);
}

@Override
public ResourcesEndpointRetryPolicy retryPolicy() {
return this.retryPolicy;
}

public static class Builder {
private URI endpoint;
private Map<String, String> additionalHeaders = new HashMap<>();
private Duration customTimeout;
private ResourcesEndpointRetryPolicy retryPolicy;

public Builder endpoint(URI endpoint) {
this.endpoint = Validate.paramNotNull(endpoint, "endpoint");
Expand All @@ -80,8 +90,13 @@ public Builder connectionTimeout(Duration timeout) {
return this;
}

public Builder retryPolicy(ResourcesEndpointRetryPolicy retryPolicy) {
this.retryPolicy = retryPolicy;
return this;
}

public StaticResourcesEndpointProvider build() {
return new StaticResourcesEndpointProvider(endpoint, additionalHeaders, customTimeout);
return new StaticResourcesEndpointProvider(endpoint, additionalHeaders, customTimeout, retryPolicy);
}
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -17,13 +17,15 @@

import static com.github.tomakehurst.wiremock.client.WireMock.aResponse;
import static com.github.tomakehurst.wiremock.client.WireMock.equalTo;
import static com.github.tomakehurst.wiremock.client.WireMock.exactly;
import static com.github.tomakehurst.wiremock.client.WireMock.get;
import static com.github.tomakehurst.wiremock.client.WireMock.getRequestedFor;
import static com.github.tomakehurst.wiremock.client.WireMock.put;
import static com.github.tomakehurst.wiremock.client.WireMock.putRequestedFor;
import static com.github.tomakehurst.wiremock.client.WireMock.stubFor;
import static com.github.tomakehurst.wiremock.client.WireMock.urlPathEqualTo;
import static com.github.tomakehurst.wiremock.core.WireMockConfiguration.wireMockConfig;
import static com.github.tomakehurst.wiremock.stubbing.Scenario.STARTED;
import static java.time.temporal.ChronoUnit.HOURS;
import static java.time.temporal.ChronoUnit.MINUTES;
import static java.time.temporal.ChronoUnit.SECONDS;
Expand Down Expand Up @@ -55,6 +57,7 @@
import org.junit.jupiter.api.extension.RegisterExtension;
import org.junit.jupiter.params.ParameterizedTest;
import org.junit.jupiter.params.provider.ValueSource;
import org.mockito.Mockito;
import software.amazon.awssdk.core.SdkSystemSetting;
import software.amazon.awssdk.core.exception.SdkClientException;
import software.amazon.awssdk.core.util.SdkUserAgent;
Expand Down Expand Up @@ -596,13 +599,74 @@ void imdsCallFrequencyIsLimited() {
}
}

@Test
void testErrorWhileCacheIsStale_shouldRecover() {
AdjustableClock clock = new AdjustableClock();

Instant now = Instant.now();
Instant expiration = now.plus(Duration.ofHours(6));

String successfulCredentialsResponse =
"{"
+ "\"AccessKeyId\":\"ACCESS_KEY_ID\","
+ "\"SecretAccessKey\":\"SECRET_ACCESS_KEY\","
+ "\"Expiration\":\"" + DateUtils.formatIso8601Date(expiration) + '"'
+ "}";

String staleResponse =
"{"
+ "\"AccessKeyId\":\"ACCESS_KEY_ID_2\","
+ "\"SecretAccessKey\":\"SECRET_ACCESS_KEY_2\","
+ "\"Expiration\":\"" + DateUtils.formatIso8601Date(Instant.now()) + '"'
+ "}";


Duration staleTime = Duration.ofMinutes(5);
AwsCredentialsProvider provider = credentialsProviderWithClock(clock, staleTime);

// cache expiration with expiration = 6 hours
clock.time = now;
stubSecureCredentialsResponse(aResponse().withBody(successfulCredentialsResponse));
AwsCredentials validCreds = provider.resolveCredentials();

// failure while cache is stale
clock.time = expiration.minus(staleTime.minus(Duration.ofMinutes(2)));
stubTokenFetchErrorResponse(aResponse().withFixedDelay(2000).withBody(STUB_CREDENTIALS), 500);
stubSecureCredentialsResponse(aResponse().withBody(staleResponse));
AwsCredentials refreshedWhileStale = provider.resolveCredentials();

assertThat(refreshedWhileStale).isNotEqualTo(validCreds);
assertThat(refreshedWhileStale.secretAccessKey()).isEqualTo("SECRET_ACCESS_KEY_2");
}

@Test
void shouldNotRetry_whenSucceeds() {
stubSecureCredentialsResponse(aResponse().withBody(STUB_CREDENTIALS));
InstanceProfileCredentialsProvider provider = InstanceProfileCredentialsProvider.builder().build();
AwsCredentials credentials = provider.resolveCredentials();
assertThat(credentials.accessKeyId()).isEqualTo("ACCESS_KEY_ID");
assertThat(credentials.secretAccessKey()).isEqualTo("SECRET_ACCESS_KEY");
assertThat(credentials.providerName()).isPresent().contains("InstanceProfileCredentialsProvider");
verifyImdsCallWithToken();
WireMock.verify(exactly(1), getRequestedFor(urlPathEqualTo(CREDENTIALS_RESOURCE_PATH + "some-profile")));
}

private AwsCredentialsProvider credentialsProviderWithClock(Clock clock) {
InstanceProfileCredentialsProvider.BuilderImpl builder =
(InstanceProfileCredentialsProvider.BuilderImpl) InstanceProfileCredentialsProvider.builder();
builder.clock(clock);
return builder.build();
}

private AwsCredentialsProvider credentialsProviderWithClock(Clock clock, Duration staleTime) {
InstanceProfileCredentialsProvider.BuilderImpl builder =
(InstanceProfileCredentialsProvider.BuilderImpl) InstanceProfileCredentialsProvider.builder();
builder.clock(clock);
builder.staleTime(staleTime);
return builder.build();
}


private static class AdjustableClock extends Clock {
private Instant time;

Expand Down
6 changes: 6 additions & 0 deletions core/auth/src/test/resources/log4j2.properties
Original file line number Diff line number Diff line change
Expand Up @@ -36,3 +36,9 @@ rootLogger.appenderRef.stdout.ref = ConsoleAppender
#
#logger.netty.name = io.netty.handler.logging
#logger.netty.level = debug

#logger.cache.name = software.amazon.awssdk.utils.cache.CachedSupplier
#logger.cache.level = DEBUG

#logger.instance.name = software.amazon.awssdk.auth.credentials.InstanceProfileCredentialsProvider
#logger.instance.level = DEBUG

0 comments on commit 3c74396

Please sign in to comment.