Adding functionality to config preferred authschemeProvider

.changes/next-release/feature-AWSSDKforJavav2-8e1c19d.json

@@ -0,0 +1,6 @@
+{
+    "type": "feature",
+    "category": "AWS SDK for Java v2",
+    "contributor": "",
+    "description": "Added ability to configure preferred authentication schemes when multiple auth options are available."
+}

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

@@ -28,6 +28,7 @@
 import software.amazon.awssdk.codegen.poet.auth.scheme.EndpointAwareAuthSchemeParamsSpec;
 import software.amazon.awssdk.codegen.poet.auth.scheme.EndpointBasedAuthSchemeProviderSpec;
 import software.amazon.awssdk.codegen.poet.auth.scheme.ModelBasedAuthSchemeProviderSpec;
+import software.amazon.awssdk.codegen.poet.auth.scheme.PreferredAuthSchemeProviderSpec;
 
 public final class AuthSchemeGeneratorTasks extends BaseGeneratorTasks {
     private final GeneratorTaskParams generatorTaskParams;
@@ -45,6 +46,7 @@ protected List<GeneratorTask> createTasks() {
         tasks.add(generateProviderInterface());
         tasks.add(generateDefaultParamsImpl());
         tasks.add(generateModelBasedProvider());
+        tasks.add(generatePreferenceProvider());
         tasks.add(generateAuthSchemeInterceptor());
         if (authSchemeSpecUtils.useEndpointBasedAuthProvider()) {
             tasks.add(generateEndpointBasedProvider());
@@ -69,6 +71,10 @@ private GeneratorTask generateModelBasedProvider() {
         return new PoetGeneratorTask(authSchemeInternalDir(), model.getFileHeader(), new ModelBasedAuthSchemeProviderSpec(model));
     }
 
+    private GeneratorTask generatePreferenceProvider() {
+        return new PoetGeneratorTask(authSchemeInternalDir(), model.getFileHeader(), new PreferredAuthSchemeProviderSpec(model));
+    }
+
     private GeneratorTask generateEndpointBasedProvider() {
         return new PoetGeneratorTask(authSchemeInternalDir(), model.getFileHeader(),
                                      new EndpointBasedAuthSchemeProviderSpec(model));

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

@@ -17,12 +17,16 @@
 
 import com.squareup.javapoet.ClassName;
 import com.squareup.javapoet.CodeBlock;
+import com.squareup.javapoet.FieldSpec;
 import com.squareup.javapoet.MethodSpec;
 import com.squareup.javapoet.ParameterizedTypeName;
 import com.squareup.javapoet.TypeName;
 import com.squareup.javapoet.TypeSpec;
+import java.util.ArrayList;
+import java.util.List;
 import java.util.function.Consumer;
 import javax.lang.model.element.Modifier;
+import software.amazon.awssdk.annotations.SdkInternalApi;
 import software.amazon.awssdk.annotations.SdkPublicApi;
 import software.amazon.awssdk.codegen.model.intermediate.IntermediateModel;
 import software.amazon.awssdk.codegen.poet.ClassSpec;
@@ -54,6 +58,9 @@ public TypeSpec poetSpec() {
                         .addMethod(resolveAuthSchemeMethod())
                         .addMethod(resolveAuthSchemeConsumerBuilderMethod())
                         .addMethod(defaultProviderMethod())
+                        .addMethod(staticBuilderMethodSpec())
+                        .addType(builderInterfaceSpec())
+                        .addType(builderClassSpec())
                         .build();
     }
 
@@ -104,4 +111,74 @@ private CodeBlock interfaceJavadoc() {
 
         return b.build();
     }
+
+    private MethodSpec staticBuilderMethodSpec() {
+        return MethodSpec.methodBuilder("builder")
+                         .addJavadoc("Create a builder for the auth scheme provider.")
+                         .addModifiers(Modifier.PUBLIC, Modifier.STATIC)
+                         .returns(className().nestedClass("Builder"))
+                         .addStatement("return new $T()", ClassName.get(className().packageName(),
+                                                                      className().simpleName(),
+                                                                      className().simpleName() + "Builder")
+                         )
+                         .build();
+    }
+
+
+    private TypeSpec builderInterfaceSpec() {
+        return TypeSpec.interfaceBuilder("Builder")
+                       .addModifiers(Modifier.PUBLIC, Modifier.STATIC)
+                       .addMethod(MethodSpec.methodBuilder("build")
+                                            .addModifiers(Modifier.PUBLIC, Modifier.ABSTRACT)
+                                            .addJavadoc("Returns a {@link $T} object that is created from the "
+                                                        + "properties that have been set on the builder.",
+                                                        className())
+                                            .returns(className())
+                                            .build())
+
+                       .addMethod(MethodSpec.methodBuilder("withPreferredAuthSchemes")
+                                            .addModifiers(Modifier.PUBLIC, Modifier.ABSTRACT)
+                                            .addJavadoc("Set the preferred auth schemes in order of preference.")
+                                            .returns(className().nestedClass("Builder"))
+                                            .addParameter(
+                                                ParameterizedTypeName.get(List.class, String.class),
+                                                "authSchemePreference"
+                                            )
+                                            .build())
+                       .build();
+    }
+
+    private TypeSpec builderClassSpec() {
+        return TypeSpec.classBuilder(authSchemeSpecUtils.authSchemeProviderBuilderName())
+                       .addAnnotation(SdkInternalApi.class)
+                       .addModifiers(Modifier.PUBLIC, Modifier.FINAL, Modifier.STATIC)
+                       .addSuperinterface(className().nestedClass("Builder"))
+                       .addField(
+                           FieldSpec
+                               .builder(ParameterizedTypeName.get(List.class, String.class), "authSchemePreference")
+                               .addModifiers(Modifier.PRIVATE)
+                               .build())
+                       .addMethod(
+                           MethodSpec
+                               .methodBuilder("withPreferredAuthSchemes").addAnnotation(Override.class)
+                               .addModifiers(Modifier.PUBLIC)
+                               .addParameter(
+                                   ParameterizedTypeName.get(List.class, String.class),
+                                   "authSchemePreference"
+                               )
+                               .returns(className().nestedClass("Builder"))
+                               .addStatement("this.authSchemePreference = new $T<>(authSchemePreference)", ArrayList.class)
+                               .addStatement("return this")
+                               .build())
+                       .addMethod(
+                           MethodSpec
+                               .methodBuilder("build").addAnnotation(Override.class)
+                               .addModifiers(Modifier.PUBLIC)
+                               .returns(className())
+                               .addStatement("return new $T(defaultProvider(), authSchemePreference)",
+                                             authSchemeSpecUtils.preferredAuthSchemeProviderName())
+                               .build())
+                       .build();
+    }
 }
+

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

@@ -97,6 +97,15 @@ public ClassName modeledAuthSchemeProviderName() {
         return ClassName.get(internalPackage(), "Modeled" + providerInterfaceName().simpleName());
     }
 
+    public ClassName preferredAuthSchemeProviderName() {
+        return ClassName.get(internalPackage(), "Preferred" + providerInterfaceName().simpleName());
+    }
+
+    public ClassName authSchemeProviderBuilderName() {
+        return ClassName.get(basePackage(),
+                             intermediateModel.getMetadata().getServiceName() + "AuthSchemeProviderBuilder");
+    }
+
     public ClassName authSchemeInterceptor() {
         return ClassName.get(internalPackage(), intermediateModel.getMetadata().getServiceName() + "AuthSchemeInterceptor");
     }

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

@@ -0,0 +1,108 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License").
+ * You may not use this file except in compliance with the License.
+ * A copy of the License is located at
+ *
+ *  http://aws.amazon.com/apache2.0
+ *
+ * or in the "license" file accompanying this file. This file is distributed
+ * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ * express or implied. See the License for the specific language governing
+ * permissions and limitations under the License.
+ */
+
+package software.amazon.awssdk.codegen.poet.auth.scheme;
+
+import com.squareup.javapoet.ClassName;
+import com.squareup.javapoet.MethodSpec;
+import com.squareup.javapoet.ParameterizedTypeName;
+import com.squareup.javapoet.TypeSpec;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+import javax.lang.model.element.Modifier;
+import software.amazon.awssdk.annotations.SdkInternalApi;
+import software.amazon.awssdk.codegen.model.intermediate.IntermediateModel;
+import software.amazon.awssdk.codegen.poet.ClassSpec;
+import software.amazon.awssdk.codegen.poet.PoetUtils;
+import software.amazon.awssdk.utils.CollectionUtils;
+
+public class PreferredAuthSchemeProviderSpec implements ClassSpec {
+    private final AuthSchemeSpecUtils authSchemeSpecUtils;
+
+    public PreferredAuthSchemeProviderSpec(IntermediateModel intermediateModel) {
+        this.authSchemeSpecUtils = new AuthSchemeSpecUtils(intermediateModel);
+    }
+
+    @Override
+    public ClassName className() {
+        return authSchemeSpecUtils.preferredAuthSchemeProviderName();
+    }
+
+    @Override
+    public TypeSpec poetSpec() {
+        return PoetUtils.createClassBuilder(className())
+                        .addModifiers(Modifier.PUBLIC, Modifier.FINAL)
+                        .addAnnotation(SdkInternalApi.class)
+                        .addField(
+                            authSchemeSpecUtils.providerInterfaceName(), "delegate",
+                            Modifier.PRIVATE, Modifier.FINAL)
+                        .addField(
+                            ParameterizedTypeName.get(List.class, String.class), "authSchemePreference",
+                            Modifier.PRIVATE, Modifier.FINAL)
+                        .addSuperinterface(authSchemeSpecUtils.providerInterfaceName())
+                        .addMethod(constructor())
+                        .addMethod(resolveAuthSchemeMethod())
+                        .build();
+    }
+
+    private MethodSpec constructor() {
+        return MethodSpec
+            .constructorBuilder()
+            .addModifiers(Modifier.PUBLIC)
+            .addParameter(authSchemeSpecUtils.providerInterfaceName(), "delegate")
+            .addParameter(ParameterizedTypeName.get(List.class, String.class), "authSchemePreference")
+            .addStatement("this.delegate = delegate")
+            .addStatement("this.authSchemePreference = authSchemePreference != null ? authSchemePreference "
+                          + ": $T.emptyList()",
+                          Collections.class)
+            .build();
+    }
+
+    private MethodSpec resolveAuthSchemeMethod() {
+        MethodSpec.Builder b = MethodSpec.methodBuilder("resolveAuthScheme")
+                                         .addModifiers(Modifier.PUBLIC)
+                                         .addAnnotation(Override.class)
+                                         .returns(authSchemeSpecUtils.resolverReturnType())
+                                         .addParameter(authSchemeSpecUtils.parametersInterfaceName(), "params");
+        b.addJavadoc("Resolve the auth schemes based on the given set of parameters.");
+        b.addStatement("$T candidateAuthSchemes = delegate.resolveAuthScheme(params)",
+                       authSchemeSpecUtils.resolverReturnType());
+        b.beginControlFlow("if ($T.isNullOrEmpty(authSchemePreference))", CollectionUtils.class)
+            .addStatement("return candidateAuthSchemes")
+            .endControlFlow();
+
+        b.addStatement("$T authSchemes = new $T<>()", authSchemeSpecUtils.resolverReturnType(), ArrayList.class);
+
+        b.beginControlFlow("authSchemePreference.forEach(preferredSchemeId -> ");
+
+        b.beginControlFlow("candidateAuthSchemes.stream().filter(candidate -> ");
+        b.addStatement("String candidateSchemeName = candidate.schemeId().contains(\"#\") ? " +
+                       "candidate.schemeId().split(\"#\")[1] : candidate.schemeId()");
+        b.addStatement("return candidateSchemeName.equals(preferredSchemeId)");
+        b.endControlFlow(").findFirst().ifPresent(authSchemes::add)");
+        b.endControlFlow(")");
+
+        b.beginControlFlow("candidateAuthSchemes.forEach(candidate -> ")
+            .beginControlFlow("if (!authSchemes.contains(candidate))")
+            .addStatement("authSchemes.add(candidate)")
+            .endControlFlow()
+            .endControlFlow(")");
+
+        b.addStatement("return authSchemes");
+        return b.build();
+    }
+
+}

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

@@ -43,6 +43,7 @@
 import software.amazon.awssdk.auth.signer.Aws4Signer;
 import software.amazon.awssdk.auth.token.credentials.aws.DefaultAwsTokenProvider;
 import software.amazon.awssdk.auth.token.signer.aws.BearerTokenSigner;
+import software.amazon.awssdk.awscore.auth.AuthSchemePreferenceProvider;
 import software.amazon.awssdk.awscore.client.builder.AwsDefaultClientBuilder;
 import software.amazon.awssdk.awscore.client.config.AwsClientOption;
 import software.amazon.awssdk.awscore.endpoint.AwsClientEndpointProvider;
@@ -277,7 +278,7 @@ private MethodSpec mergeServiceDefaultsMethod() {
         builder.addCode(".option($T.ENDPOINT_PROVIDER, defaultEndpointProvider())", SdkClientOption.class);
 
         if (authSchemeSpecUtils.useSraAuth()) {
-            builder.addCode(".option($T.AUTH_SCHEME_PROVIDER, defaultAuthSchemeProvider())", SdkClientOption.class);
+            builder.addCode(".option($T.AUTH_SCHEME_PROVIDER, defaultAuthSchemeProvider(config))", SdkClientOption.class);
             builder.addCode(".option($T.AUTH_SCHEMES, authSchemes())", SdkClientOption.class);
         } else {
             if (defaultAwsAuthSignerMethod().isPresent()) {
@@ -442,7 +443,7 @@ private MethodSpec finalizeServiceConfigurationMethod() {
             // serviceConfigBuilder; the service configuration classes (e.g. S3Configuration) return primitive booleans that
             // have a default when not present.
             builder.addStatement("builder.option($T.DUALSTACK_ENDPOINT_ENABLED, serviceConfigBuilder.dualstackEnabled())",
-                            AwsClientOption.class);
+                                 AwsClientOption.class);
         }
 
         if (model.getCustomizationConfig().getServiceConfig().hasFipsProperty()) {
@@ -452,14 +453,14 @@ private MethodSpec finalizeServiceConfigurationMethod() {
 
         if (model.getEndpointOperation().isPresent()) {
             builder.addStatement("builder.option($T.ENDPOINT_DISCOVERY_ENABLED, endpointDiscoveryEnabled)\n",
-                            SdkClientOption.class);
+                                 SdkClientOption.class);
         }
 
 
         if (StringUtils.isNotBlank(model.getCustomizationConfig().getCustomRetryStrategy())) {
             builder.addStatement("builder.option($1T.RETRY_STRATEGY, $2T.resolveRetryStrategy(config))",
-                            SdkClientOption.class,
-                            PoetUtils.classNameFromFqcn(model.getCustomizationConfig().getCustomRetryStrategy()));
+                                 SdkClientOption.class,
+                                 PoetUtils.classNameFromFqcn(model.getCustomizationConfig().getCustomRetryStrategy()));
         }
 
         if (StringUtils.isNotBlank(model.getCustomizationConfig().getCustomRetryPolicy())) {
@@ -485,7 +486,7 @@ private MethodSpec finalizeServiceConfigurationMethod() {
 
         if (endpointParamsKnowledgeIndex.hasAccountIdEndpointModeBuiltIn()) {
             builder.addStatement("builder.option($T.$L, resolveAccountIdEndpointMode(config))",
-                            AwsClientOption.class, model.getNamingStrategy().getEnumValueName("accountIdEndpointMode"));
+                                 AwsClientOption.class, model.getNamingStrategy().getEnumValueName("accountIdEndpointMode"));
         }
 
         String serviceNameForEnvVar = model.getNamingStrategy().getServiceNameForEnvironmentVariables();
@@ -829,7 +830,19 @@ private MethodSpec sigv4aSigningRegionSetMethod() {
     private MethodSpec defaultAuthSchemeProviderMethod() {
         return MethodSpec.methodBuilder("defaultAuthSchemeProvider")
                          .addModifiers(PRIVATE)
+                         .addParameter(SdkClientConfiguration.class, "config")
                          .returns(authSchemeSpecUtils.providerInterfaceName())
+                         .addCode("$T authSchemePreferenceProvider = "
+                                  + "$T.builder()",
+                                  AuthSchemePreferenceProvider.class, AuthSchemePreferenceProvider.class)
+                         .addCode(".profileFile(config.option($T.PROFILE_FILE_SUPPLIER))", SdkClientOption.class)
+                         .addCode(".profileName(config.option($T.PROFILE_NAME))", SdkClientOption.class)
+                         .addStatement(".build()")
+                         .addStatement("List<String> preferences = authSchemePreferenceProvider.resolveAuthSchemePreference()")
+                         .beginControlFlow("if(preferences != null && !preferences.isEmpty())")
+                         .addStatement("return $T.builder().withPreferredAuthSchemes(preferences).build()",
+                                       authSchemeSpecUtils.providerInterfaceName())
+                         .endControlFlow()
                          .addStatement("return $T.defaultProvider()", authSchemeSpecUtils.providerInterfaceName())
                          .build();
     }
@@ -965,10 +978,10 @@ private MethodSpec internalPluginsMethod() {
         List<String> internalPlugins = model.getCustomizationConfig().getInternalPlugins();
         if (internalPlugins.isEmpty()) {
             return builder.addStatement("return $T.emptyList()", Collections.class)
-                .build();
+                          .build();
         }
 
-        builder.addStatement("$T internalPlugins = new $T<>()", parameterizedTypeName,  ArrayList.class);
+        builder.addStatement("$T internalPlugins = new $T<>()", parameterizedTypeName, ArrayList.class);
 
         for (String internalPlugin : internalPlugins) {
             String arguments = internalPluginNewArguments(internalPlugin);