Codegen changes for AuthSchemeInterceptorSpecs to update the RegionSet from Execution Attributes

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

@@ -50,6 +50,7 @@
 import software.amazon.awssdk.core.internal.util.MetricUtils;
 import software.amazon.awssdk.core.metrics.CoreMetric;
 import software.amazon.awssdk.endpoints.EndpointProvider;
+import software.amazon.awssdk.http.auth.aws.signer.RegionSet;
 import software.amazon.awssdk.http.auth.spi.scheme.AuthScheme;
 import software.amazon.awssdk.http.auth.spi.scheme.AuthSchemeOption;
 import software.amazon.awssdk.http.auth.spi.signer.HttpSigner;
@@ -148,20 +149,17 @@ private MethodSpec generateAuthSchemeParams() {
         if (!authSchemeSpecUtils.useEndpointBasedAuthProvider()) {
             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("return $T.builder()"
-                                     + ".operation(operation)"
-                                     + ".region(region)"
-                                     + ".build()",
-                                     authSchemeSpecUtils.parametersInterfaceName());
-            } else {
-                builder.addStatement("return $T.builder()"
-                                     + ".operation(operation)"
-                                     + ".build()",
-                                     authSchemeSpecUtils.parametersInterfaceName());
+                builder.addStatement("builder.region(region)");
             }
+            addRegionSet(builder);
+            builder.addStatement("return builder.build()");
             return builder.build();
         }
 
@@ -198,6 +196,7 @@ 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.
         builder.addStatement("return builder.build()");
         return builder.build();
     }
@@ -449,4 +448,23 @@ private TypeName toTypeName(Object valueType) {
         }
         return result;
     }
+
+    private void addRegionSet(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,
+                RegionSet.class
+            );
+
+            builder.addStatement("builder.regionSet(regionSet)");
+        }
+    }
 }

codegen/src/test/java/software/amazon/awssdk/codegen/poet/auth/scheme/AuthSchemeSpecTest.java

@@ -208,6 +208,12 @@ static List<TestCase> parameters() {
                     .classSpecProvider(DefaultAuthSchemeParamsSpec::new)
                     .caseName("ops-auth-sigv4a-value")
                     .outputFileSuffix("default-params")
+                    .build(),
+            TestCase.builder()
+                    .modelProvider(ClientTestModels::opsWithSigv4a)
+                    .classSpecProvider(AuthSchemeInterceptorSpec::new)
+                    .caseName("ops-auth-sigv4a-value")
+                    .outputFileSuffix("interceptor")
                     .build()
         );
     }

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

@@ -0,0 +1,159 @@
+package software.amazon.awssdk.services.database.auth.scheme.internal;
+
+import java.time.Duration;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.CompletableFuture;
+import java.util.function.Supplier;
+import java.util.stream.Collectors;
+import software.amazon.awssdk.annotations.Generated;
+import software.amazon.awssdk.annotations.SdkInternalApi;
+import software.amazon.awssdk.awscore.AwsExecutionAttribute;
+import software.amazon.awssdk.core.SdkRequest;
+import software.amazon.awssdk.core.SelectedAuthScheme;
+import software.amazon.awssdk.core.exception.SdkException;
+import software.amazon.awssdk.core.interceptor.Context;
+import software.amazon.awssdk.core.interceptor.ExecutionAttributes;
+import software.amazon.awssdk.core.interceptor.ExecutionInterceptor;
+import software.amazon.awssdk.core.interceptor.SdkExecutionAttribute;
+import software.amazon.awssdk.core.interceptor.SdkInternalExecutionAttribute;
+import software.amazon.awssdk.core.internal.util.MetricUtils;
+import software.amazon.awssdk.core.metrics.CoreMetric;
+import software.amazon.awssdk.http.auth.aws.signer.RegionSet;
+import software.amazon.awssdk.http.auth.spi.scheme.AuthScheme;
+import software.amazon.awssdk.http.auth.spi.scheme.AuthSchemeOption;
+import software.amazon.awssdk.http.auth.spi.signer.HttpSigner;
+import software.amazon.awssdk.identity.spi.AwsCredentialsIdentity;
+import software.amazon.awssdk.identity.spi.Identity;
+import software.amazon.awssdk.identity.spi.IdentityProvider;
+import software.amazon.awssdk.identity.spi.IdentityProviders;
+import software.amazon.awssdk.identity.spi.ResolveIdentityRequest;
+import software.amazon.awssdk.identity.spi.TokenIdentity;
+import software.amazon.awssdk.metrics.MetricCollector;
+import software.amazon.awssdk.metrics.SdkMetric;
+import software.amazon.awssdk.regions.Region;
+import software.amazon.awssdk.services.database.auth.scheme.DatabaseAuthSchemeParams;
+import software.amazon.awssdk.services.database.auth.scheme.DatabaseAuthSchemeProvider;
+import software.amazon.awssdk.utils.Logger;
+import software.amazon.awssdk.utils.Validate;
+
+@Generated("software.amazon.awssdk:codegen")
+@SdkInternalApi
+public final class DatabaseAuthSchemeInterceptor implements ExecutionInterceptor {
+    private static Logger LOG = Logger.loggerFor(DatabaseAuthSchemeInterceptor.class);
+
+    @Override
+    public void beforeExecution(Context.BeforeExecution context, ExecutionAttributes executionAttributes) {
+        List<AuthSchemeOption> authOptions = resolveAuthOptions(context, executionAttributes);
+        SelectedAuthScheme<? extends Identity> selectedAuthScheme = selectAuthScheme(authOptions, executionAttributes);
+        putSelectedAuthScheme(executionAttributes, selectedAuthScheme);
+    }
+
+    private List<AuthSchemeOption> resolveAuthOptions(Context.BeforeExecution context, ExecutionAttributes executionAttributes) {
+        DatabaseAuthSchemeProvider authSchemeProvider = Validate.isInstanceOf(DatabaseAuthSchemeProvider.class,
+                                                                              executionAttributes.getAttribute(SdkInternalExecutionAttribute.AUTH_SCHEME_RESOLVER),
+                                                                              "Expected an instance of DatabaseAuthSchemeProvider");
+        DatabaseAuthSchemeParams params = authSchemeParams(context.request(), executionAttributes);
+        return authSchemeProvider.resolveAuthScheme(params);
+    }
+
+    private SelectedAuthScheme<? extends Identity> selectAuthScheme(List<AuthSchemeOption> authOptions,
+                                                                    ExecutionAttributes executionAttributes) {
+        MetricCollector metricCollector = executionAttributes.getAttribute(SdkExecutionAttribute.API_CALL_METRIC_COLLECTOR);
+        Map<String, AuthScheme<?>> authSchemes = executionAttributes.getAttribute(SdkInternalExecutionAttribute.AUTH_SCHEMES);
+        IdentityProviders identityProviders = executionAttributes.getAttribute(SdkInternalExecutionAttribute.IDENTITY_PROVIDERS);
+        List<Supplier<String>> discardedReasons = new ArrayList<>();
+        for (AuthSchemeOption authOption : authOptions) {
+            AuthScheme<?> authScheme = authSchemes.get(authOption.schemeId());
+            SelectedAuthScheme<? extends Identity> selectedAuthScheme = trySelectAuthScheme(authOption, authScheme,
+                                                                                            identityProviders, discardedReasons, metricCollector, executionAttributes);
+            if (selectedAuthScheme != null) {
+                if (!discardedReasons.isEmpty()) {
+                    LOG.debug(() -> String.format("%s auth will be used, discarded: '%s'", authOption.schemeId(),
+                                                  discardedReasons.stream().map(Supplier::get).collect(Collectors.joining(", "))));
+                }
+                return selectedAuthScheme;
+            }
+        }
+        throw SdkException
+            .builder()
+            .message(
+                "Failed to determine how to authenticate the user: "
+                + discardedReasons.stream().map(Supplier::get).collect(Collectors.joining(", "))).build();
+    }
+
+    private DatabaseAuthSchemeParams authSchemeParams(SdkRequest request, ExecutionAttributes executionAttributes) {
+        String operation = executionAttributes.getAttribute(SdkExecutionAttribute.OPERATION_NAME);
+        DatabaseAuthSchemeParams.Builder builder = DatabaseAuthSchemeParams.builder().operation(operation);
+        Region region = executionAttributes.getAttribute(AwsExecutionAttribute.AWS_REGION);
+        builder.region(region);
+        RegionSet regionSet = executionAttributes.getOptionalAttribute(AwsExecutionAttribute.AWS_SIGV4A_SIGNING_REGION_SET)
+                                                 .filter(regions -> !regions.isEmpty()).map(regions -> RegionSet.create(String.join(", ", regions)))
+                                                 .orElseGet(() -> {
+                                                     Region fallbackRegion = executionAttributes.getAttribute(AwsExecutionAttribute.AWS_REGION);
+                                                     return fallbackRegion != null ? RegionSet.create(fallbackRegion.toString()) : null;
+                                                 });
+        ;
+        builder.regionSet(regionSet);
+        return builder.build();
+    }
+
+    private <T extends Identity> SelectedAuthScheme<T> trySelectAuthScheme(AuthSchemeOption authOption, AuthScheme<T> authScheme,
+                                                                           IdentityProviders identityProviders, List<Supplier<String>> discardedReasons, MetricCollector metricCollector,
+                                                                           ExecutionAttributes executionAttributes) {
+        if (authScheme == null) {
+            discardedReasons.add(() -> String.format("'%s' is not enabled for this request.", authOption.schemeId()));
+            return null;
+        }
+        IdentityProvider<T> identityProvider = authScheme.identityProvider(identityProviders);
+        if (identityProvider == null) {
+            discardedReasons
+                .add(() -> String.format("'%s' does not have an identity provider configured.", authOption.schemeId()));
+            return null;
+        }
+        HttpSigner<T> signer;
+        try {
+            signer = authScheme.signer();
+        } catch (RuntimeException e) {
+            discardedReasons.add(() -> String.format("'%s' signer could not be retrieved: %s", authOption.schemeId(),
+                                                     e.getMessage()));
+            return null;
+        }
+        ResolveIdentityRequest.Builder identityRequestBuilder = ResolveIdentityRequest.builder();
+        authOption.forEachIdentityProperty(identityRequestBuilder::putProperty);
+        CompletableFuture<? extends T> identity;
+        SdkMetric<Duration> metric = getIdentityMetric(identityProvider);
+        if (metric == null) {
+            identity = identityProvider.resolveIdentity(identityRequestBuilder.build());
+        } else {
+            identity = MetricUtils.reportDuration(() -> identityProvider.resolveIdentity(identityRequestBuilder.build()),
+                                                  metricCollector, metric);
+        }
+        return new SelectedAuthScheme<>(identity, signer, authOption);
+    }
+
+    private SdkMetric<Duration> getIdentityMetric(IdentityProvider<?> identityProvider) {
+        Class<?> identityType = identityProvider.identityType();
+        if (identityType == AwsCredentialsIdentity.class) {
+            return CoreMetric.CREDENTIALS_FETCH_DURATION;
+        }
+        if (identityType == TokenIdentity.class) {
+            return CoreMetric.TOKEN_FETCH_DURATION;
+        }
+        return null;
+    }
+
+    private <T extends Identity> void putSelectedAuthScheme(ExecutionAttributes attributes,
+                                                            SelectedAuthScheme<T> selectedAuthScheme) {
+        SelectedAuthScheme<?> existingAuthScheme = attributes.getAttribute(SdkInternalExecutionAttribute.SELECTED_AUTH_SCHEME);
+        if (existingAuthScheme != null) {
+            AuthSchemeOption.Builder selectedOption = selectedAuthScheme.authSchemeOption().toBuilder();
+            existingAuthScheme.authSchemeOption().forEachIdentityProperty(selectedOption::putIdentityPropertyIfAbsent);
+            existingAuthScheme.authSchemeOption().forEachSignerProperty(selectedOption::putSignerPropertyIfAbsent);
+            selectedAuthScheme = new SelectedAuthScheme<>(selectedAuthScheme.identity(), selectedAuthScheme.signer(),
+                                                          selectedOption.build());
+        }
+        attributes.putAttribute(SdkInternalExecutionAttribute.SELECTED_AUTH_SCHEME, selectedAuthScheme);
+    }
+}

codegen/src/test/resources/software/amazon/awssdk/codegen/poet/auth/scheme/query-auth-scheme-interceptor.java

@@ -84,10 +84,11 @@ private SelectedAuthScheme<? extends Identity> selectAuthScheme(List<AuthSchemeO
 
     private QueryAuthSchemeParams authSchemeParams(SdkRequest request, ExecutionAttributes executionAttributes) {
         String operation = executionAttributes.getAttribute(SdkExecutionAttribute.OPERATION_NAME);
+        QueryAuthSchemeParams.Builder builder = QueryAuthSchemeParams.builder().operation(operation);
         Region region = executionAttributes.getAttribute(AwsExecutionAttribute.AWS_REGION);
-        return QueryAuthSchemeParams.builder().operation(operation).region(region).build();
+        builder.region(region);
+        return builder.build();
     }
-
     private <T extends Identity> SelectedAuthScheme<T> trySelectAuthScheme(AuthSchemeOption authOption, AuthScheme<T> authScheme,
                                                                            IdentityProviders identityProviders, List<Supplier<String>> discardedReasons, MetricCollector metricCollector,
                                                                            ExecutionAttributes executionAttributes) {

codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/c2j/ops-with-auth-sigv4a-value/service-2.json

@@ -6,7 +6,7 @@
         "globalEndpoint": "database-service.amazonaws.com",
         "protocol": "rest-json",
         "serviceAbbreviation": "Database Service",
-        "serviceFullName": "Some Service That Uses AWS Database Protocol",
+        "serviceFullName": "Some Service That Uses AWS Database Protocol With Sigv4a",
         "serviceId": "Database Service",
         "signingName": "database-service",
         "signatureVersion": "v4",

test/codegen-generated-classes-test/src/main/resources/codegen-resources/multiauth/customization.config

@@ -0,0 +1,4 @@
+{
+    "skipEndpointTestGeneration": true,
+     "useMultiAuth": true
+}