aws · gosar · May 25, 2023 · May 23, 2023 · May 23, 2023
diff --git a/...en/src/main/java/software/amazon/awssdk/codegen/poet/auth/scheme/AuthSchemeSpecUtils.java b/...en/src/main/java/software/amazon/awssdk/codegen/poet/auth/scheme/AuthSchemeSpecUtils.java
@@ -23,10 +23,10 @@
 import software.amazon.awssdk.codegen.utils.AuthUtils;
 import software.amazon.awssdk.http.auth.spi.HttpAuthOption;
 
-final class AuthSchemeSpecUtils {
+public final class AuthSchemeSpecUtils {
     private final IntermediateModel intermediateModel;
 
-    AuthSchemeSpecUtils(IntermediateModel intermediateModel) {
+    public AuthSchemeSpecUtils(IntermediateModel intermediateModel) {
         this.intermediateModel = intermediateModel;
     }
 

diff --git a/...gen/src/main/java/software/amazon/awssdk/codegen/poet/builder/BaseClientBuilderClass.java b/...gen/src/main/java/software/amazon/awssdk/codegen/poet/builder/BaseClientBuilderClass.java
@@ -47,6 +47,7 @@
 import software.amazon.awssdk.codegen.model.service.ClientContextParam;
 import software.amazon.awssdk.codegen.poet.ClassSpec;
 import software.amazon.awssdk.codegen.poet.PoetUtils;
+import software.amazon.awssdk.codegen.poet.auth.scheme.AuthSchemeSpecUtils;
 import software.amazon.awssdk.codegen.poet.rules.EndpointRulesSpecUtils;
 import software.amazon.awssdk.codegen.utils.AuthUtils;
 import software.amazon.awssdk.core.client.config.SdkAdvancedClientOption;
@@ -73,13 +74,15 @@ public class BaseClientBuilderClass implements ClassSpec {
     private final ClassName builderClassName;
     private final String basePackage;
     private final EndpointRulesSpecUtils endpointRulesSpecUtils;
+    private final AuthSchemeSpecUtils authSchemeSpecUtils;
 
     public BaseClientBuilderClass(IntermediateModel model) {
         this.model = model;
         this.basePackage = model.getMetadata().getFullClientPackageName();
         this.builderInterfaceName = ClassName.get(basePackage, model.getMetadata().getBaseBuilderInterface());
         this.builderClassName = ClassName.get(basePackage, model.getMetadata().getBaseBuilder());
         this.endpointRulesSpecUtils = new EndpointRulesSpecUtils(model);
+        this.authSchemeSpecUtils = new AuthSchemeSpecUtils(model);
     }
 
     @Override
@@ -115,6 +118,9 @@ public TypeSpec poetSpec() {
         builder.addMethod(signingNameMethod());
         builder.addMethod(defaultEndpointProviderMethod());
 
+        builder.addMethod(authSchemeProviderMethod());
+        builder.addMethod(defaultAuthSchemeProviderMethod());
+
         if (hasClientContextParams()) {
             model.getClientContextParams().forEach((n, m) -> {
                 builder.addMethod(clientContextParamSetter(n, m));
@@ -192,7 +198,7 @@ private MethodSpec mergeServiceDefaultsMethod() {
                                                .addCode("return config.merge(c -> c");
 
         builder.addCode(".option($T.ENDPOINT_PROVIDER, defaultEndpointProvider())", SdkClientOption.class);
-
+        builder.addCode(".option($T.AUTH_SCHEME_PROVIDER, defaultAuthSchemeProvider())", SdkClientOption.class);
 
         if (defaultAwsAuthSignerMethod().isPresent()) {
             builder.addCode(".option($T.SIGNER, defaultSigner())\n", SdkAdvancedClientOption.class);
@@ -580,6 +586,25 @@ private MethodSpec defaultEndpointProviderMethod() {
                          .build();
     }
 
+    private MethodSpec authSchemeProviderMethod() {
+        return MethodSpec.methodBuilder("authSchemeProvider")
+                         .addModifiers(Modifier.PUBLIC)
+                         .returns(TypeVariableName.get("B"))
+                         .addParameter(authSchemeSpecUtils.providerInterfaceName(), "authSchemeProvider")
+                         .addStatement("clientConfiguration.option($T.AUTH_SCHEME_PROVIDER, authSchemeProvider)",
+                                       SdkClientOption.class)
+                         .addStatement("return thisBuilder()")
+                         .build();
+    }
+
+    private MethodSpec defaultAuthSchemeProviderMethod() {
+        return MethodSpec.methodBuilder("defaultAuthSchemeProvider")
+                         .addModifiers(PRIVATE)
+                         .returns(authSchemeSpecUtils.providerInterfaceName())
+                         .addStatement("return $T.defaultProvider()", authSchemeSpecUtils.providerInterfaceName())
+                         .build();
+    }
+
     private MethodSpec clientContextParamSetter(String name, ClientContextParam param) {
         String setterName = endpointRulesSpecUtils.paramMethodName(name);
         String keyName = model.getNamingStrategy().getEnumValueName(name);

diff --git a/...src/main/java/software/amazon/awssdk/codegen/poet/builder/BaseClientBuilderInterface.java b/...src/main/java/software/amazon/awssdk/codegen/poet/builder/BaseClientBuilderInterface.java
@@ -34,6 +34,7 @@
 import software.amazon.awssdk.codegen.model.service.ClientContextParam;
 import software.amazon.awssdk.codegen.poet.ClassSpec;
 import software.amazon.awssdk.codegen.poet.PoetUtils;
+import software.amazon.awssdk.codegen.poet.auth.scheme.AuthSchemeSpecUtils;
 import software.amazon.awssdk.codegen.poet.rules.EndpointRulesSpecUtils;
 import software.amazon.awssdk.codegen.utils.AuthUtils;
 import software.amazon.awssdk.core.client.config.SdkAdvancedClientOption;
@@ -49,12 +50,14 @@ public class BaseClientBuilderInterface implements ClassSpec {
     private final String basePackage;
     private final ClassName builderInterfaceName;
     private final EndpointRulesSpecUtils endpointRulesSpecUtils;
+    private final AuthSchemeSpecUtils authSchemeSpecUtils;
 
     public BaseClientBuilderInterface(IntermediateModel model) {
         this.model = model;
         this.basePackage = model.getMetadata().getFullClientPackageName();
         this.builderInterfaceName = ClassName.get(basePackage, model.getMetadata().getBaseBuilderInterface());
         this.endpointRulesSpecUtils = new EndpointRulesSpecUtils(model);
+        this.authSchemeSpecUtils = new AuthSchemeSpecUtils(model);
     }
 
     @Override
@@ -79,6 +82,7 @@ public TypeSpec poetSpec() {
         }
 
         builder.addMethod(endpointProviderMethod());
+        builder.addMethod(authSchemeProviderMethod());
 
         if (hasClientContextParams()) {
             model.getClientContextParams().forEach((n, m) -> {
@@ -155,6 +159,19 @@ private MethodSpec endpointProviderMethod() {
                          .build();
     }
 
+    private MethodSpec authSchemeProviderMethod() {
+        return MethodSpec.methodBuilder("authSchemeProvider")
+                         .addModifiers(Modifier.PUBLIC, Modifier.DEFAULT)
+                         .addParameter(authSchemeSpecUtils.providerInterfaceName(), "authSchemeProvider")
+                         .addJavadoc("Set the {@link $T} implementation that will be used by the client to resolve the "
+                                     + "auth scheme for each request. This is optional; if none is provided a "
+                                     + "default implementation will be used the SDK.",
+                                     authSchemeSpecUtils.providerInterfaceName())
+                         .returns(TypeVariableName.get("B"))
+                         .addStatement("throw new $T()", UnsupportedOperationException.class)
+                         .build();
+    }
+
     private MethodSpec clientContextParamSetter(String name, ClientContextParam param) {
         String setterName = Utils.unCapitalize(CodegenNamingUtils.pascalCase(name));
         TypeName type = endpointRulesSpecUtils.toJavaType(param.getType());

diff --git a/...es/software/amazon/awssdk/codegen/poet/builder/test-bearer-auth-client-builder-class.java b/...es/software/amazon/awssdk/codegen/poet/builder/test-bearer-auth-client-builder-class.java
@@ -16,6 +16,7 @@
 import software.amazon.awssdk.core.signer.Signer;
 import software.amazon.awssdk.identity.spi.IdentityProvider;
 import software.amazon.awssdk.identity.spi.TokenIdentity;
+import software.amazon.awssdk.services.json.auth.scheme.JsonAuthSchemeProvider;
 import software.amazon.awssdk.services.json.endpoints.JsonEndpointProvider;
 import software.amazon.awssdk.services.json.endpoints.internal.JsonEndpointAuthSchemeInterceptor;
 import software.amazon.awssdk.services.json.endpoints.internal.JsonRequestSetEndpointInterceptor;
@@ -42,6 +43,7 @@ protected final String serviceName() {
     @Override
     protected final SdkClientConfiguration mergeServiceDefaults(SdkClientConfiguration config) {
         return config.merge(c -> c.option(SdkClientOption.ENDPOINT_PROVIDER, defaultEndpointProvider())
+                                  .option(SdkClientOption.AUTH_SCHEME_PROVIDER, defaultAuthSchemeProvider())
                                   .option(SdkClientOption.CRC32_FROM_COMPRESSED_DATA_ENABLED, false)
                                   .option(AwsClientOption.TOKEN_IDENTITY_PROVIDER, defaultTokenProvider())
                                   .option(SdkAdvancedClientOption.TOKEN_SIGNER, defaultTokenSigner()));
@@ -72,6 +74,15 @@ private JsonEndpointProvider defaultEndpointProvider() {
         return JsonEndpointProvider.defaultProvider();
     }
 
+    public B authSchemeProvider(JsonAuthSchemeProvider authSchemeProvider) {
+        clientConfiguration.option(SdkClientOption.AUTH_SCHEME_PROVIDER, authSchemeProvider);
+        return thisBuilder();
+    }
+
+    private JsonAuthSchemeProvider defaultAuthSchemeProvider() {
+        return JsonAuthSchemeProvider.defaultProvider();
+    }
+
     private IdentityProvider<? extends TokenIdentity> defaultTokenProvider() {
         return DefaultAwsTokenProvider.create();
     }

diff --git a/...oftware/amazon/awssdk/codegen/poet/builder/test-bearer-auth-client-builder-interface.java b/...oftware/amazon/awssdk/codegen/poet/builder/test-bearer-auth-client-builder-interface.java
@@ -5,6 +5,7 @@
 import software.amazon.awssdk.awscore.client.builder.AwsClientBuilder;
 import software.amazon.awssdk.identity.spi.IdentityProvider;
 import software.amazon.awssdk.identity.spi.TokenIdentity;
+import software.amazon.awssdk.services.json.auth.scheme.JsonAuthSchemeProvider;
 import software.amazon.awssdk.services.json.endpoints.JsonEndpointProvider;
 
 /**
@@ -21,6 +22,14 @@ default B endpointProvider(JsonEndpointProvider endpointProvider) {
         throw new UnsupportedOperationException();
     }
 
+    /**
+     * Set the {@link JsonAuthSchemeProvider} implementation that will be used by the client to resolve the auth scheme for
+     * each request. This is optional; if none is provided a default implementation will be used the SDK.
+     */
+    default B authSchemeProvider(JsonAuthSchemeProvider authSchemeProvider) {
+        throw new UnsupportedOperationException();
+    }
+
     /**
      * Set the token provider to use for bearer token authorization. This is optional, if none is provided, the SDK will
      * use {@link software.amazon.awssdk.auth.token.credentials.aws.DefaultAwsTokenProvider}.

diff --git a/...test/resources/software/amazon/awssdk/codegen/poet/builder/test-client-builder-class.java b/...test/resources/software/amazon/awssdk/codegen/poet/builder/test-client-builder-class.java
@@ -19,6 +19,7 @@
 import software.amazon.awssdk.core.signer.Signer;
 import software.amazon.awssdk.identity.spi.IdentityProvider;
 import software.amazon.awssdk.identity.spi.TokenIdentity;
+import software.amazon.awssdk.services.json.auth.scheme.JsonAuthSchemeProvider;
 import software.amazon.awssdk.services.json.endpoints.JsonClientContextParams;
 import software.amazon.awssdk.services.json.endpoints.JsonEndpointProvider;
 import software.amazon.awssdk.services.json.endpoints.internal.JsonEndpointAuthSchemeInterceptor;
@@ -47,6 +48,7 @@ protected final String serviceName() {
     @Override
     protected final SdkClientConfiguration mergeServiceDefaults(SdkClientConfiguration config) {
         return config.merge(c -> c.option(SdkClientOption.ENDPOINT_PROVIDER, defaultEndpointProvider())
+                                  .option(SdkClientOption.AUTH_SCHEME_PROVIDER, defaultAuthSchemeProvider())
                                   .option(SdkAdvancedClientOption.SIGNER, defaultSigner())
                                   .option(SdkClientOption.CRC32_FROM_COMPRESSED_DATA_ENABLED, false)
                                   .option(SdkClientOption.SERVICE_CONFIGURATION, ServiceConfiguration.builder().build())
@@ -143,6 +145,15 @@ private JsonEndpointProvider defaultEndpointProvider() {
         return JsonEndpointProvider.defaultProvider();
     }
 
+    public B authSchemeProvider(JsonAuthSchemeProvider authSchemeProvider) {
+        clientConfiguration.option(SdkClientOption.AUTH_SCHEME_PROVIDER, authSchemeProvider);
+        return thisBuilder();
+    }
+
+    private JsonAuthSchemeProvider defaultAuthSchemeProvider() {
+        return JsonAuthSchemeProvider.defaultProvider();
+    }
+
     public B serviceConfiguration(ServiceConfiguration serviceConfiguration) {
         clientConfiguration.option(SdkClientOption.SERVICE_CONFIGURATION, serviceConfiguration);
         return thisBuilder();

diff --git a/.../resources/software/amazon/awssdk/codegen/poet/builder/test-client-builder-interface.java b/.../resources/software/amazon/awssdk/codegen/poet/builder/test-client-builder-interface.java
@@ -6,6 +6,7 @@
 import software.amazon.awssdk.awscore.client.builder.AwsClientBuilder;
 import software.amazon.awssdk.identity.spi.IdentityProvider;
 import software.amazon.awssdk.identity.spi.TokenIdentity;
+import software.amazon.awssdk.services.json.auth.scheme.JsonAuthSchemeProvider;
 import software.amazon.awssdk.services.json.endpoints.JsonEndpointProvider;
 
 /**
@@ -28,6 +29,14 @@ default B endpointProvider(JsonEndpointProvider endpointProvider) {
         throw new UnsupportedOperationException();
     }
 
+    /**
+     * Set the {@link JsonAuthSchemeProvider} implementation that will be used by the client to resolve the auth scheme for
+     * each request. This is optional; if none is provided a default implementation will be used the SDK.
+     */
+    default B authSchemeProvider(JsonAuthSchemeProvider authSchemeProvider) {
+        throw new UnsupportedOperationException();
+    }
+
     /**
      * Set the token provider to use for bearer token authorization. This is optional, if none is provided, the SDK will
      * use {@link software.amazon.awssdk.auth.token.credentials.aws.DefaultAwsTokenProvider}.

diff --git a/...tware/amazon/awssdk/codegen/poet/builder/test-client-builder-internal-defaults-class.java b/...tware/amazon/awssdk/codegen/poet/builder/test-client-builder-internal-defaults-class.java
@@ -13,6 +13,7 @@
 import software.amazon.awssdk.core.interceptor.ExecutionInterceptor;
 import software.amazon.awssdk.core.retry.RetryMode;
 import software.amazon.awssdk.core.signer.Signer;
+import software.amazon.awssdk.services.json.auth.scheme.JsonAuthSchemeProvider;
 import software.amazon.awssdk.services.json.endpoints.JsonEndpointProvider;
 import software.amazon.awssdk.services.json.endpoints.internal.JsonEndpointAuthSchemeInterceptor;
 import software.amazon.awssdk.services.json.endpoints.internal.JsonRequestSetEndpointInterceptor;
@@ -39,6 +40,7 @@ protected final String serviceName() {
     @Override
     protected final SdkClientConfiguration mergeServiceDefaults(SdkClientConfiguration config) {
         return config.merge(c -> c.option(SdkClientOption.ENDPOINT_PROVIDER, defaultEndpointProvider())
+                                  .option(SdkClientOption.AUTH_SCHEME_PROVIDER, defaultAuthSchemeProvider())
                                   .option(SdkAdvancedClientOption.SIGNER, defaultSigner())
                                   .option(SdkClientOption.CRC32_FROM_COMPRESSED_DATA_ENABLED, false));
     }
@@ -80,6 +82,15 @@ private JsonEndpointProvider defaultEndpointProvider() {
         return JsonEndpointProvider.defaultProvider();
     }
 
+    public B authSchemeProvider(JsonAuthSchemeProvider authSchemeProvider) {
+        clientConfiguration.option(SdkClientOption.AUTH_SCHEME_PROVIDER, authSchemeProvider);
+        return thisBuilder();
+    }
+
+    private JsonAuthSchemeProvider defaultAuthSchemeProvider() {
+        return JsonAuthSchemeProvider.defaultProvider();
+    }
+
     protected static void validateClientOptions(SdkClientConfiguration c) {
         Validate.notNull(c.option(SdkAdvancedClientOption.SIGNER),
                          "The 'overrideConfiguration.advancedOption[SIGNER]' must be configured in the client builder.");

diff --git a/...esources/software/amazon/awssdk/codegen/poet/builder/test-query-client-builder-class.java b/...esources/software/amazon/awssdk/codegen/poet/builder/test-query-client-builder-class.java
@@ -19,6 +19,7 @@
 import software.amazon.awssdk.identity.spi.IdentityProvider;
 import software.amazon.awssdk.identity.spi.TokenIdentity;
 import software.amazon.awssdk.protocols.query.interceptor.QueryParametersToBodyInterceptor;
+import software.amazon.awssdk.services.query.auth.scheme.QueryAuthSchemeProvider;
 import software.amazon.awssdk.services.query.endpoints.QueryClientContextParams;
 import software.amazon.awssdk.services.query.endpoints.QueryEndpointProvider;
 import software.amazon.awssdk.services.query.endpoints.internal.QueryEndpointAuthSchemeInterceptor;
@@ -46,6 +47,7 @@ protected final String serviceName() {
     @Override
     protected final SdkClientConfiguration mergeServiceDefaults(SdkClientConfiguration config) {
         return config.merge(c -> c.option(SdkClientOption.ENDPOINT_PROVIDER, defaultEndpointProvider())
+                                  .option(SdkClientOption.AUTH_SCHEME_PROVIDER, defaultAuthSchemeProvider())
                                   .option(SdkAdvancedClientOption.SIGNER, defaultSigner())
                                   .option(SdkClientOption.CRC32_FROM_COMPRESSED_DATA_ENABLED, false)
                                   .option(AwsClientOption.TOKEN_IDENTITY_PROVIDER, defaultTokenProvider())
@@ -84,6 +86,15 @@ private QueryEndpointProvider defaultEndpointProvider() {
         return QueryEndpointProvider.defaultProvider();
     }
 
+    public B authSchemeProvider(QueryAuthSchemeProvider authSchemeProvider) {
+        clientConfiguration.option(SdkClientOption.AUTH_SCHEME_PROVIDER, authSchemeProvider);
+        return thisBuilder();
+    }
+
+    private QueryAuthSchemeProvider defaultAuthSchemeProvider() {
+        return QueryAuthSchemeProvider.defaultProvider();
+    }
+
     public B booleanContextParam(Boolean booleanContextParam) {
         clientContextParams.put(QueryClientContextParams.BOOLEAN_CONTEXT_PARAM, booleanContextParam);
         return thisBuilder();

diff --git a/core/sdk-core/pom.xml b/core/sdk-core/pom.xml
@@ -51,6 +51,11 @@
             <artifactId>endpoints-spi</artifactId>
             <version>${awsjavasdk.version}</version>
         </dependency>
+        <dependency>
+            <groupId>software.amazon.awssdk</groupId>
+            <artifactId>http-auth-spi</artifactId>
+            <version>${awsjavasdk.version}</version>
+        </dependency>
         <dependency>
             <groupId>software.amazon.awssdk</groupId>
             <artifactId>utils</artifactId>

diff --git a/core/sdk-core/src/main/java/software/amazon/awssdk/core/client/config/SdkClientOption.java b/core/sdk-core/src/main/java/software/amazon/awssdk/core/client/config/SdkClientOption.java
@@ -31,6 +31,7 @@
 import software.amazon.awssdk.endpoints.EndpointProvider;
 import software.amazon.awssdk.http.SdkHttpClient;
 import software.amazon.awssdk.http.async.SdkAsyncHttpClient;
+import software.amazon.awssdk.http.auth.spi.AuthSchemeProvider;
 import software.amazon.awssdk.metrics.MetricPublisher;
 import software.amazon.awssdk.profiles.ProfileFile;
 import software.amazon.awssdk.utils.AttributeMap;
@@ -184,6 +185,12 @@ public final class SdkClientOption<T> extends ClientOption<T> {
      */
     public static final SdkClientOption<EndpointProvider> ENDPOINT_PROVIDER = new SdkClientOption<>(EndpointProvider.class);
 
+    /**
+     * The {@link AuthSchemeProvider} configured on the client.
+     */
+    public static final SdkClientOption<AuthSchemeProvider> AUTH_SCHEME_PROVIDER =
+        new SdkClientOption<>(AuthSchemeProvider.class);
+
     /**
      * The container for any client contexts parameters set on the client.
      */