Skip to content

Commit 34a6beb

Browse files
alextwoodsRanVaknindavidh44L-Applin
authored
Add Environment Token support (#6130)
* Prototype implementation of Auth Scheme Preference * Add tracking of explictly set token provider * Use generated PreferredAuthSchemeProvider to wrap/delegate * Add generic-service environment token provider + customization config (support for bedrock) * Use generated PreferredAuthSchemeProvider to wrap/delegate * Include tokenProvider in service config when bearer is on the model * Support sourcing token from jvm settings + env variable. * Refactor + use namingStrategy + add tests * Set business metric using an interceptor * Add ability to override token provider on request * Adding functionality to config preferred authschemeProvider * adding test coverage * fix formatting checkstyle * Added changelog * Fix checkstyle on prefered auth scheme provider * Add validation of service+customization * Refactor env token customizaiton logic + add more tests * Testing and cleanup * Adding test coverage * Use SdkSystemSetting for both env and system * Add changelog * Add profiles to service pom * Minor cleanups * Fix test * Fix protocol test dependencies * Remove dependency on profiles (use option get in generated code instead) * Fix checkstyle * Move AuthSchemePreferenceProvider out of internal module * More checkstyle fixes * Refactor and cleanup - move anon classes to full codegen classes. * Update docs * Additional codegen tests * Update core/aws-core/src/main/java/software/amazon/awssdk/awscore/internal/AwsExecutionContextBuilder.java Co-authored-by: David Ho <70000000+davidh44@users.noreply.github.com> * Add codegen tset for base client builder w/ env bearer token * Fixing comments, adding fixture file, renaming provider to resolver, adding coverage * Add async client test * Move metric interceptor logic from beforeExecute to beforeMarshall * Move metrics logic into auth scheme interceptor * Remove unused import * Add codegen test for preferred auth scheme provider * Check for empty string * Add more documentation * Revert "Add codegen test for preferred auth scheme provider" This reverts commit 141b9d6. * Replace authprovider builder with overridden defaultProvider method * Update core/aws-core/src/main/java/software/amazon/awssdk/awscore/auth/AuthSchemePreferenceResolver.java Co-authored-by: Olivier L Applin <olapplin@amazon.com> * Update core/aws-core/src/main/java/software/amazon/awssdk/awscore/auth/AuthSchemePreferenceResolver.java Co-authored-by: Olivier L Applin <olapplin@amazon.com> * Fix import from suggested changes * move test to aws-core module * Update codegen tests * Fix checkstyle * Fix checkstyle * Improve test coverage * Fix docs --------- Co-authored-by: Ran Vaknin <RanVaknin@users.noreply.github.com> Co-authored-by: David Ho <70000000+davidh44@users.noreply.github.com> Co-authored-by: RanVaknin <50976344+RanVaknin@users.noreply.github.com> Co-authored-by: Olivier L Applin <olapplin@amazon.com>
1 parent 992948a commit 34a6beb

File tree

53 files changed

+1986
-339
lines changed

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

53 files changed

+1986
-339
lines changed

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

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,6 @@
1+
{
2+
"type": "feature",
3+
"category": "AWS SDK for Java v2",
4+
"contributor": "",
5+
"description": "Adds support for configuring bearer auth using a token sourced from the environment for services with the `enableEnvironmentBearerToken` customization flag."
6+
}

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

Lines changed: 14 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -15,11 +15,12 @@
1515

1616
package software.amazon.awssdk.codegen.emitters.tasks;
1717

18-
import java.util.Arrays;
18+
import java.util.ArrayList;
1919
import java.util.List;
2020
import software.amazon.awssdk.codegen.emitters.GeneratorTask;
2121
import software.amazon.awssdk.codegen.emitters.GeneratorTaskParams;
2222
import software.amazon.awssdk.codegen.emitters.PoetGeneratorTask;
23+
import software.amazon.awssdk.codegen.poet.client.EnvironmentTokenSystemSettingsClass;
2324
import software.amazon.awssdk.codegen.poet.client.SdkClientOptions;
2425
import software.amazon.awssdk.codegen.poet.common.UserAgentUtilsSpec;
2526

@@ -33,7 +34,13 @@ public CommonInternalGeneratorTasks(GeneratorTaskParams params) {
3334

3435
@Override
3536
protected List<GeneratorTask> createTasks() throws Exception {
36-
return Arrays.asList(createClientOptionTask(), createUserAgentTask());
37+
List<GeneratorTask> tasks = new ArrayList<>();
38+
tasks.add(createClientOptionTask());
39+
tasks.add(createUserAgentTask());
40+
if (params.getModel().getCustomizationConfig().isEnableEnvironmentBearerToken()) {
41+
tasks.add(createEnvironmentTokenSystemSettingTask());
42+
}
43+
return tasks;
3744
}
3845

3946
private PoetGeneratorTask createClientOptionTask() {
@@ -46,6 +53,11 @@ private PoetGeneratorTask createUserAgentTask() {
4653
new UserAgentUtilsSpec(params.getModel()));
4754
}
4855

56+
private GeneratorTask createEnvironmentTokenSystemSettingTask() {
57+
return new PoetGeneratorTask(clientOptionsDir(), params.getModel().getFileHeader(),
58+
new EnvironmentTokenSystemSettingsClass(params.getModel()));
59+
}
60+
4961
private String clientOptionsDir() {
5062
return params.getPathProvider().getClientInternalDirectory();
5163
}

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

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -350,6 +350,13 @@ public class CustomizationConfig {
350350
*/
351351
private boolean enableFastUnmarshaller;
352352

353+
/**
354+
* A boolean flag to indicate if support for configuring a bearer token sourced from the environment should be added to the
355+
* generated service. When enabled, the generated client will use bearer auth with the token sourced from the
356+
* `AWS_BEARER_TOKEN_[SigningName]` environment variable.
357+
*/
358+
private boolean enableEnvironmentBearerToken = false;
359+
353360
private CustomizationConfig() {
354361
}
355362

@@ -924,4 +931,12 @@ public boolean getEnableFastUnmarshaller() {
924931
public void setEnableFastUnmarshaller(boolean enableFastUnmarshaller) {
925932
this.enableFastUnmarshaller = enableFastUnmarshaller;
926933
}
934+
935+
public boolean isEnableEnvironmentBearerToken() {
936+
return enableEnvironmentBearerToken;
937+
}
938+
939+
public void setEnableEnvironmentBearerToken(boolean enableEnvironmentBearerToken) {
940+
this.enableEnvironmentBearerToken = enableEnvironmentBearerToken;
941+
}
927942
}

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

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -432,6 +432,22 @@ private boolean isDisallowedNameForShape(String name, Shape parentShape) {
432432
}
433433
}
434434

435+
@Override
436+
public String getSigningName() {
437+
return Optional.ofNullable(serviceModel.getMetadata().getSigningName())
438+
.orElseGet(() -> serviceModel.getMetadata().getEndpointPrefix());
439+
}
440+
441+
@Override
442+
public String getSigningNameForEnvironmentVariables() {
443+
return screamCase(getSigningName());
444+
}
445+
446+
@Override
447+
public String getSigningNameForSystemProperties() {
448+
return pascalCase(getSigningName());
449+
}
450+
435451
@Override
436452
public void validateCustomerVisibleNaming(IntermediateModel trimmedModel) {
437453
Metadata metadata = trimmedModel.getMetadata();

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

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -200,6 +200,21 @@ public interface NamingStrategy {
200200
*/
201201
String getExistenceCheckMethodName(String memberName, Shape parentShape);
202202

203+
/**
204+
* Retrieve the service's signing name that should be used based on the model.
205+
*/
206+
String getSigningName();
207+
208+
/**
209+
* Retrieve the service's signing name that should be used for environment variables.
210+
*/
211+
String getSigningNameForEnvironmentVariables();
212+
213+
/**
214+
* Retrieve the service's signing name that should be used for system properties.
215+
*/
216+
String getSigningNameForSystemProperties();
217+
203218
/**
204219
* Verify the customer-visible naming in the provided intermediate model will compile and is idiomatic to Java.
205220
*/

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

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -79,6 +79,10 @@ public ClassName getUserAgentClass() {
7979
return ClassName.get(model.getMetadata().getFullClientInternalPackageName(), "UserAgentUtils");
8080
}
8181

82+
public ClassName getEnvironmentTokenSystemSettingsClass() {
83+
return ClassName.get(model.getMetadata().getFullClientInternalPackageName(), "EnvironmentTokenSystemSettings");
84+
}
85+
8286
/**
8387
* @param operationName Name of the operation
8488
* @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‎

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -49,8 +49,10 @@
4949
import software.amazon.awssdk.core.interceptor.SdkInternalExecutionAttribute;
5050
import software.amazon.awssdk.core.internal.util.MetricUtils;
5151
import software.amazon.awssdk.core.metrics.CoreMetric;
52+
import software.amazon.awssdk.core.useragent.BusinessMetricFeatureId;
5253
import software.amazon.awssdk.endpoints.EndpointProvider;
5354
import software.amazon.awssdk.http.auth.aws.signer.RegionSet;
55+
import software.amazon.awssdk.http.auth.scheme.BearerAuthScheme;
5456
import software.amazon.awssdk.http.auth.spi.scheme.AuthScheme;
5557
import software.amazon.awssdk.http.auth.spi.scheme.AuthSchemeOption;
5658
import software.amazon.awssdk.http.auth.spi.signer.HttpSigner;
@@ -70,8 +72,10 @@
7072
public final class AuthSchemeInterceptorSpec implements ClassSpec {
7173
private final AuthSchemeSpecUtils authSchemeSpecUtils;
7274
private final EndpointRulesSpecUtils endpointRulesSpecUtils;
75+
private final IntermediateModel intermediateModel;
7376

7477
public AuthSchemeInterceptorSpec(IntermediateModel intermediateModel) {
78+
this.intermediateModel = intermediateModel;
7579
this.authSchemeSpecUtils = new AuthSchemeSpecUtils(intermediateModel);
7680
this.endpointRulesSpecUtils = new EndpointRulesSpecUtils(intermediateModel);
7781
}
@@ -99,9 +103,42 @@ public TypeSpec poetSpec() {
99103
.addMethod(generateTrySelectAuthScheme())
100104
.addMethod(generateGetIdentityMetric())
101105
.addMethod(putSelectedAuthSchemeMethodSpec());
106+
if (intermediateModel.getCustomizationConfig().isEnableEnvironmentBearerToken()) {
107+
builder.addMethod(generateEnvironmentTokenMetric());
108+
}
102109
return builder.build();
103110
}
104111

112+
private MethodSpec generateEnvironmentTokenMetric() {
113+
return MethodSpec
114+
.methodBuilder("recordEnvironmentTokenBusinessMetric")
115+
.addModifiers(Modifier.PRIVATE)
116+
.addTypeVariable(TypeVariableName.get("T", Identity.class))
117+
.addParameter(ParameterSpec.builder(
118+
ParameterizedTypeName.get(ClassName.get(SelectedAuthScheme.class),
119+
TypeVariableName.get("T")),
120+
"selectedAuthScheme").build())
121+
.addParameter(ExecutionAttributes.class, "executionAttributes")
122+
.addStatement("$T tokenFromEnv = executionAttributes.getAttribute($T.TOKEN_CONFIGURED_FROM_ENV)",
123+
String.class, SdkInternalExecutionAttribute.class)
124+
.beginControlFlow("if (selectedAuthScheme != null && selectedAuthScheme.authSchemeOption().schemeId().equals($T"
125+
+ ".SCHEME_ID) && selectedAuthScheme.identity().isDone())", BearerAuthScheme.class)
126+
.beginControlFlow("if (selectedAuthScheme.identity().getNow(null) instanceof $T)", TokenIdentity.class)
127+
128+
.addStatement("$T configuredToken = ($T) selectedAuthScheme.identity().getNow(null)",
129+
TokenIdentity.class, TokenIdentity.class)
130+
.beginControlFlow("if (configuredToken.token().equals(tokenFromEnv))")
131+
.addStatement("executionAttributes.getAttribute($T.BUSINESS_METRICS)"
132+
+ ".addMetric($T.BEARER_SERVICE_ENV_VARS.value())",
133+
SdkInternalExecutionAttribute.class, BusinessMetricFeatureId.class)
134+
.endControlFlow()
135+
.endControlFlow()
136+
.endControlFlow()
137+
.build();
138+
139+
140+
}
141+
105142
private MethodSpec generateBeforeExecution() {
106143
MethodSpec.Builder builder = MethodSpec.methodBuilder("beforeExecution")
107144
.addAnnotation(Override.class)
@@ -116,6 +153,11 @@ private MethodSpec generateBeforeExecution() {
116153
.addStatement("$T selectedAuthScheme = selectAuthScheme(authOptions, executionAttributes)",
117154
wildcardSelectedAuthScheme())
118155
.addStatement("putSelectedAuthScheme(executionAttributes, selectedAuthScheme)");
156+
157+
if (intermediateModel.getCustomizationConfig().isEnableEnvironmentBearerToken()) {
158+
builder.addStatement("recordEnvironmentTokenBusinessMetric(selectedAuthScheme, "
159+
+ "executionAttributes)");
160+
}
119161
return builder.build();
120162
}
121163

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

Lines changed: 54 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,7 @@
4141
import software.amazon.awssdk.annotations.SdkInternalApi;
4242
import software.amazon.awssdk.auth.credentials.TokenUtils;
4343
import software.amazon.awssdk.auth.signer.Aws4Signer;
44+
import software.amazon.awssdk.auth.token.credentials.StaticTokenProvider;
4445
import software.amazon.awssdk.auth.token.credentials.aws.DefaultAwsTokenProvider;
4546
import software.amazon.awssdk.auth.token.signer.aws.BearerTokenSigner;
4647
import software.amazon.awssdk.awscore.auth.AuthSchemePreferenceResolver;
@@ -53,6 +54,7 @@
5354
import software.amazon.awssdk.codegen.model.service.AuthType;
5455
import software.amazon.awssdk.codegen.model.service.ClientContextParam;
5556
import software.amazon.awssdk.codegen.poet.ClassSpec;
57+
import software.amazon.awssdk.codegen.poet.PoetExtension;
5658
import software.amazon.awssdk.codegen.poet.PoetUtils;
5759
import software.amazon.awssdk.codegen.poet.auth.scheme.AuthSchemeSpecUtils;
5860
import software.amazon.awssdk.codegen.poet.auth.scheme.ModelAuthSchemeClassesKnowledgeIndex;
@@ -71,7 +73,9 @@
7173
import software.amazon.awssdk.core.client.config.SdkClientOption;
7274
import software.amazon.awssdk.core.endpointdiscovery.providers.DefaultEndpointDiscoveryProviderChain;
7375
import software.amazon.awssdk.core.interceptor.ClasspathInterceptorChainFactory;
76+
import software.amazon.awssdk.core.interceptor.ExecutionAttributes;
7477
import software.amazon.awssdk.core.interceptor.ExecutionInterceptor;
78+
import software.amazon.awssdk.core.interceptor.SdkInternalExecutionAttribute;
7579
import software.amazon.awssdk.core.retry.RetryMode;
7680
import software.amazon.awssdk.core.signer.Signer;
7781
import software.amazon.awssdk.http.Protocol;
@@ -102,6 +106,8 @@ public class BaseClientBuilderClass implements ClassSpec {
102106
private final AuthSchemeSpecUtils authSchemeSpecUtils;
103107
private final ServiceClientConfigurationUtils configurationUtils;
104108
private final EndpointParamsKnowledgeIndex endpointParamsKnowledgeIndex;
109+
private final PoetExtension poetExtensions;
110+
105111

106112
public BaseClientBuilderClass(IntermediateModel model) {
107113
this.model = model;
@@ -112,6 +118,7 @@ public BaseClientBuilderClass(IntermediateModel model) {
112118
this.authSchemeSpecUtils = new AuthSchemeSpecUtils(model);
113119
this.configurationUtils = new ServiceClientConfigurationUtils(model);
114120
this.endpointParamsKnowledgeIndex = EndpointParamsKnowledgeIndex.of(model);
121+
this.poetExtensions = new PoetExtension(model);
115122
}
116123

117124
@Override
@@ -266,24 +273,24 @@ private MethodSpec serviceNameMethod() {
266273
}
267274

268275
private MethodSpec mergeServiceDefaultsMethod() {
269-
boolean crc32FromCompressedDataEnabled = model.getCustomizationConfig().isCalculateCrc32FromCompressedData();
270-
271276
MethodSpec.Builder builder = MethodSpec.methodBuilder("mergeServiceDefaults")
272277
.addAnnotation(Override.class)
273278
.addModifiers(PROTECTED, FINAL)
274279
.returns(SdkClientConfiguration.class)
275-
.addParameter(SdkClientConfiguration.class, "config")
276-
.addCode("return config.merge(c -> c");
280+
.addParameter(SdkClientConfiguration.class, "config");
277281

278-
builder.addCode(".option($T.ENDPOINT_PROVIDER, defaultEndpointProvider())", SdkClientOption.class);
282+
boolean crc32FromCompressedDataEnabled = model.getCustomizationConfig().isCalculateCrc32FromCompressedData();
283+
284+
builder.beginControlFlow("return config.merge(c -> ");
285+
builder.addCode("c.option($T.ENDPOINT_PROVIDER, defaultEndpointProvider())", SdkClientOption.class);
279286

280287
if (authSchemeSpecUtils.useSraAuth()) {
281-
builder.addCode(".option($T.AUTH_SCHEME_PROVIDER, defaultAuthSchemeProvider(config))", SdkClientOption.class);
282-
builder.addCode(".option($T.AUTH_SCHEMES, authSchemes())", SdkClientOption.class);
283-
} else {
284-
if (defaultAwsAuthSignerMethod().isPresent()) {
285-
builder.addCode(".option($T.SIGNER, defaultSigner())\n", SdkAdvancedClientOption.class);
288+
if (!model.getCustomizationConfig().isEnableEnvironmentBearerToken()) {
289+
builder.addCode(".option($T.AUTH_SCHEME_PROVIDER, defaultAuthSchemeProvider(config))", SdkClientOption.class);
286290
}
291+
builder.addCode(".option($T.AUTH_SCHEMES, authSchemes())", SdkClientOption.class);
292+
} else if (defaultAwsAuthSignerMethod().isPresent()) {
293+
builder.addCode(".option($T.SIGNER, defaultSigner())\n", SdkAdvancedClientOption.class);
287294
}
288295
builder.addCode(".option($T.CRC32_FROM_COMPRESSED_DATA_ENABLED, $L)\n",
289296
SdkClientOption.class, crc32FromCompressedDataEnabled);
@@ -302,11 +309,47 @@ private MethodSpec mergeServiceDefaultsMethod() {
302309
builder.addCode(".option($T.TOKEN_SIGNER, defaultTokenSigner())", SdkAdvancedClientOption.class);
303310
}
304311
}
312+
builder.addStatement("");
305313

306-
builder.addCode(");");
314+
if (model.getCustomizationConfig().isEnableEnvironmentBearerToken()) {
315+
configureEnvironmentBearerToken(builder);
316+
}
317+
builder.endControlFlow(")");
307318
return builder.build();
308319
}
309320

321+
private void configureEnvironmentBearerToken(MethodSpec.Builder builder) {
322+
if (!authSchemeSpecUtils.useSraAuth()) {
323+
throw new IllegalStateException("The enableEnvironmentBearerToken customization requires SRA Auth.");
324+
}
325+
if (!AuthUtils.usesBearerAuth(model)) {
326+
throw new IllegalStateException("The enableEnvironmentBearerToken customization requires the service to model and "
327+
+ "support smithy.api#httpBearerAuth.");
328+
}
329+
330+
builder.addStatement("$T tokenFromEnv = new $T().getStringValue()",
331+
ParameterizedTypeName.get(Optional.class, String.class),
332+
poetExtensions.getEnvironmentTokenSystemSettingsClass());
333+
334+
builder
335+
.beginControlFlow("if (tokenFromEnv.isPresent() && config.option($T.AUTH_SCHEME_PROVIDER) == null && config.option($T"
336+
+ ".TOKEN_IDENTITY_PROVIDER) == null)",
337+
SdkClientOption.class, AwsClientOption.class)
338+
.addStatement("c.option($T.AUTH_SCHEME_PROVIDER, $T.defaultProvider($T.singletonList($S)))",
339+
SdkClientOption.class, authSchemeSpecUtils.providerInterfaceName(), Collections.class,
340+
"httpBearerAuth")
341+
.addStatement("c.option($T.TOKEN_IDENTITY_PROVIDER, $T.create(tokenFromEnv::get))",
342+
AwsClientOption.class, StaticTokenProvider.class)
343+
.addStatement("c.option($T.EXECUTION_ATTRIBUTES, "
344+
+ "$T.builder().put($T.TOKEN_CONFIGURED_FROM_ENV, tokenFromEnv.get()).build())",
345+
SdkClientOption.class, ExecutionAttributes.class, SdkInternalExecutionAttribute.class)
346+
.endControlFlow()
347+
.beginControlFlow("else")
348+
.addStatement("c.option($T.AUTH_SCHEME_PROVIDER, defaultAuthSchemeProvider(config))", SdkClientOption.class)
349+
.endControlFlow();
350+
351+
}
352+
310353
private Optional<MethodSpec> mergeInternalDefaultsMethod() {
311354
String userAgent = model.getCustomizationConfig().getUserAgent();
312355
RetryMode defaultRetryMode = model.getCustomizationConfig().getDefaultRetryMode();

0 commit comments

Comments
 (0)