Add Environment Token support

.changes/next-release/feature-AWSSDKforJavav2-1932d6a.json

@@ -0,0 +1,6 @@
+{
+    "type": "feature",
+    "category": "AWS SDK for Java v2",
+    "contributor": "",
+    "description": "Adds support for configuring bearer auth using a token sourced from the environment for services with the `enableEnvironmentBearerToken` customization flag."
+}

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

@@ -15,11 +15,12 @@
 
 package software.amazon.awssdk.codegen.emitters.tasks;
 
-import java.util.Arrays;
+import java.util.ArrayList;
 import java.util.List;
 import software.amazon.awssdk.codegen.emitters.GeneratorTask;
 import software.amazon.awssdk.codegen.emitters.GeneratorTaskParams;
 import software.amazon.awssdk.codegen.emitters.PoetGeneratorTask;
+import software.amazon.awssdk.codegen.poet.client.EnvironmentTokenSystemSettingsClass;
 import software.amazon.awssdk.codegen.poet.client.SdkClientOptions;
 import software.amazon.awssdk.codegen.poet.common.UserAgentUtilsSpec;
 
@@ -33,7 +34,13 @@ public CommonInternalGeneratorTasks(GeneratorTaskParams params) {
 
     @Override
     protected List<GeneratorTask> createTasks() throws Exception {
-        return Arrays.asList(createClientOptionTask(), createUserAgentTask());
+        List<GeneratorTask> tasks = new ArrayList<>();
+        tasks.add(createClientOptionTask());
+        tasks.add(createUserAgentTask());
+        if (params.getModel().getCustomizationConfig().isEnableEnvironmentBearerToken()) {
+            tasks.add(createEnvironmentTokenSystemSettingTask());
+        }
+        return tasks;
     }
 
     private PoetGeneratorTask createClientOptionTask() {
@@ -46,6 +53,11 @@ private PoetGeneratorTask createUserAgentTask() {
                                      new UserAgentUtilsSpec(params.getModel()));
     }
 
+    private GeneratorTask createEnvironmentTokenSystemSettingTask() {
+        return new PoetGeneratorTask(clientOptionsDir(), params.getModel().getFileHeader(),
+                                     new EnvironmentTokenSystemSettingsClass(params.getModel()));
+    }
+
     private String clientOptionsDir() {
         return params.getPathProvider().getClientInternalDirectory();
     }

codegen/src/main/java/software/amazon/awssdk/codegen/model/config/customization/CustomizationConfig.java

@@ -350,6 +350,13 @@ public class CustomizationConfig {
      */
     private boolean enableFastUnmarshaller;
 
+    /**
+     * A boolean flag to indicate if support for configuring a bearer token sourced from the environment should be added to the
+     * generated service. When enabled, the generated client will use bearer auth with the token sourced from the
+     * `AWS_BEARER_TOKEN_[SigningName]` environment variable.
+     */
+    private boolean enableEnvironmentBearerToken = false;
+
     private CustomizationConfig() {
     }
 
@@ -924,4 +931,12 @@ public boolean getEnableFastUnmarshaller() {
     public void setEnableFastUnmarshaller(boolean enableFastUnmarshaller) {
         this.enableFastUnmarshaller = enableFastUnmarshaller;
     }
+
+    public boolean isEnableEnvironmentBearerToken() {
+        return enableEnvironmentBearerToken;
+    }
+
+    public void setEnableEnvironmentBearerToken(boolean enableEnvironmentBearerToken) {
+        this.enableEnvironmentBearerToken = enableEnvironmentBearerToken;
+    }
 }

codegen/src/main/java/software/amazon/awssdk/codegen/naming/DefaultNamingStrategy.java

@@ -432,6 +432,22 @@ private boolean isDisallowedNameForShape(String name, Shape parentShape) {
         }
     }
 
+    @Override
+    public String getSigningName() {
+        return Optional.ofNullable(serviceModel.getMetadata().getSigningName())
+                       .orElseGet(() -> serviceModel.getMetadata().getEndpointPrefix());
+    }
+
+    @Override
+    public String getSigningNameForEnvironmentVariables() {
+        return screamCase(getSigningName());
+    }
+
+    @Override
+    public String getSigningNameForSystemProperties() {
+        return pascalCase(getSigningName());
+    }
+
     @Override
     public void validateCustomerVisibleNaming(IntermediateModel trimmedModel) {
         Metadata metadata = trimmedModel.getMetadata();

codegen/src/main/java/software/amazon/awssdk/codegen/naming/NamingStrategy.java

@@ -200,6 +200,21 @@ public interface NamingStrategy {
      */
     String getExistenceCheckMethodName(String memberName, Shape parentShape);
 
+    /**
+     * Retrieve the service's signing name that should be used based on the model.
+     */
+    String getSigningName();
+
+    /**
+     * Retrieve the service's signing name that should be used for environment variables.
+     */
+    String getSigningNameForEnvironmentVariables();
+
+    /**
+     * Retrieve the service's signing name that should be used for system properties.
+     */
+    String getSigningNameForSystemProperties();
+
     /**
      * Verify the customer-visible naming in the provided intermediate model will compile and is idiomatic to Java.
      */

codegen/src/main/java/software/amazon/awssdk/codegen/poet/PoetExtension.java

@@ -79,6 +79,10 @@ public ClassName getUserAgentClass() {
         return ClassName.get(model.getMetadata().getFullClientInternalPackageName(), "UserAgentUtils");
     }
 
+    public ClassName getEnvironmentTokenSystemSettingsClass() {
+        return ClassName.get(model.getMetadata().getFullClientInternalPackageName(), "EnvironmentTokenSystemSettings");
+    }
+
     /**
      * @param operationName Name of the operation
      * @return A Poet {@link ClassName} for the response type of a paginated operation in the base service package.

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

@@ -49,8 +49,10 @@
 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.core.useragent.BusinessMetricFeatureId;
 import software.amazon.awssdk.endpoints.EndpointProvider;
 import software.amazon.awssdk.http.auth.aws.signer.RegionSet;
+import software.amazon.awssdk.http.auth.scheme.BearerAuthScheme;
 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;
@@ -70,8 +72,10 @@
 public final class AuthSchemeInterceptorSpec implements ClassSpec {
     private final AuthSchemeSpecUtils authSchemeSpecUtils;
     private final EndpointRulesSpecUtils endpointRulesSpecUtils;
+    private final IntermediateModel intermediateModel;
 
     public AuthSchemeInterceptorSpec(IntermediateModel intermediateModel) {
+        this.intermediateModel = intermediateModel;
         this.authSchemeSpecUtils = new AuthSchemeSpecUtils(intermediateModel);
         this.endpointRulesSpecUtils = new EndpointRulesSpecUtils(intermediateModel);
     }
@@ -99,9 +103,42 @@ public TypeSpec poetSpec() {
                .addMethod(generateTrySelectAuthScheme())
                .addMethod(generateGetIdentityMetric())
                .addMethod(putSelectedAuthSchemeMethodSpec());
+        if (intermediateModel.getCustomizationConfig().isEnableEnvironmentBearerToken()) {
+            builder.addMethod(generateEnvironmentTokenMetric());
+        }
         return builder.build();
     }
 
+    private MethodSpec generateEnvironmentTokenMetric() {
+        return MethodSpec
+            .methodBuilder("recordEnvironmentTokenBusinessMetric")
+            .addModifiers(Modifier.PRIVATE)
+            .addTypeVariable(TypeVariableName.get("T", Identity.class))
+            .addParameter(ParameterSpec.builder(
+                ParameterizedTypeName.get(ClassName.get(SelectedAuthScheme.class),
+                                          TypeVariableName.get("T")),
+                "selectedAuthScheme").build())
+            .addParameter(ExecutionAttributes.class, "executionAttributes")
+            .addStatement("$T tokenFromEnv = executionAttributes.getAttribute($T.TOKEN_CONFIGURED_FROM_ENV)",
+                          String.class, SdkInternalExecutionAttribute.class)
+            .beginControlFlow("if (selectedAuthScheme != null && selectedAuthScheme.authSchemeOption().schemeId().equals($T"
+                              + ".SCHEME_ID) && selectedAuthScheme.identity().isDone())", BearerAuthScheme.class)
+            .beginControlFlow("if (selectedAuthScheme.identity().getNow(null) instanceof $T)", TokenIdentity.class)
+
+            .addStatement("$T configuredToken = ($T) selectedAuthScheme.identity().getNow(null)",
+                          TokenIdentity.class, TokenIdentity.class)
+            .beginControlFlow("if (configuredToken.token().equals(tokenFromEnv))")
+            .addStatement("executionAttributes.getAttribute($T.BUSINESS_METRICS)"
+                          + ".addMetric($T.BEARER_SERVICE_ENV_VARS.value())",
+                          SdkInternalExecutionAttribute.class, BusinessMetricFeatureId.class)
+            .endControlFlow()
+            .endControlFlow()
+            .endControlFlow()
+            .build();
+
+
+    }
+
     private MethodSpec generateBeforeExecution() {
         MethodSpec.Builder builder = MethodSpec.methodBuilder("beforeExecution")
                                                .addAnnotation(Override.class)
@@ -116,6 +153,11 @@ private MethodSpec generateBeforeExecution() {
                .addStatement("$T selectedAuthScheme = selectAuthScheme(authOptions, executionAttributes)",
                              wildcardSelectedAuthScheme())
                .addStatement("putSelectedAuthScheme(executionAttributes, selectedAuthScheme)");
+
+        if (intermediateModel.getCustomizationConfig().isEnableEnvironmentBearerToken()) {
+            builder.addStatement("recordEnvironmentTokenBusinessMetric(selectedAuthScheme, "
+                                 + "executionAttributes)");
+        }
         return builder.build();
     }
 

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

@@ -41,6 +41,7 @@
 import software.amazon.awssdk.annotations.SdkInternalApi;
 import software.amazon.awssdk.auth.credentials.TokenUtils;
 import software.amazon.awssdk.auth.signer.Aws4Signer;
+import software.amazon.awssdk.auth.token.credentials.StaticTokenProvider;
 import software.amazon.awssdk.auth.token.credentials.aws.DefaultAwsTokenProvider;
 import software.amazon.awssdk.auth.token.signer.aws.BearerTokenSigner;
 import software.amazon.awssdk.awscore.auth.AuthSchemePreferenceResolver;
@@ -53,6 +54,7 @@
 import software.amazon.awssdk.codegen.model.service.AuthType;
 import software.amazon.awssdk.codegen.model.service.ClientContextParam;
 import software.amazon.awssdk.codegen.poet.ClassSpec;
+import software.amazon.awssdk.codegen.poet.PoetExtension;
 import software.amazon.awssdk.codegen.poet.PoetUtils;
 import software.amazon.awssdk.codegen.poet.auth.scheme.AuthSchemeSpecUtils;
 import software.amazon.awssdk.codegen.poet.auth.scheme.ModelAuthSchemeClassesKnowledgeIndex;
@@ -71,7 +73,9 @@
 import software.amazon.awssdk.core.client.config.SdkClientOption;
 import software.amazon.awssdk.core.endpointdiscovery.providers.DefaultEndpointDiscoveryProviderChain;
 import software.amazon.awssdk.core.interceptor.ClasspathInterceptorChainFactory;
+import software.amazon.awssdk.core.interceptor.ExecutionAttributes;
 import software.amazon.awssdk.core.interceptor.ExecutionInterceptor;
+import software.amazon.awssdk.core.interceptor.SdkInternalExecutionAttribute;
 import software.amazon.awssdk.core.retry.RetryMode;
 import software.amazon.awssdk.core.signer.Signer;
 import software.amazon.awssdk.http.Protocol;
@@ -102,6 +106,8 @@ public class BaseClientBuilderClass implements ClassSpec {
     private final AuthSchemeSpecUtils authSchemeSpecUtils;
     private final ServiceClientConfigurationUtils configurationUtils;
     private final EndpointParamsKnowledgeIndex endpointParamsKnowledgeIndex;
+    private final PoetExtension poetExtensions;
+
 
     public BaseClientBuilderClass(IntermediateModel model) {
         this.model = model;
@@ -112,6 +118,7 @@ public BaseClientBuilderClass(IntermediateModel model) {
         this.authSchemeSpecUtils = new AuthSchemeSpecUtils(model);
         this.configurationUtils = new ServiceClientConfigurationUtils(model);
         this.endpointParamsKnowledgeIndex = EndpointParamsKnowledgeIndex.of(model);
+        this.poetExtensions = new PoetExtension(model);
     }
 
     @Override
@@ -266,24 +273,24 @@ private MethodSpec serviceNameMethod() {
     }
 
     private MethodSpec mergeServiceDefaultsMethod() {
-        boolean crc32FromCompressedDataEnabled = model.getCustomizationConfig().isCalculateCrc32FromCompressedData();
-
         MethodSpec.Builder builder = MethodSpec.methodBuilder("mergeServiceDefaults")
                                                .addAnnotation(Override.class)
                                                .addModifiers(PROTECTED, FINAL)
                                                .returns(SdkClientConfiguration.class)
-                                               .addParameter(SdkClientConfiguration.class, "config")
-                                               .addCode("return config.merge(c -> c");
+                                               .addParameter(SdkClientConfiguration.class, "config");
 
-        builder.addCode(".option($T.ENDPOINT_PROVIDER, defaultEndpointProvider())", SdkClientOption.class);
+        boolean crc32FromCompressedDataEnabled = model.getCustomizationConfig().isCalculateCrc32FromCompressedData();
+
+        builder.beginControlFlow("return config.merge(c -> ");
+        builder.addCode("c.option($T.ENDPOINT_PROVIDER, defaultEndpointProvider())", SdkClientOption.class);
 
         if (authSchemeSpecUtils.useSraAuth()) {
-            builder.addCode(".option($T.AUTH_SCHEME_PROVIDER, defaultAuthSchemeProvider(config))", SdkClientOption.class);
-            builder.addCode(".option($T.AUTH_SCHEMES, authSchemes())", SdkClientOption.class);
-        } else {
-            if (defaultAwsAuthSignerMethod().isPresent()) {
-                builder.addCode(".option($T.SIGNER, defaultSigner())\n", SdkAdvancedClientOption.class);
+            if (!model.getCustomizationConfig().isEnableEnvironmentBearerToken()) {
+                builder.addCode(".option($T.AUTH_SCHEME_PROVIDER, defaultAuthSchemeProvider(config))", SdkClientOption.class);
             }
+            builder.addCode(".option($T.AUTH_SCHEMES, authSchemes())", SdkClientOption.class);
+        } else if (defaultAwsAuthSignerMethod().isPresent()) {
+            builder.addCode(".option($T.SIGNER, defaultSigner())\n", SdkAdvancedClientOption.class);
         }
         builder.addCode(".option($T.CRC32_FROM_COMPRESSED_DATA_ENABLED, $L)\n",
                         SdkClientOption.class, crc32FromCompressedDataEnabled);
@@ -302,11 +309,47 @@ private MethodSpec mergeServiceDefaultsMethod() {
                 builder.addCode(".option($T.TOKEN_SIGNER, defaultTokenSigner())", SdkAdvancedClientOption.class);
             }
         }
+        builder.addStatement("");
 
-        builder.addCode(");");
+        if (model.getCustomizationConfig().isEnableEnvironmentBearerToken()) {
+            configureEnvironmentBearerToken(builder);
+        }
+        builder.endControlFlow(")");
         return builder.build();
     }
 
+    private void configureEnvironmentBearerToken(MethodSpec.Builder builder) {
+        if (!authSchemeSpecUtils.useSraAuth()) {
+            throw new IllegalStateException("The enableEnvironmentBearerToken customization requires SRA Auth.");
+        }
+        if (!AuthUtils.usesBearerAuth(model)) {
+            throw new IllegalStateException("The enableEnvironmentBearerToken customization requires the service to model and "
+                                            + "support smithy.api#httpBearerAuth.");
+        }
+
+        builder.addStatement("$T tokenFromEnv = new $T().getStringValue()",
+                             ParameterizedTypeName.get(Optional.class, String.class),
+                             poetExtensions.getEnvironmentTokenSystemSettingsClass());
+
+        builder
+            .beginControlFlow("if (tokenFromEnv.isPresent() && config.option($T.AUTH_SCHEME_PROVIDER) == null && config.option($T"
+                              + ".TOKEN_IDENTITY_PROVIDER) == null)",
+                              SdkClientOption.class, AwsClientOption.class)
+            .addStatement("c.option($T.AUTH_SCHEME_PROVIDER, $T.defaultProvider($T.singletonList($S)))",
+                          SdkClientOption.class, authSchemeSpecUtils.providerInterfaceName(), Collections.class,
+                          "httpBearerAuth")
+            .addStatement("c.option($T.TOKEN_IDENTITY_PROVIDER, $T.create(tokenFromEnv::get))",
+                          AwsClientOption.class, StaticTokenProvider.class)
+            .addStatement("c.option($T.EXECUTION_ATTRIBUTES, "
+                           + "$T.builder().put($T.TOKEN_CONFIGURED_FROM_ENV, tokenFromEnv.get()).build())",
+                          SdkClientOption.class, ExecutionAttributes.class, SdkInternalExecutionAttribute.class)
+            .endControlFlow()
+            .beginControlFlow("else")
+            .addStatement("c.option($T.AUTH_SCHEME_PROVIDER, defaultAuthSchemeProvider(config))", SdkClientOption.class)
+            .endControlFlow();
+
+    }
+
     private Optional<MethodSpec> mergeInternalDefaultsMethod() {
         String userAgent = model.getCustomizationConfig().getUserAgent();
         RetryMode defaultRetryMode = model.getCustomizationConfig().getDefaultRetryMode();