From 6b1dac5469581b9f43f45f87e4422589b2b3bc77 Mon Sep 17 00:00:00 2001 From: Michael Dowling Date: Wed, 11 Mar 2026 15:43:12 -0500 Subject: [PATCH] Use endpoint auth hints with resolved scheme --- .../java/client/core/ClientPipeline.java | 32 ++++++++- .../java/client/core/ClientPipelineTest.java | 71 +++++++++++++++++++ 2 files changed, 102 insertions(+), 1 deletion(-) diff --git a/client/client-api/src/main/java/software/amazon/smithy/java/client/core/ClientPipeline.java b/client/client-api/src/main/java/software/amazon/smithy/java/client/core/ClientPipeline.java index bbf36ce07..f05f76ad3 100644 --- a/client/client-api/src/main/java/software/amazon/smithy/java/client/core/ClientPipeline.java +++ b/client/client-api/src/main/java/software/amazon/smithy/java/client/core/ClientPipeline.java @@ -189,10 +189,12 @@ private O afterIden var identity = identityResult.unwrap(); call.context.put(CallContext.IDENTITY, identity); - // TODO: what to do with supportedAuthSchemes of an endpoint? Endpoint endpoint = resolveEndpoint(call); call.context.put(CallContext.ENDPOINT, endpoint); + // Augment signer properties with endpoint auth scheme overrides if present. + resolvedAuthScheme = applyEndpointAuthSchemeOverrides(endpoint, resolvedAuthScheme); + RequestT req = protocol.setServiceEndpoint(requestHook.request(), endpoint); var signResult = resolvedAuthScheme.sign(req); req = signResult.signedRequest(); @@ -311,6 +313,34 @@ private Endpoint re return call.endpointResolver.resolveEndpoint(request); } + @SuppressWarnings("unchecked") + private ResolvedScheme applyEndpointAuthSchemeOverrides( + Endpoint endpoint, + ResolvedScheme resolvedScheme + ) { + var endpointAuthSchemes = endpoint.authSchemes(); + if (!endpointAuthSchemes.isEmpty()) { + var schemeId = resolvedScheme.authScheme().schemeId().toString(); + for (var endpointAuthScheme : endpointAuthSchemes) { + if (schemeId.equals(endpointAuthScheme.authSchemeId())) { + var overrides = endpointAuthScheme.properties(); + if (overrides.isEmpty()) { + return resolvedScheme; + } + // Apply the found overrides for the auth scheme. + var merged = Context.create(); + resolvedScheme.signerProperties().copyTo(merged); + for (var key : overrides) { + merged.put((Context.Key) key, endpointAuthScheme.property(key)); + } + return new ResolvedScheme<>(merged, resolvedScheme.authScheme(), resolvedScheme.identity()); + } + } + } + + return resolvedScheme; + } + private O deserialize( ClientCall call, RequestT request, 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 8de1e64fe..33dea73e0 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 @@ -13,17 +13,29 @@ import java.time.Duration; import java.util.ArrayList; +import java.util.List; import java.util.Map; +import java.util.concurrent.atomic.AtomicReference; import org.junit.jupiter.api.Assertions; 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.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.endpoint.Endpoint; +import software.amazon.smithy.java.client.core.endpoint.EndpointAuthScheme; import software.amazon.smithy.java.client.core.endpoint.EndpointResolver; 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; +import software.amazon.smithy.java.context.Context; import software.amazon.smithy.java.core.serde.document.Document; import software.amazon.smithy.java.dynamicclient.DynamicClient; +import software.amazon.smithy.java.http.api.HttpRequest; import software.amazon.smithy.java.http.api.HttpResponse; import software.amazon.smithy.java.io.datastream.DataStream; import software.amazon.smithy.java.retries.api.AcquireInitialTokenRequest; @@ -188,6 +200,65 @@ public Builder toBuilder() { assertThat(calls, contains("Acquire", "Refresh", "Success: 1")); } + @Test + public void endpointAuthSchemeOverridesAugmentSignerProperties() { + var service = ShapeId.from("smithy.example#Sprockets"); + var testSchemeId = ShapeId.from("smithy.test#testAuth"); + var TEST_KEY = Context.key("test-signing-override"); + var capturedProperties = new AtomicReference(); + + // Auth scheme with a signer that captures the properties it receives. + var testScheme = AuthScheme.of( + testSchemeId, + HttpRequest.class, + Identity.class, + (request, identity, properties) -> { + capturedProperties.set(properties); + return new SignResult<>(request); + }); + + // Endpoint resolver that returns an endpoint with an auth scheme override. + EndpointResolver endpointResolver = params -> Endpoint.builder() + .uri("https://example.com") + .addAuthScheme( + EndpointAuthScheme.builder() + .authSchemeId(testSchemeId.toString()) + .putProperty(TEST_KEY, "overridden-value") + .build()) + .build(); + + var mockQueue = new MockQueue() + .enqueue(HttpResponse.builder() + .statusCode(200) + .body(DataStream.ofString("{\"id\":\"1\"}")) + .build()); + var mock = MockPlugin.builder().addQueue(mockQueue).build(); + + var client = DynamicClient.builder() + .serviceId(service) + .model(MODEL) + .addPlugin(mock) + .endpointResolver(endpointResolver) + .authSchemeResolver(params -> List.of(new AuthSchemeOption(testSchemeId))) + .putSupportedAuthSchemes(testScheme) + .addIdentityResolver(new IdentityResolver<>() { + @Override + public IdentityResult resolveIdentity(Context requestProperties) { + return IdentityResult.of(new Identity() {}); + } + + @Override + public Class identityType() { + return Identity.class; + } + }) + .build(); + + client.call("GetSprocket", Document.ofObject(Map.of("id", "1"))); + + assertThat(capturedProperties.get().get(TEST_KEY), equalTo("overridden-value")); + } + private static final class Token implements RetryToken { int retry;