Remove legacy endpoint auth scheme logic and add putAuthScheme to AWS client builder.

codegen/src/main/java/software/amazon/awssdk/codegen/emitters/tasks/EndpointProviderTasks.java

@@ -27,7 +27,6 @@
 import software.amazon.awssdk.codegen.model.service.ClientContextParam;
 import software.amazon.awssdk.codegen.poet.rules.ClientContextParamsClassSpec;
 import software.amazon.awssdk.codegen.poet.rules.DefaultPartitionDataProviderSpec;
-import software.amazon.awssdk.codegen.poet.rules.EndpointAuthSchemeInterceptorClassSpec;
 import software.amazon.awssdk.codegen.poet.rules.EndpointParametersClassSpec;
 import software.amazon.awssdk.codegen.poet.rules.EndpointProviderInterfaceSpec;
 import software.amazon.awssdk.codegen.poet.rules.EndpointProviderSpec;
@@ -85,9 +84,7 @@ private GeneratorTask generateDefaultPartitionsProvider() {
     private Collection<GeneratorTask> generateInterceptors() {
         return Arrays.asList(
             new PoetGeneratorTask(endpointRulesInternalDir(), model.getFileHeader(), new EndpointResolverInterceptorSpec(model)),
-            new PoetGeneratorTask(endpointRulesInternalDir(), model.getFileHeader(), new RequestEndpointInterceptorSpec(model)),
-            new PoetGeneratorTask(endpointRulesInternalDir(), model.getFileHeader(),
-                                  new EndpointAuthSchemeInterceptorClassSpec(model)));
+            new PoetGeneratorTask(endpointRulesInternalDir(), model.getFileHeader(), new RequestEndpointInterceptorSpec(model)));
     }
 
     private GeneratorTask generateClientTests() {

codegen/src/main/java/software/amazon/awssdk/codegen/poet/auth/scheme/EndpointBasedAuthSchemeProviderSpec.java

@@ -43,6 +43,7 @@
 import software.amazon.awssdk.http.auth.aws.AwsV4aAuthScheme;
 import software.amazon.awssdk.http.auth.aws.AwsV4aHttpSigner;
 import software.amazon.awssdk.http.auth.spi.AuthSchemeOption;
+import software.amazon.awssdk.utils.CompletableFutureUtils;
 import software.amazon.awssdk.utils.Validate;
 
 public class EndpointBasedAuthSchemeProviderSpec implements ClassSpec {
@@ -123,7 +124,8 @@ private MethodSpec resolveAuthSchemeMethod() {
             }
         });
         spec.addStatement(".build()");
-        spec.addStatement("$T endpoint = DELEGATE.resolveEndpoint(endpointParameters).join()", Endpoint.class);
+        spec.addStatement("$T endpoint = $T.joinLikeSync(DELEGATE.resolveEndpoint(endpointParameters))",
+                          Endpoint.class, CompletableFutureUtils.class);
         spec.addStatement("$T authSchemes = endpoint.attribute($T.AUTH_SCHEMES)",
                           ParameterizedTypeName.get(List.class, EndpointAuthScheme.class), AwsEndpointAttribute.class);
         spec.beginControlFlow("if (authSchemes == null)");

codegen/src/main/java/software/amazon/awssdk/codegen/poet/builder/BaseClientBuilderClass.java

@@ -69,6 +69,9 @@
 import software.amazon.awssdk.utils.internal.CodegenNamingUtils;
 
 public class BaseClientBuilderClass implements ClassSpec {
+    private static final ParameterizedTypeName GENERIC_AUTH_SCHEME_TYPE =
+        ParameterizedTypeName.get(ClassName.get(AuthScheme.class), WildcardTypeName.subtypeOf(Object.class));
+
     private final IntermediateModel model;
     private final ClassName builderInterfaceName;
     private final ClassName builderClassName;
@@ -107,6 +110,14 @@ public TypeSpec poetSpec() {
                                       .build());
         }
 
+        builder.addField(FieldSpec.builder(ParameterizedTypeName.get(ClassName.get(Map.class),
+                                                                     ClassName.get(String.class),
+                                                                     GENERIC_AUTH_SCHEME_TYPE),
+                                           "additionalAuthSchemes")
+                                  .addModifiers(PRIVATE, FINAL)
+                                  .initializer("new $T<>()", HashMap.class)
+                                  .build());
+
         builder.addMethod(serviceEndpointPrefixMethod());
         builder.addMethod(serviceNameMethod());
         builder.addMethod(mergeServiceDefaultsMethod());
@@ -119,6 +130,7 @@ public TypeSpec poetSpec() {
 
         builder.addMethod(authSchemeProviderMethod());
         builder.addMethod(defaultAuthSchemeProviderMethod());
+        builder.addMethod(putAuthSchemeMethod());
 
         if (hasClientContextParams()) {
             model.getClientContextParams().forEach((n, m) -> {
@@ -140,7 +152,7 @@ public TypeSpec poetSpec() {
         if (AuthUtils.usesBearerAuth(model)) {
             builder.addMethod(defaultBearerTokenProviderMethod());
         }
-        builder.addMethod(defaultAuthSchemesMethod());
+        builder.addMethod(authSchemesMethod());
 
         addServiceHttpConfigIfNeeded(builder, model);
 
@@ -195,7 +207,7 @@ private MethodSpec mergeServiceDefaultsMethod() {
 
         builder.addCode(".option($T.ENDPOINT_PROVIDER, defaultEndpointProvider())", SdkClientOption.class);
         builder.addCode(".option($T.AUTH_SCHEME_PROVIDER, defaultAuthSchemeProvider())", SdkClientOption.class);
-        builder.addCode(".option($T.AUTH_SCHEMES, defaultAuthSchemes())", SdkClientOption.class);
+        builder.addCode(".option($T.AUTH_SCHEMES, authSchemes())", SdkClientOption.class);
 
         builder.addCode(".option($T.CRC32_FROM_COMPRESSED_DATA_ENABLED, $L)\n",
                         SdkClientOption.class, crc32FromCompressedDataEnabled);
@@ -259,14 +271,8 @@ private MethodSpec finalizeServiceConfigurationMethod() {
 
         List<ClassName> builtInInterceptors = new ArrayList<>();
 
-        // TODO(sra-identity-and-auth): Skip S3 for now as there are several tests failing that need to be fixed somewhere else.
-        boolean isS3 = "S3".equals(model.getMetadata().getServiceName())
-            || "S3Control".equals(model.getMetadata().getServiceName());
-        if (!isS3) {
-            builtInInterceptors.add(authSchemeSpecUtils.authSchemeInterceptor());
-        }
+        builtInInterceptors.add(authSchemeSpecUtils.authSchemeInterceptor());
         builtInInterceptors.add(endpointRulesSpecUtils.resolverInterceptorName());
-        builtInInterceptors.add(endpointRulesSpecUtils.authSchemesInterceptorName());
         builtInInterceptors.add(endpointRulesSpecUtils.requestModifierInterceptorName());
 
         for (String interceptor : model.getCustomizationConfig().getInterceptors()) {
@@ -592,6 +598,17 @@ private MethodSpec defaultAuthSchemeProviderMethod() {
                          .build();
     }
 
+    private MethodSpec putAuthSchemeMethod() {
+        return MethodSpec.methodBuilder("putAuthScheme")
+                         .addAnnotation(Override.class)
+                         .addModifiers(Modifier.PUBLIC)
+                         .returns(TypeVariableName.get("B"))
+                         .addParameter(GENERIC_AUTH_SCHEME_TYPE, "authScheme")
+                         .addStatement("additionalAuthSchemes.put(authScheme.schemeId(), authScheme)")
+                         .addStatement("return thisBuilder()")
+                         .build();
+    }
+
     private MethodSpec clientContextParamSetter(String name, ClientContextParam param) {
         String setterName = endpointRulesSpecUtils.paramMethodName(name);
         String keyName = model.getNamingStrategy().getEnumValueName(name);
@@ -616,12 +633,12 @@ private MethodSpec defaultBearerTokenProviderMethod() {
                          .build();
     }
 
-    private MethodSpec defaultAuthSchemesMethod() {
+    private MethodSpec authSchemesMethod() {
         TypeName returns = ParameterizedTypeName.get(ClassName.get(Map.class), ClassName.get(String.class),
                                                      ParameterizedTypeName.get(ClassName.get(AuthScheme.class),
                                                                                WildcardTypeName.subtypeOf(Object.class)));
 
-        MethodSpec.Builder builder = MethodSpec.methodBuilder("defaultAuthSchemes")
+        MethodSpec.Builder builder = MethodSpec.methodBuilder("authSchemes")
                                                .addModifiers(PRIVATE)
                                                .returns(returns);
 
@@ -631,12 +648,14 @@ private MethodSpec defaultAuthSchemesMethod() {
             return builder.build();
         }
 
-        builder.addStatement("$T schemes = new $T<>($L)", returns, HashMap.class, concreteAuthSchemeClasses.size());
+        builder.addStatement("$T schemes = new $T<>($L + this.additionalAuthSchemes.size())",
+                             returns, HashMap.class, concreteAuthSchemeClasses.size());
         for (Class<?> concreteAuthScheme : concreteAuthSchemeClasses) {
             String instanceVariable = CodegenNamingUtils.lowercaseFirstChar(concreteAuthScheme.getSimpleName());
             builder.addStatement("$1T $2L = $1T.create()", concreteAuthScheme, instanceVariable);
             builder.addStatement("schemes.put($1N.schemeId(), $1N)", instanceVariable);
         }
+        builder.addStatement("schemes.putAll(this.additionalAuthSchemes)");
         builder.addStatement("return $T.unmodifiableMap(schemes)", Collections.class);
         return builder.build();
     }