diff --git a/auth-api/src/main/java/software/amazon/smithy/java/auth/api/identity/CachingIdentityResolver.java b/auth-api/src/main/java/software/amazon/smithy/java/auth/api/identity/CachingIdentityResolver.java index f16416305f..d1715aa7f0 100644 --- a/auth-api/src/main/java/software/amazon/smithy/java/auth/api/identity/CachingIdentityResolver.java +++ b/auth-api/src/main/java/software/amazon/smithy/java/auth/api/identity/CachingIdentityResolver.java @@ -9,79 +9,90 @@ import java.time.Duration; import java.time.Instant; import java.util.Objects; -import java.util.concurrent.CountDownLatch; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ExecutionException; import java.util.concurrent.Executors; +import java.util.concurrent.RejectedExecutionException; import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.ScheduledFuture; +import java.util.concurrent.ThreadLocalRandom; import java.util.concurrent.TimeUnit; -import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.locks.ReentrantLock; +import java.util.function.BiPredicate; import software.amazon.smithy.java.context.Context; -/** * An {@link IdentityResolver} that caches the result of a delegate resolver and refreshes it asynchronously in the - * background before expiration. +/** + * An {@link IdentityResolver} that caches identities and refreshes them using advisory and mandatory refresh + * windows. * - *

Behavior: - *

+ *

Only one refresh can run at a time. Advisory refreshes run in the background while callers return cached + * identities. Mandatory refreshes block callers and share the result of the in-flight refresh. When static stability + * is enabled, failed refreshes retain and return cached identities, including expired identities, and rate-limit + * subsequent attempts with a jittered backoff. * - *

This class is thread-safe. At most one refresh runs at a time (enforced by an {@link AtomicBoolean}). - * Callers never block except on cold start. - * - * @param the identity type. + * @param identity type. */ public final class CachingIdentityResolver implements IdentityResolver, AutoCloseable { private static final System.Logger LOGGER = System.getLogger(CachingIdentityResolver.class.getName()); + private static final Duration DEFAULT_MANDATORY_WINDOW = Duration.ofMinutes(1); + private static final Duration DEFAULT_BACKOFF_MIN = Duration.ofMinutes(5); + private static final Duration DEFAULT_BACKOFF_MAX = Duration.ofMinutes(10); private final IdentityResolver delegate; - private final Duration prefetchBuffer; + private final Duration configuredAdvisoryWindow; + private final Duration mandatoryRefreshWindow; private final boolean allowExpiredCredentials; - private final Duration staleRefreshDelay; + private final Duration refreshBackoffMin; + private final Duration refreshBackoffMax; private final Clock clock; private final ScheduledExecutorService executor; private final boolean ownsExecutor; - private final AtomicBoolean refreshing = new AtomicBoolean(false); - private volatile CountDownLatch coldStartLatch = new CountDownLatch(1); - - private volatile CachedValue cached; - private volatile ScheduledFuture scheduledRefresh; + private final boolean closeDelegate; + private final boolean proactiveRefresh; + private final BiPredicate identityMatcher; + private final ReentrantLock lock = new ReentrantLock(); + + private CachedValue cached; + private CompletableFuture> inFlight; + private ScheduledFuture scheduledRefresh; + private Instant nextRefreshAllowedAt; + private Instant nextRefreshAfterSuccessAt; + private RefreshOutcome lastFailure; + private boolean refreshRequired; + private boolean closed; private CachingIdentityResolver(Builder builder) { this.delegate = Objects.requireNonNull(builder.delegate, "delegate"); - this.prefetchBuffer = builder.prefetchBuffer; + this.configuredAdvisoryWindow = builder.advisoryRefreshWindow; + this.mandatoryRefreshWindow = builder.mandatoryRefreshWindow; this.allowExpiredCredentials = builder.allowExpiredCredentials; - this.staleRefreshDelay = builder.staleRefreshDelay; + this.refreshBackoffMin = builder.refreshBackoffMin; + this.refreshBackoffMax = builder.refreshBackoffMax; this.clock = builder.clock; - - if (builder.executor != null) { - this.executor = builder.executor; - this.ownsExecutor = false; - } else { - this.executor = Executors.newSingleThreadScheduledExecutor(r -> { - Thread t = new Thread(r, "smithy-identity-cache-refresh"); - t.setDaemon(true); - return t; + this.closeDelegate = builder.closeDelegate; + this.proactiveRefresh = builder.proactiveRefresh; + this.identityMatcher = builder.identityMatcher; + + if (builder.executor == null) { + this.executor = Executors.newSingleThreadScheduledExecutor(runnable -> { + Thread thread = new Thread(runnable, "smithy-identity-cache-refresh"); + thread.setDaemon(true); + return thread; }); this.ownsExecutor = true; + } else { + this.executor = builder.executor; + this.ownsExecutor = false; } } /** - * Create a builder. + * Creates a caching resolver builder. * - * @param delegate the underlying resolver to cache. - * @param identity type. - * @return a new builder. + * @param delegate resolver that obtains identities from the source. + * @param identity type. + * @return builder. */ public static Builder builder(IdentityResolver delegate) { return new Builder<>(delegate); @@ -89,27 +100,53 @@ public static Builder builder(IdentityResolver delega @Override public IdentityResult resolveIdentity(Context requestProperties) { - CachedValue current = cached; - - // Cold start: first caller triggers refresh, others wait. - if (current == null) { - return coldStart(requestProperties); - } + CompletableFuture> refresh; + boolean performRefresh = false; + IdentityResult advisoryResult = null; - // Cache is fresh — return immediately. - if (!isInPrefetchWindow(current) && !isExpired(current)) { - return current.result; + lock.lock(); + try { + Instant now = clock.instant(); + if (cached == null) { + refresh = inFlight; + if (refresh == null) { + refresh = new CompletableFuture<>(); + inFlight = refresh; + performRefresh = true; + } + } else if (!refreshNeeded(cached, now)) { + return cached.result; + } else if (refreshRateLimited(now)) { + return resolveOutcome(fallbackOutcome(cached, lastFailure, now)); + } else if (!mandatoryRefreshNeeded(cached, now)) { + refresh = inFlight; + if (refresh != null) { + return cached.result; + } + refresh = new CompletableFuture<>(); + inFlight = refresh; + performRefresh = true; + advisoryResult = cached.result; + } else { + refresh = inFlight; + if (refresh == null) { + refresh = new CompletableFuture<>(); + inFlight = refresh; + performRefresh = true; + } + } + } finally { + lock.unlock(); } - // Cache is in prefetch window or expired. Kick off async refresh if not already running. - triggerAsyncRefresh(requestProperties); - - // If expired and strict mode, we can't return stale — block for the refresh. - if (isExpired(current) && !allowExpiredCredentials) { - return blockForRefresh(current, requestProperties); + if (performRefresh) { + if (advisoryResult != null) { + executeAdvisoryRefresh(requestProperties, refresh, advisoryResult); + return advisoryResult; + } + executeRefresh(requestProperties, refresh); } - - return current.result; + return awaitRefresh(refresh); } @Override @@ -118,244 +155,547 @@ public Class identityType() { } @Override - public void invalidate() { - cached = null; - coldStartLatch = new CountDownLatch(1); - cancelScheduledRefresh(); + public void invalidate(I rejectedIdentity) { + // Invalidation must not make the request path wait for refresh lifecycle state. + if (!lock.tryLock()) { + return; + } + try { + // An active refresh supersedes this rejection and will determine the next cached identity. + if (inFlight != null) { + return; + } + if (cached == null || !identityMatcher.test(cached.identity, rejectedIdentity)) { + return; + } + refreshRequired = true; + nextRefreshAfterSuccessAt = null; + cancelScheduledRefreshLocked(); + if (nextRefreshAllowedAt != null && clock.instant().isBefore(nextRefreshAllowedAt)) { + scheduleRefreshLocked(nextRefreshAllowedAt); + } + } finally { + lock.unlock(); + } } @Override public void close() { - cancelScheduledRefresh(); + lock.lock(); + try { + if (closed) { + return; + } + closed = true; + cancelScheduledRefreshLocked(); + } finally { + lock.unlock(); + } + if (closeDelegate && delegate instanceof AutoCloseable closeable) { + try { + closeable.close(); + } catch (Exception error) { + LOGGER.log(System.Logger.Level.WARNING, "Failed to close identity resolver", error); + } + } if (ownsExecutor) { executor.shutdownNow(); } } - private IdentityResult coldStart(Context requestProperties) { - if (refreshing.compareAndSet(false, true)) { + private void executeRefresh(Context requestProperties, CompletableFuture> refresh) { + try { + RefreshAttempt attempt = callCredentialSource(requestProperties); + RefreshOutcome callerOutcome; + + lock.lock(); try { - return doRefresh(requestProperties); + Instant now = clock.instant(); + if (attempt.identity != null && isFresh(attempt.identity, now)) { + cached = createCachedValue(attempt.identity, now, requestProperties); + nextRefreshAllowedAt = null; + lastFailure = null; + refreshRequired = false; + nextRefreshAfterSuccessAt = firstRefreshAt(cached, now); + scheduleRefreshLocked(nextRefreshAfterSuccessAt); + callerOutcome = RefreshOutcome.result(cached.result); + } else { + RefreshOutcome failure = attempt.failureOutcome(this); + nextRefreshAfterSuccessAt = null; + lastFailure = failure; + if (attempt.nonRecoverable) { + // Non-recoverable failures are not backed off; cached state forces the next call to retry. + nextRefreshAllowedAt = null; + refreshRequired = cached != null; + cancelScheduledRefreshLocked(); + callerOutcome = failure; + } else if (cached == null) { + nextRefreshAllowedAt = null; + cancelScheduledRefreshLocked(); + callerOutcome = failure; + } else { + Duration backoff = refreshBackoff(); + nextRefreshAllowedAt = now.plus(backoff); + boolean useCachedCredentials = allowExpiredCredentials || !isExpired(cached, now); + logRefreshFailure(attempt.failureDescription(), backoff, useCachedCredentials); + callerOutcome = fallbackOutcome(cached, failure, now); + scheduleRefreshLocked(nextRefreshAllowedAt); + } + } + if (inFlight == refresh) { + inFlight = null; + } } finally { - refreshing.set(false); - coldStartLatch.countDown(); + lock.unlock(); } - } - // Another thread is doing the cold start — wait for it. - try { - coldStartLatch.await(); - } catch (InterruptedException e) { - Thread.currentThread().interrupt(); - return IdentityResult.ofError(getClass(), "Interrupted waiting for initial credential resolution"); + refresh.complete(callerOutcome); + } catch (RuntimeException | Error failure) { + failRefresh(refresh, failure); + throw failure; } - - CachedValue result = cached; - return result != null ? result.result : IdentityResult.ofError(getClass(), "Failed to resolve credentials"); } - private void triggerAsyncRefresh(Context requestProperties) { - if (refreshing.compareAndSet(false, true)) { - executor.submit(() -> { - try { - doRefresh(requestProperties); - } finally { - refreshing.set(false); - } - }); + private void failRefresh(CompletableFuture> refresh, Throwable failure) { + lock.lock(); + try { + if (inFlight == refresh) { + inFlight = null; + } + } finally { + lock.unlock(); } + refresh.completeExceptionally(failure); } - private IdentityResult blockForRefresh(CachedValue current, Context requestProperties) { - // Strict mode: cache is expired. Try one synchronous refresh. - if (refreshing.compareAndSet(false, true)) { + private void executeAdvisoryRefresh( + Context requestProperties, + CompletableFuture> refresh, + IdentityResult cachedResult + ) { + Context refreshProperties = Context.unmodifiableCopy(requestProperties); + try { + executor.execute(() -> executeRefresh(refreshProperties, refresh)); + } catch (RejectedExecutionException error) { + lock.lock(); try { - IdentityResult result = doRefresh(requestProperties); - // If doRefresh returned the stale cached value (shouldn't in strict mode), check again. - CachedValue latest = cached; - if (latest != current) { - return latest.result; + if (inFlight == refresh) { + inFlight = null; } - return result; } finally { - refreshing.set(false); + lock.unlock(); } + refresh.complete(RefreshOutcome.result(cachedResult)); } + } - // Another thread is refreshing — wait briefly then check. + private RefreshAttempt callCredentialSource(Context requestProperties) { try { - Thread.sleep(50); - } catch (InterruptedException e) { - Thread.currentThread().interrupt(); + IdentityResult result = delegate.resolveIdentity(requestProperties); + if (result == null) { + return RefreshAttempt.resultFailure( + IdentityResult.ofError(getClass(), "Credential source returned no result")); + } + if (result.identity() != null) { + return RefreshAttempt.success(result.identity()); + } + return RefreshAttempt.resultFailure(result); + } catch (NonRecoverableIdentityException error) { + return RefreshAttempt.nonRecoverable(error); + } catch (RuntimeException error) { + return RefreshAttempt.exceptionFailure(error); } + } - CachedValue latest = cached; - if (latest != current && latest != null && !isExpired(latest)) { - return latest.result; + private CachedValue createCachedValue(I identity, Instant obtainedAt, Context requestProperties) { + Context refreshProperties = Context.unmodifiableCopy(requestProperties); + Instant expiration = identity.expirationTime(); + if (expiration == null) { + return new CachedValue<>(identity, null, null, refreshProperties); } - // Still expired — return error. - return IdentityResult.ofError(getClass(), "Credentials are expired and refresh failed"); + Duration lifetime = Duration.between(obtainedAt, expiration); + Duration advisoryWindow = configuredAdvisoryWindow == null + ? defaultAdvisoryRefreshWindow(lifetime) + : configuredAdvisoryWindow; + Instant advisoryAt = expiration.minus(advisoryWindow); + Duration effectiveMandatoryWindow = mandatoryRefreshWindow.compareTo(advisoryWindow) > 0 + ? advisoryWindow + : mandatoryRefreshWindow; + return new CachedValue<>( + identity, + advisoryAt, + expiration.minus(effectiveMandatoryWindow), + refreshProperties); } - private IdentityResult doRefresh(Context requestProperties) { - CachedValue current = cached; + private boolean refreshNeeded(CachedValue value, Instant now) { + boolean advisoryRefreshNeeded = reached(now, value.advisoryRefreshAt) + && (nextRefreshAfterSuccessAt == null || reached(now, nextRefreshAfterSuccessAt)); + return refreshRequired + || advisoryRefreshNeeded + || reached(now, value.mandatoryRefreshAt) + || isExpired(value, now); + } + + private boolean mandatoryRefreshNeeded(CachedValue value, Instant now) { + return refreshRequired || reached(now, value.mandatoryRefreshAt) || isExpired(value, now); + } + + private boolean refreshRateLimited(Instant now) { + return nextRefreshAllowedAt != null && now.isBefore(nextRefreshAllowedAt); + } + + private RefreshOutcome fallbackOutcome( + CachedValue value, + RefreshOutcome failure, + Instant now + ) { + if (allowExpiredCredentials || !isExpired(value, now)) { + return RefreshOutcome.result(value.result); + } + if (failure != null) { + return failure; + } + return RefreshOutcome.result(IdentityResult.ofError( + getClass(), + "Credentials are expired and credential refresh is rate limited")); + } - // Stale delay: don't hammer the source (only in static stability mode). - if (allowExpiredCredentials && current != null - && current.nextRefreshAfter != null - && clock.instant().isBefore(current.nextRefreshAfter)) { - return current.result; + private IdentityResult resolveOutcome(RefreshOutcome outcome) { + if (outcome.exception != null) { + throw outcome.exception; } + return outcome.result; + } - IdentityResult result; + private IdentityResult awaitRefresh(CompletableFuture> refresh) { try { - result = delegate.resolveIdentity(requestProperties); - } catch (RuntimeException e) { - LOGGER.log(System.Logger.Level.WARNING, "Credential refresh failed", e); - if (current != null && allowExpiredCredentials) { - current.nextRefreshAfter = clock.instant().plus(jitteredStaleDelay()); - return current.result; + return resolveOutcome(refresh.get()); + } catch (InterruptedException error) { + Thread.currentThread().interrupt(); + return IdentityResult.ofError(getClass(), "Interrupted waiting for credential refresh"); + } catch (ExecutionException error) { + Throwable cause = error.getCause(); + if (cause instanceof RuntimeException runtimeException) { + throw runtimeException; + } + if (cause instanceof Error executionError) { + throw executionError; } - throw e; + throw new IllegalStateException("Unexpected credential refresh failure", cause); } + } - if (result.identity() != null) { - CachedValue newCached = new CachedValue<>(result.identity()); - cached = newCached; - scheduleNextRefresh(newCached, requestProperties); - return newCached.result; + private void scheduleRefreshLocked(Instant refreshAt) { + cancelScheduledRefreshLocked(); + if (!proactiveRefresh || closed || cached == null || refreshAt == null) { + return; } - // Delegate returned an error. - if (current != null && allowExpiredCredentials) { - current.nextRefreshAfter = clock.instant().plus(jitteredStaleDelay()); - return current.result; + long delayMillis = Math.max(0, Duration.between(clock.instant(), refreshAt).toMillis()); + try { + scheduledRefresh = executor.schedule(this::runScheduledRefresh, delayMillis, TimeUnit.MILLISECONDS); + } catch (RejectedExecutionException ignored) { + scheduledRefresh = null; } - - return result; } - private void scheduleNextRefresh(CachedValue value, Context requestProperties) { - cancelScheduledRefresh(); - Instant expiration = value.identity.expirationTime(); - if (expiration == null) { - return; + private void runScheduledRefresh() { + CompletableFuture> refresh = null; + Context refreshProperties = null; + lock.lock(); + try { + scheduledRefresh = null; + if (closed || cached == null || inFlight != null) { + return; + } + + Instant now = clock.instant(); + if (refreshRateLimited(now)) { + scheduleRefreshLocked(nextRefreshAllowedAt); + return; + } + if (!refreshNeeded(cached, now)) { + scheduleRefreshLocked(firstRefreshAt(cached, now)); + return; + } + + refresh = new CompletableFuture<>(); + inFlight = refresh; + refreshProperties = cached.refreshProperties; + } finally { + lock.unlock(); } - Instant refreshAt = expiration.minus(prefetchBuffer); - long delayMillis = Duration.between(clock.instant(), refreshAt).toMillis(); - if (delayMillis <= 0) { - // Already in prefetch window; refresh was just done. - return; + if (refresh != null) { + executeRefresh(refreshProperties, refresh); } + } - scheduledRefresh = executor.schedule(() -> { - if (refreshing.compareAndSet(false, true)) { - try { - doRefresh(requestProperties); - } finally { - refreshing.set(false); - } - } - }, delayMillis, TimeUnit.MILLISECONDS); + private Instant firstRefreshAt(CachedValue value, Instant now) { + if (value.advisoryRefreshAt != null && now.isBefore(value.advisoryRefreshAt)) { + return value.advisoryRefreshAt; + } + if (value.mandatoryRefreshAt != null && now.isBefore(value.mandatoryRefreshAt)) { + return value.mandatoryRefreshAt; + } + return null; } - private void cancelScheduledRefresh() { - ScheduledFuture f = scheduledRefresh; - if (f != null) { - f.cancel(false); + private void cancelScheduledRefreshLocked() { + if (scheduledRefresh != null) { + scheduledRefresh.cancel(false); scheduledRefresh = null; } } - private boolean isInPrefetchWindow(CachedValue value) { + private Duration refreshBackoff() { + long minMillis = refreshBackoffMin.toMillis(); + long maxMillis = refreshBackoffMax.toMillis(); + if (minMillis == maxMillis) { + return refreshBackoffMin; + } + return Duration.ofMillis(ThreadLocalRandom.current().nextLong(minMillis, maxMillis + 1)); + } + + private void logRefreshFailure(String error, Duration backoff, boolean useCachedCredentials) { + LOGGER.log( + System.Logger.Level.WARNING, + refreshFailureMessage(error, backoff, useCachedCredentials)); + } + + static String refreshFailureMessage(String error, Duration backoff, boolean useCachedCredentials) { + String cachedCredentialsMessage = useCachedCredentials + ? "The SDK will continue using cached credentials." + : "The SDK will not use the expired cached credentials."; + return "Credential refresh failed: " + error + ". " + cachedCredentialsMessage + + " A refresh of these credentials will be attempted again after " + + backoff.toSeconds() + " seconds."; + } + + private boolean isFresh(I identity, Instant now) { + Instant expiration = identity.expirationTime(); + return expiration == null || expiration.isAfter(now); + } + + private static boolean isExpired(CachedValue value, Instant now) { Instant expiration = value.identity.expirationTime(); - return expiration != null && clock.instant().isAfter(expiration.minus(prefetchBuffer)); + return expiration != null && !now.isBefore(expiration); } - private boolean isExpired(CachedValue value) { - Instant exp = value.identity.expirationTime(); - return exp != null && clock.instant().isAfter(exp); + private static boolean reached(Instant now, Instant threshold) { + return threshold != null && !now.isBefore(threshold); } - private Duration jitteredStaleDelay() { - long baseMillis = staleRefreshDelay.toMillis(); - long jitter = (long) (Math.random() * baseMillis); - return Duration.ofMillis(baseMillis + jitter); + static Duration defaultAdvisoryRefreshWindow(Duration lifetime) { + if (lifetime.compareTo(Duration.ofMinutes(20)) <= 0) { + return Duration.ofMinutes(5); + } + if (lifetime.compareTo(Duration.ofMinutes(90)) < 0) { + return Duration.ofMinutes(15); + } + return Duration.ofMinutes(60); } private static final class CachedValue { final I identity; final IdentityResult result; - volatile Instant nextRefreshAfter; - - CachedValue(I identity) { + final Instant advisoryRefreshAt; + final Instant mandatoryRefreshAt; + final Context refreshProperties; + + CachedValue( + I identity, + Instant advisoryRefreshAt, + Instant mandatoryRefreshAt, + Context refreshProperties + ) { this.identity = identity; this.result = IdentityResult.of(identity); + this.advisoryRefreshAt = advisoryRefreshAt; + this.mandatoryRefreshAt = mandatoryRefreshAt; + this.refreshProperties = refreshProperties; } } - /** - * Builder for {@link CachingIdentityResolver}. - */ + private static final class RefreshOutcome { + final IdentityResult result; + final RuntimeException exception; + + private RefreshOutcome(IdentityResult result, RuntimeException exception) { + this.result = result; + this.exception = exception; + } + + static RefreshOutcome result(IdentityResult result) { + return new RefreshOutcome<>(Objects.requireNonNull(result), null); + } + + static RefreshOutcome exception(RuntimeException exception) { + return new RefreshOutcome<>(null, Objects.requireNonNull(exception)); + } + } + + private static final class RefreshAttempt { + final I identity; + final IdentityResult failureResult; + final RuntimeException failureException; + final boolean nonRecoverable; + + private RefreshAttempt( + I identity, + IdentityResult failureResult, + RuntimeException failureException, + boolean nonRecoverable + ) { + this.identity = identity; + this.failureResult = failureResult; + this.failureException = failureException; + this.nonRecoverable = nonRecoverable; + } + + static RefreshAttempt success(I identity) { + return new RefreshAttempt<>(Objects.requireNonNull(identity), null, null, false); + } + + static RefreshAttempt resultFailure(IdentityResult result) { + return new RefreshAttempt<>(null, Objects.requireNonNull(result), null, false); + } + + static RefreshAttempt exceptionFailure(RuntimeException exception) { + return new RefreshAttempt<>(null, null, Objects.requireNonNull(exception), false); + } + + static RefreshAttempt nonRecoverable(NonRecoverableIdentityException exception) { + return new RefreshAttempt<>(null, null, exception, true); + } + + RefreshOutcome failureOutcome(CachingIdentityResolver owner) { + if (identity != null) { + return RefreshOutcome.result(IdentityResult.ofError( + owner.getClass(), + "Credential source returned credentials that are already expired")); + } + if (failureException != null) { + return RefreshOutcome.exception(failureException); + } + return RefreshOutcome.result(failureResult); + } + + String failureDescription() { + if (identity != null) { + return "credential source returned credentials that are already expired"; + } + if (failureException != null) { + return failureException.toString(); + } + return failureResult.error(); + } + } + + /** Builder for {@link CachingIdentityResolver}. */ public static final class Builder { private final IdentityResolver delegate; - private Duration prefetchBuffer = Duration.ofMinutes(5); - private boolean allowExpiredCredentials = false; - private Duration staleRefreshDelay = Duration.ofMinutes(5); + private Duration advisoryRefreshWindow; + private Duration mandatoryRefreshWindow = DEFAULT_MANDATORY_WINDOW; + private boolean allowExpiredCredentials; + private Duration refreshBackoffMin = DEFAULT_BACKOFF_MIN; + private Duration refreshBackoffMax = DEFAULT_BACKOFF_MAX; private Clock clock = Clock.systemUTC(); private ScheduledExecutorService executor; + private boolean closeDelegate; + private boolean proactiveRefresh = true; + private BiPredicate identityMatcher = Objects::equals; private Builder(IdentityResolver delegate) { - this.delegate = delegate; + this.delegate = Objects.requireNonNull(delegate, "delegate"); } /** - * How far before expiration to trigger a background refresh. Default: 5 minutes. + * Sets an explicit advisory refresh window. By default the window is derived from credential lifetime. */ - public Builder prefetchBuffer(Duration prefetchBuffer) { - this.prefetchBuffer = Objects.requireNonNull(prefetchBuffer); + public Builder prefetchBuffer(Duration advisoryRefreshWindow) { + this.advisoryRefreshWindow = requirePositive(advisoryRefreshWindow, "advisoryRefreshWindow"); return this; } - /** - * When {@code true}, expired credentials are returned instead of failing. Enables - * AWS Static Stability behavior. Default: {@code false}. - */ + /** Sets the mandatory blocking refresh window. Default: 1 minute. */ + public Builder mandatoryRefreshWindow(Duration mandatoryRefreshWindow) { + this.mandatoryRefreshWindow = requirePositive(mandatoryRefreshWindow, "mandatoryRefreshWindow"); + return this; + } + + /** Enables returning cached expired identities after refresh failure. */ public Builder allowExpiredCredentials(boolean allowExpiredCredentials) { this.allowExpiredCredentials = allowExpiredCredentials; return this; } /** - * Base delay before retrying refresh when credentials are expired and refresh failed. - * Actual delay is jittered up to 2x this value. Default: 5 minutes. + * Sets the base stale refresh delay. The actual delay is uniformly selected from this value to twice this + * value. The default is 5-10 minutes. */ public Builder staleRefreshDelay(Duration staleRefreshDelay) { - this.staleRefreshDelay = Objects.requireNonNull(staleRefreshDelay); + this.refreshBackoffMin = requirePositive(staleRefreshDelay, "staleRefreshDelay"); + this.refreshBackoffMax = staleRefreshDelay.multipliedBy(2); return this; } - /** - * Clock for time comparisons. Default: {@link Clock#systemUTC()}. - */ + /** Sets the clock used for refresh decisions. */ public Builder clock(Clock clock) { - this.clock = Objects.requireNonNull(clock); + this.clock = Objects.requireNonNull(clock, "clock"); + return this; + } + + /** Sets the executor used for proactive and advisory refresh work. */ + public Builder executor(ScheduledExecutorService executor) { + this.executor = Objects.requireNonNull(executor, "executor"); return this; } /** - * Executor for background refresh tasks. If not set, a single daemon thread is created - * and owned by this resolver (shut down on {@link CachingIdentityResolver#close()}). + * Sets whether closing the cache also closes an {@link AutoCloseable} delegate. Disabled by default. + * + *

Enable this only when the cache owns the delegate rather than wrapping a caller-supplied resolver. */ - public Builder executor(ScheduledExecutorService executor) { - this.executor = Objects.requireNonNull(executor); + public Builder closeDelegate(boolean closeDelegate) { + this.closeDelegate = closeDelegate; return this; } + /** Enables or disables proactive scheduled refresh. Enabled by default. */ + public Builder proactiveRefresh(boolean proactiveRefresh) { + this.proactiveRefresh = proactiveRefresh; + return this; + } + + /** Sets how an invalidated identity is matched against the currently cached identity. */ + public Builder identityMatcher(BiPredicate identityMatcher) { + this.identityMatcher = Objects.requireNonNull(identityMatcher, "identityMatcher"); + return this; + } + + Builder refreshBackoff(Duration minimum, Duration maximum) { + this.refreshBackoffMin = requirePositive(minimum, "refreshBackoffMin"); + this.refreshBackoffMax = requirePositive(maximum, "refreshBackoffMax"); + return this; + } + + /** Builds the caching resolver. */ public CachingIdentityResolver build() { + if (refreshBackoffMin.compareTo(refreshBackoffMax) > 0) { + throw new IllegalArgumentException("refreshBackoffMin must not exceed refreshBackoffMax"); + } return new CachingIdentityResolver<>(this); } + + private static Duration requirePositive(Duration duration, String name) { + Objects.requireNonNull(duration, name); + if (duration.isZero() || duration.isNegative()) { + throw new IllegalArgumentException(name + " must be positive"); + } + return duration; + } } } diff --git a/auth-api/src/main/java/software/amazon/smithy/java/auth/api/identity/IdentityResolver.java b/auth-api/src/main/java/software/amazon/smithy/java/auth/api/identity/IdentityResolver.java index 11e22ae63e..b9d381219a 100644 --- a/auth-api/src/main/java/software/amazon/smithy/java/auth/api/identity/IdentityResolver.java +++ b/auth-api/src/main/java/software/amazon/smithy/java/auth/api/identity/IdentityResolver.java @@ -31,16 +31,17 @@ public interface IdentityResolver { Class identityType(); /** - * Invalidate any cached identity, forcing the next call to {@link #resolveIdentity(Context)} to fetch fresh - * credentials from the underlying source. + * Signals that a rejected identity should be invalidated if it is still current. * *

This is typically called by an interceptor when a service returns an expired- or invalid-credential error - * (e.g., {@code ExpiredTokenException}), indicating that the currently cached identity is no longer valid. + * (e.g., {@code ExpiredTokenException}), indicating that the identity used for the request is no longer valid. * - *

The default implementation is a no-op. Caching resolvers (such as {@link CachingIdentityResolver}) override - * this to clear their cache. + *

The default implementation is a no-op. Refreshable caching resolvers such as + * {@link CachingIdentityResolver} can require a refresh while retaining the cached identity as a fallback. + * + * @param rejectedIdentity identity used to sign the rejected request. */ - default void invalidate() {} + default void invalidate(IdentityT rejectedIdentity) {} /** * Combines multiple identity resolvers with the same identity type into a single resolver. diff --git a/auth-api/src/main/java/software/amazon/smithy/java/auth/api/identity/IdentityResolverChain.java b/auth-api/src/main/java/software/amazon/smithy/java/auth/api/identity/IdentityResolverChain.java index 2d737d9d05..4991caa099 100644 --- a/auth-api/src/main/java/software/amazon/smithy/java/auth/api/identity/IdentityResolverChain.java +++ b/auth-api/src/main/java/software/amazon/smithy/java/auth/api/identity/IdentityResolverChain.java @@ -39,4 +39,11 @@ public IdentityResult resolveIdentity(Context requestProperties) { } return IdentityResult.ofError(IdentityResolverChain.class, "Attempted resolvers: " + errors); } + + @Override + public void invalidate(IdentityT rejectedIdentity) { + for (var resolver : resolvers) { + resolver.invalidate(rejectedIdentity); + } + } } diff --git a/auth-api/src/main/java/software/amazon/smithy/java/auth/api/identity/NonRecoverableIdentityException.java b/auth-api/src/main/java/software/amazon/smithy/java/auth/api/identity/NonRecoverableIdentityException.java new file mode 100644 index 0000000000..8ee915e52e --- /dev/null +++ b/auth-api/src/main/java/software/amazon/smithy/java/auth/api/identity/NonRecoverableIdentityException.java @@ -0,0 +1,23 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package software.amazon.smithy.java.auth.api.identity; + +/** + * Signals that refreshing an identity is not expected to succeed without external action. + * + *

Caching resolvers propagate this exception from on-demand resolutions and do not apply refresh backoff. + * Background refreshes retain cached identities and retry on the next on-demand resolution. + */ +public class NonRecoverableIdentityException extends RuntimeException { + + public NonRecoverableIdentityException(String message) { + super(message); + } + + public NonRecoverableIdentityException(String message, Throwable cause) { + super(message, cause); + } +} diff --git a/auth-api/src/test/java/software/amazon/smithy/java/auth/api/identity/CachingIdentityResolverTest.java b/auth-api/src/test/java/software/amazon/smithy/java/auth/api/identity/CachingIdentityResolverTest.java index a36516f3aa..0db3d756f1 100644 --- a/auth-api/src/test/java/software/amazon/smithy/java/auth/api/identity/CachingIdentityResolverTest.java +++ b/auth-api/src/test/java/software/amazon/smithy/java/auth/api/identity/CachingIdentityResolverTest.java @@ -6,28 +6,42 @@ package software.amazon.smithy.java.auth.api.identity; import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertNotNull; import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTimeoutPreemptively; import static org.junit.jupiter.api.Assertions.assertTrue; import java.time.Clock; import java.time.Duration; import java.time.Instant; import java.time.ZoneId; +import java.time.ZoneOffset; +import java.util.ArrayDeque; +import java.util.ArrayList; +import java.util.Deque; +import java.util.List; import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; +import java.util.concurrent.Future; import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicInteger; import java.util.concurrent.atomic.AtomicReference; +import java.util.function.BooleanSupplier; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.Test; import software.amazon.smithy.java.context.Context; class CachingIdentityResolverTest { + private static final Instant BASE = Instant.parse("2026-01-01T00:00:00Z"); + private final ScheduledExecutorService executor = Executors.newSingleThreadScheduledExecutor(); - private CachingIdentityResolver resolver; + private CachingIdentityResolver resolver; @AfterEach void tearDown() { @@ -39,65 +53,461 @@ void tearDown() { @Test void coldStartBlocksAndCachesResult() { - var identity = new TestIdentity("cached", Instant.now().plusSeconds(3600)); - var delegate = new CountingResolver(identity); - resolver = CachingIdentityResolver.builder(delegate) - .executor(executor) - .build(); + var delegate = new QueueResolver(identity("cached", BASE.plusSeconds(3600))); + resolver = resolver(delegate, new MutableClock(BASE), true); - IdentityResult result = resolve(); - assertNotNull(result.identity()); - assertEquals("cached", result.identity().value); - assertEquals(1, delegate.callCount.get()); + assertEquals("cached", resolve().identity().value()); + assertEquals("cached", resolve().identity().value()); + assertEquals(1, delegate.calls.get()); + } - // Second call returns cached — no delegate invocation. - IdentityResult result2 = resolve(); - assertEquals("cached", result2.identity().value); - assertEquals(1, delegate.callCount.get()); + @Test + void failedColdStartReturnsSourceErrorAndRetriesOnNextCall() { + var delegate = new QueueResolver( + error("source unavailable"), + identity("fresh", BASE.plusSeconds(3600))); + resolver = resolver(delegate, new MutableClock(BASE), true); + + IdentityResult failed = resolve(); + assertNull(failed.identity()); + assertEquals("source unavailable", failed.error()); + + assertEquals("fresh", resolve().identity().value()); + assertEquals(2, delegate.calls.get()); } @Test - void backgroundRefreshHappensBeforeExpiry() throws InterruptedException { - Instant now = Instant.now(); - // Credentials expire in 2 seconds, prefetch buffer is 1 second → refresh at T+1s. - var delegate = new CountingResolver(expiringIdentity(now.plusSeconds(2))); + void nonRecoverableColdStartRaisesSourceErrorAndRetriesOnNextCall() { + var delegate = new QueueResolver( + new NonRecoverableIdentityException("reauthenticate"), + identity("fresh", BASE.plusSeconds(3600))); + resolver = resolver(delegate, new MutableClock(BASE), true); + + var error = assertThrows(NonRecoverableIdentityException.class, this::resolve); + assertEquals("reauthenticate", error.getMessage()); + + assertEquals("fresh", resolve().identity().value()); + assertEquals(2, delegate.calls.get()); + } + + @Test + void concurrentColdStartUsesOneSourceCall() throws Exception { + var delegate = new BlockingResolver(identity("initial", BASE.plusSeconds(3600)), true); + resolver = resolver(delegate, new MutableClock(BASE), true); + ExecutorService callers = Executors.newFixedThreadPool(8); + try { + List>> results = new ArrayList<>(); + for (int i = 0; i < 8; i++) { + results.add(callers.submit(this::resolve)); + } + + assertTrue(delegate.started.await(5, TimeUnit.SECONDS)); + assertEquals(1, delegate.calls.get()); + delegate.release.countDown(); + + for (var result : results) { + assertEquals("initial", result.get(5, TimeUnit.SECONDS).identity().value()); + } + assertEquals(1, delegate.calls.get()); + } finally { + callers.shutdownNow(); + } + } + + @Test + void interruptedColdStartWaiterReturnsAnError() throws Exception { + var delegate = new BlockingResolver(identity("initial", BASE.plusSeconds(3600)), true); + resolver = resolver(delegate, new MutableClock(BASE), true); + ExecutorService sourceCaller = Executors.newSingleThreadExecutor(); + var waitingResult = new AtomicReference>(); + var interruptRestored = new AtomicBoolean(); + + try { + Future> sourceResult = sourceCaller.submit(this::resolve); + assertTrue(delegate.started.await(5, TimeUnit.SECONDS)); + + Thread waiter = new Thread(() -> { + waitingResult.set(resolve()); + interruptRestored.set(Thread.currentThread().isInterrupted()); + }); + waiter.start(); + await(() -> waiter.getState() == Thread.State.WAITING); + + waiter.interrupt(); + waiter.join(TimeUnit.SECONDS.toMillis(5)); + + assertFalse(waiter.isAlive()); + assertNull(waitingResult.get().identity()); + assertTrue(waitingResult.get().error().contains("Interrupted waiting for credential refresh")); + assertTrue(interruptRestored.get()); + + delegate.release.countDown(); + assertEquals("initial", sourceResult.get(5, TimeUnit.SECONDS).identity().value()); + } finally { + delegate.release.countDown(); + sourceCaller.shutdownNow(); + } + } + + @Test + void unexpectedRefreshFailureDoesNotStrandFutureCallers() { + var delegate = new QueueResolver( + identity("failed", BASE.plusSeconds(3600)), + identity("fresh", BASE.plusSeconds(3600))); + resolver = resolver(delegate, new FailOnceClock(BASE, 2), true); + + assertThrows(IllegalStateException.class, this::resolve); + + assertTimeoutPreemptively( + Duration.ofSeconds(1), + () -> assertEquals("fresh", resolve().identity().value())); + assertEquals(2, delegate.calls.get()); + } + + @Test + void advisoryRefreshDoesNotBlockCallers() throws Exception { + var clock = new MutableClock(BASE); + var initial = identity("initial", BASE.plus(Duration.ofHours(2))); + var refreshed = identity("refreshed", BASE.plus(Duration.ofHours(3))); + var delegate = new BlockingRefreshResolver(initial, refreshed); + resolver = resolver(delegate, clock, true); + + assertEquals("initial", resolve().identity().value()); + clock.advance(Duration.ofMinutes(61)); + + assertEquals("initial", resolve().identity().value()); + assertTrue(delegate.refreshStarted.await(5, TimeUnit.SECONDS)); + for (int i = 0; i < 10; i++) { + assertEquals("initial", resolve().identity().value()); + } + assertEquals(2, delegate.calls.get()); + + delegate.releaseRefresh.countDown(); + await(() -> "refreshed".equals(resolve().identity().value())); + assertEquals(2, delegate.calls.get()); + } + + @Test + void advisoryRefreshFailureUsesCachedCredentialsAndBacksOff() throws Exception { + var clock = new MutableClock(BASE); + var delegate = new QueueResolver( + identity("cached", BASE.plus(Duration.ofHours(2))), + error("source unavailable"), + identity("fresh", BASE.plus(Duration.ofHours(3)))); + resolver = resolver(delegate, clock, true); + + assertEquals("cached", resolve().identity().value()); + clock.advance(Duration.ofMinutes(61)); + assertEquals("cached", resolve().identity().value()); + await(() -> delegate.calls.get() == 2); + assertEquals(2, delegate.calls.get()); + + clock.advance(Duration.ofSeconds(299)); + assertEquals("cached", resolve().identity().value()); + assertEquals(2, delegate.calls.get()); + + clock.advance(Duration.ofSeconds(1)); + await(() -> "fresh".equals(resolve().identity().value())); + assertEquals(3, delegate.calls.get()); + } + + @Test + void invalidationDuringInFlightAdvisoryRefreshIsIgnored() throws Exception { + var clock = new MutableClock(BASE); + var initial = identity("initial", BASE.plus(Duration.ofHours(2))); + var refreshed = identity("refreshed", BASE.plus(Duration.ofHours(3))); + var delegate = new BlockingRefreshResolver(initial, refreshed); + resolver = resolver(delegate, clock, true); + + assertEquals("initial", resolve().identity().value()); + clock.advance(Duration.ofMinutes(61)); + + assertEquals("initial", resolve().identity().value()); + assertTrue(delegate.refreshStarted.await(5, TimeUnit.SECONDS)); + + resolver.invalidate(initial); + assertEquals("initial", resolve().identity().value()); + assertEquals(2, delegate.calls.get()); + + delegate.releaseRefresh.countDown(); + await(() -> "refreshed".equals(resolve().identity().value())); + assertEquals(2, delegate.calls.get()); + } + + @Test + void ignoredInvalidationDoesNotForceAnotherRefreshWhenSourceReturnsSameIdentity() throws Exception { + var clock = new MutableClock(BASE); + var initial = identity("same", BASE.plus(Duration.ofHours(2))); + var refreshed = identity("same", BASE.plus(Duration.ofHours(3))); + var delegate = new BlockingRefreshResolver(initial, refreshed); resolver = CachingIdentityResolver.builder(delegate) .executor(executor) - .prefetchBuffer(Duration.ofSeconds(1)) + .clock(clock) + .allowExpiredCredentials(true) + .proactiveRefresh(false) + .identityMatcher((cached, rejected) -> cached.value().equals(rejected.value())) .build(); - // Cold start. + assertEquals("same", resolve().identity().value()); + clock.advance(Duration.ofMinutes(61)); + assertEquals("same", resolve().identity().value()); + assertTrue(delegate.refreshStarted.await(5, TimeUnit.SECONDS)); + + resolver.invalidate(initial); + delegate.releaseRefresh.countDown(); + + await(() -> refreshed.expirationTime().equals(resolve().identity().expirationTime())); + assertEquals(2, delegate.calls.get()); + } + + @Test + void mandatoryRefreshBlocksConcurrentCallersOnOneSourceCall() throws Exception { + var clock = new MutableClock(BASE); + var initial = identity("initial", BASE.plusSeconds(3600)); + var refreshed = identity("refreshed", BASE.plusSeconds(7200)); + var delegate = new BlockingRefreshResolver(initial, refreshed); + resolver = resolver(delegate, clock, true); + assertNotNull(resolve().identity()); + clock.advance(Duration.ofSeconds(3590)); + + ExecutorService callers = Executors.newFixedThreadPool(6); + try { + List>> results = new ArrayList<>(); + for (int i = 0; i < 6; i++) { + results.add(callers.submit(this::resolve)); + } + assertTrue(delegate.refreshStarted.await(5, TimeUnit.SECONDS)); + assertEquals(2, delegate.calls.get()); + for (var result : results) { + assertTrue(!result.isDone()); + } + + delegate.releaseRefresh.countDown(); + for (var result : results) { + assertEquals("refreshed", result.get(5, TimeUnit.SECONDS).identity().value()); + } + assertEquals(2, delegate.calls.get()); + } finally { + callers.shutdownNow(); + } + } + + @Test + void failedRefreshUsesCachedCredentialsAndBacksOff() { + var clock = new MutableClock(BASE); + var delegate = new QueueResolver( + identity("cached", BASE.plusSeconds(60)), + error("source unavailable"), + identity("fresh", BASE.plusSeconds(7200))); + resolver = resolver(delegate, clock, true); + + assertEquals("cached", resolve().identity().value()); + clock.advance(Duration.ofSeconds(61)); + assertEquals("cached", resolve().identity().value()); + assertEquals(2, delegate.calls.get()); + + clock.advance(Duration.ofSeconds(299)); + assertEquals("cached", resolve().identity().value()); + assertEquals(2, delegate.calls.get()); + + clock.advance(Duration.ofSeconds(1)); + assertEquals("fresh", resolve().identity().value()); + assertEquals(3, delegate.calls.get()); + } + + @Test + void invalidationPreservesCachedCredentialsAndBackoff() { + var clock = new MutableClock(BASE); + var cached = identity("cached", BASE.plusSeconds(60)); + var delegate = new QueueResolver( + cached, + error("source unavailable"), + identity("fresh", BASE.plusSeconds(7200))); + resolver = resolver(delegate, clock, true); + + resolve(); + clock.advance(Duration.ofSeconds(61)); resolve(); - assertEquals(1, delegate.callCount.get()); + resolver.invalidate(cached); + + assertEquals("cached", resolve().identity().value()); + assertEquals(2, delegate.calls.get()); + clock.advance(Duration.ofMinutes(5)); + assertEquals("fresh", resolve().identity().value()); + assertEquals(3, delegate.calls.get()); + } - // Wait for background refresh to fire (should happen ~1s from now). - Thread.sleep(1500); - assertTrue(delegate.callCount.get() >= 2, "Expected background refresh, got " + delegate.callCount.get()); + @Test + void matchingInvalidationForcesMandatoryRefresh() { + var clock = new MutableClock(BASE); + var initial = identity("initial", BASE.plusSeconds(3600)); + var delegate = new QueueResolver(initial, identity("fresh", BASE.plusSeconds(7200))); + resolver = resolver(delegate, clock, true); + + assertEquals("initial", resolve().identity().value()); + resolver.invalidate(new TestIdentity("other", initial.expirationTime())); + assertEquals("initial", resolve().identity().value()); + assertEquals(1, delegate.calls.get()); + + resolver.invalidate(initial); + assertEquals("fresh", resolve().identity().value()); + assertEquals(2, delegate.calls.get()); } @Test - void nonExpiringIdentityIsCachedIndefinitely() throws InterruptedException { - var delegate = new CountingResolver(new TestIdentity("permanent", null)); + void concurrentInvalidationDoesNotBlock() throws Exception { + var initial = identity("initial", BASE.plusSeconds(3600)); + var delegate = new QueueResolver(initial); + var matcherStarted = new CountDownLatch(1); + var releaseMatcher = new CountDownLatch(1); + var blockFirstMatcher = new AtomicBoolean(true); resolver = CachingIdentityResolver.builder(delegate) .executor(executor) + .clock(new MutableClock(BASE)) + .identityMatcher((cached, rejected) -> { + if (blockFirstMatcher.compareAndSet(true, false)) { + matcherStarted.countDown(); + awaitLatch(releaseMatcher); + } + return cached.value().equals(rejected.value()); + }) .build(); + assertEquals("initial", resolve().identity().value()); + + ExecutorService callers = Executors.newFixedThreadPool(2); + Future firstInvalidation = callers.submit(() -> resolver.invalidate(initial)); + try { + assertTrue(matcherStarted.await(5, TimeUnit.SECONDS)); + + Future concurrentInvalidation = callers.submit(() -> resolver.invalidate(initial)); + concurrentInvalidation.get(1, TimeUnit.SECONDS); + assertFalse(firstInvalidation.isDone()); + } finally { + releaseMatcher.countDown(); + firstInvalidation.get(5, TimeUnit.SECONDS); + callers.shutdownNow(); + } + } + + @Test + void staleSourceCredentialsAreRefreshFailures() { + var clock = new MutableClock(BASE); + var delegate = new QueueResolver( + identity("cached", BASE.plusSeconds(60)), + identity("already-expired", BASE.minusSeconds(1))); + resolver = resolver(delegate, clock, true); resolve(); - Thread.sleep(200); + clock.advance(Duration.ofSeconds(61)); + + assertEquals("cached", resolve().identity().value()); + assertEquals(2, delegate.calls.get()); + assertEquals("cached", resolve().identity().value()); + assertEquals(2, delegate.calls.get()); + } + + @Test + void staleCredentialsOnColdStartReturnAnError() { + var clock = new MutableClock(BASE); + var delegate = new QueueResolver(identity("expired", BASE)); + resolver = resolver(delegate, clock, true); + + IdentityResult result = resolve(); + + assertNull(result.identity()); + assertTrue(result.error().contains("already expired")); + } + + @Test + void strictModeDoesNotReturnExpiredCredentialsDuringBackoff() { + var clock = new MutableClock(BASE); + var delegate = new QueueResolver( + identity("cached", BASE.plusSeconds(60)), + error("source unavailable")); + resolver = resolver(delegate, clock, false); + + resolve(); + clock.advance(Duration.ofSeconds(61)); + IdentityResult failed = resolve(); + assertNull(failed.identity()); + assertEquals("source unavailable", failed.error()); + + IdentityResult rateLimited = resolve(); + assertNull(rateLimited.identity()); + assertEquals("source unavailable", rateLimited.error()); + assertEquals(2, delegate.calls.get()); + } + + @Test + void strictModeRefreshFailureMessageDoesNotClaimExpiredCredentialsAreUsed() { + assertEquals( + "Credential refresh failed: source unavailable. " + + "The SDK will not use the expired cached credentials. " + + "A refresh of these credentials will be attempted again after 300 seconds.", + CachingIdentityResolver.refreshFailureMessage( + "source unavailable", + Duration.ofMinutes(5), + false)); + } + + @Test + void advisoryNonRecoverableFailureIsRaisedOnNextResolutionWithoutBackoff() throws Exception { + var clock = new MutableClock(BASE); + var delegate = new QueueResolver( + identity("cached", BASE.plus(Duration.ofHours(2))), + new NonRecoverableIdentityException("reauthenticate"), + new NonRecoverableIdentityException("reauthenticate"), + identity("fresh", BASE.plus(Duration.ofHours(3)))); + resolver = resolver(delegate, clock, true); + resolve(); - assertEquals(1, delegate.callCount.get()); + clock.advance(Duration.ofMinutes(61)); + assertEquals("cached", resolve().identity().value()); + await(() -> delegate.calls.get() == 2); + + var error = assertThrows(NonRecoverableIdentityException.class, this::resolve); + assertEquals("reauthenticate", error.getMessage()); + assertEquals(3, delegate.calls.get()); + + assertEquals("fresh", resolve().identity().value()); + assertEquals(4, delegate.calls.get()); } @Test - void allowExpiredCredentialsReturnsStaleOnFailure() { - AtomicInteger calls = new AtomicInteger(0); + void nonRecoverableFailureIsNotBackedOff() { + var clock = new MutableClock(BASE); + var delegate = new QueueResolver( + identity("cached", BASE.plusSeconds(60)), + new NonRecoverableIdentityException("reauthenticate"), + identity("fresh", BASE.plusSeconds(7200))); + resolver = resolver(delegate, clock, true); + + resolve(); + clock.advance(Duration.ofSeconds(61)); + var error = assertThrows(NonRecoverableIdentityException.class, this::resolve); + assertEquals("reauthenticate", error.getMessage()); + + assertEquals("fresh", resolve().identity().value()); + assertEquals(3, delegate.calls.get()); + } + + @Test + void proactiveNonRecoverableFailureIsRetriedOnDemandWithoutBackoff() throws Exception { + var calls = new AtomicInteger(); + var proactiveAttempted = new CountDownLatch(1); + Instant now = Instant.now(); IdentityResolver delegate = new IdentityResolver<>() { @Override - public IdentityResult resolveIdentity(Context ctx) { - if (calls.incrementAndGet() == 1) { - return IdentityResult.of(expiringIdentity(Instant.now().minusSeconds(10))); - } - throw new RuntimeException("refresh failed"); + public IdentityResult resolveIdentity(Context requestProperties) { + return switch (calls.incrementAndGet()) { + case 1 -> IdentityResult.of(identity("initial", now.plusMillis(300))); + case 2 -> { + proactiveAttempted.countDown(); + throw new NonRecoverableIdentityException("reauthenticate"); + } + case 3 -> throw new NonRecoverableIdentityException("reauthenticate"); + default -> IdentityResult.of(identity("fresh", now.plusSeconds(3600))); + }; } @Override @@ -105,37 +515,118 @@ public Class identityType() { return TestIdentity.class; } }; + resolver = CachingIdentityResolver.builder(delegate) + .executor(executor) + .clock(Clock.systemUTC()) + .allowExpiredCredentials(true) + .prefetchBuffer(Duration.ofMillis(200)) + .mandatoryRefreshWindow(Duration.ofMillis(50)) + .build(); + + assertEquals("initial", resolve().identity().value()); + assertTrue(proactiveAttempted.await(5, TimeUnit.SECONDS)); + + var observed = new AtomicReference(); + await(() -> { + if (observed.get() != null) { + return true; + } + try { + resolve(); + return false; + } catch (NonRecoverableIdentityException error) { + observed.set(error); + return true; + } + }); + assertEquals("reauthenticate", observed.get().getMessage()); + assertEquals(3, calls.get()); + + assertEquals("fresh", resolve().identity().value()); + assertEquals(4, calls.get()); + } + + @Test + void computesAdvisoryWindowFromCredentialLifetime() { + assertEquals( + Duration.ofMinutes(5), + CachingIdentityResolver.defaultAdvisoryRefreshWindow(Duration.ofMinutes(20))); + assertEquals( + Duration.ofMinutes(15), + CachingIdentityResolver.defaultAdvisoryRefreshWindow(Duration.ofMinutes(21))); + assertEquals( + Duration.ofMinutes(15), + CachingIdentityResolver.defaultAdvisoryRefreshWindow(Duration.ofMinutes(89))); + assertEquals( + Duration.ofMinutes(60), + CachingIdentityResolver.defaultAdvisoryRefreshWindow(Duration.ofMinutes(90))); + } - // Use a fixed clock that's past expiration so the cache is immediately stale. - Clock pastClock = Clock.fixed(Instant.now().plusSeconds(60), ZoneId.of("UTC")); + @Test + void configuredAdvisoryWindowOverridesLifetimeDefault() throws Exception { + var clock = new MutableClock(BASE); + var delegate = new QueueResolver( + identity("initial", BASE.plus(Duration.ofHours(1))), + identity("fresh", BASE.plus(Duration.ofHours(2)))); resolver = CachingIdentityResolver.builder(delegate) .executor(executor) + .clock(clock) .allowExpiredCredentials(true) - .clock(pastClock) - .prefetchBuffer(Duration.ofSeconds(1)) + .prefetchBuffer(Duration.ofMinutes(10)) + .proactiveRefresh(false) .build(); - // Cold start succeeds (returns expired identity). - IdentityResult result = resolve(); - assertNotNull(result.identity()); + assertEquals("initial", resolve().identity().value()); + clock.advance(Duration.ofMinutes(49)); + assertEquals("initial", resolve().identity().value()); + assertEquals(1, delegate.calls.get()); + + clock.advance(Duration.ofMinutes(1)); + assertEquals("initial", resolve().identity().value()); + await(() -> "fresh".equals(resolve().identity().value())); + assertEquals(2, delegate.calls.get()); + } + + @Test + void waitsForMandatoryWindowWhenAdvisoryWindowEqualsLifetime() { + var clock = new MutableClock(BASE); + var delegate = new QueueResolver( + identity("initial", BASE.plusSeconds(60)), + identity("fresh", BASE.plusSeconds(3600))); + resolver = CachingIdentityResolver.builder(delegate) + .executor(executor) + .clock(clock) + .allowExpiredCredentials(true) + .prefetchBuffer(Duration.ofSeconds(60)) + .mandatoryRefreshWindow(Duration.ofSeconds(30)) + .proactiveRefresh(false) + .build(); + + assertEquals("initial", resolve().identity().value()); + assertEquals("initial", resolve().identity().value()); + assertEquals(1, delegate.calls.get()); - // Trigger refresh — it fails, but we still get the stale value. - IdentityResult result2 = resolve(); - assertNotNull(result2.identity()); + clock.advance(Duration.ofSeconds(29)); + assertEquals("initial", resolve().identity().value()); + assertEquals(1, delegate.calls.get()); + + clock.advance(Duration.ofSeconds(1)); + assertEquals("fresh", resolve().identity().value()); + assertEquals(2, delegate.calls.get()); } @Test - void strictModeReturnsErrorWhenExpiredAndRefreshFails() { - Instant expiration = Instant.now().plusMillis(50); // Expires very soon. - var identity = new TestIdentity("expiring", expiration); - AtomicInteger calls = new AtomicInteger(0); + void proactiveRefreshWaitsForMandatoryWindowWhenAdvisoryWindowIsAlreadyOpen() throws Exception { + var calls = new AtomicInteger(); + var refreshed = new CountDownLatch(1); IdentityResolver delegate = new IdentityResolver<>() { @Override - public IdentityResult resolveIdentity(Context ctx) { + public IdentityResult resolveIdentity(Context requestProperties) { if (calls.incrementAndGet() == 1) { - return IdentityResult.of(identity); + return IdentityResult.of(identity("initial", Instant.now().plusMillis(500))); } - return IdentityResult.ofError(getClass(), "no creds"); + refreshed.countDown(); + return IdentityResult.of(identity("refreshed", Instant.now().plusSeconds(3600))); } @Override @@ -143,60 +634,106 @@ public Class identityType() { return TestIdentity.class; } }; - resolver = CachingIdentityResolver.builder(delegate) .executor(executor) - .allowExpiredCredentials(false) - .prefetchBuffer(Duration.ofMillis(10)) + .clock(Clock.systemUTC()) + .allowExpiredCredentials(true) + .prefetchBuffer(Duration.ofSeconds(1)) + .mandatoryRefreshWindow(Duration.ofMillis(100)) .build(); - // Cold start succeeds. - IdentityResult first = resolve(); - assertNotNull(first.identity()); + assertEquals("initial", resolve().identity().value()); + assertTrue(!refreshed.await(100, TimeUnit.MILLISECONDS)); + assertTrue(refreshed.await(5, TimeUnit.SECONDS)); + assertEquals(2, calls.get()); + assertEquals("refreshed", resolve().identity().value()); + } - // Wait for expiration. - try { - Thread.sleep(100); - } catch (InterruptedException e) { - Thread.currentThread().interrupt(); - } + @Test + void mandatoryWindowIsClampedToComputedAdvisoryWindow() { + var clock = new MutableClock(BASE); + var delegate = new QueueResolver( + identity("initial", BASE.plus(Duration.ofMinutes(20))), + identity("fresh", BASE.plus(Duration.ofHours(1)))); + resolver = CachingIdentityResolver.builder(delegate) + .executor(executor) + .clock(clock) + .allowExpiredCredentials(true) + .mandatoryRefreshWindow(Duration.ofMinutes(10)) + .proactiveRefresh(false) + .build(); - // Now expired + strict → blocks for retry → delegate returns error. - IdentityResult result = resolve(); - assertNull(result.identity()); - assertNotNull(result.error()); + assertEquals("initial", resolve().identity().value()); + clock.advance(Duration.ofMinutes(10)); + assertEquals("initial", resolve().identity().value()); + assertEquals(1, delegate.calls.get()); + + clock.advance(Duration.ofMinutes(5)); + assertEquals("fresh", resolve().identity().value()); + assertEquals(2, delegate.calls.get()); } @Test - void invalidateForcesNextCallToRefresh() { - var delegate = new CountingResolver(expiringIdentity(Instant.now().plusSeconds(3600))); + void mandatoryWindowIsClampedToConfiguredAdvisoryWindow() { + var clock = new MutableClock(BASE); + var delegate = new QueueResolver( + identity("initial", BASE.plus(Duration.ofMinutes(10))), + identity("fresh", BASE.plus(Duration.ofHours(1)))); resolver = CachingIdentityResolver.builder(delegate) .executor(executor) + .clock(clock) + .allowExpiredCredentials(true) + .prefetchBuffer(Duration.ofSeconds(30)) + .proactiveRefresh(false) .build(); - resolve(); - assertEquals(1, delegate.callCount.get()); + assertEquals("initial", resolve().identity().value()); + clock.advance(Duration.ofMinutes(9).plusSeconds(29)); + assertEquals("initial", resolve().identity().value()); + assertEquals(1, delegate.calls.get()); - resolver.invalidate(); - resolve(); - assertEquals(2, delegate.callCount.get()); + clock.advance(Duration.ofSeconds(1)); + assertEquals("fresh", resolve().identity().value()); + assertEquals(2, delegate.calls.get()); } @Test - void concurrentColdStartOnlyCallsDelegateOnce() throws InterruptedException { - CountDownLatch startGate = new CountDownLatch(1); - var delegate = new IdentityResolver() { - final AtomicInteger callCount = new AtomicInteger(0); + void proactiveRefreshRetriesAfterBackoff() throws Exception { + Instant now = Instant.now(); + var delegate = new QueueResolver( + identity("initial", now.plusMillis(300)), + error("source unavailable"), + identity("fresh", now.plusSeconds(3600))); + resolver = CachingIdentityResolver.builder(delegate) + .executor(executor) + .clock(Clock.systemUTC()) + .allowExpiredCredentials(true) + .prefetchBuffer(Duration.ofMillis(200)) + .mandatoryRefreshWindow(Duration.ofMillis(50)) + .refreshBackoff(Duration.ofMillis(50), Duration.ofMillis(50)) + .build(); + + assertEquals("initial", resolve().identity().value()); + await(() -> delegate.calls.get() >= 3); + assertEquals("fresh", resolve().identity().value()); + } + @Test + void proactiveRefreshPreservesResolutionContext() throws Exception { + var property = Context.key("refresh property"); + var refreshedProperty = new AtomicReference(); + var refreshed = new CountDownLatch(1); + var calls = new AtomicInteger(); + Instant now = Instant.now(); + IdentityResolver delegate = new IdentityResolver<>() { @Override - public IdentityResult resolveIdentity(Context ctx) { - callCount.incrementAndGet(); - try { - Thread.sleep(100); // Simulate slow first fetch. - } catch (InterruptedException e) { - Thread.currentThread().interrupt(); + public IdentityResult resolveIdentity(Context requestProperties) { + if (calls.incrementAndGet() == 1) { + return IdentityResult.of(identity("initial", now.plusMillis(300))); } - return IdentityResult.of(new TestIdentity("shared", Instant.now().plusSeconds(3600))); + refreshedProperty.set(requestProperties.get(property)); + refreshed.countDown(); + return IdentityResult.of(identity("refreshed", now.plusSeconds(3600))); } @Override @@ -204,58 +741,217 @@ public Class identityType() { return TestIdentity.class; } }; + resolver = CachingIdentityResolver.builder(delegate) + .executor(executor) + .clock(Clock.systemUTC()) + .prefetchBuffer(Duration.ofMillis(200)) + .mandatoryRefreshWindow(Duration.ofMillis(50)) + .build(); + var properties = Context.create(); + properties.put(property, "expected"); + + assertEquals("initial", resolver.resolveIdentity(properties).identity().value()); + assertTrue(refreshed.await(5, TimeUnit.SECONDS)); + assertEquals("expected", refreshedProperty.get()); + } + @Test + void nonExpiringIdentityIsCachedUntilInvalidated() { + var clock = new MutableClock(BASE); + var first = identity("first", null); + var delegate = new QueueResolver(first, identity("second", null)); + resolver = resolver(delegate, clock, true); + + assertEquals("first", resolve().identity().value()); + clock.advance(Duration.ofDays(30)); + assertEquals("first", resolve().identity().value()); + resolver.invalidate(first); + assertEquals("second", resolve().identity().value()); + } + + @Test + void doesNotCloseCallerSuppliedDelegateByDefault() { + var delegate = new CloseableResolver(); resolver = CachingIdentityResolver.builder(delegate) .executor(executor) .build(); - int threadCount = 10; - CountDownLatch done = new CountDownLatch(threadCount); - AtomicReference firstValue = new AtomicReference<>(); + resolver.close(); - for (int i = 0; i < threadCount; i++) { - new Thread(() -> { - try { - startGate.await(); - } catch (InterruptedException e) { - return; - } - @SuppressWarnings("unchecked") - var r = (CachingIdentityResolver) resolver; - IdentityResult result = r.resolveIdentity(Context.empty()); - firstValue.compareAndSet(null, result.identity().value); - done.countDown(); - }).start(); - } + assertFalse(delegate.closed.get()); + } + + @Test + void closesOwnedDelegateWhenConfigured() { + var delegate = new CloseableResolver(); + resolver = CachingIdentityResolver.builder(delegate) + .executor(executor) + .closeDelegate(true) + .build(); - startGate.countDown(); - assertTrue(done.await(5, TimeUnit.SECONDS)); - assertEquals(1, delegate.callCount.get()); - assertEquals("shared", firstValue.get()); + resolver.close(); + + assertTrue(delegate.closed.get()); + } + + private CachingIdentityResolver resolver( + IdentityResolver delegate, + Clock clock, + boolean allowExpired + ) { + return CachingIdentityResolver.builder(delegate) + .executor(executor) + .clock(clock) + .allowExpiredCredentials(allowExpired) + .proactiveRefresh(false) + .refreshBackoff(Duration.ofMinutes(5), Duration.ofMinutes(5)) + .build(); } - @SuppressWarnings("unchecked") private IdentityResult resolve() { - return ((CachingIdentityResolver) resolver).resolveIdentity(Context.empty()); + return resolver.resolveIdentity(Context.empty()); + } + + private static TestIdentity identity(String value, Instant expiration) { + return new TestIdentity(value, expiration); + } + + private static IdentityResult error(String message) { + return IdentityResult.ofError(CachingIdentityResolverTest.class, message); + } + + private static void await(BooleanSupplier condition) throws Exception { + long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(5); + while (!condition.getAsBoolean() && System.nanoTime() < deadline) { + Thread.sleep(10); + } + assertTrue(condition.getAsBoolean(), "condition was not satisfied before timeout"); + } + + private record TestIdentity(String value, Instant expirationTime) implements Identity {} + + private static final class MutableClock extends Clock { + private final AtomicReference now; + + MutableClock(Instant now) { + this.now = new AtomicReference<>(now); + } + + void advance(Duration duration) { + now.updateAndGet(current -> current.plus(duration)); + } + + @Override + public ZoneId getZone() { + return ZoneOffset.UTC; + } + + @Override + public Clock withZone(ZoneId zone) { + return this; + } + + @Override + public Instant instant() { + return now.get(); + } } - private static TestIdentity expiringIdentity(Instant expiration) { - return new TestIdentity("id-" + System.nanoTime(), expiration); + private static final class FailOnceClock extends Clock { + private final Instant now; + private final int failureCall; + private final AtomicInteger calls = new AtomicInteger(); + + FailOnceClock(Instant now, int failureCall) { + this.now = now; + this.failureCall = failureCall; + } + + @Override + public ZoneId getZone() { + return ZoneOffset.UTC; + } + + @Override + public Clock withZone(ZoneId zone) { + return this; + } + + @Override + public Instant instant() { + if (calls.incrementAndGet() == failureCall) { + throw new IllegalStateException("clock failure"); + } + return now; + } } - record TestIdentity(String value, Instant expirationTime) implements Identity {} + private static final class QueueResolver implements IdentityResolver { + private final Deque responses = new ArrayDeque<>(); + private final AtomicInteger calls = new AtomicInteger(); + + QueueResolver(Object... responses) { + this.responses.addAll(List.of(responses)); + } + + @Override + @SuppressWarnings("unchecked") + public synchronized IdentityResult resolveIdentity(Context requestProperties) { + calls.incrementAndGet(); + Object response = responses.removeFirst(); + if (response instanceof RuntimeException error) { + throw error; + } + if (response instanceof IdentityResult result) { + return (IdentityResult) result; + } + return IdentityResult.of((TestIdentity) response); + } + + @Override + public Class identityType() { + return TestIdentity.class; + } + } - static class CountingResolver implements IdentityResolver { - final AtomicInteger callCount = new AtomicInteger(0); + private static final class CloseableResolver implements IdentityResolver, AutoCloseable { + private final AtomicBoolean closed = new AtomicBoolean(); + + @Override + public IdentityResult resolveIdentity(Context requestProperties) { + return IdentityResult.of(identity("identity", null)); + } + + @Override + public Class identityType() { + return TestIdentity.class; + } + + @Override + public void close() { + closed.set(true); + } + } + + private static final class BlockingResolver implements IdentityResolver { private final TestIdentity identity; + private final boolean block; + private final AtomicInteger calls = new AtomicInteger(); + private final CountDownLatch started = new CountDownLatch(1); + private final CountDownLatch release = new CountDownLatch(1); - CountingResolver(TestIdentity identity) { + BlockingResolver(TestIdentity identity, boolean block) { this.identity = identity; + this.block = block; } @Override - public IdentityResult resolveIdentity(Context ctx) { - callCount.incrementAndGet(); + public IdentityResult resolveIdentity(Context requestProperties) { + calls.incrementAndGet(); + started.countDown(); + if (block) { + awaitLatch(release); + } return IdentityResult.of(identity); } @@ -264,4 +960,41 @@ public Class identityType() { return TestIdentity.class; } } + + private static final class BlockingRefreshResolver implements IdentityResolver { + private final TestIdentity initial; + private final TestIdentity refreshed; + private final AtomicInteger calls = new AtomicInteger(); + private final CountDownLatch refreshStarted = new CountDownLatch(1); + private final CountDownLatch releaseRefresh = new CountDownLatch(1); + + BlockingRefreshResolver(TestIdentity initial, TestIdentity refreshed) { + this.initial = initial; + this.refreshed = refreshed; + } + + @Override + public IdentityResult resolveIdentity(Context requestProperties) { + if (calls.incrementAndGet() == 1) { + return IdentityResult.of(initial); + } + refreshStarted.countDown(); + awaitLatch(releaseRefresh); + return IdentityResult.of(refreshed); + } + + @Override + public Class identityType() { + return TestIdentity.class; + } + } + + private static void awaitLatch(CountDownLatch latch) { + try { + latch.await(); + } catch (InterruptedException error) { + Thread.currentThread().interrupt(); + throw new RuntimeException(error); + } + } } diff --git a/aws/aws-credential-chain/src/main/java/software/amazon/smithy/java/aws/credentials/chain/AwsCredentialCaching.java b/aws/aws-credential-chain/src/main/java/software/amazon/smithy/java/aws/credentials/chain/AwsCredentialCaching.java new file mode 100644 index 0000000000..3d9f11e2c0 --- /dev/null +++ b/aws/aws-credential-chain/src/main/java/software/amazon/smithy/java/aws/credentials/chain/AwsCredentialCaching.java @@ -0,0 +1,55 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package software.amazon.smithy.java.aws.credentials.chain; + +import java.util.Objects; +import java.util.concurrent.ScheduledExecutorService; +import software.amazon.smithy.java.auth.api.identity.CachingIdentityResolver; +import software.amazon.smithy.java.auth.api.identity.IdentityResolver; +import software.amazon.smithy.java.aws.auth.api.identity.AwsCredentialsIdentity; +import software.amazon.smithy.utils.SmithyInternalApi; + +/** Applies the AWS credential refresh policy to credential-source resolvers. */ +@SmithyInternalApi +public final class AwsCredentialCaching { + + private AwsCredentialCaching() {} + + /** Wraps an AWS-managed credential source with caching and static stability. */ + public static IdentityResolver staticallyStable( + IdentityResolver delegate, + ScheduledExecutorService executor + ) { + return builder(delegate, executor) + .allowExpiredCredentials(true) + .build(); + } + + /** Wraps an opaque credential source with caching but without expired-credential fallback. */ + public static IdentityResolver cachingOnly( + IdentityResolver delegate, + ScheduledExecutorService executor + ) { + return builder(delegate, executor) + .allowExpiredCredentials(false) + .build(); + } + + private static CachingIdentityResolver.Builder builder( + IdentityResolver delegate, + ScheduledExecutorService executor + ) { + var builder = CachingIdentityResolver.builder(Objects.requireNonNull(delegate, "delegate")) + .closeDelegate(true) + .identityMatcher((cached, rejected) -> Objects.equals( + cached.accessKeyId(), + rejected.accessKeyId())); + if (executor != null) { + builder.executor(executor); + } + return builder; + } +} diff --git a/aws/aws-credential-chain/src/main/java/software/amazon/smithy/java/aws/credentials/chain/ChainSetup.java b/aws/aws-credential-chain/src/main/java/software/amazon/smithy/java/aws/credentials/chain/ChainSetup.java index 029b68bf05..713a3315cd 100644 --- a/aws/aws-credential-chain/src/main/java/software/amazon/smithy/java/aws/credentials/chain/ChainSetup.java +++ b/aws/aws-credential-chain/src/main/java/software/amazon/smithy/java/aws/credentials/chain/ChainSetup.java @@ -6,7 +6,9 @@ package software.amazon.smithy.java.aws.credentials.chain; import java.util.ArrayList; +import java.util.EnumSet; import java.util.List; +import java.util.Objects; import java.util.Set; import java.util.concurrent.ScheduledExecutorService; import java.util.function.Function; @@ -40,6 +42,7 @@ public final class ChainSetup { private final String regionOverride; private final Context properties; private final List resolvers = new ArrayList<>(); + private final Set detectedSlots = EnumSet.noneOf(StandardProvider.class); private final Function envFn; private AwsProfileFile profileFile; private AwsProfile profile; @@ -110,6 +113,18 @@ public String getenv(String name) { return envFn.apply(name); } + /** + * Returns the environment lookup function configured for this assembly. + * + *

Resolvers that need environment values after assembly can retain this function without retaining the + * mutable {@code ChainSetup}. + * + * @return the configured environment variable lookup function. + */ + public Function environment() { + return envFn; + } + /** * Returns the environment variable lookup function this setup was built with, so the chain can perform slot * detection ({@link StandardProvider#isDetected(Function)}) against the same environment it was assembled @@ -172,15 +187,47 @@ public void setProfile(AwsProfile profile) { this.profile = profile; } + /** + * Records that configuration for a standard provider slot was detected during assembly. + * + *

This is used by providers such as {@code SHARED_CONFIG}, which detects profile-based + * sources on behalf of optional provider modules. + * + * @param slot detected standard provider slot. + */ + public void markDetected(StandardProvider slot) { + detectedSlots.add(Objects.requireNonNull(slot, "slot")); + } + + Set detectedSlots() { + return detectedSlots; + } + /** * Registers a resolver at the current provider's position. Assembly continues after * this call. May be called multiple times to register multiple resolvers that stack * at this position. * + *

If the resolver is {@link AutoCloseable}, ownership transfers to the assembled chain. It is closed when the + * chain is closed or if assembly fails after registration. + * * @param resolver the identity resolver to register. */ public void addResolver(IdentityResolver resolver) { - resolvers.add(new NamedResolver(currentProvider.name(), currentProvider.featureIds(), resolver)); + addResolver(resolver, currentProvider.featureIds()); + } + + /** + * Registers a resolver with source-specific feature IDs at the current provider's position. + * + *

If the resolver is {@link AutoCloseable}, ownership transfers to the assembled chain. It is closed when the + * chain is closed or if assembly fails after registration. + * + * @param resolver resolver to register. + * @param featureIds feature IDs emitted when this resolver succeeds. + */ + public void addResolver(IdentityResolver resolver, Set featureIds) { + resolvers.add(new NamedResolver(currentProvider.name(), featureIds, resolver)); } /** @@ -189,10 +236,26 @@ public void addResolver(IdentityResolver resolver) { * environment variables contain a complete set of credentials, or a profile explicitly * configures assume-role). * + *

If the resolver is {@link AutoCloseable}, ownership transfers to the assembled chain. It is closed when the + * chain is closed or if assembly fails after registration. + * * @param resolver the identity resolver to register. */ public void addTerminalResolver(IdentityResolver resolver) { - resolvers.add(new NamedResolver(currentProvider.name(), currentProvider.featureIds(), resolver)); + addTerminalResolver(resolver, currentProvider.featureIds()); + } + + /** + * Registers a terminal resolver with source-specific feature IDs. + * + *

If the resolver is {@link AutoCloseable}, ownership transfers to the assembled chain. It is closed when the + * chain is closed or if assembly fails after registration. + * + * @param resolver resolver to register. + * @param featureIds feature IDs emitted when this resolver succeeds. + */ + public void addTerminalResolver(IdentityResolver resolver, Set featureIds) { + resolvers.add(new NamedResolver(currentProvider.name(), featureIds, resolver)); this.terminal = true; } @@ -232,7 +295,13 @@ boolean isTerminal() { * @param featureIds the feature IDs to emit on successful resolution. * @param resolver the identity resolver. */ - public record NamedResolver(String name, Set featureIds, IdentityResolver resolver) {} + public record NamedResolver(String name, Set featureIds, IdentityResolver resolver) { + public NamedResolver { + Objects.requireNonNull(name, "name"); + featureIds = Set.copyOf(featureIds); + Objects.requireNonNull(resolver, "resolver"); + } + } /** * Builder for {@link ChainSetup}. diff --git a/aws/aws-credential-chain/src/main/java/software/amazon/smithy/java/aws/credentials/chain/IdentityChain.java b/aws/aws-credential-chain/src/main/java/software/amazon/smithy/java/aws/credentials/chain/IdentityChain.java index 825efcda4e..1040ba4ebe 100644 --- a/aws/aws-credential-chain/src/main/java/software/amazon/smithy/java/aws/credentials/chain/IdentityChain.java +++ b/aws/aws-credential-chain/src/main/java/software/amazon/smithy/java/aws/credentials/chain/IdentityChain.java @@ -8,6 +8,7 @@ import java.util.ArrayList; import java.util.Collections; import java.util.HashSet; +import java.util.IdentityHashMap; import java.util.List; import java.util.ServiceLoader; import java.util.Set; @@ -58,6 +59,7 @@ public final class IdentityChain implements IdentityResolver private final Class identityType; private final List resolvers; private final Set claimedSlots; + private final Set detectedSlots; private final Function envFn; private final ScheduledExecutorService executor; @@ -65,12 +67,14 @@ private IdentityChain( Class identityType, List resolvers, Set claimedSlots, + Set detectedSlots, Function envFn, ScheduledExecutorService executor ) { this.identityType = identityType; this.resolvers = resolvers; this.claimedSlots = claimedSlots; + this.detectedSlots = detectedSlots; this.envFn = envFn; this.executor = executor; } @@ -116,7 +120,8 @@ private static ScheduledExecutorService defaultExecutor() { * Create an identity chain by discovering providers via ServiceLoader. * * @param identityType Identity type to resolve. - * @param ex Executor used for background resolution. + * @param ex Executor used for background resolution. Ownership transfers to the returned chain. If assembly + * fails, the executor is shut down before the failure is propagated. * @return the assembled chain. * @throws IllegalStateException if two providers claim the same standard slot. */ @@ -137,7 +142,8 @@ public static IdentityChain create(Class identityType * how a client's configured region flows into credential resolution. * * @param identityType Identity type to resolve. - * @param ex Executor used for background resolution. + * @param ex Executor used for background resolution. Ownership transfers to the returned chain. If assembly + * fails, the executor is shut down before the failure is propagated. * @param profileFile Already-parsed profile file to use, or {@code null} to load from the default locations. * @param regionOverride Region for service-calling providers to use, or {@code null} to resolve it normally. * @return the assembled chain. @@ -149,15 +155,22 @@ public static IdentityChain create( AwsProfileFile profileFile, String regionOverride ) { - List registrations = new ArrayList<>(); - for (ChainIdentityProvider r : ServiceLoader.load(ChainIdentityProvider.class)) { - registrations.add(r); + List registrations; + ChainSetup setup; + try { + registrations = new ArrayList<>(); + for (ChainIdentityProvider r : ServiceLoader.load(ChainIdentityProvider.class)) { + registrations.add(r); + } + setup = ChainSetup.builder() + .executor(ex) + .profileFile(profileFile) + .regionOverride(regionOverride) + .build(); + } catch (RuntimeException | Error failure) { + shutdownExecutor(ex, failure); + throw failure; } - ChainSetup setup = ChainSetup.builder() - .executor(ex) - .profileFile(profileFile) - .regionOverride(regionOverride) - .build(); return assemble(identityType, registrations, ex, setup); } @@ -179,112 +192,125 @@ static IdentityChain assemble( ScheduledExecutorService executor, ChainSetup setup ) { - // Check for duplicate names. - Set seenNames = new HashSet<>(); - for (ChainIdentityProvider r : registrations) { - if (!seenNames.add(r.name())) { - throw new IllegalStateException("Duplicate credential provider registration name: '" + r.name() + "'"); + try { + // Check for duplicate names. + Set seenNames = new HashSet<>(); + for (ChainIdentityProvider r : registrations) { + if (!seenNames.add(r.name())) { + throw new IllegalStateException( + "Duplicate credential provider registration name: '" + r.name() + "'"); + } } - } - // Sort providers by ordering constraint (enum order for Standard, relative for Before/After). - List sorted = sortByOrdering(registrations); + // Sort providers by ordering constraint (enum order for Standard, relative for Before/After). + List sorted = sortByOrdering(registrations); - // Call setup() on each provider in sorted order. - for (ChainIdentityProvider provider : sorted) { - setup.setCurrentProvider(provider); - provider.setup(identityType, setup); - if (setup.isTerminal()) { - break; + // Call setup() on each provider in sorted order. + for (ChainIdentityProvider provider : sorted) { + setup.setCurrentProvider(provider); + provider.setup(identityType, setup); + if (setup.isTerminal()) { + break; + } } - } - var ordered = setup.resolvers(); + var ordered = setup.resolvers(); + validateResolverTypes(identityType, ordered); - if (LOGGER.isDebugEnabled()) { - LOGGER.debug("Assembled identity chain: {}", - ordered.stream().map(ChainSetup.NamedResolver::name).collect(Collectors.joining(", "))); - } + if (LOGGER.isDebugEnabled()) { + LOGGER.debug("Assembled identity chain: {}", + ordered.stream().map(ChainSetup.NamedResolver::name).collect(Collectors.joining(", "))); + } - // Warn about detected-but-unclaimed slots. - Set claimed = new HashSet<>(); - for (var nr : ordered) { - for (ChainIdentityProvider p : sorted) { - if (p.name().equals(nr.name()) - && p.ordering() instanceof OrderingConstraint.Standard(StandardProvider s)) { - claimed.add(s); + // A discovered registration claims its module's slot even when it does not produce a + // resolver for this identity type or configuration. + Set claimed = new HashSet<>(); + for (ChainIdentityProvider provider : sorted) { + if (provider.ordering() instanceof OrderingConstraint.Standard(StandardProvider slot)) { + claimed.add(slot); } } + Set detected = Set.copyOf(setup.detectedSlots()); + warnDetectedButUnclaimed(claimed, detected, setup.envFn()); + return new IdentityChain<>(identityType, + Collections.unmodifiableList(ordered), + Collections.unmodifiableSet(claimed), + detected, + setup.envFn(), + executor); + } catch (RuntimeException | Error failure) { + closeResolvers(setup.resolvers(), failure); + shutdownExecutor(executor, failure); + throw failure; } - warnDetectedButUnclaimed(claimed, setup.envFn()); - return new IdentityChain<>(identityType, - Collections.unmodifiableList(ordered), - Collections.unmodifiableSet(claimed), - setup.envFn(), - executor); } - private static List sortByOrdering(List providers) { - // Separate into standard-slot providers and relative providers. - List standards = new ArrayList<>(); - List befores = new ArrayList<>(); - List afters = new ArrayList<>(); - Set seenSlots = new HashSet<>(); + private static void validateResolverTypes( + Class identityType, + List resolvers + ) { + for (var namedResolver : resolvers) { + Class resolverType = namedResolver.resolver().identityType(); + if (!identityType.isAssignableFrom(resolverType)) { + throw new IllegalStateException( + "Credential provider '" + namedResolver.name() + "' registered a resolver for " + + resolverType.getName() + " when the chain requires " + identityType.getName()); + } + } + } + private static List sortByOrdering(List providers) { + ChainIdentityProvider[] standards = new ChainIdentityProvider[StandardProvider.values().length]; for (ChainIdentityProvider p : providers) { - switch (p.ordering()) { - case OrderingConstraint.Standard(StandardProvider slot) -> { - if (!seenSlots.add(slot)) { - throw new IllegalStateException("Two providers claim the same standard slot '" - + slot + "': check provider '" + p.name() + "'"); - } - standards.add(p); + if (p.ordering() instanceof OrderingConstraint.Standard(StandardProvider slot)) { + if (standards[slot.ordinal()] != null) { + throw new IllegalStateException("Two providers claim the same standard slot '" + + slot + "': check provider '" + p.name() + "'"); } - case OrderingConstraint.Before b -> befores.add(p); - case OrderingConstraint.After a -> afters.add(p); + standards[slot.ordinal()] = p; } } - // Sort standards by enum ordinal. - standards.sort((a, b) -> { - var slotA = ((OrderingConstraint.Standard) a.ordering()).slot(); - var slotB = ((OrderingConstraint.Standard) b.ordering()).slot(); - return slotA.compareTo(slotB); - }); - - // Build final list: insert Before/After relative to their referenced slot's position. - List result = new ArrayList<>(standards); - for (ChainIdentityProvider p : befores) { - var slot = ((OrderingConstraint.Before) p.ordering()).slot(); - int idx = indexOfSlot(result, slot); - result.add(idx, p); - } - for (ChainIdentityProvider p : afters) { - var slot = ((OrderingConstraint.After) p.ordering()).slot(); - int idx = indexOfSlot(result, slot); - int insertAt = Math.min(idx + 1, result.size()); - result.add(insertAt, p); + List result = new ArrayList<>(providers.size()); + for (StandardProvider slot : StandardProvider.values()) { + addRelativeProviders(result, providers, slot, true); + ChainIdentityProvider standard = standards[slot.ordinal()]; + if (standard != null) { + result.add(standard); + } + addRelativeProviders(result, providers, slot, false); } return result; } - private static int indexOfSlot(List list, StandardProvider slot) { - for (int i = 0; i < list.size(); i++) { - if (list.get(i).ordering() instanceof OrderingConstraint.Standard(StandardProvider s) && s == slot) { - return i; - } - // If slot not found, find where it would be by enum order. - if (list.get(i).ordering() instanceof OrderingConstraint.Standard(StandardProvider s) - && s.ordinal() > slot.ordinal()) { - return i; + private static void addRelativeProviders( + List result, + List providers, + StandardProvider slot, + boolean before + ) { + for (ChainIdentityProvider provider : providers) { + if (before + && provider.ordering() instanceof OrderingConstraint.Before(StandardProvider target) + && target == slot) { + result.add(provider); + } else if (!before + && provider.ordering() instanceof OrderingConstraint.After(StandardProvider target) + && target == slot) { + result.add(provider); } } - return list.size(); } - private static void warnDetectedButUnclaimed(Set claimed, Function envFn) { + private static void warnDetectedButUnclaimed( + Set claimed, + Set detected, + Function envFn + ) { for (StandardProvider slot : StandardProvider.values()) { - if (slot.moduleSuggestion() != null && !claimed.contains(slot) && slot.isDetected(envFn)) { + if (slot.moduleSuggestion() != null + && !claimed.contains(slot) + && isDetected(slot, detected, envFn)) { LOGGER.warn("{} credentials detected but no provider is registered for the '{}' slot. " + "Add '{}' to your dependencies.", slot.name(), @@ -353,7 +379,7 @@ private void emitDiagnostics(Context requestProperties, ChainResolutionDiagnosti private List detectedButMissingModules() { List suggestions = new ArrayList<>(); for (StandardProvider slot : StandardProvider.values()) { - if (slot.moduleSuggestion() != null && slot.isDetected(envFn) && !isClaimed(slot)) { + if (slot.moduleSuggestion() != null && isDetected(slot) && !isClaimed(slot)) { suggestions.add(slot.moduleSuggestion()); } } @@ -363,7 +389,7 @@ private List detectedButMissingModules() { private String detectedButMissingHints() { StringBuilder hints = new StringBuilder(); for (StandardProvider slot : StandardProvider.values()) { - if (slot.moduleSuggestion() != null && slot.isDetected(envFn) && !isClaimed(slot)) { + if (slot.moduleSuggestion() != null && isDetected(slot) && !isClaimed(slot)) { hints.append(" Detected ") .append(slot.name()) .append(" credentials; add '") @@ -378,6 +404,18 @@ private boolean isClaimed(StandardProvider slot) { return claimedSlots.contains(slot); } + private boolean isDetected(StandardProvider slot) { + return isDetected(slot, detectedSlots, envFn); + } + + private static boolean isDetected( + StandardProvider slot, + Set detectedSlots, + Function envFn + ) { + return detectedSlots.contains(slot) || slot.isDetected(envFn); + } + /** * @return the ordered list of provider names in this chain. */ @@ -395,14 +433,60 @@ public Class identityType() { } @Override - public void invalidate() { + public void invalidate(I rejectedIdentity) { for (var nr : resolvers) { - nr.resolver().invalidate(); + invalidateResolver(nr.resolver(), rejectedIdentity); + } + } + + private static void invalidateResolver( + IdentityResolver resolver, + Identity rejectedIdentity + ) { + if (resolver.identityType().isInstance(rejectedIdentity)) { + resolver.invalidate(resolver.identityType().cast(rejectedIdentity)); } } @Override public void close() { - executor.shutdownNow(); + closeResolvers(resolvers, null); + shutdownExecutor(executor, null); + } + + private static void closeResolvers(List resolvers, Throwable priorFailure) { + Set> closed = Collections.newSetFromMap(new IdentityHashMap<>()); + for (var namedResolver : resolvers) { + IdentityResolver resolver = namedResolver.resolver(); + if (closed.add(resolver) && resolver instanceof AutoCloseable closeable) { + try { + closeable.close(); + } catch (Exception error) { + if (priorFailure == null) { + LOGGER.warn("Failed to close credential resolver {}: {}", namedResolver.name(), error); + } else { + priorFailure.addSuppressed(error); + } + } catch (Error error) { + if (priorFailure == null) { + throw error; + } + priorFailure.addSuppressed(error); + } + } + } + } + + private static void shutdownExecutor(ScheduledExecutorService executor, Throwable priorFailure) { + if (executor != null) { + try { + executor.shutdownNow(); + } catch (RuntimeException | Error error) { + if (priorFailure == null) { + throw error; + } + priorFailure.addSuppressed(error); + } + } } } diff --git a/aws/aws-credential-chain/src/main/java/software/amazon/smithy/java/aws/credentials/chain/StandardProvider.java b/aws/aws-credential-chain/src/main/java/software/amazon/smithy/java/aws/credentials/chain/StandardProvider.java index ffc7ec100f..556f2b7667 100644 --- a/aws/aws-credential-chain/src/main/java/software/amazon/smithy/java/aws/credentials/chain/StandardProvider.java +++ b/aws/aws-credential-chain/src/main/java/software/amazon/smithy/java/aws/credentials/chain/StandardProvider.java @@ -175,8 +175,8 @@ public boolean isDetected(Function env) { /** * Profile-based credential process ({@code credential_process}). * - *

Invokes an external process on each resolution. The command string is captured at - * assembly time from the active profile. + *

Invokes an external process when the credential cache needs to be populated or refreshed. + * The command string is captured at assembly time from the active profile. */ PROFILE_CREDENTIAL_PROCESS(null) { @Override diff --git a/aws/aws-credential-chain/src/main/java/software/amazon/smithy/java/aws/credentials/chain/config/CredentialProcessHandler.java b/aws/aws-credential-chain/src/main/java/software/amazon/smithy/java/aws/credentials/chain/config/CredentialProcessHandler.java index 52b33fa3c0..1cace88af1 100644 --- a/aws/aws-credential-chain/src/main/java/software/amazon/smithy/java/aws/credentials/chain/config/CredentialProcessHandler.java +++ b/aws/aws-credential-chain/src/main/java/software/amazon/smithy/java/aws/credentials/chain/config/CredentialProcessHandler.java @@ -20,6 +20,7 @@ import software.amazon.smithy.java.aws.auth.api.identity.AwsCredentialsIdentity; import software.amazon.smithy.java.aws.config.AwsConfigCredentialSource; import software.amazon.smithy.java.aws.config.AwsProfile; +import software.amazon.smithy.java.aws.credentials.chain.AwsCredentialCaching; import software.amazon.smithy.java.aws.credentials.chain.ChainIdentityProvider; import software.amazon.smithy.java.aws.credentials.chain.ChainSetup; import software.amazon.smithy.java.aws.credentials.chain.CredentialFeatureId; @@ -34,9 +35,9 @@ * Resolves credentials by invoking an external process specified by {@code credential_process} * in the active AWS profile. * - *

The command string is captured at assembly time and does not change after construction. - * Each call to {@code resolveIdentity()} re-executes the process to obtain fresh credentials, - * so expiring credentials returned by the process are naturally refreshed without caching. + *

The command string is captured at assembly time and does not change after construction. Returned credentials + * are cached and refreshed before expiration. Because the process is an opaque customer-provided source, refresh + * failures do not enable expired-credential fallback. * *

This provider registers terminally: once a profile declares {@code credential_process}, the chain commits * to it, and a process failure is returned as an error rather than falling through to a lower-priority provider @@ -80,7 +81,8 @@ public void setup(Class identityType, ChainSetup setup) { } for (AwsConfigCredentialSource source : profile.credentialSources()) { if (source instanceof AwsConfigCredentialSource.CredentialProcess(String commandLine)) { - setup.addTerminalResolver(new Resolver(commandLine)); + setup.addTerminalResolver( + AwsCredentialCaching.cachingOnly(new Resolver(commandLine), setup.executor())); return; } } @@ -167,7 +169,9 @@ private static IdentityResult parseOutput(String json) { try { expiration = Instant.parse(expirationStr); } catch (DateTimeParseException e) { - LOGGER.warn("credential_process returned unparseable Expiration: {}", expirationStr); + return IdentityResult.ofError( + CredentialProcessHandler.class, + "credential_process returned unparseable Expiration: " + expirationStr); } } return IdentityResult.of(AwsCredentialsIdentity.create( diff --git a/aws/aws-credential-chain/src/main/java/software/amazon/smithy/java/aws/credentials/chain/config/SharedConfigProvider.java b/aws/aws-credential-chain/src/main/java/software/amazon/smithy/java/aws/credentials/chain/config/SharedConfigProvider.java index 2bd70e0233..611ea8f856 100644 --- a/aws/aws-credential-chain/src/main/java/software/amazon/smithy/java/aws/credentials/chain/config/SharedConfigProvider.java +++ b/aws/aws-credential-chain/src/main/java/software/amazon/smithy/java/aws/credentials/chain/config/SharedConfigProvider.java @@ -6,6 +6,7 @@ package software.amazon.smithy.java.aws.credentials.chain.config; import software.amazon.smithy.java.auth.api.identity.Identity; +import software.amazon.smithy.java.aws.config.AwsConfigCredentialSource; import software.amazon.smithy.java.aws.config.AwsProfileFile; import software.amazon.smithy.java.aws.credentials.chain.ChainIdentityProvider; import software.amazon.smithy.java.aws.credentials.chain.ChainSetup; @@ -46,6 +47,29 @@ public void setup(Class identityType, ChainSetup setup) { } else { setup.setProfile(profileFile.activeProfile(setup::getenv)); } + markDetectedProfileSources(setup); } } + + private static void markDetectedProfileSources(ChainSetup setup) { + if (setup.profile() == null) { + return; + } + for (AwsConfigCredentialSource source : setup.profile().credentialSources()) { + setup.markDetected(slot(source)); + } + } + + private static StandardProvider slot(AwsConfigCredentialSource source) { + return switch (source) { + case AwsConfigCredentialSource.StaticKeys ignored -> StandardProvider.PROFILE_STATIC_KEYS; + case AwsConfigCredentialSource.SessionKeys ignored -> StandardProvider.PROFILE_SESSION_KEYS; + case AwsConfigCredentialSource.AssumeRole ignored -> StandardProvider.PROFILE_ASSUME_ROLE; + case AwsConfigCredentialSource.WebIdentityToken ignored -> StandardProvider.PROFILE_WEB_IDENTITY; + case AwsConfigCredentialSource.SsoSession ignored -> StandardProvider.PROFILE_SSO_SESSION; + case AwsConfigCredentialSource.LegacySso ignored -> StandardProvider.PROFILE_LEGACY_SSO; + case AwsConfigCredentialSource.CredentialProcess ignored -> StandardProvider.PROFILE_CREDENTIAL_PROCESS; + case AwsConfigCredentialSource.LoginSession ignored -> StandardProvider.PROFILE_LOGIN; + }; + } } diff --git a/aws/aws-credential-chain/src/test/java/software/amazon/smithy/java/aws/credentials/chain/CredentialProcessTerminalTest.java b/aws/aws-credential-chain/src/test/java/software/amazon/smithy/java/aws/credentials/chain/CredentialProcessTerminalTest.java index 7dbc7facbf..9406366e65 100644 --- a/aws/aws-credential-chain/src/test/java/software/amazon/smithy/java/aws/credentials/chain/CredentialProcessTerminalTest.java +++ b/aws/aws-credential-chain/src/test/java/software/amazon/smithy/java/aws/credentials/chain/CredentialProcessTerminalTest.java @@ -5,6 +5,7 @@ package software.amazon.smithy.java.aws.credentials.chain; +import static org.junit.jupiter.api.Assertions.assertInstanceOf; import static org.junit.jupiter.api.Assertions.assertTrue; import java.io.IOException; @@ -12,6 +13,7 @@ import java.nio.file.Path; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.io.TempDir; +import software.amazon.smithy.java.auth.api.identity.CachingIdentityResolver; import software.amazon.smithy.java.aws.auth.api.identity.AwsCredentialsIdentity; import software.amazon.smithy.java.aws.config.AwsProfileFile; import software.amazon.smithy.java.aws.credentials.chain.config.CredentialProcessHandler; @@ -31,6 +33,7 @@ void claimsTerminalWhenProfileDeclaresCredentialProcess(@TempDir Path tmp) throw setup.setCurrentProvider(handler); handler.setup(AwsCredentialsIdentity.class, setup); + assertInstanceOf(CachingIdentityResolver.class, setup.resolvers().getFirst().resolver()); assertTrue(setup.isTerminal(), "A profile declaring credential_process must claim the chain terminally so a process failure " + "stops resolution instead of falling through to a lower-priority provider."); diff --git a/aws/aws-credential-chain/src/test/java/software/amazon/smithy/java/aws/credentials/chain/IdentityChainTest.java b/aws/aws-credential-chain/src/test/java/software/amazon/smithy/java/aws/credentials/chain/IdentityChainTest.java index 6bc5ea0012..1b51cfe14e 100644 --- a/aws/aws-credential-chain/src/test/java/software/amazon/smithy/java/aws/credentials/chain/IdentityChainTest.java +++ b/aws/aws-credential-chain/src/test/java/software/amazon/smithy/java/aws/credentials/chain/IdentityChainTest.java @@ -11,15 +11,57 @@ import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; +import java.io.IOException; +import java.nio.file.Files; +import java.nio.file.Path; import java.util.List; +import java.util.concurrent.Executors; +import java.util.concurrent.atomic.AtomicInteger; import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; import software.amazon.smithy.java.auth.api.identity.Identity; import software.amazon.smithy.java.auth.api.identity.IdentityResolver; import software.amazon.smithy.java.auth.api.identity.IdentityResult; +import software.amazon.smithy.java.auth.api.identity.TokenIdentity; import software.amazon.smithy.java.aws.auth.api.identity.AwsCredentialsIdentity; +import software.amazon.smithy.java.aws.config.AwsProfileFile; +import software.amazon.smithy.java.aws.credentials.chain.config.SharedConfigProvider; import software.amazon.smithy.java.context.Context; class IdentityChainTest { + @Test + void awsCachingMatchesInvalidationByAccessKeyId() { + var calls = new AtomicInteger(); + IdentityResolver delegate = new IdentityResolver<>() { + @Override + public IdentityResult resolveIdentity(Context requestProperties) { + int call = calls.incrementAndGet(); + return IdentityResult.of(AwsCredentialsIdentity.create( + call == 1 ? "AKID" : "REFRESHED", + "secret-" + call, + null, + null)); + } + + @Override + public Class identityType() { + return AwsCredentialsIdentity.class; + } + }; + var executor = Executors.newSingleThreadScheduledExecutor(); + var resolver = AwsCredentialCaching.staticallyStable(delegate, executor); + try { + assertEquals("AKID", resolver.resolveIdentity(Context.empty()).identity().accessKeyId()); + + resolver.invalidate(AwsCredentialsIdentity.create("AKID", "different-secret", null, null)); + + assertEquals("REFRESHED", resolver.resolveIdentity(Context.empty()).identity().accessKeyId()); + assertEquals(2, calls.get()); + } finally { + executor.shutdownNow(); + } + } + @Test void standardProvidersAreOrderedByEnumOrder() { var chain = IdentityChain.assemble(AwsCredentialsIdentity.class, @@ -38,6 +80,204 @@ void standardProvidersAreOrderedByEnumOrder() { assertEquals(List.of("env", "profile", "imds"), chain.providerNames()); } + @Test + void rejectsResolverWithIncompatibleIdentityType() { + IdentityResolver tokenResolver = new IdentityResolver<>() { + @Override + public IdentityResult resolveIdentity(Context requestProperties) { + return IdentityResult.of(TokenIdentity.create("token")); + } + + @Override + public Class identityType() { + return TokenIdentity.class; + } + }; + + var error = assertThrows(IllegalStateException.class, + () -> IdentityChain.assemble( + AwsCredentialsIdentity.class, + List.of(registration( + "wrong-type", + new OrderingConstraint.Standard(StandardProvider.ENVIRONMENT), + tokenResolver)), + null)); + + assertTrue(error.getMessage().contains("wrong-type")); + assertTrue(error.getMessage().contains(TokenIdentity.class.getName())); + assertTrue(error.getMessage().contains(AwsCredentialsIdentity.class.getName())); + } + + @Test + void failedTypeValidationClosesResolverAndExecutor() { + var closes = new AtomicInteger(); + var resolver = new CloseableResolver<>( + TokenIdentity.class, + TokenIdentity.create("token"), + closes); + var executor = Executors.newSingleThreadScheduledExecutor(); + try { + assertThrows( + IllegalStateException.class, + () -> IdentityChain.assemble( + AwsCredentialsIdentity.class, + List.of(registration( + "wrong-type", + new OrderingConstraint.Standard(StandardProvider.ENVIRONMENT), + resolver)), + executor)); + + assertEquals(1, closes.get()); + assertTrue(executor.isShutdown()); + } finally { + executor.shutdownNow(); + } + } + + @Test + void providerSetupFailureClosesPreviouslyRegisteredResolversAndExecutor() { + var closes = new AtomicInteger(); + var resolver = new CloseableResolver<>( + AwsCredentialsIdentity.class, + AwsCredentialsIdentity.create("AK", "SK"), + closes); + ChainIdentityProvider failingProvider = new ChainIdentityProvider() { + @Override + public String name() { + return "failing"; + } + + @Override + public OrderingConstraint ordering() { + return new OrderingConstraint.Standard(StandardProvider.SHARED_CONFIG); + } + + @Override + public void setup(Class identityType, ChainSetup setup) { + throw new IllegalStateException("setup failed"); + } + }; + var executor = Executors.newSingleThreadScheduledExecutor(); + try { + var error = assertThrows( + IllegalStateException.class, + () -> IdentityChain.assemble( + AwsCredentialsIdentity.class, + List.of( + registration( + "first", + new OrderingConstraint.Standard(StandardProvider.ENVIRONMENT), + resolver), + failingProvider), + executor)); + + assertEquals("setup failed", error.getMessage()); + assertEquals(1, closes.get()); + assertTrue(executor.isShutdown()); + } finally { + executor.shutdownNow(); + } + } + + @Test + void closesResolverRegisteredMoreThanOnceOnlyOnce() { + var closes = new AtomicInteger(); + var resolver = new CloseableResolver<>( + AwsCredentialsIdentity.class, + AwsCredentialsIdentity.create("AK", "SK"), + closes); + ChainIdentityProvider provider = new ChainIdentityProvider() { + @Override + public String name() { + return "duplicate-resolver"; + } + + @Override + public OrderingConstraint ordering() { + return new OrderingConstraint.Standard(StandardProvider.ENVIRONMENT); + } + + @Override + public void setup(Class identityType, ChainSetup setup) { + setup.addResolver(resolver); + setup.addResolver(resolver); + } + }; + var chain = IdentityChain.assemble( + AwsCredentialsIdentity.class, + List.of(provider), + null); + + chain.close(); + + assertEquals(1, closes.get()); + } + + @Test + void awsCredentialCacheClosesItsOwnedDelegate() throws Exception { + var closes = new AtomicInteger(); + var delegate = new CloseableResolver<>( + AwsCredentialsIdentity.class, + AwsCredentialsIdentity.create("AK", "SK"), + closes); + var executor = Executors.newSingleThreadScheduledExecutor(); + var cache = AwsCredentialCaching.staticallyStable(delegate, executor); + try { + ((AutoCloseable) cache).close(); + + assertEquals(1, closes.get()); + } finally { + executor.shutdownNow(); + } + } + + @Test + void acceptsResolverWithCompatibleIdentitySubtype() { + var identity = new TestCredentials("AK", "SK"); + var chain = IdentityChain.assemble( + AwsCredentialsIdentity.class, + List.of(registration( + "subtype", + new OrderingConstraint.Standard(StandardProvider.ENVIRONMENT), + IdentityResolver.of(identity))), + null); + + IdentityResult result = chain.resolveIdentity(Context.empty()); + + assertEquals(identity, result.identity()); + } + + @Test + void invalidatesOnlyResolverCompatibleWithRejectedIdentity() { + var firstInvalidations = new AtomicInteger(); + var secondInvalidations = new AtomicInteger(); + var chain = IdentityChain.assemble( + AwsCredentialsIdentity.class, + List.of( + registration( + "first", + new OrderingConstraint.Standard(StandardProvider.ENVIRONMENT), + trackingResolver( + TestCredentials.class, + new TestCredentials("AK", "SK"), + firstInvalidations)), + registration( + "second", + new OrderingConstraint.Standard(StandardProvider.SHARED_CONFIG), + trackingResolver( + OtherCredentials.class, + new OtherCredentials("OTHER_AK", "OTHER_SK"), + secondInvalidations))), + null); + + AwsCredentialsIdentity rejectedIdentity = chain.resolveIdentity(Context.empty()).identity(); + assertNotNull(rejectedIdentity); + chain.invalidate(rejectedIdentity); + + assertEquals(1, firstInvalidations.get()); + assertEquals(0, secondInvalidations.get()); + } + @Test void firstSuccessfulProviderWins() { var chain = IdentityChain.assemble(AwsCredentialsIdentity.class, @@ -138,6 +378,50 @@ void relativeToUnclaimedSlotAppendsAtEnd() { assertEquals(List.of("env", "custom"), chain.providerNames()); } + @Test + void afterUnclaimedSlotPrecedesNextClaimedSlot() { + var chain = IdentityChain.assemble(AwsCredentialsIdentity.class, + List.of( + registration("env", + new OrderingConstraint.Standard(StandardProvider.ENVIRONMENT), + errorResolver("env")), + registration("imds", + new OrderingConstraint.Standard(StandardProvider.EC2_INSTANCE_METADATA), + errorResolver("imds")), + registration("custom", + new OrderingConstraint.After(StandardProvider.PROFILE_SSO_SESSION), + errorResolver("custom"))), + null); + + assertEquals(List.of("env", "custom", "imds"), chain.providerNames()); + } + + @Test + void relativeProvidersPreserveInsertionOrder() { + var chain = IdentityChain.assemble(AwsCredentialsIdentity.class, + List.of( + registration("before-a", + new OrderingConstraint.Before(StandardProvider.ENVIRONMENT), + errorResolver("before-a")), + registration("before-b", + new OrderingConstraint.Before(StandardProvider.ENVIRONMENT), + errorResolver("before-b")), + registration("env", + new OrderingConstraint.Standard(StandardProvider.ENVIRONMENT), + errorResolver("env")), + registration("after-a", + new OrderingConstraint.After(StandardProvider.ENVIRONMENT), + errorResolver("after-a")), + registration("after-b", + new OrderingConstraint.After(StandardProvider.ENVIRONMENT), + errorResolver("after-b"))), + null); + + assertEquals( + List.of("before-a", "before-b", "env", "after-a", "after-b"), + chain.providerNames()); + } + @Test void duplicateNameThrows() { assertThrows(IllegalStateException.class, @@ -235,10 +519,79 @@ void detectedSlotClaimedByProviderProducesNoMissingModuleHint() { } } + @Test + void registeredProviderWithoutResolverProducesNoMissingModuleHint() { + ChainIdentityProvider registeredEcsProvider = new ChainIdentityProvider() { + @Override + public String name() { + return "Ecs"; + } + + @Override + public OrderingConstraint ordering() { + return new OrderingConstraint.Standard(StandardProvider.ECS_CONTAINER); + } + + @Override + public void setup(Class identityType, ChainSetup setup) { + // This installed module does not support the requested identity type. + } + }; + var setup = ChainSetup.builder() + .env(name -> name.equals("AWS_CONTAINER_CREDENTIALS_FULL_URI") ? "http://localhost" : null) + .build(); + var chain = IdentityChain.assemble( + AwsCredentialsIdentity.class, + List.of(registeredEcsProvider), + null, + setup); + + var diagnostics = new ChainResolutionDiagnostics[1]; + var context = Context.create(); + context.put(IdentityChain.DIAGNOSTICS, value -> diagnostics[0] = value); + var result = chain.resolveIdentity(context); + + assertTrue(result.error().contains("No credential providers were discovered")); + assertTrue(diagnostics[0].moduleSuggestions().isEmpty()); + assertTrue(!result.error().contains("aws-credentials-ecs")); + } + + @Test + void profileSourceProducesMissingModuleSuggestion(@TempDir Path tempDir) throws IOException { + Path config = tempDir.resolve("config"); + Files.writeString(config, """ + [default] + role_arn = arn:aws:iam::123456789012:role/Foo + source_profile = source + + [profile source] + aws_access_key_id = AKID + aws_secret_access_key = SECRET + """); + var profileFile = AwsProfileFile.builder().configFile(config).credentialsFile(null).build(); + var setup = ChainSetup.builder().profileFile(profileFile).env(name -> null).build(); + var chain = IdentityChain.assemble( + AwsCredentialsIdentity.class, + List.of(new SharedConfigProvider()), + null, + setup); + var diagnostics = new ChainResolutionDiagnostics[1]; + var context = Context.create(); + context.put(IdentityChain.DIAGNOSTICS, value -> diagnostics[0] = value); + + var result = chain.resolveIdentity(context); + + assertNull(result.identity()); + assertEquals( + List.of("software.amazon.smithy.java:aws-credentials-sts"), + diagnostics[0].moduleSuggestions()); + assertTrue(result.error().contains("aws-credentials-sts")); + } + private static ChainIdentityProvider registration( String name, OrderingConstraint ordering, - IdentityResolver resolver + IdentityResolver resolver ) { return new ChainIdentityProvider() { public String name() { @@ -280,4 +633,59 @@ public Class identityType() { } }; } + + private static IdentityResolver trackingResolver( + Class identityType, + I identity, + AtomicInteger invalidations + ) { + return new IdentityResolver<>() { + @Override + public IdentityResult resolveIdentity(Context requestProperties) { + return IdentityResult.of(identity); + } + + @Override + public Class identityType() { + return identityType; + } + + @Override + public void invalidate(I rejectedIdentity) { + invalidations.incrementAndGet(); + } + }; + } + + private static final class CloseableResolver + implements IdentityResolver, AutoCloseable { + private final Class identityType; + private final IdentityResult result; + private final AtomicInteger closes; + + CloseableResolver(Class identityType, I identity, AtomicInteger closes) { + this.identityType = identityType; + this.result = IdentityResult.of(identity); + this.closes = closes; + } + + @Override + public IdentityResult resolveIdentity(Context requestProperties) { + return result; + } + + @Override + public Class identityType() { + return identityType; + } + + @Override + public void close() { + closes.incrementAndGet(); + } + } + + private record TestCredentials(String accessKeyId, String secretAccessKey) implements AwsCredentialsIdentity {} + + private record OtherCredentials(String accessKeyId, String secretAccessKey) implements AwsCredentialsIdentity {} } diff --git a/aws/aws-credential-chain/src/test/java/software/amazon/smithy/java/aws/credentials/chain/ModularCredentialChainSuiteTest.java b/aws/aws-credential-chain/src/test/java/software/amazon/smithy/java/aws/credentials/chain/ModularCredentialChainSuiteTest.java index 4b710bec30..b0d28faedf 100644 --- a/aws/aws-credential-chain/src/test/java/software/amazon/smithy/java/aws/credentials/chain/ModularCredentialChainSuiteTest.java +++ b/aws/aws-credential-chain/src/test/java/software/amazon/smithy/java/aws/credentials/chain/ModularCredentialChainSuiteTest.java @@ -8,6 +8,7 @@ import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertNotNull; import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertSame; import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; import static org.junit.jupiter.api.Assertions.fail; @@ -24,7 +25,6 @@ import java.util.List; import java.util.Map; import java.util.Set; -import java.util.concurrent.atomic.AtomicInteger; import java.util.function.Function; import org.junit.jupiter.api.Assumptions; import org.junit.jupiter.api.DynamicTest; @@ -178,6 +178,8 @@ private void runCase(Document testCase) throws IOException { // Assemble the provider list: standard-slot stubs for installed modules, plus any custom providers. List providers = new ArrayList<>(); + List invalidatedProviders = new ArrayList<>(); + List invalidatedIdentities = new ArrayList<>(); Set installedSlots = installedSlots(input); for (StandardProvider slot : installedSlots) { providers.add(standardStub(slot, input, env)); @@ -185,7 +187,7 @@ private void runCase(Document testCase) throws IOException { Document customProviders = input.getMember("customProviders"); if (customProviders != null) { for (Document custom : customProviders.asList()) { - providers.add(customStub(custom)); + providers.add(customStub(custom, invalidatedProviders, invalidatedIdentities)); } } @@ -220,6 +222,20 @@ private void runCase(Document testCase) throws IOException { IdentityChain chain = IdentityChain.assemble(identityType, providers, null, setup); + Document invalidationIdentity = input.getMember("invalidationIdentity"); + if (invalidationIdentity != null) { + var rejectedIdentity = awsIdentity(invalidationIdentity); + invalidateChain(chain, rejectedIdentity); + assertEquals(stringList(expected.getMember("invalidatedProviders")), + invalidatedProviders, + "invalidatedProviders"); + assertEquals(invalidatedProviders.size(), invalidatedIdentities.size(), "invalidated identity count"); + for (Identity invalidatedIdentity : invalidatedIdentities) { + assertSame(rejectedIdentity, invalidatedIdentity, "chain must propagate the same identity instance"); + } + return; + } + var context = Context.create(); // Seed the feature-ID sink the way the client pipeline does, so the chain can record the winning provider's // feature IDs for the User-Agent. @@ -232,13 +248,6 @@ private void runCase(Document testCase) throws IOException { IdentityResult result; try { result = chain.resolveIdentity(context); - - // Invalidation: resolve, invalidate, resolve again; the second resolve is the asserted one. - Document invalidateAgain = input.getMember("invalidateAndResolveAgain"); - if (invalidateAgain != null && invalidateAgain.asBoolean()) { - chain.invalidate(); - result = chain.resolveIdentity(context); - } } catch (RuntimeException thrown) { if (expected.getMember("resolutionError") != null) { // The chain stopped at the throwing provider; assert the providers tried up to that point. @@ -329,7 +338,8 @@ private void slotSetup(StandardProvider slot, Document input, Function { - /* profile is pre-set on the setup by the runner */ } + markDetectedProfileSources(setup); + } case PROFILE_STATIC_KEYS -> profileSlot(setup, AwsConfigCredentialSource.StaticKeys.class, src -> staticResolver(src.accessKeyId(), src.secretAccessKey(), null, src.accountId())); @@ -341,7 +351,10 @@ private void slotSetup(StandardProvider slot, Document input, Function profileSlot(setup, AwsConfigCredentialSource.AssumeRole.class, - src -> responseResolver(input.getMember("stsResponse"))); + src -> responseResolver(input.getMember("stsResponse")), + src -> Set.of( + new CredentialFeatureId(src.sourceProfile() != null ? "o" : "p"), + new CredentialFeatureId("i"))); case PROFILE_WEB_IDENTITY -> profileSlot(setup, AwsConfigCredentialSource.WebIdentityToken.class, src -> responseResolver(input.getMember("stsResponse"))); @@ -365,20 +378,58 @@ private void profileSlot( ChainSetup setup, Class sourceType, Function> resolverFactory + ) { + profileSlot(setup, sourceType, resolverFactory, null); + } + + private void profileSlot( + ChainSetup setup, + Class sourceType, + Function> resolverFactory, + Function> featureIds ) { if (setup.profile() == null) { return; } for (AwsConfigCredentialSource source : setup.profile().credentialSources()) { if (sourceType.isInstance(source)) { - setup.addTerminalResolver(resolverFactory.apply(sourceType.cast(source))); + T typedSource = sourceType.cast(source); + IdentityResolver resolver = resolverFactory.apply(typedSource); + if (featureIds == null) { + setup.addTerminalResolver(resolver); + } else { + setup.addTerminalResolver(resolver, featureIds.apply(typedSource)); + } return; } } } + private void markDetectedProfileSources(ChainSetup setup) { + if (setup.profile() == null) { + return; + } + for (AwsConfigCredentialSource source : setup.profile().credentialSources()) { + setup.markDetected(switch (source) { + case AwsConfigCredentialSource.StaticKeys ignored -> StandardProvider.PROFILE_STATIC_KEYS; + case AwsConfigCredentialSource.SessionKeys ignored -> StandardProvider.PROFILE_SESSION_KEYS; + case AwsConfigCredentialSource.AssumeRole ignored -> StandardProvider.PROFILE_ASSUME_ROLE; + case AwsConfigCredentialSource.WebIdentityToken ignored -> StandardProvider.PROFILE_WEB_IDENTITY; + case AwsConfigCredentialSource.SsoSession ignored -> StandardProvider.PROFILE_SSO_SESSION; + case AwsConfigCredentialSource.LegacySso ignored -> StandardProvider.PROFILE_LEGACY_SSO; + case AwsConfigCredentialSource.CredentialProcess ignored -> + StandardProvider.PROFILE_CREDENTIAL_PROCESS; + case AwsConfigCredentialSource.LoginSession ignored -> StandardProvider.PROFILE_LOGIN; + }); + } + } + /** A custom (third-party) provider, exercising Before/After ordering and terminal/non-terminal semantics. */ - private ChainIdentityProvider customStub(Document custom) { + private ChainIdentityProvider customStub( + Document custom, + List invalidatedProviders, + List invalidatedIdentities + ) { String name = custom.getMember("name").asString(); OrderingConstraint ordering = ordering(custom.getMember("ordering")); boolean supportsType = boolOrDefault(custom.getMember("supportsIdentityType"), true); @@ -401,7 +452,11 @@ public void setup(Class identityType, ChainSetup setup) { if (!supportsType || !present) { return; } - IdentityResolver resolver = perAttemptResolver(response); + IdentityResolver resolver = recordingResolver( + responseResolver(response), + name, + invalidatedProviders, + invalidatedIdentities); if (terminal) { setup.addTerminalResolver(resolver); } else { @@ -432,34 +487,39 @@ private IdentityResolver responseResolver(Document response) { return IdentityResolver.of(awsIdentity(response)); } - /** - * Build a resolver whose response may be an array of per-attempt responses (for invalidate-then-resolve cases). - * Each call to {@code resolveIdentity} advances to the next element; a single object is returned every time. - */ - private IdentityResolver perAttemptResolver(Document response) { - if (response == null || !response.isType(ShapeType.LIST)) { - return responseResolver(response); - } - List> perAttempt = new ArrayList<>(); - for (Document attempt : response.asList()) { - perAttempt.add(responseResolver(attempt)); - } - AtomicInteger index = new AtomicInteger(0); - return new IdentityResolver() { + private IdentityResolver recordingResolver( + IdentityResolver delegate, + String name, + List invalidatedProviders, + List invalidatedIdentities + ) { + return new IdentityResolver<>() { + @Override - @SuppressWarnings("unchecked") - public IdentityResult resolveIdentity(Context ctx) { - int i = Math.min(index.getAndIncrement(), perAttempt.size() - 1); - return (IdentityResult) perAttempt.get(i).resolveIdentity(ctx); + public IdentityResult resolveIdentity(Context requestProperties) { + return delegate.resolveIdentity(requestProperties); } @Override - public Class identityType() { - return Identity.class; + public Class identityType() { + return delegate.identityType(); + } + + @Override + public void invalidate(I rejectedIdentity) { + invalidatedProviders.add(name); + invalidatedIdentities.add(rejectedIdentity); } }; } + private static void invalidateChain( + IdentityChain chain, + Identity rejectedIdentity + ) { + chain.invalidate(chain.identityType().cast(rejectedIdentity)); + } + private IdentityResolver errorResolver(String message) { // A resolver that raises a thrown exception stops the chain; the suite distinguishes this from an error // result via a message that reads like an unrecoverable failure. Circular source_profile is the case that diff --git a/aws/aws-credential-chain/src/test/java/software/amazon/smithy/java/aws/credentials/chain/config/CredentialProcessHandlerTest.java b/aws/aws-credential-chain/src/test/java/software/amazon/smithy/java/aws/credentials/chain/config/CredentialProcessHandlerTest.java index 4ff2edb7d6..039ca68c57 100644 --- a/aws/aws-credential-chain/src/test/java/software/amazon/smithy/java/aws/credentials/chain/config/CredentialProcessHandlerTest.java +++ b/aws/aws-credential-chain/src/test/java/software/amazon/smithy/java/aws/credentials/chain/config/CredentialProcessHandlerTest.java @@ -61,6 +61,22 @@ void processWithExpirationParsesTimestamp(@TempDir Path tmp) throws IOException assertEquals("2099-01-01T00:00:00Z", id.expirationTime().toString()); } + @Test + void processWithMalformedExpirationReturnsError(@TempDir Path tmp) throws IOException { + Path script = writeScript(tmp, + """ + #!/bin/sh + echo '{"Version": 1, "AccessKeyId": "AK", "SecretAccessKey": "SK", "Expiration": "invalid"}' + """); + + AwsConfigCredentialSource.CredentialProcess source = + new AwsConfigCredentialSource.CredentialProcess(script.toString()); + IdentityResult result = createFromProfileResult(source); + + assertNull(result.identity()); + assertTrue(result.error().contains("unparseable Expiration")); + } + @Test void processWithoutSessionTokenReturnsBasicCredentials(@TempDir Path tmp) throws IOException { Path script = writeScript(tmp, """ diff --git a/aws/aws-credential-chain/src/test/resources/software/amazon/smithy/java/aws/credentials/chain/modular-credential-chain-tests-schema.json b/aws/aws-credential-chain/src/test/resources/software/amazon/smithy/java/aws/credentials/chain/modular-credential-chain-tests-schema.json index 0c588f0c91..0114d7278d 100644 --- a/aws/aws-credential-chain/src/test/resources/software/amazon/smithy/java/aws/credentials/chain/modular-credential-chain-tests-schema.json +++ b/aws/aws-credential-chain/src/test/resources/software/amazon/smithy/java/aws/credentials/chain/modular-credential-chain-tests-schema.json @@ -377,15 +377,33 @@ } } }, - "invalidateAndResolveAgain": { - "type": "boolean", - "description": "When true, the harness resolves once, calls invalidate() on the chain, then resolves a second time. Source responses whose value is an array are interpreted per-attempt (index 0 for the first resolve, index 1 for the second). 'expected' describes the result of the SECOND resolve. Defaults to false." + "invalidationIdentity": { + "type": "object", + "description": "When present, the harness calls invalidate() on the assembled chain with this identity and does not perform resolution. Used only to validate propagation; cache behavior is outside this suite.", + "contents": { + "accessKeyId": { + "type": "string", + "description": "The rejected AWS access key ID" + }, + "secretAccessKey": { + "type": "string", + "description": "The rejected AWS secret access key" + }, + "sessionToken": { + "type": "string", + "description": "The rejected AWS session token" + }, + "accountId": { + "type": "string", + "description": "The rejected AWS account ID" + } + } } } }, "expected": { "type": "object", - "description": "The expected outcome. Exactly one of 'resolved', 'assemblyError', or 'resolutionError' is populated. 'assembledOrder', 'selectedProvider', and 'featureIds' are optional assertions layered on top.", + "description": "The expected outcome. Exactly one of 'resolved', 'assemblyError', 'resolutionError', or 'invalidatedProviders' is populated. 'assembledOrder', 'selectedProvider', and 'featureIds' are optional assertions layered on top.", "contents": { "resolved": { "type": "object", @@ -425,6 +443,10 @@ "type": "array", "description": "The canonical names of the resolvers in the assembled chain, in order. Optional assertion used by ordering tests. Assembly stops after a terminal resolver, so this list ends at the first terminal provider." }, + "invalidatedProviders": { + "type": "array", + "description": "The canonical names of every assembled child resolver whose declared identity type accepted invalidationIdentity, in propagation order. Used only by invalidation propagation cases." + }, "assemblyError": { "type": "object", "description": "Populated when client construction MUST fail during chain assembly (e.g. duplicate name, duplicate Standard slot claim).", diff --git a/aws/aws-credential-chain/src/test/resources/software/amazon/smithy/java/aws/credentials/chain/modular-credential-chain-tests.json b/aws/aws-credential-chain/src/test/resources/software/amazon/smithy/java/aws/credentials/chain/modular-credential-chain-tests.json index b21ee9b053..2354c6f88e 100644 --- a/aws/aws-credential-chain/src/test/resources/software/amazon/smithy/java/aws/credentials/chain/modular-credential-chain-tests.json +++ b/aws/aws-credential-chain/src/test/resources/software/amazon/smithy/java/aws/credentials/chain/modular-credential-chain-tests.json @@ -544,81 +544,34 @@ }, { "id": "23", - "documentation": "Invalidation does not change the source with terminal providers: a terminal custom provider resolves, invalidate() is called, and the second resolve still uses the same provider with refreshed values.", + "documentation": "Invalidation propagation: the chain passes the rejected identity unchanged to every assembled child resolver whose declared identity type accepts it. Cache behavior is outside this suite.", "input": { "installedModules": ["core"], - "invalidateAndResolveAgain": true, + "invalidationIdentity": { + "accessKeyId": "1", + "secretAccessKey": "2", + "sessionToken": "3", + "accountId": "4" + }, "customProviders": [ { - "name": "CachingVault", + "name": "RecorderA", "ordering": { "type": "Before", "slot": "ENVIRONMENT" }, "present": true, - "terminal": true, - "response": [ - { - "accessKeyId": "1", - "secretAccessKey": "2", - "sessionToken": "3", - "accountId": "4" - }, - { - "accessKeyId": "5", - "secretAccessKey": "6", - "sessionToken": "7", - "accountId": "8" - } - ] - } - ] - }, - "expected": { - "resolved": { - "accessKeyId": "5", - "secretAccessKey": "6", - "sessionToken": "7", - "accountId": "8" - }, - "selectedProvider": "CachingVault" - } - }, - { - "id": "24", - "documentation": "Invalidation can change source with a non-terminal custom provider: it resolves on the first attempt, becomes unavailable after invalidate(), and a lower-priority terminal provider resolves on the second attempt.", - "input": { - "installedModules": ["core", "imds"], - "invalidateAndResolveAgain": true, - "customProviders": [ + "terminal": false, + "response": { "error": "not resolved in this case" } + }, { - "name": "BestEffort", - "ordering": { "type": "Before", "slot": "EC2_INSTANCE_METADATA" }, + "name": "RecorderB", + "ordering": { "type": "After", "slot": "ENVIRONMENT" }, "present": true, "terminal": false, - "response": [ - { - "accessKeyId": "1", - "secretAccessKey": "2", - "sessionToken": "3", - "accountId": "4" - }, - { "error": "Vault unavailable" } - ] + "response": { "error": "not resolved in this case" } } - ], - "imdsResponse": { - "accessKeyId": "5", - "secretAccessKey": "6", - "sessionToken": "7", - "accountId": "8" - } + ] }, "expected": { - "resolved": { - "accessKeyId": "5", - "secretAccessKey": "6", - "sessionToken": "7", - "accountId": "8" - }, - "selectedProvider": "Ec2InstanceMetadata" + "invalidatedProviders": ["RecorderA", "RecorderB"] } }, { diff --git a/aws/aws-credentials-imds/src/main/java/software/amazon/smithy/java/aws/credentials/imds/ImdsCredentialProvider.java b/aws/aws-credentials-imds/src/main/java/software/amazon/smithy/java/aws/credentials/imds/ImdsCredentialProvider.java index 88aa43bc79..e6098c59c8 100644 --- a/aws/aws-credentials-imds/src/main/java/software/amazon/smithy/java/aws/credentials/imds/ImdsCredentialProvider.java +++ b/aws/aws-credentials-imds/src/main/java/software/amazon/smithy/java/aws/credentials/imds/ImdsCredentialProvider.java @@ -12,12 +12,12 @@ import java.time.format.DateTimeParseException; import java.util.Set; import java.util.function.Function; -import software.amazon.smithy.java.auth.api.identity.CachingIdentityResolver; import software.amazon.smithy.java.auth.api.identity.Identity; import software.amazon.smithy.java.auth.api.identity.IdentityResult; import software.amazon.smithy.java.aws.auth.api.identity.AwsCredentialsIdentity; import software.amazon.smithy.java.aws.auth.api.identity.AwsCredentialsResolver; import software.amazon.smithy.java.aws.config.AwsProfile; +import software.amazon.smithy.java.aws.credentials.chain.AwsCredentialCaching; import software.amazon.smithy.java.aws.credentials.chain.ChainIdentityProvider; import software.amazon.smithy.java.aws.credentials.chain.ChainSetup; import software.amazon.smithy.java.aws.credentials.chain.CredentialFeatureId; @@ -75,10 +75,7 @@ public void setup(Class identityType, ChainSetup setup) { ImdsClient client = new ImdsClient(endpoint); AwsCredentialsResolver delegate = ctx -> fetchAndParse(client, profileName); - setup.addResolver(CachingIdentityResolver.builder(delegate) - .executor(setup.executor()) - .allowExpiredCredentials(true) - .build()); + setup.addResolver(AwsCredentialCaching.staticallyStable(delegate, setup.executor())); } private static IdentityResult fetchAndParse(ImdsClient client, String profileName) { diff --git a/aws/aws-credentials-sts/build.gradle.kts b/aws/aws-credentials-sts/build.gradle.kts index 3831d580c3..e692b103f7 100644 --- a/aws/aws-credentials-sts/build.gradle.kts +++ b/aws/aws-credentials-sts/build.gradle.kts @@ -52,6 +52,7 @@ dependencies { implementation(project(":aws:client:aws-client-awsquery")) implementation(project(":codecs:json-codec", configuration = "shadow")) implementation(project(":logging")) + runtimeOnly(project(":aws:aws-sigv4")) stsModel("software.amazon.api.models:sts:1.0.7") { isTransitive = false } testImplementation(project(":client:client-mock-plugin")) testImplementation(project(":http:http-api")) diff --git a/aws/aws-credentials-sts/src/main/java/software/amazon/smithy/java/aws/credentials/sts/EnvWebIdentityProvider.java b/aws/aws-credentials-sts/src/main/java/software/amazon/smithy/java/aws/credentials/sts/EnvWebIdentityProvider.java index 927c30b24f..64028fef4e 100644 --- a/aws/aws-credentials-sts/src/main/java/software/amazon/smithy/java/aws/credentials/sts/EnvWebIdentityProvider.java +++ b/aws/aws-credentials-sts/src/main/java/software/amazon/smithy/java/aws/credentials/sts/EnvWebIdentityProvider.java @@ -9,6 +9,7 @@ import software.amazon.smithy.java.auth.api.identity.Identity; import software.amazon.smithy.java.aws.auth.api.identity.AwsCredentialsIdentity; import software.amazon.smithy.java.aws.config.AwsConfigCredentialSource; +import software.amazon.smithy.java.aws.credentials.chain.AwsCredentialCaching; import software.amazon.smithy.java.aws.credentials.chain.ChainIdentityProvider; import software.amazon.smithy.java.aws.credentials.chain.ChainSetup; import software.amazon.smithy.java.aws.credentials.chain.CredentialFeatureId; @@ -53,7 +54,8 @@ public void setup(Class identityType, ChainSetup setup) { String sessionName = setup.getenv("AWS_ROLE_SESSION_NAME"); var wit = new AwsConfigCredentialSource.WebIdentityToken(roleArn, tokenFile, sessionName, null); var endpoint = StsEndpointConfig.resolve(null, setup); - setup.addTerminalResolver(new StsWebIdentityResolver(wit, StsClientFactory.createNoAuth(endpoint))); + var resolver = new StsWebIdentityResolver(wit, StsClientFactory.createNoAuth(endpoint)); + setup.addTerminalResolver(AwsCredentialCaching.staticallyStable(resolver, setup.executor())); } } } diff --git a/aws/aws-credentials-sts/src/main/java/software/amazon/smithy/java/aws/credentials/sts/ProfileAssumeRoleProvider.java b/aws/aws-credentials-sts/src/main/java/software/amazon/smithy/java/aws/credentials/sts/ProfileAssumeRoleProvider.java index 8b5e8cbf5d..19e3c78ba8 100644 --- a/aws/aws-credentials-sts/src/main/java/software/amazon/smithy/java/aws/credentials/sts/ProfileAssumeRoleProvider.java +++ b/aws/aws-credentials-sts/src/main/java/software/amazon/smithy/java/aws/credentials/sts/ProfileAssumeRoleProvider.java @@ -9,6 +9,7 @@ import software.amazon.smithy.java.auth.api.identity.Identity; import software.amazon.smithy.java.aws.auth.api.identity.AwsCredentialsIdentity; import software.amazon.smithy.java.aws.config.AwsConfigCredentialSource; +import software.amazon.smithy.java.aws.credentials.chain.AwsCredentialCaching; import software.amazon.smithy.java.aws.credentials.chain.ChainIdentityProvider; import software.amazon.smithy.java.aws.credentials.chain.ChainSetup; import software.amazon.smithy.java.aws.credentials.chain.CredentialFeatureId; @@ -24,9 +25,13 @@ */ public final class ProfileAssumeRoleProvider implements ChainIdentityProvider { - private static final Set FEATURE_IDS = Set.of( + private static final CredentialFeatureId STS_FEATURE_ID = new CredentialFeatureId("i"); + private static final Set SOURCE_PROFILE_FEATURE_IDS = Set.of( new CredentialFeatureId("o"), - new CredentialFeatureId("i")); + STS_FEATURE_ID); + private static final Set CREDENTIAL_SOURCE_FEATURE_IDS = Set.of( + new CredentialFeatureId("p"), + STS_FEATURE_ID); @Override public String name() { @@ -38,11 +43,6 @@ public OrderingConstraint ordering() { return new OrderingConstraint.Standard(StandardProvider.PROFILE_ASSUME_ROLE); } - @Override - public Set featureIds() { - return FEATURE_IDS; - } - @Override public void setup(Class identityType, ChainSetup setup) { if (identityType != AwsCredentialsIdentity.class || setup.profile() == null) { @@ -51,10 +51,29 @@ public void setup(Class identityType, ChainSetup setup) { for (AwsConfigCredentialSource source : setup.profile().credentialSources()) { if (source instanceof AwsConfigCredentialSource.AssumeRole ar) { - var endpoint = StsEndpointConfig.resolve(ar.region(), setup); - setup.addTerminalResolver(new StsAssumeRoleResolver(ar, setup.profileFile(), endpoint)); + var resolver = createResolver(ar, setup); + setup.addTerminalResolver( + AwsCredentialCaching.staticallyStable(resolver, setup.executor()), + featureIds(ar)); return; } } } + + private static Set featureIds(AwsConfigCredentialSource.AssumeRole source) { + return source.sourceProfile() != null ? SOURCE_PROFILE_FEATURE_IDS : CREDENTIAL_SOURCE_FEATURE_IDS; + } + + static StsAssumeRoleResolver createResolver( + AwsConfigCredentialSource.AssumeRole source, + ChainSetup setup + ) { + var endpoint = StsEndpointConfig.resolve(source.region(), setup); + return new StsAssumeRoleResolver( + source, + setup.profileFile(), + endpoint, + setup.executor(), + setup.environment()); + } } diff --git a/aws/aws-credentials-sts/src/main/java/software/amazon/smithy/java/aws/credentials/sts/ProfileWebIdentityProvider.java b/aws/aws-credentials-sts/src/main/java/software/amazon/smithy/java/aws/credentials/sts/ProfileWebIdentityProvider.java index ac6d57975c..39c4bc12a8 100644 --- a/aws/aws-credentials-sts/src/main/java/software/amazon/smithy/java/aws/credentials/sts/ProfileWebIdentityProvider.java +++ b/aws/aws-credentials-sts/src/main/java/software/amazon/smithy/java/aws/credentials/sts/ProfileWebIdentityProvider.java @@ -9,6 +9,7 @@ import software.amazon.smithy.java.auth.api.identity.Identity; import software.amazon.smithy.java.aws.auth.api.identity.AwsCredentialsIdentity; import software.amazon.smithy.java.aws.config.AwsConfigCredentialSource; +import software.amazon.smithy.java.aws.credentials.chain.AwsCredentialCaching; import software.amazon.smithy.java.aws.credentials.chain.ChainIdentityProvider; import software.amazon.smithy.java.aws.credentials.chain.ChainSetup; import software.amazon.smithy.java.aws.credentials.chain.CredentialFeatureId; @@ -49,7 +50,8 @@ public void setup(Class identityType, ChainSetup setup) { for (AwsConfigCredentialSource source : setup.profile().credentialSources()) { if (source instanceof AwsConfigCredentialSource.WebIdentityToken wit) { var endpoint = StsEndpointConfig.resolve(wit.region(), setup); - setup.addTerminalResolver(new StsWebIdentityResolver(wit, StsClientFactory.createNoAuth(endpoint))); + var resolver = new StsWebIdentityResolver(wit, StsClientFactory.createNoAuth(endpoint)); + setup.addTerminalResolver(AwsCredentialCaching.staticallyStable(resolver, setup.executor())); return; } } diff --git a/aws/aws-credentials-sts/src/main/java/software/amazon/smithy/java/aws/credentials/sts/StsAssumeRoleResolver.java b/aws/aws-credentials-sts/src/main/java/software/amazon/smithy/java/aws/credentials/sts/StsAssumeRoleResolver.java index 99f2565161..f2fa2b6bb3 100644 --- a/aws/aws-credentials-sts/src/main/java/software/amazon/smithy/java/aws/credentials/sts/StsAssumeRoleResolver.java +++ b/aws/aws-credentials-sts/src/main/java/software/amazon/smithy/java/aws/credentials/sts/StsAssumeRoleResolver.java @@ -9,12 +9,15 @@ import java.util.HashSet; import java.util.Map; import java.util.Set; +import java.util.concurrent.ScheduledExecutorService; +import java.util.function.Function; import software.amazon.smithy.java.auth.api.identity.IdentityResolver; import software.amazon.smithy.java.auth.api.identity.IdentityResult; import software.amazon.smithy.java.aws.auth.api.identity.AwsCredentialsIdentity; import software.amazon.smithy.java.aws.config.AwsConfigCredentialSource; import software.amazon.smithy.java.aws.config.AwsProfile; import software.amazon.smithy.java.aws.config.AwsProfileFile; +import software.amazon.smithy.java.aws.credentials.chain.AwsCredentialCaching; import software.amazon.smithy.java.aws.credentials.chain.ChainSetup; import software.amazon.smithy.java.aws.credentials.imds.ImdsCredentialProvider; import software.amazon.smithy.java.context.Context; @@ -27,20 +30,40 @@ *

Handles recursive source_profile resolution with cycle detection per the * Assume Role SEP. */ -final class StsAssumeRoleResolver implements IdentityResolver { +final class StsAssumeRoleResolver implements IdentityResolver, AutoCloseable { private final AwsConfigCredentialSource.AssumeRole source; private final AwsProfileFile profileFile; private final StsEndpointConfig endpoint; + private final ScheduledExecutorService executor; + private final Set sourceProfilePath; + private final Function environment; + private volatile IdentityResolver sourceResolver; StsAssumeRoleResolver( AwsConfigCredentialSource.AssumeRole source, AwsProfileFile profileFile, - StsEndpointConfig endpoint + StsEndpointConfig endpoint, + ScheduledExecutorService executor, + Function environment + ) { + this(source, profileFile, endpoint, executor, Set.of(), environment); + } + + private StsAssumeRoleResolver( + AwsConfigCredentialSource.AssumeRole source, + AwsProfileFile profileFile, + StsEndpointConfig endpoint, + ScheduledExecutorService executor, + Set sourceProfilePath, + Function environment ) { this.source = source; this.profileFile = profileFile; this.endpoint = endpoint; + this.executor = executor; + this.sourceProfilePath = Set.copyOf(sourceProfilePath); + this.environment = environment; } @Override @@ -55,23 +78,40 @@ StsEndpointConfig endpoint() { @Override public IdentityResult resolveIdentity(Context ctx) { - AwsCredentialsIdentity sourceCredentials = resolveSourceCredentials(source, new HashSet<>()); - return callAssumeRole(sourceCredentials, source.roleArn(), source.externalId()); + var sourceResolver = sourceResolver(); + AwsCredentialsIdentity sourceCredentials = sourceResolver.resolveIdentity(ctx).unwrap(); + return callAssumeRole(sourceResolver, sourceCredentials, source.roleArn(), source.externalId()); } - private AwsCredentialsIdentity resolveSourceCredentials( + IdentityResolver sourceResolver() { + IdentityResolver current = sourceResolver; + if (current != null) { + return current; + } + synchronized (this) { + if (sourceResolver == null) { + sourceResolver = createSourceResolver(source, new HashSet<>(sourceProfilePath)); + } + return sourceResolver; + } + } + + private IdentityResolver createSourceResolver( AwsConfigCredentialSource.AssumeRole ar, Set visited ) { if (ar.sourceProfile() != null) { - return resolveFromSourceProfile(ar.sourceProfile(), visited); + return createSourceProfileResolver(ar.sourceProfile(), visited); } else if (ar.credentialSource() != null) { - return resolveFromCredentialSource(ar.credentialSource()); + return createCredentialSourceResolver(ar.credentialSource()); } throw new IllegalStateException("Profile with role_arn must have either source_profile or credential_source"); } - private AwsCredentialsIdentity resolveFromSourceProfile(String profileName, Set visited) { + private IdentityResolver createSourceProfileResolver( + String profileName, + Set visited + ) { if (!visited.add(profileName)) { throw new IllegalStateException("Circular source_profile reference detected: " + visited); } else if (profileFile == null) { @@ -86,68 +126,82 @@ private AwsCredentialsIdentity resolveFromSourceProfile(String profileName, Set< // Per the Assume Role SEP: terminate at static credentials for (AwsConfigCredentialSource src : sourceProfile.credentialSources()) { if (src instanceof AwsConfigCredentialSource.StaticKeys(String accessKeyId, String secretAccessKey, String accountId)) { - return AwsCredentialsIdentity.create( + return IdentityResolver.of(AwsCredentialsIdentity.create( accessKeyId, secretAccessKey, null, null, - accountId); + accountId)); } else if (src instanceof AwsConfigCredentialSource.SessionKeys(String accessKeyId, String secretAccessKey, String sessionToken, String accountId)) { - return AwsCredentialsIdentity.create( + return IdentityResolver.of(AwsCredentialsIdentity.create( accessKeyId, secretAccessKey, sessionToken, null, - accountId); + accountId)); } else if (src instanceof AwsConfigCredentialSource.AssumeRole nested) { - // Recursive: resolve source creds for the nested role, then assume it - AwsCredentialsIdentity nestedSource = resolveSourceCredentials(nested, visited); - return callAssumeRole(nestedSource, nested.roleArn(), nested.externalId()).unwrap(); + var nestedResolver = + new StsAssumeRoleResolver(nested, profileFile, endpoint, executor, visited, environment); + return AwsCredentialCaching.staticallyStable(nestedResolver, executor); } } throw new IllegalStateException("Source profile '" + profileName + "' has no resolvable credential source"); } - private AwsCredentialsIdentity resolveFromCredentialSource(String credentialSource) { + private IdentityResolver createCredentialSourceResolver(String credentialSource) { return switch (credentialSource) { case "Environment" -> { String ak = getRequireEnv("AWS_ACCESS_KEY_ID"); String sk = getRequireEnv("AWS_SECRET_ACCESS_KEY"); - String st = System.getenv("AWS_SESSION_TOKEN"); - yield AwsCredentialsIdentity.create(ak, sk, st, null, System.getenv("AWS_ACCOUNT_ID")); + String st = environment.apply("AWS_SESSION_TOKEN"); + yield IdentityResolver.of( + AwsCredentialsIdentity.create(ak, sk, st, null, environment.apply("AWS_ACCOUNT_ID"))); } case "Ec2InstanceMetadata" -> { - // Create a temporary ChainSetup to let ImdsCredentialProvider register its resolver - var tempSetup = ChainSetup.builder().build(); - new ImdsCredentialProvider().setup(AwsCredentialsIdentity.class, tempSetup); + var tempSetup = ChainSetup.builder().executor(executor).env(environment).build(); + var provider = new ImdsCredentialProvider(); + tempSetup.setCurrentProvider(provider); + provider.setup(AwsCredentialsIdentity.class, tempSetup); var resolvers = tempSetup.resolvers(); if (resolvers.isEmpty()) { throw new IllegalStateException("IMDS credential provider did not produce a resolver"); } @SuppressWarnings("unchecked") - var r = (IdentityResolver) resolvers.getFirst().resolver(); - yield r.resolveIdentity(Context.create()).unwrap(); + var resolver = (IdentityResolver) resolvers.getFirst().resolver(); + yield resolver; } default -> throw new IllegalStateException("Unsupported credential_source: " + credentialSource); }; } - private static String getRequireEnv(String name) { - var result = System.getenv(name); + private String getRequireEnv(String name) { + var result = environment.apply(name); if (result == null) { throw new IllegalStateException("credential_source=Environment but " + name + " not set"); } return result; } + @Override + public void close() { + IdentityResolver current = sourceResolver; + if (current instanceof AutoCloseable closeable) { + try { + closeable.close(); + } catch (Exception error) { + throw new IllegalStateException("Failed to close source credential resolver", error); + } + } + } + private IdentityResult callAssumeRole( + IdentityResolver retainedSourceResolver, AwsCredentialsIdentity sourceCredentials, String roleArn, String externalId ) { - // Create a static resolver for the source credentials - var sourceResolver = createSourceResolver(sourceCredentials); + var sourceResolver = createSourceResolver(sourceCredentials, retainedSourceResolver); try (DynamicClient client = StsClientFactory.create(sourceResolver, endpoint)) { // ExternalId is optional; Map.of rejects null values, so only include it when present. @@ -161,7 +215,10 @@ private IdentityResult callAssumeRole( } } - private static IdentityResolver createSourceResolver(AwsCredentialsIdentity creds) { + static IdentityResolver createSourceResolver( + AwsCredentialsIdentity creds, + IdentityResolver retainedSourceResolver + ) { IdentityResult sourceResult = IdentityResult.of(creds); return new IdentityResolver<>() { @Override @@ -173,6 +230,11 @@ public IdentityResult resolveIdentity(Context c) { public Class identityType() { return AwsCredentialsIdentity.class; } + + @Override + public void invalidate(AwsCredentialsIdentity rejectedIdentity) { + retainedSourceResolver.invalidate(rejectedIdentity); + } }; } } diff --git a/aws/aws-credentials-sts/src/main/java/software/amazon/smithy/java/aws/credentials/sts/StsClientFactory.java b/aws/aws-credentials-sts/src/main/java/software/amazon/smithy/java/aws/credentials/sts/StsClientFactory.java index 79ad1cd591..e0b237f749 100644 --- a/aws/aws-credentials-sts/src/main/java/software/amazon/smithy/java/aws/credentials/sts/StsClientFactory.java +++ b/aws/aws-credentials-sts/src/main/java/software/amazon/smithy/java/aws/credentials/sts/StsClientFactory.java @@ -9,6 +9,7 @@ import software.amazon.smithy.java.auth.api.identity.IdentityResolver; import software.amazon.smithy.java.aws.auth.api.identity.AwsCredentialsIdentity; import software.amazon.smithy.java.aws.client.awsquery.AwsQueryClientProtocol; +import software.amazon.smithy.java.aws.client.core.AwsCredentialChainPlugin; import software.amazon.smithy.java.aws.client.core.settings.RegionSetting; import software.amazon.smithy.java.aws.client.core.settings.StsEndpointSettings; import software.amazon.smithy.java.client.core.auth.scheme.AuthSchemeResolver; @@ -51,7 +52,8 @@ private static DynamicClient.Builder baseBuilder(StsEndpointConfig endpoint) { var builder = DynamicClient.builder() .model(model()) .serviceId(STS_SERVICE) - .protocol(new AwsQueryClientProtocol(STS_SERVICE, STS_VERSION)); + .protocol(new AwsQueryClientProtocol(STS_SERVICE, STS_VERSION)) + .addInterceptor(AwsCredentialChainPlugin.invalidationInterceptor()); if (endpoint.region() != null) { builder.putConfig(RegionSetting.REGION, endpoint.region()); } diff --git a/aws/aws-credentials-sts/src/main/java/software/amazon/smithy/java/aws/credentials/sts/StsWebIdentityResolver.java b/aws/aws-credentials-sts/src/main/java/software/amazon/smithy/java/aws/credentials/sts/StsWebIdentityResolver.java index 193f076795..23b45f6714 100644 --- a/aws/aws-credentials-sts/src/main/java/software/amazon/smithy/java/aws/credentials/sts/StsWebIdentityResolver.java +++ b/aws/aws-credentials-sts/src/main/java/software/amazon/smithy/java/aws/credentials/sts/StsWebIdentityResolver.java @@ -11,11 +11,13 @@ import java.nio.file.Path; import java.time.Instant; import java.util.Map; +import java.util.Set; import software.amazon.smithy.java.auth.api.identity.IdentityResolver; import software.amazon.smithy.java.auth.api.identity.IdentityResult; import software.amazon.smithy.java.aws.auth.api.identity.AwsCredentialsIdentity; import software.amazon.smithy.java.aws.config.AwsConfigCredentialSource; import software.amazon.smithy.java.context.Context; +import software.amazon.smithy.java.core.error.ModeledException; import software.amazon.smithy.java.core.serde.document.Document; import software.amazon.smithy.java.dynamicclient.DynamicClient; @@ -25,14 +27,30 @@ *

This operation does not require source credentials (the web identity token is the * authentication), so the STS client is configured with no auth. */ -final class StsWebIdentityResolver implements IdentityResolver { +final class StsWebIdentityResolver implements IdentityResolver, AutoCloseable { + + private static final Set EXPIRED_TOKEN_ERRORS = Set.of("ExpiredToken", "ExpiredTokenException"); private final AwsConfigCredentialSource.WebIdentityToken source; - private final DynamicClient client; + private final WebIdentityCall call; + private final AutoCloseable closeable; StsWebIdentityResolver(AwsConfigCredentialSource.WebIdentityToken source, DynamicClient client) { + this(source, input -> client.call("AssumeRoleWithWebIdentity", input), client); + } + + StsWebIdentityResolver(AwsConfigCredentialSource.WebIdentityToken source, WebIdentityCall call) { + this(source, call, () -> {}); + } + + private StsWebIdentityResolver( + AwsConfigCredentialSource.WebIdentityToken source, + WebIdentityCall call, + AutoCloseable closeable + ) { this.source = source; - this.client = client; + this.call = call; + this.closeable = closeable; } @Override @@ -43,19 +61,14 @@ public Class identityType() { @Override public IdentityResult resolveIdentity(Context ctx) { try { - var token = Files.readString(Path.of(source.webIdentityTokenFile()), StandardCharsets.UTF_8).trim(); - var sessionName = source.roleSessionName() != null - ? source.roleSessionName() - : "smithy-java-" + System.currentTimeMillis(); - - return parseCredentials(client.call("AssumeRoleWithWebIdentity", - Map.of( - "RoleArn", - source.roleArn(), - "RoleSessionName", - sessionName, - "WebIdentityToken", - token))); + try { + return callWithCurrentToken(); + } catch (RuntimeException error) { + if (!isExpiredToken(error)) { + throw error; + } + return callWithCurrentToken(); + } } catch (RuntimeException | IOException e) { var msg = String.format("Failed to assume role with web identity (role=%s): %s", source.roleArn(), @@ -64,6 +77,42 @@ public IdentityResult resolveIdentity(Context ctx) { } } + private IdentityResult callWithCurrentToken() throws IOException { + var token = Files.readString(Path.of(source.webIdentityTokenFile()), StandardCharsets.UTF_8).trim(); + var sessionName = source.roleSessionName() != null + ? source.roleSessionName() + : "smithy-java-" + System.currentTimeMillis(); + + return parseCredentials(call.call(Map.of( + "RoleArn", + source.roleArn(), + "RoleSessionName", + sessionName, + "WebIdentityToken", + token))); + } + + private static boolean isExpiredToken(RuntimeException error) { + return error instanceof ModeledException modeled + && EXPIRED_TOKEN_ERRORS.contains(modeled.schema().id().getName()); + } + + @Override + public void close() { + try { + closeable.close(); + } catch (RuntimeException error) { + throw error; + } catch (Exception error) { + throw new IllegalStateException("Failed to close STS client", error); + } + } + + @FunctionalInterface + interface WebIdentityCall { + Document call(Map input); + } + static IdentityResult parseCredentials(Document output) { Document creds = output.getMember("Credentials"); String accessKeyId = creds.getMember("AccessKeyId").asString(); diff --git a/aws/aws-credentials-sts/src/test/java/software/amazon/smithy/java/aws/credentials/sts/EnvWebIdentityProviderTest.java b/aws/aws-credentials-sts/src/test/java/software/amazon/smithy/java/aws/credentials/sts/EnvWebIdentityProviderTest.java index 73fce50aa6..856f99f709 100644 --- a/aws/aws-credentials-sts/src/test/java/software/amazon/smithy/java/aws/credentials/sts/EnvWebIdentityProviderTest.java +++ b/aws/aws-credentials-sts/src/test/java/software/amazon/smithy/java/aws/credentials/sts/EnvWebIdentityProviderTest.java @@ -6,6 +6,7 @@ package software.amazon.smithy.java.aws.credentials.sts; import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertInstanceOf; import java.io.IOException; import java.nio.file.Files; @@ -13,6 +14,7 @@ import java.util.Map; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.io.TempDir; +import software.amazon.smithy.java.auth.api.identity.CachingIdentityResolver; import software.amazon.smithy.java.aws.auth.api.identity.AwsCredentialsIdentity; import software.amazon.smithy.java.aws.credentials.chain.ChainSetup; @@ -37,6 +39,7 @@ void registersWhenEnvVarsSet(@TempDir Path tmp) throws IOException { provider.setup(AwsCredentialsIdentity.class, setup); assertEquals(1, setup.resolvers().size()); + assertInstanceOf(CachingIdentityResolver.class, setup.resolvers().getFirst().resolver()); } @Test diff --git a/aws/aws-credentials-sts/src/test/java/software/amazon/smithy/java/aws/credentials/sts/ProfileAssumeRoleProviderTest.java b/aws/aws-credentials-sts/src/test/java/software/amazon/smithy/java/aws/credentials/sts/ProfileAssumeRoleProviderTest.java index 0c78507ac2..6d58e9015e 100644 --- a/aws/aws-credentials-sts/src/test/java/software/amazon/smithy/java/aws/credentials/sts/ProfileAssumeRoleProviderTest.java +++ b/aws/aws-credentials-sts/src/test/java/software/amazon/smithy/java/aws/credentials/sts/ProfileAssumeRoleProviderTest.java @@ -6,18 +6,22 @@ package software.amazon.smithy.java.aws.credentials.sts; import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertInstanceOf; import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; import java.io.IOException; import java.nio.file.Files; import java.nio.file.Path; +import java.util.Set; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.io.TempDir; +import software.amazon.smithy.java.auth.api.identity.CachingIdentityResolver; import software.amazon.smithy.java.aws.auth.api.identity.AwsCredentialsIdentity; import software.amazon.smithy.java.aws.config.AwsConfigCredentialSource; import software.amazon.smithy.java.aws.config.AwsProfileFile; import software.amazon.smithy.java.aws.credentials.chain.ChainSetup; +import software.amazon.smithy.java.aws.credentials.chain.CredentialFeatureId; import software.amazon.smithy.java.context.Context; class ProfileAssumeRoleProviderTest { @@ -62,6 +66,32 @@ void registersWhenProfileHasRoleArn(@TempDir Path tmp) throws IOException { provider.setup(AwsCredentialsIdentity.class, setup); assertEquals(1, setup.resolvers().size()); + assertEquals( + Set.of(new CredentialFeatureId("o"), new CredentialFeatureId("i")), + setup.resolvers().getFirst().featureIds()); + } + + @Test + void credentialSourceUsesCredentialSourceFeatureId(@TempDir Path tmp) throws IOException { + Path config = tmp.resolve("config"); + Files.writeString(config, """ + [default] + role_arn = arn:aws:iam::123456789:role/RoleA + credential_source = Environment + """); + + var profileFile = AwsProfileFile.builder().configFile(config).credentialsFile(null).build(); + var setup = ChainSetup.builder().build(); + setup.setProfileFile(profileFile); + setup.setProfile(profileFile.activeProfile(k -> null)); + var provider = new ProfileAssumeRoleProvider(); + setup.setCurrentProvider(provider); + + provider.setup(AwsCredentialsIdentity.class, setup); + + assertEquals( + Set.of(new CredentialFeatureId("p"), new CredentialFeatureId("i")), + setup.resolvers().getFirst().featureIds()); } @Test @@ -88,7 +118,9 @@ void regionOverrideFlowsIntoStsEndpoint(@TempDir Path tmp) throws IOException { provider.setup(AwsCredentialsIdentity.class, setup); assertEquals(1, setup.resolvers().size()); - var resolver = (StsAssumeRoleResolver) setup.resolvers().get(0).resolver(); + assertInstanceOf(CachingIdentityResolver.class, setup.resolvers().get(0).resolver()); + var source = (AwsConfigCredentialSource.AssumeRole) setup.profile().credentialSources().getFirst(); + var resolver = ProfileAssumeRoleProvider.createResolver(source, setup); assertEquals("eu-central-1", resolver.endpoint().region()); } @@ -127,7 +159,7 @@ void detectsCircularSourceProfile(@TempDir Path tmp) throws IOException { var profileFile = AwsProfileFile.builder().configFile(config).credentialsFile(null).build(); var source = assumeRole("arn:aws:iam::123456789:role/RoleA", "B", null); - var resolver = new StsAssumeRoleResolver(source, profileFile, TEST_ENDPOINT); + var resolver = new StsAssumeRoleResolver(source, profileFile, TEST_ENDPOINT, null, name -> null); var ex = assertThrows(RuntimeException.class, () -> resolver.resolveIdentity(Context.create())); assertTrue(ex.getMessage().contains("Circular") || ex.getCause().getMessage().contains("Circular")); @@ -144,7 +176,7 @@ void failsWhenSourceProfileNotFound(@TempDir Path tmp) throws IOException { var profileFile = AwsProfileFile.builder().configFile(config).credentialsFile(null).build(); var source = assumeRole("arn:aws:iam::123456789:role/RoleA", "nonexistent", null); - var resolver = new StsAssumeRoleResolver(source, profileFile, TEST_ENDPOINT); + var resolver = new StsAssumeRoleResolver(source, profileFile, TEST_ENDPOINT, null, name -> null); var ex = assertThrows(RuntimeException.class, () -> resolver.resolveIdentity(Context.create())); assertTrue(ex.getMessage().contains("nonexistent") || ex.getCause().getMessage().contains("nonexistent")); @@ -153,7 +185,7 @@ void failsWhenSourceProfileNotFound(@TempDir Path tmp) throws IOException { @Test void failsWithUnsupportedCredentialSource() { var source = assumeRole("arn:aws:iam::123456789:role/RoleA", null, "CustomUnsupportedProvider"); - var resolver = new StsAssumeRoleResolver(source, null, TEST_ENDPOINT); + var resolver = new StsAssumeRoleResolver(source, null, TEST_ENDPOINT, null, name -> null); var ex = assertThrows(RuntimeException.class, () -> resolver.resolveIdentity(Context.create())); assertTrue(ex.getMessage().contains("Unsupported") || ex.getCause().getMessage().contains("Unsupported")); diff --git a/aws/aws-credentials-sts/src/test/java/software/amazon/smithy/java/aws/credentials/sts/ProfileWebIdentityProviderTest.java b/aws/aws-credentials-sts/src/test/java/software/amazon/smithy/java/aws/credentials/sts/ProfileWebIdentityProviderTest.java index 5df3c67095..8638e99f5d 100644 --- a/aws/aws-credentials-sts/src/test/java/software/amazon/smithy/java/aws/credentials/sts/ProfileWebIdentityProviderTest.java +++ b/aws/aws-credentials-sts/src/test/java/software/amazon/smithy/java/aws/credentials/sts/ProfileWebIdentityProviderTest.java @@ -6,12 +6,14 @@ package software.amazon.smithy.java.aws.credentials.sts; import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertInstanceOf; import java.io.IOException; import java.nio.file.Files; import java.nio.file.Path; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.io.TempDir; +import software.amazon.smithy.java.auth.api.identity.CachingIdentityResolver; import software.amazon.smithy.java.aws.auth.api.identity.AwsCredentialsIdentity; import software.amazon.smithy.java.aws.config.AwsProfileFile; import software.amazon.smithy.java.aws.credentials.chain.ChainSetup; @@ -36,6 +38,7 @@ void registersWhenProfileHasTokenFile(@TempDir Path tmp) throws IOException { provider.setup(AwsCredentialsIdentity.class, setup); assertEquals(1, setup.resolvers().size()); + assertInstanceOf(CachingIdentityResolver.class, setup.resolvers().getFirst().resolver()); } @Test diff --git a/aws/aws-credentials-sts/src/test/java/software/amazon/smithy/java/aws/credentials/sts/StsAssumeRoleResolverTest.java b/aws/aws-credentials-sts/src/test/java/software/amazon/smithy/java/aws/credentials/sts/StsAssumeRoleResolverTest.java index 18f923acec..34d28e30cd 100644 --- a/aws/aws-credentials-sts/src/test/java/software/amazon/smithy/java/aws/credentials/sts/StsAssumeRoleResolverTest.java +++ b/aws/aws-credentials-sts/src/test/java/software/amazon/smithy/java/aws/credentials/sts/StsAssumeRoleResolverTest.java @@ -5,16 +5,27 @@ package software.amazon.smithy.java.aws.credentials.sts; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertInstanceOf; +import static org.junit.jupiter.api.Assertions.assertSame; import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; import java.io.IOException; import java.nio.file.Files; import java.nio.file.Path; +import java.util.Map; +import java.util.concurrent.atomic.AtomicInteger; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.io.TempDir; +import software.amazon.smithy.java.auth.api.identity.CachingIdentityResolver; +import software.amazon.smithy.java.auth.api.identity.IdentityResolver; +import software.amazon.smithy.java.auth.api.identity.IdentityResult; +import software.amazon.smithy.java.aws.auth.api.identity.AwsCredentialsIdentity; import software.amazon.smithy.java.aws.config.AwsConfigCredentialSource; import software.amazon.smithy.java.aws.config.AwsProfileFile; import software.amazon.smithy.java.context.Context; +import software.amazon.smithy.model.shapes.ShapeId; class StsAssumeRoleResolverTest { @@ -51,7 +62,7 @@ void resolvesSourceCredsAndAttemptsStsCall(@TempDir Path tmp) throws IOException var profileFile = AwsProfileFile.builder().configFile(config).credentialsFile(null).build(); var source = assumeRole("arn:aws:iam::123456789:role/RoleA", "src", null); - var resolver = new StsAssumeRoleResolver(source, profileFile, TEST_ENDPOINT); + var resolver = new StsAssumeRoleResolver(source, profileFile, TEST_ENDPOINT, null, name -> null); // Source creds resolve, STS call fails (no real endpoint) — that's expected assertThrows(RuntimeException.class, () -> resolver.resolveIdentity(Context.create())); @@ -73,7 +84,7 @@ void resolvesSessionKeysFromSourceProfile(@TempDir Path tmp) throws IOException var profileFile = AwsProfileFile.builder().configFile(config).credentialsFile(null).build(); var source = assumeRole("arn:aws:iam::123456789:role/RoleA", "src", null); - var resolver = new StsAssumeRoleResolver(source, profileFile, TEST_ENDPOINT); + var resolver = new StsAssumeRoleResolver(source, profileFile, TEST_ENDPOINT, null, name -> null); // Session keys resolve, STS call fails — expected assertThrows(RuntimeException.class, () -> resolver.resolveIdentity(Context.create())); @@ -94,7 +105,7 @@ void detectsCircularSourceProfile(@TempDir Path tmp) throws IOException { var profileFile = AwsProfileFile.builder().configFile(config).credentialsFile(null).build(); var source = assumeRole("arn:aws:iam::123456789:role/RoleA", "B", null); - var resolver = new StsAssumeRoleResolver(source, profileFile, TEST_ENDPOINT); + var resolver = new StsAssumeRoleResolver(source, profileFile, TEST_ENDPOINT, null, name -> null); assertThrows(RuntimeException.class, () -> resolver.resolveIdentity(Context.create())); } @@ -110,7 +121,7 @@ void failsWhenSourceProfileNotFound(@TempDir Path tmp) throws IOException { var profileFile = AwsProfileFile.builder().configFile(config).credentialsFile(null).build(); var source = assumeRole("arn:aws:iam::123456789:role/RoleA", "nonexistent", null); - var resolver = new StsAssumeRoleResolver(source, profileFile, TEST_ENDPOINT); + var resolver = new StsAssumeRoleResolver(source, profileFile, TEST_ENDPOINT, null, name -> null); assertThrows(RuntimeException.class, () -> resolver.resolveIdentity(Context.create())); } @@ -118,7 +129,7 @@ void failsWhenSourceProfileNotFound(@TempDir Path tmp) throws IOException { @Test void failsWithUnsupportedCredentialSource() { var source = assumeRole("arn:aws:iam::123456789:role/RoleA", null, "CustomUnsupportedProvider"); - var resolver = new StsAssumeRoleResolver(source, null, TEST_ENDPOINT); + var resolver = new StsAssumeRoleResolver(source, null, TEST_ENDPOINT, null, name -> null); assertThrows(RuntimeException.class, () -> resolver.resolveIdentity(Context.create())); } @@ -126,7 +137,7 @@ void failsWithUnsupportedCredentialSource() { @Test void failsWhenNeitherSourceProfileNorCredentialSource() { var source = assumeRole("arn:aws:iam::123456789:role/RoleA", null, null); - var resolver = new StsAssumeRoleResolver(source, null, TEST_ENDPOINT); + var resolver = new StsAssumeRoleResolver(source, null, TEST_ENDPOINT, null, name -> null); assertThrows(RuntimeException.class, () -> resolver.resolveIdentity(Context.create())); } @@ -145,7 +156,7 @@ void failsWhenSourceProfileHasNoCredentialSources(@TempDir Path tmp) throws IOEx var profileFile = AwsProfileFile.builder().configFile(config).credentialsFile(null).build(); var source = assumeRole("arn:aws:iam::123456789:role/RoleA", "empty", null); - var resolver = new StsAssumeRoleResolver(source, profileFile, TEST_ENDPOINT); + var resolver = new StsAssumeRoleResolver(source, profileFile, TEST_ENDPOINT, null, name -> null); assertThrows(RuntimeException.class, () -> resolver.resolveIdentity(Context.create())); } @@ -153,7 +164,7 @@ void failsWhenSourceProfileHasNoCredentialSources(@TempDir Path tmp) throws IOEx @Test void failsWhenProfileFileIsNull() { var source = assumeRole("arn:aws:iam::123456789:role/RoleA", "src", null); - var resolver = new StsAssumeRoleResolver(source, null, TEST_ENDPOINT); + var resolver = new StsAssumeRoleResolver(source, null, TEST_ENDPOINT, null, name -> null); assertThrows(RuntimeException.class, () -> resolver.resolveIdentity(Context.create())); } @@ -177,20 +188,102 @@ void chainedAssumeRoleResolvesNestedSourceProfile(@TempDir Path tmp) throws IOEx var profileFile = AwsProfileFile.builder().configFile(config).credentialsFile(null).build(); var source = assumeRole("arn:aws:iam::111:role/RoleA", "B", null); - var resolver = new StsAssumeRoleResolver(source, profileFile, TEST_ENDPOINT); + var resolver = new StsAssumeRoleResolver(source, profileFile, TEST_ENDPOINT, null, name -> null); // Walks A -> B -> C (static keys), then attempts STS calls which fail assertThrows(RuntimeException.class, () -> resolver.resolveIdentity(Context.create())); } @Test - void credentialSourceEnvironmentResolvesAndAttemptsSts(@TempDir Path tmp) throws IOException { - // This test requires real env vars — will fail if AWS_ACCESS_KEY_ID not set - // We test the error path (env vars not set) + void retainsAndCachesNestedAssumeRoleSource(@TempDir Path tmp) throws IOException { + Path config = tmp.resolve("config"); + Files.writeString(config, """ + [default] + role_arn = arn:aws:iam::111:role/RoleA + source_profile = B + + [profile B] + role_arn = arn:aws:iam::222:role/RoleB + source_profile = C + + [profile C] + aws_access_key_id = LEAF_AK + aws_secret_access_key = LEAF_SK + """); + + var profileFile = AwsProfileFile.builder().configFile(config).credentialsFile(null).build(); + var source = assumeRole("arn:aws:iam::111:role/RoleA", "B", null); + var resolver = new StsAssumeRoleResolver(source, profileFile, TEST_ENDPOINT, null, name -> null); + + var nestedSource = resolver.sourceResolver(); + assertSame(nestedSource, resolver.sourceResolver()); + assertInstanceOf(CachingIdentityResolver.class, nestedSource); + resolver.close(); + } + + @Test + void credentialSourceEnvironmentUsesConfiguredEnvironment() { var source = assumeRole("arn:aws:iam::123456789:role/RoleA", null, "Environment"); - var resolver = new StsAssumeRoleResolver(source, null, TEST_ENDPOINT); + var environment = Map.of( + "AWS_ACCESS_KEY_ID", + "CONFIGURED_AK", + "AWS_SECRET_ACCESS_KEY", + "CONFIGURED_SK", + "AWS_SESSION_TOKEN", + "CONFIGURED_TOKEN", + "AWS_ACCOUNT_ID", + "123456789012"); + var resolver = new StsAssumeRoleResolver(source, null, TEST_ENDPOINT, null, environment::get); + + var identity = resolver.sourceResolver().resolveIdentity(Context.empty()).identity(); + + assertEquals("CONFIGURED_AK", identity.accessKeyId()); + assertEquals("CONFIGURED_TOKEN", identity.sessionToken()); + assertEquals("123456789012", identity.accountId()); + } - // Fails because AWS_ACCESS_KEY_ID is not set in test environment - assertThrows(RuntimeException.class, () -> resolver.resolveIdentity(Context.create())); + @Test + void sourceCredentialSnapshotForwardsInvalidation() { + var invalidations = new AtomicInteger(); + IdentityResolver retainedResolver = new IdentityResolver<>() { + @Override + public IdentityResult resolveIdentity(Context requestProperties) { + throw new AssertionError("Snapshot resolver must not resolve the retained resolver again"); + } + + @Override + public Class identityType() { + return AwsCredentialsIdentity.class; + } + + @Override + public void invalidate(AwsCredentialsIdentity rejectedIdentity) { + invalidations.incrementAndGet(); + } + }; + var credentials = AwsCredentialsIdentity.create("AKID", "SECRET"); + var snapshot = StsAssumeRoleResolver.createSourceResolver(credentials, retainedResolver); + + assertSame(credentials, snapshot.resolveIdentity(Context.empty()).identity()); + snapshot.invalidate(credentials); + assertEquals(1, invalidations.get()); + } + + @Test + void signedStsClientInstallsCredentialInvalidationInterceptor() { + var credentials = AwsCredentialsIdentity.create("AKID", "SECRET"); + + try (var client = StsClientFactory.create(IdentityResolver.of(credentials), TEST_ENDPOINT)) { + assertTrue(client.config() + .supportedAuthSchemes() + .stream() + .anyMatch(scheme -> scheme.schemeId().equals(ShapeId.from("aws.auth#sigv4")))); + assertTrue(client.config() + .interceptors() + .stream() + .anyMatch(interceptor -> interceptor.getClass() + .getSimpleName() + .equals("InvalidateCredentialsInterceptor"))); + } } } diff --git a/aws/aws-credentials-sts/src/test/java/software/amazon/smithy/java/aws/credentials/sts/StsWebIdentityResolverTest.java b/aws/aws-credentials-sts/src/test/java/software/amazon/smithy/java/aws/credentials/sts/StsWebIdentityResolverTest.java index 18623ca292..a9eff065f2 100644 --- a/aws/aws-credentials-sts/src/test/java/software/amazon/smithy/java/aws/credentials/sts/StsWebIdentityResolverTest.java +++ b/aws/aws-credentials-sts/src/test/java/software/amazon/smithy/java/aws/credentials/sts/StsWebIdentityResolverTest.java @@ -5,16 +5,29 @@ package software.amazon.smithy.java.aws.credentials.sts; +import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; import java.io.IOException; +import java.io.UncheckedIOException; import java.nio.file.Files; import java.nio.file.Path; +import java.time.Instant; +import java.util.ArrayList; +import java.util.List; +import java.util.Map; +import java.util.concurrent.atomic.AtomicInteger; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.io.TempDir; import software.amazon.smithy.java.aws.config.AwsConfigCredentialSource; import software.amazon.smithy.java.context.Context; +import software.amazon.smithy.java.core.error.ModeledException; +import software.amazon.smithy.java.core.schema.Schema; +import software.amazon.smithy.java.core.serde.ShapeSerializer; +import software.amazon.smithy.java.core.serde.document.Document; +import software.amazon.smithy.model.shapes.ShapeId; +import software.amazon.smithy.model.traits.ErrorTrait; class StsWebIdentityResolverTest { @@ -51,4 +64,66 @@ void readsTokenFileAndAttemptsStsCall(@TempDir Path tmp) throws IOException { assertTrue(ex.getMessage().contains("Failed to assume role with web identity")); assertTrue(ex.getMessage().contains("arn:aws:iam::123:role/R")); } + + @Test + void reloadsTokenFileAndRetriesExpiredTokenOnce(@TempDir Path tmp) throws IOException { + Path tokenFile = tmp.resolve("token"); + Files.writeString(tokenFile, "expired-token"); + var source = new AwsConfigCredentialSource.WebIdentityToken( + "arn:aws:iam::123:role/R", + tokenFile.toString(), + "session", + null); + var calls = new AtomicInteger(); + var tokens = new ArrayList(); + var resolver = new StsWebIdentityResolver(source, input -> { + tokens.add((String) input.get("WebIdentityToken")); + if (calls.getAndIncrement() == 0) { + try { + Files.writeString(tokenFile, "fresh-token"); + } catch (IOException error) { + throw new UncheckedIOException(error); + } + throw new TestModeledException("ExpiredTokenException"); + } + return credentialResponse(); + }); + + var result = resolver.resolveIdentity(Context.empty()); + + assertEquals("AKID", result.identity().accessKeyId()); + assertEquals(2, calls.get()); + assertEquals(List.of("expired-token", "fresh-token"), tokens); + } + + private static Document credentialResponse() { + return Document.ofObject(Map.of( + "Credentials", + Map.of( + "AccessKeyId", + "AKID", + "SecretAccessKey", + "SECRET", + "SessionToken", + "TOKEN", + "Expiration", + Instant.now().plusSeconds(3600).toString()))); + } + + private static final class TestModeledException extends ModeledException { + private TestModeledException(String name) { + super(Schema.structureBuilder( + ShapeId.from("com.amazonaws.sts#" + name), + new ErrorTrait("client")).build(), + name); + } + + @Override + public void serializeMembers(ShapeSerializer serializer) {} + + @Override + public T getMemberValue(Schema member) { + return null; + } + } } diff --git a/aws/aws-sigv4-s3express/src/main/java/software/amazon/smithy/java/aws/client/auth/scheme/s3express/S3ExpressIdentityCache.java b/aws/aws-sigv4-s3express/src/main/java/software/amazon/smithy/java/aws/client/auth/scheme/s3express/S3ExpressIdentityCache.java index 8a7d64ed6a..67cc8b8819 100644 --- a/aws/aws-sigv4-s3express/src/main/java/software/amazon/smithy/java/aws/client/auth/scheme/s3express/S3ExpressIdentityCache.java +++ b/aws/aws-sigv4-s3express/src/main/java/software/amazon/smithy/java/aws/client/auth/scheme/s3express/S3ExpressIdentityCache.java @@ -167,11 +167,11 @@ void evictIdle() { } } - void invalidateAll() { + void invalidateAll(AwsCredentialsIdentity rejectedIdentity) { writeLock.lock(); try { for (var entry : entries.values()) { - entry.resolver.invalidate(); + entry.resolver.invalidate(rejectedIdentity); } } finally { writeLock.unlock(); diff --git a/aws/aws-sigv4-s3express/src/main/java/software/amazon/smithy/java/aws/client/auth/scheme/s3express/S3ExpressIdentityProvider.java b/aws/aws-sigv4-s3express/src/main/java/software/amazon/smithy/java/aws/client/auth/scheme/s3express/S3ExpressIdentityProvider.java index 3fac871af8..79b9053785 100644 --- a/aws/aws-sigv4-s3express/src/main/java/software/amazon/smithy/java/aws/client/auth/scheme/s3express/S3ExpressIdentityProvider.java +++ b/aws/aws-sigv4-s3express/src/main/java/software/amazon/smithy/java/aws/client/auth/scheme/s3express/S3ExpressIdentityProvider.java @@ -124,7 +124,7 @@ public Class identityType() { } @Override - public void invalidate() { - cache.invalidateAll(); + public void invalidate(AwsCredentialsIdentity rejectedIdentity) { + cache.invalidateAll(rejectedIdentity); } } diff --git a/aws/aws-sigv4-s3express/src/test/java/software/amazon/smithy/java/aws/client/auth/scheme/s3express/S3ExpressIdentityProviderTest.java b/aws/aws-sigv4-s3express/src/test/java/software/amazon/smithy/java/aws/client/auth/scheme/s3express/S3ExpressIdentityProviderTest.java index f25b63d2cd..d69041a4f7 100644 --- a/aws/aws-sigv4-s3express/src/test/java/software/amazon/smithy/java/aws/client/auth/scheme/s3express/S3ExpressIdentityProviderTest.java +++ b/aws/aws-sigv4-s3express/src/test/java/software/amazon/smithy/java/aws/client/auth/scheme/s3express/S3ExpressIdentityProviderTest.java @@ -33,7 +33,7 @@ void invalidateClearsCachedSessions() { var first = provider.resolveIdentity(context).identity(); var cached = provider.resolveIdentity(context).identity(); - provider.invalidate(); + provider.invalidate(first); var refreshed = provider.resolveIdentity(context).identity(); assertThat(cached, is(first)); diff --git a/aws/client/aws-client-awsquery/src/main/java/software/amazon/smithy/java/aws/client/awsquery/AwsQueryClientProtocol.java b/aws/client/aws-client-awsquery/src/main/java/software/amazon/smithy/java/aws/client/awsquery/AwsQueryClientProtocol.java index 0c67cb6343..bf0d6bced4 100644 --- a/aws/client/aws-client-awsquery/src/main/java/software/amazon/smithy/java/aws/client/awsquery/AwsQueryClientProtocol.java +++ b/aws/client/aws-client-awsquery/src/main/java/software/amazon/smithy/java/aws/client/awsquery/AwsQueryClientProtocol.java @@ -10,6 +10,7 @@ import java.util.Objects; import software.amazon.smithy.aws.traits.protocols.AwsQueryErrorTrait; import software.amazon.smithy.aws.traits.protocols.AwsQueryTrait; +import software.amazon.smithy.java.client.core.CallContext; import software.amazon.smithy.java.client.core.ClientProtocol; import software.amazon.smithy.java.client.core.ClientProtocolFactory; import software.amazon.smithy.java.client.core.ProtocolSettings; @@ -130,6 +131,7 @@ public CallException parsePayload( ) { var deserializer = codec.createDeserializer(buffer); String code = XmlUtil.parseErrorCodeName(deserializer); + context.put(CallContext.RESPONSE_ERROR_CODE, code); // First, resolve @awsQueryError custom codes ShapeBuilder builder = null; diff --git a/aws/client/aws-client-awsquery/src/main/java/software/amazon/smithy/java/aws/client/awsquery/Ec2QueryClientProtocol.java b/aws/client/aws-client-awsquery/src/main/java/software/amazon/smithy/java/aws/client/awsquery/Ec2QueryClientProtocol.java index 52c7afa0b5..9bd49f69ef 100644 --- a/aws/client/aws-client-awsquery/src/main/java/software/amazon/smithy/java/aws/client/awsquery/Ec2QueryClientProtocol.java +++ b/aws/client/aws-client-awsquery/src/main/java/software/amazon/smithy/java/aws/client/awsquery/Ec2QueryClientProtocol.java @@ -9,6 +9,7 @@ import java.util.List; import java.util.Objects; import software.amazon.smithy.aws.traits.protocols.Ec2QueryTrait; +import software.amazon.smithy.java.client.core.CallContext; import software.amazon.smithy.java.client.core.ClientProtocol; import software.amazon.smithy.java.client.core.ClientProtocolFactory; import software.amazon.smithy.java.client.core.ProtocolSettings; @@ -134,6 +135,7 @@ public CallException parsePayload( ) { var deserializer = codec.createDeserializer(buffer); String code = XmlUtil.parseErrorCodeName(deserializer); + context.put(CallContext.RESPONSE_ERROR_CODE, code); var id = ShapeId.fromOptionalNamespace(serviceId.getNamespace(), code); ShapeBuilder builder = typeRegistry.createBuilder(id, ModeledException.class); diff --git a/aws/client/aws-client-core/src/main/java/software/amazon/smithy/java/aws/client/core/AwsCredentialChainPlugin.java b/aws/client/aws-client-core/src/main/java/software/amazon/smithy/java/aws/client/core/AwsCredentialChainPlugin.java index 7090de3f14..0160d30c0f 100644 --- a/aws/client/aws-client-core/src/main/java/software/amazon/smithy/java/aws/client/core/AwsCredentialChainPlugin.java +++ b/aws/client/aws-client-core/src/main/java/software/amazon/smithy/java/aws/client/core/AwsCredentialChainPlugin.java @@ -12,13 +12,15 @@ import software.amazon.smithy.java.client.core.ClientConfig; import software.amazon.smithy.java.client.core.ClientPlugin; import software.amazon.smithy.java.client.core.auth.scheme.AuthScheme; +import software.amazon.smithy.java.client.core.interceptors.ClientInterceptor; /** * A {@link ClientPlugin} that registers the AWS default credential chain on any client that uses an AWS auth scheme * (one whose {@link AuthScheme#identityClass()} is {@link AwsCredentialsIdentity}). * *

This plugin is wired into generated AWS clients by codegen. It is a no-op for clients that do not use - * AWS authentication or that already have an {@link AwsCredentialsIdentity} resolver registered. + * AWS authentication. When an AWS credentials resolver is already registered, the plugin preserves it and installs + * only the authentication-failure interceptor. * *

Users can also add it explicitly: * {@snippet lang="java" : @@ -31,6 +33,18 @@ * resolver directly instead of using this plugin. */ public final class AwsCredentialChainPlugin implements ClientPlugin { + /** + * Returns the interceptor that invalidates AWS credentials rejected by a service. + * + *

This is used by clients that provide their own credential resolver and do not use the default credential + * chain plugin. + * + * @return the AWS credential invalidation interceptor. + */ + public static ClientInterceptor invalidationInterceptor() { + return InvalidateCredentialsInterceptor.INSTANCE; + } + @Override public Phase getPluginPhase() { // Run after DEFAULTS so the client's region (and any other defaults) are populated on the config before we @@ -40,15 +54,21 @@ public Phase getPluginPhase() { @Override public void configureClient(ClientConfig.Builder config) { - if (needsAwsCredentials(config) && !hasAwsCredentialsResolver(config)) { + if (!needsAwsCredentials(config)) { + return; + } + if (!hasAwsCredentialsResolver(config)) { // Pass the client's configured region so credential providers that make a service call (STS, SSO) use // the same region as the client. Null when no region is configured, in which case the providers fall // back to the environment and profile. String region = config.context().get(RegionSetting.REGION); - var chain = IdentityChain.create(AwsCredentialsIdentity.class, null, region); - config.addIdentityResolver(chain); - config.addInterceptor(new InvalidateCredentialsInterceptor(chain)); + var resolver = new LeasedIdentityResolver<>( + AwsCredentialsIdentity.class, + () -> IdentityChain.create(AwsCredentialsIdentity.class, null, region)); + config.addIdentityResolver(resolver); + config.addOwnedResource(resolver::acquire); } + config.addInterceptor(InvalidateCredentialsInterceptor.INSTANCE); } private static boolean needsAwsCredentials(ClientConfig.Builder config) { diff --git a/aws/client/aws-client-core/src/main/java/software/amazon/smithy/java/aws/client/core/InvalidateCredentialsInterceptor.java b/aws/client/aws-client-core/src/main/java/software/amazon/smithy/java/aws/client/core/InvalidateCredentialsInterceptor.java index c426e2a26a..5356442425 100644 --- a/aws/client/aws-client-core/src/main/java/software/amazon/smithy/java/aws/client/core/InvalidateCredentialsInterceptor.java +++ b/aws/client/aws-client-core/src/main/java/software/amazon/smithy/java/aws/client/core/InvalidateCredentialsInterceptor.java @@ -6,36 +6,60 @@ package software.amazon.smithy.java.aws.client.core; import java.util.Set; +import software.amazon.smithy.java.auth.api.identity.Identity; import software.amazon.smithy.java.auth.api.identity.IdentityResolver; +import software.amazon.smithy.java.client.core.CallContext; import software.amazon.smithy.java.client.core.interceptors.ClientInterceptor; import software.amazon.smithy.java.client.core.interceptors.OutputHook; +import software.amazon.smithy.java.context.Context; +import software.amazon.smithy.java.core.error.CallException; import software.amazon.smithy.java.core.error.ModeledException; +import software.amazon.smithy.java.retries.api.RetrySafety; /** - * Interceptor that invalidates cached credentials when a service returns an expired- or - * invalid-credential error, so the stale credentials are cleared and the next request resolves - * fresh ones. + * Interceptor that notifies the request identity resolver when a service returns an expired- or + * invalid-credential error. */ final class InvalidateCredentialsInterceptor implements ClientInterceptor { - private final IdentityResolver resolver; + static final InvalidateCredentialsInterceptor INSTANCE = new InvalidateCredentialsInterceptor(); - // well-known error names. Ideally the wire would have signal so we don't need this. + // Well-known error names. Ideally the wire would have a signal so we do not need this. private static final Set EXPIRED_NAMES = Set.of("ExpiredToken", "InvalidToken"); - InvalidateCredentialsInterceptor(IdentityResolver resolver) { - this.resolver = resolver; - } + private InvalidateCredentialsInterceptor() {} @Override public void readAfterAttempt(OutputHook hook, RuntimeException error) { - if (shouldInvalidate(error)) { - resolver.invalidate(); + invalidate(hook.context(), error); + } + + void invalidate(Context context, RuntimeException error) { + if (!shouldInvalidate(context, error)) { + return; + } + + if (error instanceof CallException callException) { + callException.isRetrySafe(RetrySafety.NO); } + + var identity = context.get(CallContext.IDENTITY); + var resolver = context.get(CallContext.IDENTITY_RESOLVER); + if (identity == null || resolver == null || !resolver.identityType().isInstance(identity)) { + return; + } + invalidate(resolver, identity); } - private static boolean shouldInvalidate(RuntimeException error) { - // Check for well-known expired- or invalid-credential error names. + private static void invalidate(IdentityResolver resolver, Identity rejectedIdentity) { + resolver.invalidate(resolver.identityType().cast(rejectedIdentity)); + } + + private static boolean shouldInvalidate(Context context, RuntimeException error) { + String errorCode = context.get(CallContext.RESPONSE_ERROR_CODE); + if (errorCode != null && EXPIRED_NAMES.contains(errorCode)) { + return true; + } if (error instanceof ModeledException me) { var name = me.schema().id().getName(); return EXPIRED_NAMES.contains(name); diff --git a/aws/client/aws-client-core/src/main/java/software/amazon/smithy/java/aws/client/core/LeasedIdentityResolver.java b/aws/client/aws-client-core/src/main/java/software/amazon/smithy/java/aws/client/core/LeasedIdentityResolver.java new file mode 100644 index 0000000000..f898f74222 --- /dev/null +++ b/aws/client/aws-client-core/src/main/java/software/amazon/smithy/java/aws/client/core/LeasedIdentityResolver.java @@ -0,0 +1,123 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package software.amazon.smithy.java.aws.client.core; + +import java.util.Objects; +import java.util.function.Supplier; +import software.amazon.smithy.java.auth.api.identity.Identity; +import software.amazon.smithy.java.auth.api.identity.IdentityResolver; +import software.amazon.smithy.java.auth.api.identity.IdentityResult; +import software.amazon.smithy.java.context.Context; + +/** + * Shares a lazily created resolver while leases are active and recreates it after the last lease is closed. + */ +final class LeasedIdentityResolver implements IdentityResolver { + + private final Class identityType; + private final Supplier> factory; + private IdentityResolver delegate; + private AutoCloseable closeable; + private int leases; + + LeasedIdentityResolver(Class identityType, Supplier> factory) { + this.identityType = Objects.requireNonNull(identityType, "identityType"); + this.factory = Objects.requireNonNull(factory, "factory"); + } + + synchronized AutoCloseable acquire() { + if (delegate == null) { + IdentityResolver created = Objects.requireNonNull(factory.get(), "Identity resolver must not be null"); + try { + Class createdType = created.identityType(); + if (!identityType.isAssignableFrom(createdType)) { + throw new IllegalStateException( + "Identity resolver factory returned " + createdType.getName() + + " when " + identityType.getName() + " was required"); + } + if (!(created instanceof AutoCloseable createdCloseable)) { + throw new IllegalStateException("Leased identity resolver must be AutoCloseable"); + } + delegate = created; + closeable = createdCloseable; + } catch (RuntimeException | Error error) { + closeAfterFailedAcquire(created, error); + throw error; + } + } + leases++; + return new Lease(this); + } + + @Override + public IdentityResult resolveIdentity(Context requestProperties) { + return activeDelegate().resolveIdentity(requestProperties); + } + + @Override + public Class identityType() { + return identityType; + } + + @Override + public void invalidate(I rejectedIdentity) { + IdentityResolver current; + synchronized (this) { + current = delegate; + } + if (current != null) { + current.invalidate(rejectedIdentity); + } + } + + private synchronized IdentityResolver activeDelegate() { + if (delegate == null) { + throw new IllegalStateException("Identity resolver has no active client lease"); + } + return delegate; + } + + private synchronized void release() throws Exception { + if (leases == 0) { + throw new IllegalStateException("Identity resolver lease count is already zero"); + } + if (--leases == 0) { + AutoCloseable released = closeable; + delegate = null; + closeable = null; + released.close(); + } + } + + private static void closeAfterFailedAcquire(Object created, Throwable failure) { + if (created instanceof AutoCloseable closeable) { + try { + closeable.close(); + } catch (Exception closeError) { + failure.addSuppressed(closeError); + } catch (Error closeError) { + failure.addSuppressed(closeError); + } + } + } + + private static final class Lease implements AutoCloseable { + private final LeasedIdentityResolver owner; + private boolean closed; + + Lease(LeasedIdentityResolver owner) { + this.owner = owner; + } + + @Override + public synchronized void close() throws Exception { + if (!closed) { + closed = true; + owner.release(); + } + } + } +} diff --git a/aws/client/aws-client-core/src/main/java/software/amazon/smithy/java/aws/client/core/identity/EnvironmentVariableIdentityResolver.java b/aws/client/aws-client-core/src/main/java/software/amazon/smithy/java/aws/client/core/identity/EnvironmentVariableIdentityResolver.java index cce701a73b..bb2238a826 100644 --- a/aws/client/aws-client-core/src/main/java/software/amazon/smithy/java/aws/client/core/identity/EnvironmentVariableIdentityResolver.java +++ b/aws/client/aws-client-core/src/main/java/software/amazon/smithy/java/aws/client/core/identity/EnvironmentVariableIdentityResolver.java @@ -5,6 +5,8 @@ package software.amazon.smithy.java.aws.client.core.identity; +import java.util.Objects; +import java.util.concurrent.atomic.AtomicReference; import software.amazon.smithy.java.auth.api.identity.IdentityResult; import software.amazon.smithy.java.aws.auth.api.identity.AwsCredentialsIdentity; import software.amazon.smithy.java.aws.auth.api.identity.AwsCredentialsResolver; @@ -13,8 +15,8 @@ /** * {@link AwsCredentialsResolver} implementation that loads credentials from environment variables. * - *

This resolver reads environment variables once on first access and caches the result. Use - * {@link #invalidate()} to force re-reading (e.g., in tests). + *

This resolver reads its source once on first access and caches the result. Call + * {@link #invalidate(AwsCredentialsIdentity)} with the cached identity to force re-reading (e.g., in tests). * *

Expected environment variables: *

@@ -40,23 +42,29 @@ public final class EnvironmentVariableIdentityResolver implements AwsCredentials "Could not resolve an AWS identity using the AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY environment " + "variables"); - private volatile IdentityResult cached; + private final AtomicReference> cached = new AtomicReference<>(); @Override public IdentityResult resolveIdentity(Context requestProperties) { - IdentityResult result = cached; - if (result != null) { - return result; - } + while (true) { + IdentityResult result = cached.get(); + if (result != null) { + return result; + } - result = resolve(); - cached = result; - return result; + result = resolve(); + if (cached.compareAndSet(null, result)) { + return result; + } + } } @Override - public void invalidate() { - cached = null; + public void invalidate(AwsCredentialsIdentity rejectedIdentity) { + IdentityResult current = cached.get(); + if (current != null && Objects.equals(current.identity(), rejectedIdentity)) { + cached.compareAndSet(current, null); + } } private static IdentityResult resolve() { diff --git a/aws/client/aws-client-core/src/main/java/software/amazon/smithy/java/aws/client/core/identity/SystemPropertiesIdentityResolver.java b/aws/client/aws-client-core/src/main/java/software/amazon/smithy/java/aws/client/core/identity/SystemPropertiesIdentityResolver.java index 2e42590393..317027e826 100644 --- a/aws/client/aws-client-core/src/main/java/software/amazon/smithy/java/aws/client/core/identity/SystemPropertiesIdentityResolver.java +++ b/aws/client/aws-client-core/src/main/java/software/amazon/smithy/java/aws/client/core/identity/SystemPropertiesIdentityResolver.java @@ -5,6 +5,8 @@ package software.amazon.smithy.java.aws.client.core.identity; +import java.util.Objects; +import java.util.concurrent.atomic.AtomicReference; import software.amazon.smithy.java.auth.api.identity.IdentityResult; import software.amazon.smithy.java.aws.auth.api.identity.AwsCredentialsIdentity; import software.amazon.smithy.java.aws.auth.api.identity.AwsCredentialsResolver; @@ -13,8 +15,8 @@ /** * {@link AwsCredentialsResolver} implementation that loads credentials from Java system properties. * - *

This resolver reads system properties once on first access and caches the result. Use - * {@link #invalidate()} to force re-reading (e.g., in tests). + *

This resolver reads its source once on first access and caches the result. Call + * {@link #invalidate(AwsCredentialsIdentity)} with the cached identity to force re-reading (e.g., in tests). * *

Expected system properties: *

@@ -41,23 +43,29 @@ public final class SystemPropertiesIdentityResolver implements AwsCredentialsRes SystemPropertiesIdentityResolver.class, "Could not resolve AWS identity from the aws.accessKeyId and aws.secretAccessKey system properties"); - private volatile IdentityResult cached; + private final AtomicReference> cached = new AtomicReference<>(); @Override public IdentityResult resolveIdentity(Context requestProperties) { - IdentityResult result = cached; - if (result != null) { - return result; - } + while (true) { + IdentityResult result = cached.get(); + if (result != null) { + return result; + } - result = resolve(); - cached = result; - return result; + result = resolve(); + if (cached.compareAndSet(null, result)) { + return result; + } + } } @Override - public void invalidate() { - cached = null; + public void invalidate(AwsCredentialsIdentity rejectedIdentity) { + IdentityResult current = cached.get(); + if (current != null && Objects.equals(current.identity(), rejectedIdentity)) { + cached.compareAndSet(current, null); + } } private static IdentityResult resolve() { diff --git a/aws/client/aws-client-core/src/test/java/software/amazon/smithy/java/aws/client/core/InvalidateCredentialsInterceptorTest.java b/aws/client/aws-client-core/src/test/java/software/amazon/smithy/java/aws/client/core/InvalidateCredentialsInterceptorTest.java index 26d3f489b5..b0deb38cd6 100644 --- a/aws/client/aws-client-core/src/test/java/software/amazon/smithy/java/aws/client/core/InvalidateCredentialsInterceptorTest.java +++ b/aws/client/aws-client-core/src/test/java/software/amazon/smithy/java/aws/client/core/InvalidateCredentialsInterceptorTest.java @@ -6,65 +6,144 @@ package software.amazon.smithy.java.aws.client.core; import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertSame; +import static org.junit.jupiter.api.Assertions.assertTrue; import java.util.concurrent.atomic.AtomicInteger; import org.junit.jupiter.api.Test; +import software.amazon.smithy.java.auth.api.SignResult; +import software.amazon.smithy.java.auth.api.identity.Identity; import software.amazon.smithy.java.auth.api.identity.IdentityResolver; import software.amazon.smithy.java.auth.api.identity.IdentityResult; import software.amazon.smithy.java.aws.auth.api.identity.AwsCredentialsIdentity; +import software.amazon.smithy.java.client.core.CallContext; +import software.amazon.smithy.java.client.core.ClientConfig; +import software.amazon.smithy.java.client.core.auth.scheme.AuthScheme; import software.amazon.smithy.java.context.Context; +import software.amazon.smithy.java.core.error.CallException; import software.amazon.smithy.java.core.error.ModeledException; import software.amazon.smithy.java.core.schema.Schema; import software.amazon.smithy.java.core.serde.ShapeSerializer; +import software.amazon.smithy.java.retries.api.RetrySafety; import software.amazon.smithy.model.shapes.ShapeId; class InvalidateCredentialsInterceptorTest { + private static final AwsCredentialsIdentity IDENTITY = AwsCredentialsIdentity.create("AK", "SK"); + @Test void invalidatesOnExpiredToken() { var counter = new CountingResolver(); - var interceptor = new InvalidateCredentialsInterceptor(counter); + var interceptor = InvalidateCredentialsInterceptor.INSTANCE; - interceptor.readAfterAttempt(null, credentialError("ExpiredToken")); + interceptor.invalidate(context(counter, IDENTITY), credentialError("ExpiredToken")); assertEquals(1, counter.invalidateCount.get()); + assertSame(IDENTITY, counter.invalidatedIdentity); } @Test void invalidatesOnInvalidToken() { var counter = new CountingResolver(); - var interceptor = new InvalidateCredentialsInterceptor(counter); + var interceptor = InvalidateCredentialsInterceptor.INSTANCE; + + interceptor.invalidate(context(counter, IDENTITY), credentialError("InvalidToken")); + assertEquals(1, counter.invalidateCount.get()); + } + + @Test + void invalidatesUnmodeledResponseErrorCode() { + var counter = new CountingResolver(); + var context = context(counter, IDENTITY); + context.put(CallContext.RESPONSE_ERROR_CODE, "ExpiredToken"); + + InvalidateCredentialsInterceptor.INSTANCE.invalidate(context, new CallException("unmodeled error")); - interceptor.readAfterAttempt(null, credentialError("InvalidToken")); assertEquals(1, counter.invalidateCount.get()); + assertSame(IDENTITY, counter.invalidatedIdentity); } @Test void doesNotInvalidateOnOtherModeledError() { var counter = new CountingResolver(); - var interceptor = new InvalidateCredentialsInterceptor(counter); + var interceptor = InvalidateCredentialsInterceptor.INSTANCE; - interceptor.readAfterAttempt(null, credentialError("AccessDenied")); + interceptor.invalidate(context(counter, IDENTITY), credentialError("AccessDenied")); assertEquals(0, counter.invalidateCount.get()); } @Test void doesNotInvalidateOnNonModeledError() { var counter = new CountingResolver(); - var interceptor = new InvalidateCredentialsInterceptor(counter); + var interceptor = InvalidateCredentialsInterceptor.INSTANCE; - interceptor.readAfterAttempt(null, new RuntimeException("network error")); + interceptor.invalidate(context(counter, IDENTITY), new RuntimeException("network error")); assertEquals(0, counter.invalidateCount.get()); } @Test void doesNotInvalidateOnNull() { var counter = new CountingResolver(); - var interceptor = new InvalidateCredentialsInterceptor(counter); + var interceptor = InvalidateCredentialsInterceptor.INSTANCE; + + interceptor.invalidate(context(counter, IDENTITY), null); + assertEquals(0, counter.invalidateCount.get()); + } + + @Test + void doesNotInvalidateWithoutAttemptContext() { + InvalidateCredentialsInterceptor.INSTANCE.invalidate(Context.empty(), credentialError("ExpiredToken")); + } + + @Test + void credentialErrorsAreMarkedNonRetryableWithoutAttemptContext() { + var context = Context.create(); + context.put(CallContext.RESPONSE_ERROR_CODE, "ExpiredToken"); + var error = new CallException("expired credentials"); + error.isRetrySafe(RetrySafety.YES); + + InvalidateCredentialsInterceptor.INSTANCE.invalidate(context, error); + + assertEquals(RetrySafety.NO, error.isRetrySafe()); + } + + @Test + void doesNotInvalidateWhenIdentityTypeDoesNotMatch() { + var counter = new CountingResolver(); + var context = Context.create(); + context.put(CallContext.IDENTITY, new Identity() {}); + context.put(CallContext.IDENTITY_RESOLVER, counter); + + InvalidateCredentialsInterceptor.INSTANCE.invalidate(context, credentialError("ExpiredToken")); - interceptor.readAfterAttempt(null, null); assertEquals(0, counter.invalidateCount.get()); } + @Test + void pluginInstallsInterceptorForCustomResolver() { + var resolver = new CountingResolver(); + var scheme = AuthScheme.of( + ShapeId.from("aws.auth#sigv4"), + Object.class, + AwsCredentialsIdentity.class, + (request, identity, properties) -> new SignResult<>(request)); + var config = ClientConfig.builder() + .putSupportedAuthSchemes(scheme) + .addIdentityResolver(resolver); + + new AwsCredentialChainPlugin().configureClient(config); + + assertEquals(1, config.identityResolvers().size()); + assertSame(resolver, config.identityResolvers().get(0)); + assertTrue(config.interceptors().contains(InvalidateCredentialsInterceptor.INSTANCE)); + } + + private static Context context(IdentityResolver resolver, AwsCredentialsIdentity identity) { + var context = Context.create(); + context.put(CallContext.IDENTITY, identity); + context.put(CallContext.IDENTITY_RESOLVER, resolver); + return context; + } + private static RuntimeException credentialError(String errorName) { Schema schema = Schema.createString(ShapeId.from("com.example#" + errorName)); return new ModeledException(schema, errorName + " error") { @@ -83,10 +162,11 @@ public T getMemberValue(Schema member) { private static class CountingResolver implements IdentityResolver { final AtomicInteger invalidateCount = new AtomicInteger(0); + AwsCredentialsIdentity invalidatedIdentity; @Override public IdentityResult resolveIdentity(Context ctx) { - return IdentityResult.of(AwsCredentialsIdentity.create("AK", "SK")); + return IdentityResult.of(IDENTITY); } @Override @@ -95,7 +175,8 @@ public Class identityType() { } @Override - public void invalidate() { + public void invalidate(AwsCredentialsIdentity rejectedIdentity) { + invalidatedIdentity = rejectedIdentity; invalidateCount.incrementAndGet(); } } diff --git a/aws/client/aws-client-core/src/test/java/software/amazon/smithy/java/aws/client/core/LeasedIdentityResolverTest.java b/aws/client/aws-client-core/src/test/java/software/amazon/smithy/java/aws/client/core/LeasedIdentityResolverTest.java new file mode 100644 index 0000000000..22cf8d6e7a --- /dev/null +++ b/aws/client/aws-client-core/src/test/java/software/amazon/smithy/java/aws/client/core/LeasedIdentityResolverTest.java @@ -0,0 +1,123 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package software.amazon.smithy.java.aws.client.core; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertSame; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; +import org.junit.jupiter.api.Test; +import software.amazon.smithy.java.auth.api.identity.IdentityResolver; +import software.amazon.smithy.java.auth.api.identity.IdentityResult; +import software.amazon.smithy.java.auth.api.identity.TokenIdentity; +import software.amazon.smithy.java.aws.auth.api.identity.AwsCredentialsIdentity; +import software.amazon.smithy.java.context.Context; + +class LeasedIdentityResolverTest { + + @Test + void sharesDelegateUntilLastLeaseClosesAndRecreatesItLater() throws Exception { + var creations = new AtomicInteger(); + var closes = new AtomicInteger(); + var current = new AtomicReference(); + var resolver = new LeasedIdentityResolver<>(AwsCredentialsIdentity.class, () -> { + var created = new TestResolver(creations.incrementAndGet(), closes); + current.set(created); + return created; + }); + + assertThrows(IllegalStateException.class, () -> resolver.resolveIdentity(Context.empty())); + assertEquals(0, creations.get()); + AutoCloseable first = resolver.acquire(); + AutoCloseable second = resolver.acquire(); + + assertEquals(1, creations.get()); + assertEquals("AK1", resolver.resolveIdentity(Context.empty()).identity().accessKeyId()); + var rejected = AwsCredentialsIdentity.create("AK1", "SK"); + resolver.invalidate(rejected); + assertSame(rejected, current.get().invalidated); + + first.close(); + assertEquals(0, closes.get()); + second.close(); + assertEquals(1, closes.get()); + assertThrows(IllegalStateException.class, () -> resolver.resolveIdentity(Context.empty())); + + resolver.acquire().close(); + assertEquals(2, creations.get()); + assertEquals(2, closes.get()); + } + + @Test + @SuppressWarnings({"rawtypes", "unchecked"}) + void closesResolverCreatedWithIncompatibleIdentityType() { + var wrongType = new CloseableTokenResolver(); + var resolver = new LeasedIdentityResolver<>( + AwsCredentialsIdentity.class, + () -> (IdentityResolver) wrongType); + + var error = assertThrows(IllegalStateException.class, resolver::acquire); + + assertTrue(error.getMessage().contains(TokenIdentity.class.getName())); + assertTrue(wrongType.closed.get()); + } + + private static final class TestResolver + implements IdentityResolver, AutoCloseable { + private final int generation; + private final AtomicInteger closes; + private AwsCredentialsIdentity invalidated; + + TestResolver(int generation, AtomicInteger closes) { + this.generation = generation; + this.closes = closes; + } + + @Override + public IdentityResult resolveIdentity(Context requestProperties) { + return IdentityResult.of(AwsCredentialsIdentity.create("AK" + generation, "SK")); + } + + @Override + public Class identityType() { + return AwsCredentialsIdentity.class; + } + + @Override + public void invalidate(AwsCredentialsIdentity rejectedIdentity) { + invalidated = rejectedIdentity; + } + + @Override + public void close() { + closes.incrementAndGet(); + } + } + + private static final class CloseableTokenResolver + implements IdentityResolver, AutoCloseable { + private final AtomicBoolean closed = new AtomicBoolean(); + + @Override + public IdentityResult resolveIdentity(Context requestProperties) { + return IdentityResult.of(TokenIdentity.create("token")); + } + + @Override + public Class identityType() { + return TokenIdentity.class; + } + + @Override + public void close() { + closed.set(true); + } + } +} diff --git a/aws/client/aws-client-restxml/src/main/java/software/amazon/smithy/java/aws/client/restxml/RestXmlClientProtocol.java b/aws/client/aws-client-restxml/src/main/java/software/amazon/smithy/java/aws/client/restxml/RestXmlClientProtocol.java index 6874c51501..e754cdb2fc 100644 --- a/aws/client/aws-client-restxml/src/main/java/software/amazon/smithy/java/aws/client/restxml/RestXmlClientProtocol.java +++ b/aws/client/aws-client-restxml/src/main/java/software/amazon/smithy/java/aws/client/restxml/RestXmlClientProtocol.java @@ -10,6 +10,7 @@ import software.amazon.smithy.java.aws.events.AwsEventDecoderFactory; import software.amazon.smithy.java.aws.events.AwsEventEncoderFactory; import software.amazon.smithy.java.aws.events.AwsEventFrame; +import software.amazon.smithy.java.client.core.CallContext; import software.amazon.smithy.java.client.core.ClientProtocol; import software.amazon.smithy.java.client.core.ClientProtocolFactory; import software.amazon.smithy.java.client.core.ProtocolSettings; @@ -123,6 +124,7 @@ public CallException parsePayload( ) { var deserializer = codec.createDeserializer(buffer); String code = XmlUtil.parseErrorCodeName(deserializer); + context.put(CallContext.RESPONSE_ERROR_CODE, code); var nameSpace = serviceId.getNamespace(); var id = ShapeId.fromOptionalNamespace(nameSpace, code); var builder = typeRegistry.createBuilder(id, ModeledException.class); diff --git a/client/client-core/src/main/java/software/amazon/smithy/java/client/core/CallContext.java b/client/client-core/src/main/java/software/amazon/smithy/java/client/core/CallContext.java index e14e958de1..b2918b8ade 100644 --- a/client/client-core/src/main/java/software/amazon/smithy/java/client/core/CallContext.java +++ b/client/client-core/src/main/java/software/amazon/smithy/java/client/core/CallContext.java @@ -8,6 +8,7 @@ import java.util.HashSet; import java.util.Set; import software.amazon.smithy.java.auth.api.identity.Identity; +import software.amazon.smithy.java.auth.api.identity.IdentityResolver; import software.amazon.smithy.java.context.Context; import software.amazon.smithy.java.endpoints.Endpoint; import software.amazon.smithy.java.endpoints.EndpointResolver; @@ -38,6 +39,21 @@ public final class CallContext { */ public static final Context.Key IDENTITY = Context.key("Identity of the caller"); + /** + * The read-only identity resolver that produced {@link #IDENTITY} for the current attempt. + * + *

This is a read-only value; modifying this value has no effect on a request. + */ + public static final Context.Key> IDENTITY_RESOLVER = + Context.key("Identity resolver used for the request"); + + /** + * The normalized error code returned by the service for the current attempt. + * + *

This is a read-only value; modifying this value has no effect on a request. + */ + public static final Context.Key RESPONSE_ERROR_CODE = Context.key("Response error code"); + /** * The current number of retry attempts the client has made for the current call, starting at 1. * diff --git a/client/client-core/src/main/java/software/amazon/smithy/java/client/core/Client.java b/client/client-core/src/main/java/software/amazon/smithy/java/client/core/Client.java index 54987ae63e..ff2735712e 100644 --- a/client/client-core/src/main/java/software/amazon/smithy/java/client/core/Client.java +++ b/client/client-core/src/main/java/software/amazon/smithy/java/client/core/Client.java @@ -8,6 +8,7 @@ import java.io.Closeable; import java.io.IOException; import java.io.UncheckedIOException; +import java.util.ArrayList; import java.util.List; import java.util.concurrent.CompletionException; import java.util.function.Predicate; @@ -42,24 +43,41 @@ public abstract class Client implements Closeable { private final ClientInterceptor interceptor; private final IdentityResolvers identityResolvers; private final RetryStrategy retryStrategy; + private final List ownedResources; protected Client(Builder builder) { ClientConfig.Builder configBuilder = builder.configBuilder(); - this.config = configBuilder.build(); - this.pipeline = ClientPipeline.of(config.protocol(), config.transport()); - this.interceptor = config.interceptorChain(); - this.identityResolvers = IdentityResolvers.of(config.identityResolvers()); - this.typeRegistry = typeRegistry(); - if (config.retryStrategy() != null) { - this.retryStrategy = config.retryStrategy(); - } else { - this.retryStrategy = StandardRetryStrategy.create(); - } - // Claim this strategy, if successful any other client that attempts to claim it - // will fail preventing sharing strategies among unrelated clients. - if (retryStrategy instanceof Claimable c) { - c.claim(this); + ClientConfig createdConfig = configBuilder.build(); + List createdOwnedResources = createdConfig.acquireOwnedResources(); + ClientPipeline createdPipeline; + ClientInterceptor createdInterceptor; + IdentityResolvers createdIdentityResolvers; + TypeRegistry createdTypeRegistry; + RetryStrategy createdRetryStrategy; + try { + createdPipeline = ClientPipeline.of(createdConfig.protocol(), createdConfig.transport()); + createdInterceptor = createdConfig.interceptorChain(); + createdIdentityResolvers = IdentityResolvers.of(createdConfig.identityResolvers()); + createdTypeRegistry = typeRegistry(); + createdRetryStrategy = createdConfig.retryStrategy() != null + ? createdConfig.retryStrategy() + : StandardRetryStrategy.create(); + // Claim this strategy, if successful any other client that attempts to claim it + // will fail preventing sharing strategies among unrelated clients. + if (createdRetryStrategy instanceof Claimable c) { + c.claim(this); + } + } catch (RuntimeException | Error error) { + closeResources(createdOwnedResources, error); + throw error; } + this.config = createdConfig; + this.pipeline = createdPipeline; + this.interceptor = createdInterceptor; + this.identityResolvers = createdIdentityResolvers; + this.typeRegistry = createdTypeRegistry; + this.retryStrategy = createdRetryStrategy; + this.ownedResources = createdOwnedResources; } /** @@ -77,58 +95,91 @@ protected O call( ApiOperation operation, RequestOverrideConfig overrideConfig ) { - ClientPipeline callPipeline = pipeline; - IdentityResolvers callIdentityResolvers = identityResolvers; - ClientInterceptor callInterceptor = interceptor; - - //If there is an override config first apply that before sending to interceptors. - ClientConfig callConfig = config; - if (overrideConfig != null) { - callConfig = callConfig.withRequestOverride(overrideConfig); - } - ClientConfig afterInterceptionConfig = - callInterceptor.modifyBeforeCall(new CallHook<>(operation, callConfig, input)); - if (afterInterceptionConfig != null && afterInterceptionConfig != callConfig) { + List acquiredConfigs = new ArrayList<>(); + List callOwnedResources = new ArrayList<>(); + Throwable failure = null; + try { + ClientPipeline callPipeline = pipeline; + IdentityResolvers callIdentityResolvers = identityResolvers; + ClientInterceptor callInterceptor = interceptor; + + // If there is an override config first apply that before sending to interceptors. + ClientConfig callConfig = config; if (overrideConfig != null) { - callConfig = afterInterceptionConfig.withRequestOverride(overrideConfig); - } else { - callConfig = afterInterceptionConfig; + callConfig = callConfig.withRequestOverride(overrideConfig); + acquireCallResources(callConfig, acquiredConfigs, callOwnedResources); + } + ClientConfig afterInterceptionConfig = + callInterceptor.modifyBeforeCall(new CallHook<>(operation, callConfig, input)); + if (afterInterceptionConfig != null && afterInterceptionConfig != callConfig) { + acquireCallResources(afterInterceptionConfig, acquiredConfigs, callOwnedResources); + if (overrideConfig != null) { + callConfig = afterInterceptionConfig.withRequestOverride(overrideConfig); + acquireCallResources(callConfig, acquiredConfigs, callOwnedResources); + } else { + callConfig = afterInterceptionConfig; + } } - } - // Rebuild the pipeline, resolvers, etc if the config changed. - if (callConfig != config) { - callPipeline = ClientPipeline.of(callConfig.protocol(), callConfig.transport()); - callInterceptor = callConfig.interceptorChain(); - callIdentityResolvers = IdentityResolvers.of(callConfig.identityResolvers()); - } + // Rebuild the pipeline, resolvers, etc if the config changed. + if (callConfig != config) { + callPipeline = ClientPipeline.of(callConfig.protocol(), callConfig.transport()); + callInterceptor = callConfig.interceptorChain(); + callIdentityResolvers = IdentityResolvers.of(callConfig.identityResolvers()); + } - ClientCall call = new ClientCall<>( - input, - operation, - callConfig, - callPipeline, - callInterceptor, - callIdentityResolvers, - // Compose a type registry that adds the errors this operation can encounter. - TypeRegistry.compose(operation.errorRegistry(), typeRegistry), - retryStrategy); - - // Make the running client available to interceptors that need to re-enter. - call.context.put(ClientContext.CLIENT, this); - - if (!callInterceptor.interceptCalls()) { - return callPipeline.send(call); + ClientCall call = new ClientCall<>( + input, + operation, + callConfig, + callPipeline, + callInterceptor, + callIdentityResolvers, + // Compose a type registry that adds the errors this operation can encounter. + TypeRegistry.compose(operation.errorRegistry(), typeRegistry), + retryStrategy); + + // Make the running client available to interceptors that need to re-enter. + call.context.put(ClientContext.CLIENT, this); + + if (!callInterceptor.interceptCalls()) { + return callPipeline.send(call); + } + + // Build an InputHook for interceptCall and a terminal Invoker that sends the + // (possibly input-substituted) call through the pipeline. + InputHook hook = new InputHook<>(operation, call.context, input); + return callInterceptor.interceptCall(hook, h -> { + return h.input() == call.input + ? call.pipeline.send(call) + : call.pipeline.send(new ClientCall<>(call, h.input())); + }); + } catch (RuntimeException | Error error) { + failure = error; + throw error; + } finally { + RuntimeException closeFailure = closeResources(callOwnedResources, failure); + if (failure == null && closeFailure != null) { + throw closeFailure; + } } + } - // Build an InputHook for interceptCall and a terminal Invoker that sends the - // (possibly input-substituted) call through the pipeline. - InputHook hook = new InputHook<>(operation, call.context, input); - return callInterceptor.interceptCall(hook, h -> { - return h.input() == call.input - ? call.pipeline.send(call) - : call.pipeline.send(new ClientCall<>(call, h.input())); - }); + private void acquireCallResources( + ClientConfig callConfig, + List acquiredConfigs, + List callOwnedResources + ) { + if (callConfig == config) { + return; + } + for (ClientConfig acquiredConfig : acquiredConfigs) { + if (acquiredConfig == callConfig) { + return; + } + } + callOwnedResources.addAll(callConfig.acquireOwnedResources()); + acquiredConfigs.add(callConfig); } /** @@ -150,15 +201,46 @@ public ClientConfig config() { } /** - * Closes the transport used by this client. + * Closes resources owned by this client and its transport. */ @Override public void close() { + RuntimeException failure = closeResources(ownedResources, null); try { config.transport().close(); } catch (IOException e) { - throw new UncheckedIOException(e); + failure = recordCloseFailure(failure, new UncheckedIOException(e)); + } + if (failure != null) { + throw failure; + } + } + + private static RuntimeException closeResources(List resources, Throwable priorFailure) { + RuntimeException closeFailure = null; + for (AutoCloseable resource : resources) { + try { + resource.close(); + } catch (Exception error) { + if (priorFailure == null) { + closeFailure = recordCloseFailure(closeFailure, error); + } else { + priorFailure.addSuppressed(error); + } + } + } + return closeFailure; + } + + private static RuntimeException recordCloseFailure(RuntimeException current, Exception failure) { + RuntimeException runtimeFailure = failure instanceof RuntimeException runtime + ? runtime + : new RuntimeException(failure); + if (current == null) { + return runtimeFailure; } + current.addSuppressed(runtimeFailure); + return current; } /** @@ -212,6 +294,7 @@ public B withConfiguration(ClientConfig config) { .endpointResolver(config.endpointResolver()) .authSchemeResolver(config.authSchemeResolver()) .identityResolvers(config.identityResolvers()); + config.copyOwnedResourcesTo(configBuilder); config.interceptors().forEach(configBuilder::addInterceptor); config.supportedAuthSchemes().forEach(configBuilder::putSupportedAuthSchemes); configBuilder.putAllConfig(config.context()); diff --git a/client/client-core/src/main/java/software/amazon/smithy/java/client/core/ClientConfig.java b/client/client-core/src/main/java/software/amazon/smithy/java/client/core/ClientConfig.java index b8628ba8b4..4617125b44 100644 --- a/client/client-core/src/main/java/software/amazon/smithy/java/client/core/ClientConfig.java +++ b/client/client-core/src/main/java/software/amazon/smithy/java/client/core/ClientConfig.java @@ -15,6 +15,7 @@ import java.util.Objects; import java.util.Set; import java.util.function.Predicate; +import java.util.function.Supplier; import software.amazon.smithy.java.auth.api.identity.Identity; import software.amazon.smithy.java.auth.api.identity.IdentityResolver; import software.amazon.smithy.java.client.core.auth.scheme.AuthScheme; @@ -28,6 +29,7 @@ import software.amazon.smithy.java.logging.InternalLogger; import software.amazon.smithy.java.retries.api.RetryStrategy; import software.amazon.smithy.model.shapes.ShapeId; +import software.amazon.smithy.utils.SmithyInternalApi; /** * An immutable representation of configurations of a {@link Client}. @@ -54,6 +56,7 @@ public final class ClientConfig { private final ApiService service; private final RetryStrategy retryStrategy; private final String retryScope; + private final List> ownedResources; private final Set> appliedPluginClasses; private ClientConfig(Builder builder) { @@ -99,6 +102,7 @@ private ClientConfig(Builder builder) { this.retryStrategy = builder.retryStrategy; this.retryScope = builder.retryScope; + this.ownedResources = List.copyOf(builder.ownedResources); this.context = Context.unmodifiableCopy(builder.context); this.service = Objects.requireNonNull(builder.service, "Missing required service schema"); @@ -234,6 +238,29 @@ String retryScope() { return retryScope; } + List acquireOwnedResources() { + List leases = new ArrayList<>(ownedResources.size()); + try { + for (Supplier resource : ownedResources) { + leases.add(Objects.requireNonNull(resource.get(), "Client-owned resource lease must not be null")); + } + return leases; + } catch (RuntimeException | Error error) { + for (AutoCloseable lease : leases) { + try { + lease.close(); + } catch (Exception | Error closeError) { + error.addSuppressed(closeError); + } + } + throw error; + } + } + + void copyOwnedResourcesTo(Builder builder) { + builder.ownedResources.addAll(ownedResources); + } + /** * Create a new builder to build {@link ClientConfig}. * @@ -329,6 +356,7 @@ public static final class Builder { private final Context context = Context.create(); private RetryStrategy retryStrategy; private String retryScope; + private final List> ownedResources = new ArrayList<>(); private Predicate pluginPredicate = p -> true; private final Map, ClientPlugin> plugins = new LinkedHashMap<>(); // Mutable set that tracks which plugin classes have been applied to this builder @@ -351,6 +379,7 @@ private Builder copyBuilder() { context.copyTo(builder.context); builder.retryStrategy = retryStrategy; builder.retryScope = retryScope; + builder.ownedResources.addAll(ownedResources); builder.plugins.putAll(plugins); builder.pluginPredicate = pluginPredicate; builder.appliedPluginClasses.addAll(appliedPluginClasses); @@ -434,6 +463,23 @@ public String retryScope() { return retryScope; } + /** + * Adds a resource lease factory created as part of client configuration. + * + *

The factory is invoked once for each client or request configuration that uses this configuration. The + * returned lease is closed when that client, request, or failed initialization releases its resources. + * Factories can share an underlying resource by reference counting their returned leases, and must support + * acquisition after all previously returned leases have been closed so configurations remain reusable. + * + * @param resource resource lease factory. + * @return the builder. + */ + @SmithyInternalApi + public Builder addOwnedResource(Supplier resource) { + ownedResources.add(Objects.requireNonNull(resource, "resource")); + return this; + } + /** * Set the service schema. * @@ -695,4 +741,5 @@ public ClientConfig build() { return new ClientConfig(this); } } + } diff --git a/client/client-core/src/main/java/software/amazon/smithy/java/client/core/ClientPipeline.java b/client/client-core/src/main/java/software/amazon/smithy/java/client/core/ClientPipeline.java index 1cff2656c8..012032a8b3 100644 --- a/client/client-core/src/main/java/software/amazon/smithy/java/client/core/ClientPipeline.java +++ b/client/client-core/src/main/java/software/amazon/smithy/java/client/core/ClientPipeline.java @@ -11,6 +11,7 @@ import java.util.StringJoiner; import software.amazon.smithy.java.auth.api.SignResult; import software.amazon.smithy.java.auth.api.identity.Identity; +import software.amazon.smithy.java.auth.api.identity.IdentityResolver; import software.amazon.smithy.java.auth.api.identity.IdentityResolvers; import software.amazon.smithy.java.auth.api.identity.IdentityResult; import software.amazon.smithy.java.client.core.auth.scheme.AuthScheme; @@ -153,6 +154,7 @@ private O doSendOrR RequestHook requestHook ) { var request = requestHook.request(); + call.context.put(CallContext.RESPONSE_ERROR_CODE, null); // 8.a. Interceptors: Invoke ReadBeforeAttempt. call.interceptor.readBeforeAttempt(requestHook); @@ -178,12 +180,15 @@ private O afterIden // This throws if no identity was found. var identity = identityResult.unwrap(); call.context.put(CallContext.IDENTITY, identity); + call.context.put(CallContext.IDENTITY_RESOLVER, resolvedAuthScheme.identityResolver()); Endpoint endpoint = resolveEndpoint(call); call.context.put(CallContext.ENDPOINT, endpoint); // Augment or swap the resolved auth scheme based on the endpoint's authSchemes property. resolvedAuthScheme = applyEndpointAuthSchemeOverrides(call, endpoint, resolvedAuthScheme); + call.context.put(CallContext.IDENTITY, resolvedAuthScheme.identity().unwrap()); + call.context.put(CallContext.IDENTITY_RESOLVER, resolvedAuthScheme.identityResolver()); RequestT req = protocol.setServiceEndpoint(requestHook.request(), endpoint); var signResult = resolvedAuthScheme.sign(req); @@ -272,12 +277,17 @@ private ResolvedScheme createR if (identityResolver == null) { return null; } - return new ResolvedScheme<>(signerProperties, authScheme, identityResolver.resolveIdentity(identityProperties)); + return new ResolvedScheme<>( + signerProperties, + authScheme, + identityResolver, + identityResolver.resolveIdentity(identityProperties)); } private record ResolvedScheme( Context signerProperties, AuthScheme authScheme, + IdentityResolver identityResolver, IdentityResult identity) { public SignResult sign(RequestT request) { // Throws when no identity is found. @@ -381,7 +391,11 @@ private ResolvedScheme mergeOv for (var key : overrides) { merged.put((Context.Key) key, endpointAuthScheme.property(key)); } - return new ResolvedScheme<>(merged, resolvedScheme.authScheme(), resolvedScheme.identity()); + return new ResolvedScheme<>( + merged, + resolvedScheme.authScheme(), + resolvedScheme.identityResolver(), + resolvedScheme.identity()); } private O deserialize( diff --git a/client/client-core/src/test/java/software/amazon/smithy/java/client/core/ClientPipelineTest.java b/client/client-core/src/test/java/software/amazon/smithy/java/client/core/ClientPipelineTest.java index f2403b9944..d52bbb9c88 100644 --- a/client/client-core/src/test/java/software/amazon/smithy/java/client/core/ClientPipelineTest.java +++ b/client/client-core/src/test/java/software/amazon/smithy/java/client/core/ClientPipelineTest.java @@ -23,10 +23,14 @@ import software.amazon.smithy.java.auth.api.identity.Identity; import software.amazon.smithy.java.auth.api.identity.IdentityResolver; import software.amazon.smithy.java.auth.api.identity.IdentityResult; +import software.amazon.smithy.java.auth.api.identity.TokenIdentity; import software.amazon.smithy.java.aws.client.restjson.RestJsonClientProtocol; import software.amazon.smithy.java.client.core.auth.scheme.AuthScheme; import software.amazon.smithy.java.client.core.auth.scheme.AuthSchemeOption; import software.amazon.smithy.java.client.core.auth.scheme.AuthSchemeResolver; +import software.amazon.smithy.java.client.core.interceptors.ClientInterceptor; +import software.amazon.smithy.java.client.core.interceptors.OutputHook; +import software.amazon.smithy.java.client.core.interceptors.RequestHook; import software.amazon.smithy.java.client.http.JavaHttpClientTransport; import software.amazon.smithy.java.client.http.mock.MockPlugin; import software.amazon.smithy.java.client.http.mock.MockQueue; @@ -135,12 +139,13 @@ public Builder toBuilder() { public void canRetryRequests() { var service = ShapeId.from("smithy.example#Sprockets"); var calls = new ArrayList<>(); + var responseErrorCodes = new ArrayList(); var mockQueue = new MockQueue() .enqueue( HttpResponse.create() .setStatusCode(429) - .setBody(DataStream.ofString("{\"__type\":\"InvalidSprocketId\"}")) + .setBody(DataStream.ofString("{\"__type\":\"ExpiredToken\"}")) .toUnmodifiable()) .enqueue( HttpResponse.create() @@ -153,6 +158,15 @@ public void canRetryRequests() { .serviceId(service) .model(MODEL) .addPlugin(mock) + .addInterceptor(new ClientInterceptor() { + @Override + public void readAfterAttempt( + OutputHook hook, + RuntimeException error + ) { + responseErrorCodes.add(hook.context().get(CallContext.RESPONSE_ERROR_CODE)); + } + }) .endpointResolver(EndpointResolver.staticEndpoint("https://localhost:8081")) .authSchemeResolver(AuthSchemeResolver.NO_AUTH) .retryStrategy(new RetryStrategy() { @@ -199,6 +213,9 @@ public Builder toBuilder() { assertThat(response.getMember("id").asString(), equalTo("1")); assertThat(response, instanceOf(Document.class)); assertThat(calls, contains("Acquire", "Refresh", "Success: 1")); + assertThat(responseErrorCodes.size(), equalTo(2)); + assertThat(responseErrorCodes.get(0), equalTo("ExpiredToken")); + Assertions.assertNull(responseErrorCodes.get(1)); } @Test @@ -207,6 +224,12 @@ public void endpointAuthSchemeOverridesAugmentSignerProperties() { var testSchemeId = ShapeId.from("smithy.test#testAuth"); var TEST_KEY = Context.key("test-signing-override"); var capturedProperties = new AtomicReference(); + var capturedIdentity = new AtomicReference(); + var capturedResolver = new AtomicReference>(); + var defaultIdentity = new Identity() {}; + var overrideIdentity = new Identity() {}; + var defaultResolver = identityResolver(defaultIdentity); + var overrideResolver = identityResolver(overrideIdentity); // Auth scheme with a signer that captures the properties it receives. var testScheme = AuthScheme.of( @@ -242,22 +265,24 @@ public void endpointAuthSchemeOverridesAugmentSignerProperties() { .endpointResolver(endpointResolver) .authSchemeResolver(params -> List.of(new AuthSchemeOption(testSchemeId))) .putSupportedAuthSchemes(testScheme) - .addIdentityResolver(new IdentityResolver<>() { + .addIdentityResolver(defaultResolver) + .addInterceptor(new ClientInterceptor() { @Override - public IdentityResult resolveIdentity(Context requestProperties) { - return IdentityResult.of(new Identity() {}); - } - - @Override - public Class identityType() { - return Identity.class; + public void readAfterSigning(RequestHook hook) { + capturedIdentity.set(hook.context().get(CallContext.IDENTITY)); + capturedResolver.set(hook.context().get(CallContext.IDENTITY_RESOLVER)); } }) .build(); - client.call("GetSprocket", Document.ofObject(Map.of("id", "1"))); + client.call( + "GetSprocket", + Document.ofObject(Map.of("id", "1")), + RequestOverrideConfig.builder().addIdentityResolver(overrideResolver).build()); assertThat(capturedProperties.get().get(TEST_KEY), equalTo("overridden-value")); + assertThat(capturedIdentity.get(), is(overrideIdentity)); + assertThat(capturedResolver.get(), is(overrideResolver)); } @Test @@ -278,10 +303,24 @@ public void endpointAuthSchemeSwapsToDifferentScheme() { }); var capturedProperties = new AtomicReference(); + var capturedIdentity = new AtomicReference(); + var capturedResolver = new AtomicReference>(); + var targetIdentity = TokenIdentity.create("target"); + IdentityResolver targetResolver = new IdentityResolver<>() { + @Override + public IdentityResult resolveIdentity(Context requestProperties) { + return IdentityResult.of(targetIdentity); + } + + @Override + public Class identityType() { + return TokenIdentity.class; + } + }; var targetScheme = AuthScheme.of( endpointTargetId, HttpRequest.class, - Identity.class, + TokenIdentity.class, (request, identity, properties) -> { capturedProperties.set(properties); return new SignResult<>(request); @@ -320,12 +359,21 @@ public IdentityResult resolveIdentity(Context requestProperties) { public Class identityType() { return Identity.class; } + }, targetResolver) + .addInterceptor(new ClientInterceptor() { + @Override + public void readAfterSigning(RequestHook hook) { + capturedIdentity.set(hook.context().get(CallContext.IDENTITY)); + capturedResolver.set(hook.context().get(CallContext.IDENTITY_RESOLVER)); + } }) .build(); client.call("GetSprocket", Document.ofObject(Map.of("id", "1"))); assertThat(capturedProperties.get().get(TEST_KEY), equalTo("from-endpoint")); + assertThat(capturedIdentity.get(), is(targetIdentity)); + assertThat(capturedResolver.get(), is(targetResolver)); } @Test @@ -381,6 +429,20 @@ public Class identityType() { assertThat(ex.getMessage(), Matchers.containsString(endpointTargetId.toString())); } + private static IdentityResolver identityResolver(Identity identity) { + return new IdentityResolver<>() { + @Override + public IdentityResult resolveIdentity(Context requestProperties) { + return IdentityResult.of(identity); + } + + @Override + public Class identityType() { + return Identity.class; + } + }; + } + private static final class Token implements RetryToken { int retry; diff --git a/client/client-core/src/test/java/software/amazon/smithy/java/client/core/ClientTest.java b/client/client-core/src/test/java/software/amazon/smithy/java/client/core/ClientTest.java index a4a88a9050..6a46b90a3c 100644 --- a/client/client-core/src/test/java/software/amazon/smithy/java/client/core/ClientTest.java +++ b/client/client-core/src/test/java/software/amazon/smithy/java/client/core/ClientTest.java @@ -18,6 +18,8 @@ import java.util.HashMap; import java.util.List; import java.util.Map; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.Test; import software.amazon.smithy.java.aws.client.restjson.RestJsonClientProtocol; @@ -46,6 +48,7 @@ import software.amazon.smithy.java.endpoints.EndpointContext; import software.amazon.smithy.java.endpoints.EndpointResolver; import software.amazon.smithy.java.http.api.HttpResponse; +import software.amazon.smithy.java.retries.StandardRetryStrategy; import software.amazon.smithy.model.Model; import software.amazon.smithy.model.shapes.ShapeId; @@ -110,6 +113,172 @@ public void pluginIntegrationTest() { assertThat(applied, equalTo(EXPECTED_PLUGIN_CLASSES)); } + @Test + public void closesResourcesCreatedByPlugins() { + var closed = new AtomicBoolean(); + DynamicClient client = DynamicClient.builder() + .model(MODEL) + .serviceId(SERVICE) + .protocol(new RestJsonClientProtocol(SERVICE)) + .addPlugin(config -> config.addOwnedResource(() -> () -> closed.set(true))) + .endpointResolver(EndpointResolver.staticEndpoint("http://localhost")) + .build(); + + client.close(); + + Assertions.assertTrue(closed.get()); + } + + @Test + public void closesSharedPluginResourcesAfterLastClient() { + var resource = new ReusableSharedResource(); + DynamicClient first = DynamicClient.builder() + .model(MODEL) + .serviceId(SERVICE) + .protocol(new RestJsonClientProtocol(SERVICE)) + .addPlugin(config -> config.addOwnedResource(resource::acquire)) + .endpointResolver(EndpointResolver.staticEndpoint("http://localhost")) + .build(); + DynamicClient second = DynamicClient.builder() + .model(MODEL) + .serviceId(SERVICE) + .withConfiguration(first.config()) + .build(); + + first.close(); + Assertions.assertEquals(0, resource.closes); + + second.close(); + Assertions.assertEquals(1, resource.closes); + } + + @Test + public void resourceBearingConfigurationCanBeReusedAfterLastClientCloses() { + var resource = new ReusableSharedResource(); + DynamicClient first = DynamicClient.builder() + .model(MODEL) + .serviceId(SERVICE) + .protocol(new RestJsonClientProtocol(SERVICE)) + .addPlugin(config -> config.addOwnedResource(resource::acquire)) + .endpointResolver(EndpointResolver.staticEndpoint("http://localhost")) + .build(); + ClientConfig savedConfig = first.config(); + + first.close(); + Assertions.assertEquals(1, resource.generations); + Assertions.assertEquals(1, resource.closes); + + DynamicClient second = DynamicClient.builder() + .model(MODEL) + .serviceId(SERVICE) + .withConfiguration(savedConfig) + .build(); + Assertions.assertEquals(2, resource.generations); + + second.close(); + Assertions.assertEquals(2, resource.closes); + } + + @Test + public void doesNotAcquirePluginResourcesWhenConfigurationFails() { + var pluginApplied = new AtomicBoolean(); + var acquisitions = new AtomicInteger(); + + Assertions.assertThrows( + NullPointerException.class, + () -> ClientConfig.builder() + .addPlugin(config -> { + pluginApplied.set(true); + config.addOwnedResource(() -> { + acquisitions.incrementAndGet(); + return () -> {}; + }); + }) + .build()); + + Assertions.assertTrue(pluginApplied.get()); + Assertions.assertEquals(0, acquisitions.get()); + } + + @Test + public void closesResourcesCreatedByRequestOverridePlugins() { + var queue = new MockQueue(); + queue.enqueue(HttpResponse.create().setStatusCode(200).toUnmodifiable()); + var closed = new AtomicBoolean(); + DynamicClient client = DynamicClient.builder() + .model(MODEL) + .serviceId(SERVICE) + .protocol(new RestJsonClientProtocol(SERVICE)) + .addPlugin(MockPlugin.builder().addQueue(queue).build()) + .endpointResolver(EndpointResolver.staticEndpoint("http://localhost")) + .authSchemeResolver(AuthSchemeResolver.NO_AUTH) + .build(); + var override = RequestOverrideConfig.builder() + .addPlugin(config -> config.addOwnedResource(() -> () -> closed.set(true))) + .build(); + + client.call("GetSprocket", Document.ofObject(Map.of()), override); + + Assertions.assertTrue(closed.get()); + client.close(); + } + + @Test + public void closesRequestOverrideResourcesWhenModifyBeforeCallFails() { + var closed = new AtomicBoolean(); + DynamicClient client = DynamicClient.builder() + .model(MODEL) + .serviceId(SERVICE) + .protocol(new RestJsonClientProtocol(SERVICE)) + .addPlugin(config -> config.addInterceptor(new ClientInterceptor() { + @Override + public ClientConfig modifyBeforeCall(CallHook hook) { + throw new IllegalStateException("stop"); + } + })) + .endpointResolver(EndpointResolver.staticEndpoint("http://localhost")) + .build(); + var override = RequestOverrideConfig.builder() + .addPlugin(config -> config.addOwnedResource(() -> () -> closed.set(true))) + .build(); + try { + Assertions.assertThrows( + IllegalStateException.class, + () -> client.call("GetSprocket", Document.ofObject(Map.of()), override)); + Assertions.assertTrue(closed.get()); + } finally { + client.close(); + } + } + + @Test + public void closesOwnedResourcesWhenClientInitializationFails() { + var retryStrategy = StandardRetryStrategy.create(); + DynamicClient first = DynamicClient.builder() + .model(MODEL) + .serviceId(SERVICE) + .protocol(new RestJsonClientProtocol(SERVICE)) + .retryStrategy(retryStrategy) + .endpointResolver(EndpointResolver.staticEndpoint("http://localhost")) + .build(); + var closed = new AtomicBoolean(); + try { + Assertions.assertThrows( + IllegalStateException.class, + () -> DynamicClient.builder() + .model(MODEL) + .serviceId(SERVICE) + .protocol(new RestJsonClientProtocol(SERVICE)) + .retryStrategy(retryStrategy) + .addPlugin(config -> config.addOwnedResource(() -> () -> closed.set(true))) + .endpointResolver(EndpointResolver.staticEndpoint("http://localhost")) + .build()); + Assertions.assertTrue(closed.get()); + } finally { + first.close(); + } + } + @Test public void canFilterPlugins() { List> applied = new ArrayList<>(); @@ -500,4 +669,41 @@ public void readBeforeExecution(InputHook hook) { var seen = (StructDocument) seenInputs.get(0); assertThat(seen.getMember("id").asString(), equalTo("swapped")); } + + private static final class ReusableSharedResource { + private int leases; + private int generations; + private int closes; + + synchronized AutoCloseable acquire() { + if (leases == 0) { + generations++; + } + leases++; + return new ResourceLease(this); + } + + private synchronized void release() { + if (--leases == 0) { + closes++; + } + } + } + + private static final class ResourceLease implements AutoCloseable { + private final ReusableSharedResource resource; + private boolean closed; + + ResourceLease(ReusableSharedResource resource) { + this.resource = resource; + } + + @Override + public synchronized void close() { + if (!closed) { + closed = true; + resource.release(); + } + } + } } diff --git a/client/client-http/src/main/java/software/amazon/smithy/java/client/http/AmznErrorHeaderExtractor.java b/client/client-http/src/main/java/software/amazon/smithy/java/client/http/AmznErrorHeaderExtractor.java index 51d1227b52..9c0405a2b6 100644 --- a/client/client-http/src/main/java/software/amazon/smithy/java/client/http/AmznErrorHeaderExtractor.java +++ b/client/client-http/src/main/java/software/amazon/smithy/java/client/http/AmznErrorHeaderExtractor.java @@ -22,6 +22,11 @@ public boolean hasHeader(HttpResponse response) { return response.headers().firstValue(ERROR_HEADER) != null; } + @Override + public String errorCode(HttpResponse response) { + return ErrorTypeUtils.removeNamespaceAndUri(response.headers().firstValue(ERROR_HEADER)); + } + @Override public ShapeId resolveId(HttpResponse response, String serviceNamespace, TypeRegistry registry) { var header = response.headers().firstValue(ERROR_HEADER); diff --git a/client/client-http/src/main/java/software/amazon/smithy/java/client/http/HttpErrorDeserializer.java b/client/client-http/src/main/java/software/amazon/smithy/java/client/http/HttpErrorDeserializer.java index 304074f96a..3c5920480b 100644 --- a/client/client-http/src/main/java/software/amazon/smithy/java/client/http/HttpErrorDeserializer.java +++ b/client/client-http/src/main/java/software/amazon/smithy/java/client/http/HttpErrorDeserializer.java @@ -7,6 +7,7 @@ import java.nio.ByteBuffer; import java.util.Objects; +import software.amazon.smithy.java.client.core.CallContext; import software.amazon.smithy.java.context.Context; import software.amazon.smithy.java.core.error.CallException; import software.amazon.smithy.java.core.error.ErrorFault; @@ -49,6 +50,16 @@ public interface HeaderErrorExtractor { * @return the resolved builder, or null if no builder could be found. */ ShapeId resolveId(HttpResponse response, String serviceNamespace, TypeRegistry registry); + + /** + * Extract the normalized wire error code, independent of whether it maps to a modeled error. + * + * @param response Response to inspect. + * @return the normalized error code, or null if unavailable. + */ + default String errorCode(HttpResponse response) { + return null; + } } /** @@ -144,7 +155,10 @@ default CallException parsePayload( ) { var document = codec.createDeserializer(buffer).readDocument(); var id = extractErrorType(document, serviceId.getNamespace()); - var builder = typeRegistry.createBuilder(id, ModeledException.class); + if (id != null) { + context.put(CallContext.RESPONSE_ERROR_CODE, id.getName()); + } + var builder = id == null ? null : typeRegistry.createBuilder(id, ModeledException.class); if (builder != null) { return knownErrorFactory.createErrorFromDocument( context, @@ -258,6 +272,7 @@ public CallException createError( TypeRegistry typeRegistry, HttpResponse response ) { + context.put(CallContext.RESPONSE_ERROR_CODE, null); var hasErrorHeader = headerErrorExtractor.hasHeader(response); if (hasErrorHeader) { @@ -293,6 +308,7 @@ private CallException makeErrorFromHeader( TypeRegistry typeRegistry, HttpResponse response ) { + context.put(CallContext.RESPONSE_ERROR_CODE, headerErrorExtractor.errorCode(response)); // The content can be parsed directly here rather than through an intermediate document like with __type. var id = headerErrorExtractor.resolveId(response, serviceId.getNamespace(), typeRegistry); var builder = id == null ? null : typeRegistry.createBuilder(id, ModeledException.class); diff --git a/client/client-http/src/test/java/software/amazon/smithy/java/client/http/AmznErrorHeaderExtractorTest.java b/client/client-http/src/test/java/software/amazon/smithy/java/client/http/AmznErrorHeaderExtractorTest.java index 06b8652c0a..64bff346a9 100644 --- a/client/client-http/src/test/java/software/amazon/smithy/java/client/http/AmznErrorHeaderExtractorTest.java +++ b/client/client-http/src/test/java/software/amazon/smithy/java/client/http/AmznErrorHeaderExtractorTest.java @@ -131,6 +131,19 @@ public void resolvesToServiceErrorWhenAbsoluteNotFound() { assertThat(extractor.resolveId(response, "com.foo", registry), equalTo(ShapeId.from("com.foo#Bam"))); } + @Test + public void extractsUnmodeledErrorCode() { + var extractor = new AmznErrorHeaderExtractor(); + var response = HttpResponse.create() + .setStatusCode(400) + .setHeaders(HttpHeaders.of(Map.of( + "x-amzn-errortype", + List.of("other#ExpiredToken:http://example.com")))) + .toUnmodifiable(); + + assertThat(extractor.errorCode(response), equalTo("ExpiredToken")); + } + @Test public void returnsNullWhenNoTypeFound() { var extractor = new AmznErrorHeaderExtractor(); diff --git a/client/client-http/src/test/java/software/amazon/smithy/java/client/http/HttpErrorDeserializerTest.java b/client/client-http/src/test/java/software/amazon/smithy/java/client/http/HttpErrorDeserializerTest.java index 5ef9d9c912..d0098d564b 100644 --- a/client/client-http/src/test/java/software/amazon/smithy/java/client/http/HttpErrorDeserializerTest.java +++ b/client/client-http/src/test/java/software/amazon/smithy/java/client/http/HttpErrorDeserializerTest.java @@ -9,6 +9,7 @@ import static org.hamcrest.Matchers.containsString; import static org.hamcrest.Matchers.equalTo; import static org.hamcrest.Matchers.instanceOf; +import static org.hamcrest.Matchers.nullValue; import java.util.List; import java.util.Map; @@ -16,6 +17,7 @@ import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.Arguments; import org.junit.jupiter.params.provider.MethodSource; +import software.amazon.smithy.java.client.core.CallContext; import software.amazon.smithy.java.context.Context; import software.amazon.smithy.java.core.error.CallException; import software.amazon.smithy.java.core.error.ModeledException; @@ -131,6 +133,26 @@ static List genericErrorCases() { Arguments.of(600, "", "Unknown HTTP/1.1 600 response from operation com.foo#PutFoo.")); } + @Test + public void recordsAndClearsUnmodeledPayloadErrorCode() { + var deserializer = HttpErrorDeserializer.builder() + .codec(CODEC) + .serviceId(SERVICE) + .build(); + var context = Context.create(); + var errorResponse = HttpResponse.create() + .setStatusCode(400) + .setBody(DataStream.ofString("{\"__type\":\"com.foo#ExpiredToken\"}")) + .toUnmodifiable(); + + deserializer.createError(context, OPERATION, TypeRegistry.empty(), errorResponse); + assertThat(context.get(CallContext.RESPONSE_ERROR_CODE), equalTo("ExpiredToken")); + + var emptyResponse = HttpResponse.create().setStatusCode(400).toUnmodifiable(); + deserializer.createError(context, OPERATION, TypeRegistry.empty(), emptyResponse); + assertThat(context.get(CallContext.RESPONSE_ERROR_CODE), nullValue()); + } + @Test public void deserializesIntoErrorBasedOnHeaders() { var deserializer = HttpErrorDeserializer.builder() diff --git a/client/client-metrics-otel/src/main/java/software/amazon/smithy/java/client/metrics/otel/OperationMetricsInterceptor.java b/client/client-metrics-otel/src/main/java/software/amazon/smithy/java/client/metrics/otel/OperationMetricsInterceptor.java index c702e422da..e2add0b71f 100644 --- a/client/client-metrics-otel/src/main/java/software/amazon/smithy/java/client/metrics/otel/OperationMetricsInterceptor.java +++ b/client/client-metrics-otel/src/main/java/software/amazon/smithy/java/client/metrics/otel/OperationMetricsInterceptor.java @@ -200,6 +200,11 @@ public IdentityResult resolveIdentity(Context requestProperties) { public Class identityType() { return delegate.identityType(); } + + @Override + public void invalidate(I rejectedIdentity) { + delegate.invalidate(rejectedIdentity); + } } static double elapsedSecondsSinceNs(long startNs) { diff --git a/retries/src/it/java/software/amazon/smithy/java/retries/AdaptiveRetryStrategyIntegrationTest.java b/retries/src/it/java/software/amazon/smithy/java/retries/AdaptiveRetryStrategyIntegrationTest.java index 7d13d93c56..c830c61fe4 100644 --- a/retries/src/it/java/software/amazon/smithy/java/retries/AdaptiveRetryStrategyIntegrationTest.java +++ b/retries/src/it/java/software/amazon/smithy/java/retries/AdaptiveRetryStrategyIntegrationTest.java @@ -5,6 +5,7 @@ package software.amazon.smithy.java.retries; +import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertTrue; import java.time.Duration; @@ -34,7 +35,6 @@ void rateLimiterThrottlesClientsToServerRate() throws InterruptedException { var server = new RateLimitedServer(RATE_LIMIT_MS); var successCount = new AtomicInteger(0); var throttleCount = new AtomicInteger(0); - var attemptCount = new AtomicInteger(0); var latch = new CountDownLatch(THREAD_COUNT); ExecutorService executor = Executors.newFixedThreadPool(THREAD_COUNT); @@ -42,7 +42,7 @@ void rateLimiterThrottlesClientsToServerRate() throws InterruptedException { executor.submit(() -> { try { for (var c = 0; c < CALLS_PER_THREAD; c++) { - executeWithRetry(strategy, server, successCount, throttleCount, attemptCount); + executeWithRetry(strategy, server, successCount, throttleCount); } } finally { latch.countDown(); @@ -55,27 +55,17 @@ void rateLimiterThrottlesClientsToServerRate() throws InterruptedException { var totalCalls = THREAD_COUNT * CALLS_PER_THREAD; // All calls should eventually succeed - assertTrue(successCount.get() == totalCalls, - "Expected " + totalCalls + " successes but got " + successCount.get()); + assertEquals(totalCalls, successCount.get()); // Some calls should have been throttled initially assertTrue(throttleCount.get() > 0, "Expected some throttling but got none"); - // Success rate should be greater than 80% - var totalAttempts = attemptCount.get(); - var successRate = (double) successCount.get() / totalAttempts; - assertTrue(successRate > 0.80, - String.format("Expected success rate > 80%% but got %.1f%% (%d/%d)", - successRate * 100, - successCount.get(), - totalAttempts)); } private void executeWithRetry( AdaptiveRetryStrategy strategy, RateLimitedServer server, AtomicInteger successCount, - AtomicInteger throttleCount, - AtomicInteger attemptCount + AtomicInteger throttleCount ) { var acquireResponse = strategy.acquireInitialToken(new AcquireInitialTokenRequest("test-scope")); var token = acquireResponse.token(); @@ -83,7 +73,6 @@ private void executeWithRetry( while (true) { try { - attemptCount.incrementAndGet(); server.call(); strategy.recordSuccess(new RecordSuccessRequest(token)); successCount.incrementAndGet(); @@ -96,8 +85,7 @@ private void executeWithRetry( token = refreshResponse.token(); sleep(refreshResponse.delay()); } catch (TokenAcquisitionFailedException ex) { - // Retries exhausted, count as success anyway since we'll retry the outer loop - // This shouldn't happen with maxAttempts=10 and adaptive backoff + // This shouldn't happen with maxAttempts=10 and adaptive backoff. throw new AssertionError("Retries exhausted unexpectedly", ex); } }