Add Endpoint Params to AuthSchemeParams for MultiAuth Scheme Resolution

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

@@ -48,6 +48,8 @@ protected List<GeneratorTask> createTasks() {
         tasks.add(generateAuthSchemeInterceptor());
         if (authSchemeSpecUtils.useEndpointBasedAuthProvider()) {
             tasks.add(generateEndpointBasedProvider());
+        }
+        if (authSchemeSpecUtils.useEndpointParamsInAuthScheme()) {
             tasks.add(generateEndpointAwareAuthSchemeParams());
         }
         return tasks;

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

@@ -63,6 +63,7 @@
 import software.amazon.awssdk.metrics.MetricCollector;
 import software.amazon.awssdk.metrics.SdkMetric;
 import software.amazon.awssdk.regions.Region;
+import software.amazon.awssdk.utils.CollectionUtils;
 import software.amazon.awssdk.utils.Logger;
 import software.amazon.awssdk.utils.Validate;
 
@@ -146,20 +147,23 @@ private MethodSpec generateAuthSchemeParams() {
                                                .addParameter(SdkRequest.class, "request")
                                                .addParameter(ExecutionAttributes.class, "executionAttributes");
 
-        if (!authSchemeSpecUtils.useEndpointBasedAuthProvider()) {
+        if (!authSchemeSpecUtils.useEndpointParamsInAuthScheme()) {
             builder.addStatement("$T operation = executionAttributes.getAttribute($T.OPERATION_NAME)", String.class,
                                  SdkExecutionAttribute.class);
-            builder.addStatement("$T.Builder builder = $T.builder().operation(operation)",
-                                 authSchemeSpecUtils.parametersInterfaceName(),
-                                 authSchemeSpecUtils.parametersInterfaceName());
-
             if (authSchemeSpecUtils.usesSigV4()) {
                 builder.addStatement("$T region = executionAttributes.getAttribute($T.AWS_REGION)", Region.class,
                                      AwsExecutionAttribute.class);
-                builder.addStatement("builder.region(region)");
+                builder.addStatement("return $T.builder()"
+                                     + ".operation(operation)"
+                                     + ".region(region)"
+                                     + ".build()",
+                                     authSchemeSpecUtils.parametersInterfaceName());
+            } else {
+                builder.addStatement("return $T.builder()"
+                                     + ".operation(operation)"
+                                     + ".build()",
+                                     authSchemeSpecUtils.parametersInterfaceName());
             }
-            generateSigv4aRegionSet(builder);
-            builder.addStatement("return builder.build()");
             return builder.build();
         }
 
@@ -196,7 +200,9 @@ private MethodSpec generateAuthSchemeParams() {
         builder.addStatement("(($T)builder).endpointProvider(($T)endpointProvider)", paramsBuilderClass, endpointProviderClass);
         builder.endControlFlow();
         builder.endControlFlow();
-        // TODO: Implement addRegionSet() for legacy services that resolve authentication from endpoints in one of next PRs.
+        if (authSchemeSpecUtils.hasMultiAuthSigvOrSigv4a()) {
+            generateSigv4aRegionSet(builder);
+        }
         builder.addStatement("return builder.build()");
         return builder.build();
     }
@@ -452,19 +458,13 @@ private TypeName toTypeName(Object valueType) {
     private void generateSigv4aRegionSet(MethodSpec.Builder builder) {
         if (authSchemeSpecUtils.usesSigV4a()) {
             builder.addStatement(
-                "$T regionSet = executionAttributes.getOptionalAttribute($T.AWS_SIGV4A_SIGNING_REGION_SET)\n" +
-                "    .filter(regions -> !regions.isEmpty())\n" +
-                "    .map(regions -> $T.create(String.join(\", \", regions)))\n" +
-                "    .orElseGet(() -> {\n" +
-                "        $T fallbackRegion = executionAttributes.getAttribute($T.AWS_REGION);\n" +
-                "        return fallbackRegion != null ? $T.create(fallbackRegion.toString()) : null;\n" +
-                "    });",
-                RegionSet.class, AwsExecutionAttribute.class,
-                RegionSet.class, Region.class, AwsExecutionAttribute.class,
+                "executionAttributes.getOptionalAttribute($T.AWS_SIGV4A_SIGNING_REGION_SET)\n" +
+                "                   .filter(regionSet -> !$T.isNullOrEmpty(regionSet))\n" +
+                "                   .ifPresent(nonEmptyRegionSet -> builder.regionSet($T.create(nonEmptyRegionSet)))",
+                AwsExecutionAttribute.class,
+                CollectionUtils.class,
                 RegionSet.class
             );
-
-            builder.addStatement("builder.regionSet(regionSet)");
         }
     }
 }

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

@@ -131,7 +131,7 @@ private void addAccessorMethods(TypeSpec.Builder b) {
                                   .build());
         }
 
-        if (authSchemeSpecUtils.generateEndpointBasedParams()) {
+        if (authSchemeSpecUtils.useEndpointParamsInAuthScheme()) {
             parameters().forEach((name, model) -> {
                 if (authSchemeSpecUtils.includeParam(name)) {
                     MethodSpec accessor = endpointRulesSpecUtils.parameterInterfaceAccessorMethod(name, model);
@@ -184,7 +184,7 @@ private void addBuilderSetterMethods(TypeSpec.Builder b) {
 
         }
 
-        if (authSchemeSpecUtils.generateEndpointBasedParams()) {
+        if (authSchemeSpecUtils.useEndpointParamsInAuthScheme()) {
             parameters().forEach((name, model) -> {
                 if (authSchemeSpecUtils.includeParam(name)) {
                     ClassName parametersInterfaceName = authSchemeSpecUtils.parametersInterfaceName();

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

@@ -25,8 +25,10 @@
 import java.util.Set;
 import software.amazon.awssdk.codegen.model.config.customization.CustomizationConfig;
 import software.amazon.awssdk.codegen.model.intermediate.IntermediateModel;
+import software.amazon.awssdk.codegen.model.service.AuthType;
 import software.amazon.awssdk.codegen.utils.AuthUtils;
 import software.amazon.awssdk.http.auth.spi.scheme.AuthSchemeOption;
+import software.amazon.awssdk.utils.CollectionUtils;
 
 public final class AuthSchemeSpecUtils {
     private static final Set<String> DEFAULT_AUTH_SCHEME_PARAMS = setOf("region", "operation");
@@ -147,6 +149,25 @@ public boolean includeParamForProvider(String name) {
         return true;
     }
 
+    //Multi-Auth option determined by "auth" trait on Service model or operation model.
+    public boolean hasMultiAuthSigvOrSigv4a() {
+        List<AuthType> authList = intermediateModel.getMetadata().getAuth();
+
+        return (!CollectionUtils.isNullOrEmpty(authList) &&
+                authList.stream().anyMatch(authType -> authType == AuthType.V4 || authType == AuthType.V4A))
+               ||
+               intermediateModel.getOperations()
+                                .values()
+                                .stream()
+                                .flatMap(operationModel -> operationModel.getAuth().stream())
+                                .anyMatch(authType -> authType == AuthType.V4 || authType == AuthType.V4A);
+    }
+
+    //Include Endpoint params in Auth Schemes to resolve the Endpoint for obtaining Signing properties in Multi Auth.
+    public boolean useEndpointParamsInAuthScheme() {
+        return generateEndpointBasedParams() || hasMultiAuthSigvOrSigv4a();
+    }
+
     public String serviceName() {
         return intermediateModel.getMetadata().getServiceName();
     }

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

@@ -60,7 +60,7 @@ public TypeSpec poetSpec() {
                                       .addMethod(builderMethod())
                                       .addType(builderImplSpec());
 
-        if (authSchemeSpecUtils.useEndpointBasedAuthProvider()) {
+        if (authSchemeSpecUtils.useEndpointParamsInAuthScheme()) {
             b.addSuperinterface(authSchemeSpecUtils.parametersEndpointAwareDefaultImplName());
         }
 
@@ -84,7 +84,7 @@ private MethodSpec constructor() {
             b.addStatement("this.regionSet = builder.regionSet");
         }
 
-        if (authSchemeSpecUtils.generateEndpointBasedParams()) {
+        if (authSchemeSpecUtils.useEndpointParamsInAuthScheme()) {
             parameters().forEach((name, model) -> {
                 if (authSchemeSpecUtils.includeParam(name)) {
                     String fieldName = authSchemeSpecUtils.paramMethodName(name);
@@ -122,7 +122,7 @@ private TypeSpec builderImplSpec() {
                                      .addModifiers(Modifier.PRIVATE, Modifier.STATIC, Modifier.FINAL)
                                      .addSuperinterface(authSchemeSpecUtils.parametersInterfaceBuilderInterfaceName());
 
-        if (authSchemeSpecUtils.useEndpointBasedAuthProvider()) {
+        if (authSchemeSpecUtils.useEndpointParamsInAuthScheme()) {
             b.addSuperinterface(authSchemeSpecUtils.parametersEndpointAwareDefaultImplName().nestedClass("Builder"));
         }
 
@@ -153,7 +153,7 @@ private void addBuilderConstructors(TypeSpec.Builder b) {
         if (authSchemeSpecUtils.usesSigV4a()) {
             builderFromInstance.addStatement("this.regionSet = params.regionSet");
         }
-        if (authSchemeSpecUtils.generateEndpointBasedParams()) {
+        if (authSchemeSpecUtils.useEndpointParamsInAuthScheme()) {
             parameters().forEach((name, model) -> {
                 if (authSchemeSpecUtils.includeParam(name)) {
                     builderFromInstance.addStatement("this.$1N = params.$1N", endpointRulesSpecUtils.variableName(name));
@@ -202,7 +202,7 @@ private void addFieldsAndAccessors(TypeSpec.Builder b) {
                                   .build());
         }
 
-        if (authSchemeSpecUtils.generateEndpointBasedParams()) {
+        if (authSchemeSpecUtils.useEndpointParamsInAuthScheme()) {
             parameters().forEach((name, model) -> {
                 if (authSchemeSpecUtils.includeParam(name)) {
                     b.addField(endpointRulesSpecUtils.parameterClassField(name, model));
@@ -255,7 +255,7 @@ private void addBuilderFieldsAndSetter(TypeSpec.Builder b) {
             b.addMethod(builderSetterMethod("regionSet", TypeName.get(RegionSet.class)));
         }
 
-        if (authSchemeSpecUtils.generateEndpointBasedParams()) {
+        if (authSchemeSpecUtils.useEndpointParamsInAuthScheme()) {
             parameters().forEach((name, model) -> {
                 if (authSchemeSpecUtils.includeParam(name)) {
                     b.addField(endpointRulesSpecUtils.parameterBuilderFieldSpec(name, model));

codegen/src/test/java/software/amazon/awssdk/codegen/poet/ClientTestModels.java

@@ -341,10 +341,15 @@ public static IntermediateModel opsWithSigv4a() {
         File customizationModel =
             new File(ClientTestModels.class.getResource("client/c2j/ops-with-auth-sigv4a-value/customization.config")
                                            .getFile());
+
+        File endpointRuleSetModel =
+            new File(ClientTestModels.class.getResource("client/c2j/ops-with-auth-sigv4a-value/endpoint-rule-set.json").getFile());
+
         C2jModels models = C2jModels
             .builder()
             .serviceModel(getServiceModel(serviceModel))
             .customizationConfig(getCustomizationConfig(customizationModel))
+            .endpointRuleSetModel(getEndpointRuleSet(endpointRuleSetModel))
             .build();
 
         return new IntermediateModelBuilder(models).build();

codegen/src/test/resources/software/amazon/awssdk/codegen/poet/auth/scheme/ops-auth-sigv4a-value-auth-scheme-default-params.java

@@ -5,21 +5,37 @@
 import software.amazon.awssdk.http.auth.aws.signer.RegionSet;
 import software.amazon.awssdk.regions.Region;
 import software.amazon.awssdk.services.database.auth.scheme.DatabaseAuthSchemeParams;
+import software.amazon.awssdk.services.database.endpoints.DatabaseEndpointProvider;
 import software.amazon.awssdk.utils.Validate;
 
 @Generated("software.amazon.awssdk:codegen")
 @SdkInternalApi
-public final class DefaultDatabaseAuthSchemeParams implements DatabaseAuthSchemeParams {
+public final class DefaultDatabaseAuthSchemeParams implements DatabaseAuthSchemeParams, DatabaseEndpointResolverAware {
     private final String operation;
 
     private final Region region;
 
     private final RegionSet regionSet;
 
+    private final Boolean useDualStackEndpoint;
+
+    private final Boolean useFIPSEndpoint;
+
+    private final String accountId;
+
+    private final String operationContextParam;
+
+    private final DatabaseEndpointProvider endpointProvider;
+
     private DefaultDatabaseAuthSchemeParams(Builder builder) {
         this.operation = Validate.paramNotNull(builder.operation, "operation");
         this.region = builder.region;
         this.regionSet = builder.regionSet;
+        this.useDualStackEndpoint = builder.useDualStackEndpoint;
+        this.useFIPSEndpoint = builder.useFIPSEndpoint;
+        this.accountId = builder.accountId;
+        this.operationContextParam = builder.operationContextParam;
+        this.endpointProvider = builder.endpointProvider;
     }
 
     public static DatabaseAuthSchemeParams.Builder builder() {
@@ -41,25 +57,65 @@ public RegionSet regionSet() {
         return regionSet;
     }
 
+    @Override
+    public Boolean useDualStackEndpoint() {
+        return useDualStackEndpoint;
+    }
+
+    @Override
+    public Boolean useFipsEndpoint() {
+        return useFIPSEndpoint;
+    }
+
+    @Override
+    public String accountId() {
+        return accountId;
+    }
+
+    @Override
+    public String operationContextParam() {
+        return operationContextParam;
+    }
+
+    @Override
+    public DatabaseEndpointProvider endpointProvider() {
+        return endpointProvider;
+    }
+
     @Override
     public DatabaseAuthSchemeParams.Builder toBuilder() {
         return new Builder(this);
     }
 
-    private static final class Builder implements DatabaseAuthSchemeParams.Builder {
+    private static final class Builder implements DatabaseAuthSchemeParams.Builder, DatabaseEndpointResolverAware.Builder {
         private String operation;
 
         private Region region;
 
         private RegionSet regionSet;
 
+        private Boolean useDualStackEndpoint;
+
+        private Boolean useFIPSEndpoint;
+
+        private String accountId;
+
+        private String operationContextParam;
+
+        private DatabaseEndpointProvider endpointProvider;
+
         Builder() {
         }
 
         Builder(DefaultDatabaseAuthSchemeParams params) {
             this.operation = params.operation;
             this.region = params.region;
             this.regionSet = params.regionSet;
+            this.useDualStackEndpoint = params.useDualStackEndpoint;
+            this.useFIPSEndpoint = params.useFIPSEndpoint;
+            this.accountId = params.accountId;
+            this.operationContextParam = params.operationContextParam;
+            this.endpointProvider = params.endpointProvider;
         }
 
         @Override
@@ -80,6 +136,36 @@ public Builder regionSet(RegionSet regionSet) {
             return this;
         }
 
+        @Override
+        public Builder useDualStackEndpoint(Boolean useDualStackEndpoint) {
+            this.useDualStackEndpoint = useDualStackEndpoint;
+            return this;
+        }
+
+        @Override
+        public Builder useFipsEndpoint(Boolean useFIPSEndpoint) {
+            this.useFIPSEndpoint = useFIPSEndpoint;
+            return this;
+        }
+
+        @Override
+        public Builder accountId(String accountId) {
+            this.accountId = accountId;
+            return this;
+        }
+
+        @Override
+        public Builder operationContextParam(String operationContextParam) {
+            this.operationContextParam = operationContextParam;
+            return this;
+        }
+
+        @Override
+        public Builder endpointProvider(DatabaseEndpointProvider endpointProvider) {
+            this.endpointProvider = endpointProvider;
+            return this;
+        }
+
         @Override
         public DatabaseAuthSchemeParams build() {
             return new DefaultDatabaseAuthSchemeParams(this);