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,13 @@
 
 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.EnvironmentTokenMetricsInterceptorClass;
+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 +35,14 @@ 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());
+            tasks.add(createEnvironmentTokenMetricInterceptorTask());
+        }
+        return tasks;
     }
 
     private PoetGeneratorTask createClientOptionTask() {
@@ -46,6 +55,16 @@ private PoetGeneratorTask createUserAgentTask() {
                                      new UserAgentUtilsSpec(params.getModel()));
     }
 
+    private GeneratorTask createEnvironmentTokenSystemSettingTask() {
+        return new PoetGeneratorTask(clientOptionsDir(), params.getModel().getFileHeader(),
+                                     new EnvironmentTokenSystemSettingsClass(params.getModel()));
+    }
+
+    private GeneratorTask createEnvironmentTokenMetricInterceptorTask() {
+        return new PoetGeneratorTask(clientOptionsDir(), params.getModel().getFileHeader(),
+                                     new EnvironmentTokenMetricsInterceptorClass(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,14 @@ public ClassName getUserAgentClass() {
         return ClassName.get(model.getMetadata().getFullClientInternalPackageName(), "UserAgentUtils");
     }
 
+    public ClassName getEnvironmentTokenMetricsInterceptorClass() {
+        return ClassName.get(model.getMetadata().getFullClientInternalPackageName(), "EnvironmentTokenMetricsInterceptor");
+    }
+
+    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/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.client.builder.AwsDefaultClientBuilder;
@@ -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;
@@ -102,6 +104,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 +116,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 +271,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 +307,52 @@ 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.builder()"
+                          + ".withPreferredAuthSchemes($T.singletonList($S)).build())",
+                          SdkClientOption.class, authSchemeSpecUtils.providerInterfaceName(), Collections.class,
+                          "httpBearerAuth")
+            .addStatement("c.option($T.TOKEN_IDENTITY_PROVIDER, $T.create(tokenFromEnv::get))",
+                          AwsClientOption.class, StaticTokenProvider.class)
+            .addStatement("$T interceptors = c.option($T.EXECUTION_INTERCEPTORS)",
+                          ParameterizedTypeName.get(List.class, ExecutionInterceptor.class), SdkClientOption.class)
+            .addStatement("$T envTokenMetricInterceptors = $T.singletonList(new $T(tokenFromEnv.get()))",
+                          ParameterizedTypeName.get(List.class, ExecutionInterceptor.class), Collections.class,
+                          poetExtensions.getEnvironmentTokenMetricsInterceptorClass())
+            .addStatement("c.option($T.EXECUTION_INTERCEPTORS, $T.mergeLists(interceptors, envTokenMetricInterceptors))",
+                          SdkClientOption.class, CollectionUtils.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();

codegen/src/main/java/software/amazon/awssdk/codegen/poet/client/EnvironmentTokenMetricsInterceptorClass.java

@@ -0,0 +1,94 @@
+/*
+ * 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.client;
+
+import com.squareup.javapoet.ClassName;
+import com.squareup.javapoet.MethodSpec;
+import com.squareup.javapoet.TypeSpec;
+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.PoetExtension;
+import software.amazon.awssdk.codegen.poet.PoetUtils;
+import software.amazon.awssdk.core.SelectedAuthScheme;
+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.SdkInternalExecutionAttribute;
+import software.amazon.awssdk.core.useragent.BusinessMetricFeatureId;
+import software.amazon.awssdk.http.auth.scheme.BearerAuthScheme;
+import software.amazon.awssdk.identity.spi.TokenIdentity;
+
+public class EnvironmentTokenMetricsInterceptorClass implements ClassSpec {
+    protected final PoetExtension poetExtensions;
+
+    public EnvironmentTokenMetricsInterceptorClass(IntermediateModel model) {
+        this.poetExtensions = new PoetExtension(model);
+    }
+
+
+    @Override
+    public TypeSpec poetSpec() {
+        return TypeSpec.classBuilder(className())
+                       .addModifiers(Modifier.PUBLIC)
+                       .addAnnotation(PoetUtils.generatedAnnotation())
+                       .addAnnotation(SdkInternalApi.class)
+                       .addSuperinterface(ExecutionInterceptor.class)
+                       .addField(String.class, "tokenFromEnv", Modifier.PRIVATE, Modifier.FINAL)
+                       .addMethod(constructor())
+                       .addMethod(beforeExecutionMethod())
+                       .build();
+    }
+
+    private MethodSpec constructor() {
+        return MethodSpec.constructorBuilder()
+                         .addModifiers(Modifier.PUBLIC)
+                         .addParameter(String.class, "tokenFromEnv")
+                         .addStatement("this.tokenFromEnv = tokenFromEnv")
+                         .build();
+    }
+
+    @Override
+    public ClassName className() {
+        return poetExtensions.getEnvironmentTokenMetricsInterceptorClass();
+    }
+
+    private MethodSpec beforeExecutionMethod() {
+        return MethodSpec
+            .methodBuilder("beforeExecution")
+            .addAnnotation(Override.class)
+            .addModifiers(Modifier.PUBLIC)
+            .addParameter(Context.BeforeExecution.class, "context")
+            .addParameter(ExecutionAttributes.class, "executionAttributes")
+            .addStatement("$T<?> selectedAuthScheme = executionAttributes.getAttribute($T.SELECTED_AUTH_SCHEME)",
+                          SelectedAuthScheme.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();
+    }
+}