Skip to content

Commit bfc51a7

Browse files
S-Saranya1S-Saranya1
authored
Move auth scheme resolution to pipeline stage from interceptors (#6755)
* Move auth scheme identity resolution to pipeline stage from interceptors * Address PR feedback: Added callback pattern for auth scheme options resolution Removed duplicate code Added fixture tests * Address Review comments: Move AuthSchemeResolutionStage to run before ApplyUserAgentStage and recording business metrics Simplify DefaultS3Presigner to use generated endpoint params * Remove identity provider override test from AwsExecutionContextBuilderTest, now happening in AuthSchemeResolutionStage
1 parent aca7500 commit bfc51a7

File tree

32 files changed

+1447
-358
lines changed

32 files changed

+1447
-358
lines changed

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

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -47,7 +47,6 @@ protected List<GeneratorTask> createTasks() {
4747
tasks.add(generateDefaultParamsImpl());
4848
tasks.add(generateModelBasedProvider());
4949
tasks.add(generatePreferenceProvider());
50-
tasks.add(generateAuthSchemeInterceptor());
5150
if (authSchemeSpecUtils.useEndpointBasedAuthProvider()) {
5251
tasks.add(generateEndpointBasedProvider());
5352
tasks.add(generateEndpointAwareAuthSchemeParams());

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

Lines changed: 30 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -60,14 +60,42 @@ public TypeSpec poetSpec() {
6060
.addModifiers(Modifier.PUBLIC)
6161
.addAnnotation(SdkPublicApi.class)
6262
.addJavadoc(interfaceJavadoc())
63-
.addMethod(builderMethod())
64-
.addType(builderInterfaceSpec());
63+
.addMethod(builderMethod());
64+
65+
if (authSchemeSpecUtils.generateEndpointBasedParams()) {
66+
b.addMethod(fromEndpointParamsMethod());
67+
}
68+
69+
b.addType(builderInterfaceSpec());
6570

6671
addAccessorMethods(b);
6772
addToBuilder(b);
6873
return b.build();
6974
}
7075

76+
private MethodSpec fromEndpointParamsMethod() {
77+
ClassName endpointParamsClass = endpointRulesSpecUtils.parametersClassName();
78+
MethodSpec.Builder builder = MethodSpec.methodBuilder("fromEndpointParams")
79+
.addModifiers(Modifier.PUBLIC, Modifier.STATIC)
80+
.addParameter(endpointParamsClass, "endpointParams")
81+
.returns(authSchemeSpecUtils.parametersInterfaceBuilderInterfaceName())
82+
.addJavadoc("Create a builder pre-populated with endpoint parameters.\n"
83+
+ "@param endpointParams the endpoint parameters to copy\n"
84+
+ "@return a builder with values from the endpoint parameters");
85+
86+
builder.addStatement("$T builder = builder()", authSchemeSpecUtils.parametersInterfaceBuilderInterfaceName());
87+
88+
parameters().forEach((name, model) -> {
89+
if (authSchemeSpecUtils.includeParamForProvider(name)) {
90+
String methodName = endpointRulesSpecUtils.paramMethodName(name);
91+
builder.addStatement("builder.$1N(endpointParams.$1N())", methodName);
92+
}
93+
});
94+
95+
builder.addStatement("return builder");
96+
return builder.build();
97+
}
98+
7199
private CodeBlock interfaceJavadoc() {
72100
CodeBlock.Builder b = CodeBlock.builder();
73101

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

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -370,7 +370,6 @@ private MethodSpec finalizeServiceConfigurationMethod() {
370370

371371
List<ClassName> builtInInterceptors = new ArrayList<>();
372372

373-
builtInInterceptors.add(authSchemeSpecUtils.authSchemeInterceptor());
374373
builtInInterceptors.add(endpointRulesSpecUtils.resolverInterceptorName());
375374
builtInInterceptors.add(endpointRulesSpecUtils.requestModifierInterceptorName());
376375

‎codegen/src/main/java/software/amazon/awssdk/codegen/poet/client/AsyncClientClass.java‎

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -69,10 +69,12 @@
6969
import software.amazon.awssdk.codegen.poet.PoetExtension;
7070
import software.amazon.awssdk.codegen.poet.PoetUtils;
7171
import software.amazon.awssdk.codegen.poet.StaticImport;
72+
import software.amazon.awssdk.codegen.poet.auth.scheme.AuthSchemeSpecUtils;
7273
import software.amazon.awssdk.codegen.poet.client.specs.ProtocolSpec;
7374
import software.amazon.awssdk.codegen.poet.eventstream.EventStreamUtils;
7475
import software.amazon.awssdk.codegen.poet.model.EventStreamSpecHelper;
7576
import software.amazon.awssdk.codegen.poet.model.ServiceClientConfigurationUtils;
77+
import software.amazon.awssdk.codegen.poet.rules.EndpointRulesSpecUtils;
7678
import software.amazon.awssdk.core.RequestOverrideConfiguration;
7779
import software.amazon.awssdk.core.async.AsyncResponseTransformer;
7880
import software.amazon.awssdk.core.async.AsyncResponseTransformerUtils;
@@ -100,6 +102,8 @@ public final class AsyncClientClass extends AsyncClientInterface {
100102
private final ProtocolSpec protocolSpec;
101103
private final ClassName serviceClientConfigurationClassName;
102104
private final ServiceClientConfigurationUtils configurationUtils;
105+
private final AuthSchemeSpecUtils authSchemeSpecUtils;
106+
private final EndpointRulesSpecUtils endpointRulesSpecUtils;
103107
private boolean hasScheduledExecutor;
104108

105109
public AsyncClientClass(GeneratorTaskParams dependencies) {
@@ -110,6 +114,8 @@ public AsyncClientClass(GeneratorTaskParams dependencies) {
110114
this.protocolSpec = getProtocolSpecs(poetExtensions, model);
111115
this.serviceClientConfigurationClassName = new PoetExtension(model).getServiceConfigClass();
112116
this.configurationUtils = new ServiceClientConfigurationUtils(model);
117+
this.authSchemeSpecUtils = new AuthSchemeSpecUtils(model);
118+
this.endpointRulesSpecUtils = new EndpointRulesSpecUtils(model);
113119
}
114120

115121
@Override
@@ -165,7 +171,8 @@ protected void addAdditionalMethods(TypeSpec.Builder type) {
165171
.addMethod(nameMethod())
166172
.addMethods(protocolSpec.additionalMethods())
167173
.addMethod(protocolSpec.initProtocolFactory(model))
168-
.addMethod(resolveMetricPublishersMethod());
174+
.addMethod(resolveMetricPublishersMethod())
175+
.addMethod(ClientClassUtils.resolveAuthSchemeOptionsMethod(authSchemeSpecUtils, endpointRulesSpecUtils));
169176

170177
type.addMethod(ClientClassUtils.updateRetryStrategyClientConfigurationMethod());
171178
type.addMethod(updateSdkClientConfigurationMethod(configurationUtils.serviceClientConfigurationBuilderClassName(),

‎codegen/src/main/java/software/amazon/awssdk/codegen/poet/client/ClientClassUtils.java‎

Lines changed: 136 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@
2929
import java.util.Map;
3030
import java.util.Objects;
3131
import java.util.Optional;
32+
import java.util.Set;
3233
import java.util.function.Consumer;
3334
import java.util.stream.Collectors;
3435
import javax.lang.model.element.Modifier;
@@ -45,14 +46,17 @@
4546
import software.amazon.awssdk.codegen.model.service.HostPrefixProcessor;
4647
import software.amazon.awssdk.codegen.poet.PoetExtension;
4748
import software.amazon.awssdk.codegen.poet.PoetUtils;
49+
import software.amazon.awssdk.codegen.poet.auth.scheme.AuthSchemeSpecUtils;
4850
import software.amazon.awssdk.codegen.poet.rules.EndpointRulesSpecUtils;
51+
import software.amazon.awssdk.core.SdkClient;
4952
import software.amazon.awssdk.core.SdkPlugin;
5053
import software.amazon.awssdk.core.SdkRequest;
5154
import software.amazon.awssdk.core.client.config.ClientOverrideConfiguration;
5255
import software.amazon.awssdk.core.client.config.SdkClientConfiguration;
5356
import software.amazon.awssdk.core.client.config.SdkClientOption;
5457
import software.amazon.awssdk.core.retry.RetryMode;
5558
import software.amazon.awssdk.core.signer.Signer;
59+
import software.amazon.awssdk.http.auth.spi.scheme.AuthSchemeOption;
5660
import software.amazon.awssdk.retries.api.RetryStrategy;
5761
import software.amazon.awssdk.utils.AttributeMap;
5862
import software.amazon.awssdk.utils.CollectionUtils;
@@ -346,4 +350,136 @@ public static MethodSpec updateRetryStrategyClientConfigurationMethod() {
346350
static String transformServiceId(String serviceId) {
347351
return serviceId.replace(" ", "_");
348352
}
353+
354+
static MethodSpec resolveAuthSchemeOptionsMethod(AuthSchemeSpecUtils authSchemeSpecUtils,
355+
EndpointRulesSpecUtils endpointRulesSpecUtils) {
356+
MethodSpec.Builder builder = MethodSpec.methodBuilder("resolveAuthSchemeOptions")
357+
.addModifiers(PRIVATE)
358+
.returns(ParameterizedTypeName.get(ClassName.get(List.class), ClassName.get(AuthSchemeOption.class)))
359+
.addParameter(SdkRequest.class, "request")
360+
.addParameter(String.class, "operationName")
361+
.addParameter(SdkClientConfiguration.class, "clientConfiguration");
362+
363+
ClassName providerInterface = authSchemeSpecUtils.providerInterfaceName();
364+
365+
builder.addStatement("$T authSchemeProvider = $T.isInstanceOf($T.class, "
366+
+ "clientConfiguration.option($T.AUTH_SCHEME_PROVIDER), $S)",
367+
providerInterface, Validate.class, providerInterface, SdkClientOption.class,
368+
"Expected an instance of " + authSchemeSpecUtils.providerInterfaceName().simpleName());
369+
370+
if (authSchemeSpecUtils.useEndpointBasedAuthProvider()) {
371+
addEndpointBasedAuthSchemeResolution(builder, authSchemeSpecUtils, endpointRulesSpecUtils);
372+
} else {
373+
addSimpleAuthSchemeResolution(builder, authSchemeSpecUtils);
374+
}
375+
376+
if (endpointRulesSpecUtils.isS3()) {
377+
ClassName sdkIdentityProperty = ClassName.get("software.amazon.awssdk.core.identity", "SdkIdentityProperty");
378+
builder.addStatement("$T sdkClient = clientConfiguration.option($T.SDK_CLIENT)",
379+
SdkClient.class, SdkClientOption.class);
380+
builder.addStatement("return options.stream().map(o -> o.toBuilder()"
381+
+ ".putIdentityProperty($T.SDK_CLIENT, sdkClient).build())"
382+
+ ".collect($T.toList())",
383+
sdkIdentityProperty, Collectors.class);
384+
} else {
385+
builder.addStatement("return options");
386+
}
387+
388+
return builder.build();
389+
}
390+
391+
private static void addSimpleAuthSchemeResolution(MethodSpec.Builder builder,
392+
AuthSchemeSpecUtils authSchemeSpecUtils) {
393+
ClassName paramsInterface = authSchemeSpecUtils.parametersInterfaceName();
394+
ClassName awsClientOption = ClassName.get("software.amazon.awssdk.awscore.client.config", "AwsClientOption");
395+
396+
builder.addStatement("$T.Builder paramsBuilder = $T.builder().operation(operationName)",
397+
paramsInterface, paramsInterface);
398+
399+
if (authSchemeSpecUtils.usesSigV4()) {
400+
builder.addStatement("paramsBuilder.region(clientConfiguration.option($T.AWS_REGION))", awsClientOption);
401+
}
402+
403+
if (authSchemeSpecUtils.hasSigV4aSupport()) {
404+
ClassName regionSet = ClassName.get("software.amazon.awssdk.http.auth.aws.signer", "RegionSet");
405+
builder.addStatement("$T<String> sigv4aRegionSet = clientConfiguration.option($T.AWS_SIGV4A_SIGNING_REGION_SET)",
406+
ClassName.get(Set.class), awsClientOption);
407+
builder.beginControlFlow("if (!$T.isNullOrEmpty(sigv4aRegionSet))", CollectionUtils.class);
408+
builder.addStatement("paramsBuilder.regionSet($T.create(sigv4aRegionSet))", regionSet);
409+
builder.nextControlFlow("else");
410+
builder.addStatement("paramsBuilder.regionSet($T.create(clientConfiguration.option($T.AWS_REGION).id()))",
411+
regionSet, awsClientOption);
412+
builder.endControlFlow();
413+
}
414+
415+
builder.addStatement("$T<$T> options = authSchemeProvider.resolveAuthScheme(paramsBuilder.build())",
416+
List.class, AuthSchemeOption.class);
417+
}
418+
419+
private static void addEndpointBasedAuthSchemeResolution(MethodSpec.Builder builder,
420+
AuthSchemeSpecUtils authSchemeSpecUtils,
421+
EndpointRulesSpecUtils endpointRulesSpecUtils) {
422+
ClassName paramsInterface = authSchemeSpecUtils.parametersInterfaceName();
423+
ClassName awsClientOption = ClassName.get("software.amazon.awssdk.awscore.client.config", "AwsClientOption");
424+
ClassName endpointParamsClass = endpointRulesSpecUtils.parametersClassName();
425+
ClassName resolverInterceptor = endpointRulesSpecUtils.resolverInterceptorName();
426+
ClassName executionAttributesClass = ClassName.get("software.amazon.awssdk.core.interceptor", "ExecutionAttributes");
427+
ClassName awsExecutionAttribute = ClassName.get("software.amazon.awssdk.awscore", "AwsExecutionAttribute");
428+
ClassName sdkExecutionAttribute = ClassName.get("software.amazon.awssdk.core.interceptor", "SdkExecutionAttribute");
429+
ClassName sdkInternalExecutionAttribute = ClassName.get("software.amazon.awssdk.core.interceptor",
430+
"SdkInternalExecutionAttribute");
431+
432+
builder.addStatement("$T executionAttributes = new $T()", executionAttributesClass, executionAttributesClass);
433+
builder.addStatement("executionAttributes.putAttribute($T.AWS_REGION, clientConfiguration.option($T.AWS_REGION))",
434+
awsExecutionAttribute, awsClientOption);
435+
builder.addStatement("executionAttributes.putAttribute($T.DUALSTACK_ENDPOINT_ENABLED, "
436+
+ "clientConfiguration.option($T.DUALSTACK_ENDPOINT_ENABLED))",
437+
awsExecutionAttribute, awsClientOption);
438+
builder.addStatement("executionAttributes.putAttribute($T.FIPS_ENDPOINT_ENABLED, "
439+
+ "clientConfiguration.option($T.FIPS_ENDPOINT_ENABLED))",
440+
awsExecutionAttribute, awsClientOption);
441+
builder.addStatement("executionAttributes.putAttribute($T.OPERATION_NAME, operationName)", sdkExecutionAttribute);
442+
builder.addStatement("executionAttributes.putAttribute($T.CLIENT_ENDPOINT_PROVIDER, "
443+
+ "clientConfiguration.option($T.CLIENT_ENDPOINT_PROVIDER))",
444+
sdkInternalExecutionAttribute, SdkClientOption.class);
445+
builder.addStatement("executionAttributes.putAttribute($T.CLIENT_CONTEXT_PARAMS, "
446+
+ "clientConfiguration.option($T.CLIENT_CONTEXT_PARAMS))",
447+
sdkInternalExecutionAttribute, SdkClientOption.class);
448+
449+
builder.addStatement("$T endpointParams = $T.ruleParams(request, executionAttributes)",
450+
endpointParamsClass, resolverInterceptor);
451+
452+
builder.addStatement("$T.Builder paramsBuilder = $T.builder()", paramsInterface, paramsInterface);
453+
454+
boolean regionIncluded = false;
455+
for (String paramName : endpointRulesSpecUtils.parameters().keySet()) {
456+
if (!authSchemeSpecUtils.includeParamForProvider(paramName)) {
457+
continue;
458+
}
459+
regionIncluded = regionIncluded || paramName.equalsIgnoreCase("region");
460+
String methodName = endpointRulesSpecUtils.paramMethodName(paramName);
461+
builder.addStatement("paramsBuilder.$1N(endpointParams.$1N())", methodName);
462+
}
463+
464+
builder.addStatement("paramsBuilder.operation(operationName)");
465+
466+
if (authSchemeSpecUtils.usesSigV4() && !regionIncluded) {
467+
builder.addStatement("paramsBuilder.region(clientConfiguration.option($T.AWS_REGION))", awsClientOption);
468+
}
469+
470+
ClassName paramsBuilderClass = authSchemeSpecUtils.parametersEndpointAwareDefaultImplName().nestedClass("Builder");
471+
ClassName endpointProviderInterface = endpointRulesSpecUtils.providerInterfaceName();
472+
473+
builder.beginControlFlow("if (paramsBuilder instanceof $T)", paramsBuilderClass);
474+
builder.addStatement("$T endpointProvider = clientConfiguration.option($T.ENDPOINT_PROVIDER)",
475+
ClassName.get("software.amazon.awssdk.endpoints", "EndpointProvider"), SdkClientOption.class);
476+
builder.beginControlFlow("if (endpointProvider instanceof $T)", endpointProviderInterface);
477+
builder.addStatement("(($T) paramsBuilder).endpointProvider(($T) endpointProvider)",
478+
paramsBuilderClass, endpointProviderInterface);
479+
builder.endControlFlow();
480+
builder.endControlFlow();
481+
482+
builder.addStatement("$T<$T> options = authSchemeProvider.resolveAuthScheme(paramsBuilder.build())",
483+
List.class, AuthSchemeOption.class);
484+
}
349485
}

‎codegen/src/main/java/software/amazon/awssdk/codegen/poet/client/SyncClientClass.java‎

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -54,12 +54,14 @@
5454
import software.amazon.awssdk.codegen.model.service.PreClientExecutionRequestCustomizer;
5555
import software.amazon.awssdk.codegen.poet.PoetExtension;
5656
import software.amazon.awssdk.codegen.poet.PoetUtils;
57+
import software.amazon.awssdk.codegen.poet.auth.scheme.AuthSchemeSpecUtils;
5758
import software.amazon.awssdk.codegen.poet.client.specs.Ec2ProtocolSpec;
5859
import software.amazon.awssdk.codegen.poet.client.specs.JsonProtocolSpec;
5960
import software.amazon.awssdk.codegen.poet.client.specs.ProtocolSpec;
6061
import software.amazon.awssdk.codegen.poet.client.specs.QueryProtocolSpec;
6162
import software.amazon.awssdk.codegen.poet.client.specs.XmlProtocolSpec;
6263
import software.amazon.awssdk.codegen.poet.model.ServiceClientConfigurationUtils;
64+
import software.amazon.awssdk.codegen.poet.rules.EndpointRulesSpecUtils;
6365
import software.amazon.awssdk.core.RequestOverrideConfiguration;
6466
import software.amazon.awssdk.core.client.config.SdkClientConfiguration;
6567
import software.amazon.awssdk.core.client.config.SdkClientOption;
@@ -83,6 +85,8 @@ public class SyncClientClass extends SyncClientInterface {
8385
private final ProtocolSpec protocolSpec;
8486
private final ClassName serviceClientConfigurationClassName;
8587
private final ServiceClientConfigurationUtils configurationUtils;
88+
private final AuthSchemeSpecUtils authSchemeSpecUtils;
89+
private final EndpointRulesSpecUtils endpointRulesSpecUtils;
8690

8791
public SyncClientClass(GeneratorTaskParams taskParams) {
8892
super(taskParams.getModel());
@@ -92,6 +96,8 @@ public SyncClientClass(GeneratorTaskParams taskParams) {
9296
this.protocolSpec = getProtocolSpecs(poetExtensions, model);
9397
this.serviceClientConfigurationClassName = new PoetExtension(model).getServiceConfigClass();
9498
this.configurationUtils = new ServiceClientConfigurationUtils(model);
99+
this.authSchemeSpecUtils = new AuthSchemeSpecUtils(model);
100+
this.endpointRulesSpecUtils = new EndpointRulesSpecUtils(model);
95101
}
96102

97103
@Override
@@ -133,7 +139,8 @@ protected void addAdditionalMethods(TypeSpec.Builder type) {
133139
type.addMethod(constructor())
134140
.addMethod(nameMethod())
135141
.addMethods(protocolSpec.additionalMethods())
136-
.addMethod(resolveMetricPublishersMethod());
142+
.addMethod(resolveMetricPublishersMethod())
143+
.addMethod(ClientClassUtils.resolveAuthSchemeOptionsMethod(authSchemeSpecUtils, endpointRulesSpecUtils));
137144

138145
protocolSpec.createErrorResponseHandler().ifPresent(type::addMethod);
139146
type.addMethod(ClientClassUtils.updateRetryStrategyClientConfigurationMethod());

‎codegen/src/main/java/software/amazon/awssdk/codegen/poet/client/specs/JsonProtocolSpec.java‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -221,7 +221,9 @@ public CodeBlock executionHandler(OperationModel opModel) {
221221
.add(credentialType(opModel, model))
222222
.add(".withRequestConfiguration(clientConfiguration)")
223223
.add(".withInput($L)\n", opModel.getInput().getVariableName())
224-
.add(".withMetricCollector(apiCallMetricCollector)")
224+
.add(".withMetricCollector(apiCallMetricCollector)\n")
225+
.add(".withAuthSchemeOptionsResolver(r -> resolveAuthSchemeOptions(r, $S, clientConfiguration))\n",
226+
opModel.getOperationName())
225227
.add(HttpChecksumRequiredTrait.putHttpChecksumAttribute(opModel))
226228
.add(HttpChecksumTrait.create(opModel));
227229

@@ -295,6 +297,8 @@ public CodeBlock asyncExecutionHandler(IntermediateModel intermediateModel, Oper
295297
.add(".withErrorResponseHandler(errorResponseHandler)\n")
296298
.add(".withRequestConfiguration(clientConfiguration)")
297299
.add(".withMetricCollector(apiCallMetricCollector)\n")
300+
.add(".withAuthSchemeOptionsResolver(r -> resolveAuthSchemeOptions(r, $S, clientConfiguration))\n",
301+
opModel.getOperationName())
298302
.add(hostPrefixExpression(opModel))
299303
.add(discoveredEndpoint(opModel))
300304
.add(credentialType(opModel, model))

0 commit comments

Comments
 (0)