Skip to content

Commit 40b06d0

Browse files
sugmanuesugmanue
authored
Endpoint based auth scheme resolver should honor endpoint overrides (#4838)
1 parent 0a6e10f commit 40b06d0

14 files changed

+434
-19
lines changed

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

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@
2525
import software.amazon.awssdk.codegen.poet.auth.scheme.AuthSchemeProviderSpec;
2626
import software.amazon.awssdk.codegen.poet.auth.scheme.AuthSchemeSpecUtils;
2727
import software.amazon.awssdk.codegen.poet.auth.scheme.DefaultAuthSchemeParamsSpec;
28+
import software.amazon.awssdk.codegen.poet.auth.scheme.EndpointAwareAuthSchemeParamsSpec;
2829
import software.amazon.awssdk.codegen.poet.auth.scheme.EndpointBasedAuthSchemeProviderSpec;
2930
import software.amazon.awssdk.codegen.poet.auth.scheme.ModelBasedAuthSchemeProviderSpec;
3031

@@ -47,6 +48,7 @@ protected List<GeneratorTask> createTasks() {
4748
tasks.add(generateAuthSchemeInterceptor());
4849
if (authSchemeSpecUtils.useEndpointBasedAuthProvider()) {
4950
tasks.add(generateEndpointBasedProvider());
51+
tasks.add(generateEndpointAwareAuthSchemeParams());
5052
}
5153
return tasks;
5254
}
@@ -72,6 +74,11 @@ private GeneratorTask generateEndpointBasedProvider() {
7274
new EndpointBasedAuthSchemeProviderSpec(model));
7375
}
7476

77+
private GeneratorTask generateEndpointAwareAuthSchemeParams() {
78+
return new PoetGeneratorTask(authSchemeDir(), model.getFileHeader(), new EndpointAwareAuthSchemeParamsSpec(model));
79+
80+
}
81+
7582
private GeneratorTask generateAuthSchemeInterceptor() {
7683
return new PoetGeneratorTask(authSchemeInternalDir(), model.getFileHeader(), new AuthSchemeInterceptorSpec(model));
7784
}

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

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -49,6 +49,7 @@
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.endpoints.EndpointProvider;
5253
import software.amazon.awssdk.http.auth.spi.scheme.AuthScheme;
5354
import software.amazon.awssdk.http.auth.spi.scheme.AuthSchemeOption;
5455
import software.amazon.awssdk.identity.spi.AwsCredentialsIdentity;
@@ -185,6 +186,17 @@ private MethodSpec generateAuthSchemeParams() {
185186
AwsExecutionAttribute.class);
186187
builder.addStatement("builder.region(region)");
187188
}
189+
ClassName paramsBuilderClass = authSchemeSpecUtils.parametersEndpointAwareDefaultImplName().nestedClass("Builder");
190+
builder.beginControlFlow("if (builder instanceof $T)",
191+
paramsBuilderClass);
192+
ClassName endpointProviderClass = endpointRulesSpecUtils.providerInterfaceName();
193+
builder.addStatement("$T endpointProvider = executionAttributes.getAttribute($T.ENDPOINT_PROVIDER)",
194+
EndpointProvider.class,
195+
SdkInternalExecutionAttribute.class);
196+
builder.beginControlFlow("if (endpointProvider instanceof $T)", endpointProviderClass);
197+
builder.addStatement("(($T)builder).endpointProvider(($T)endpointProvider)", paramsBuilderClass, endpointProviderClass);
198+
builder.endControlFlow();
199+
builder.endControlFlow();
188200
builder.addStatement("return builder.build()");
189201
return builder.build();
190202
}

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

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -79,6 +79,10 @@ public ClassName parametersInterfaceName() {
7979
return ClassName.get(basePackage(), intermediateModel.getMetadata().getServiceName() + "AuthSchemeParams");
8080
}
8181

82+
public ClassName parametersEndpointAwareDefaultImplName() {
83+
return ClassName.get(internalPackage(), intermediateModel.getMetadata().getServiceName() + "EndpointResolverAware");
84+
}
85+
8286
public ClassName parametersInterfaceBuilderInterfaceName() {
8387
return parametersInterfaceName().nestedClass("Builder");
8488
}

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

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -59,6 +59,10 @@ public TypeSpec poetSpec() {
5959
.addMethod(builderMethod())
6060
.addType(builderImplSpec());
6161

62+
if (authSchemeSpecUtils.useEndpointBasedAuthProvider()) {
63+
b.addSuperinterface(authSchemeSpecUtils.parametersEndpointAwareDefaultImplName());
64+
}
65+
6266
addFieldsAndAccessors(b);
6367
addToBuilder(b);
6468
return b.build();
@@ -89,6 +93,7 @@ private MethodSpec constructor() {
8993

9094
}
9195
});
96+
b.addStatement("this.endpointProvider = builder.endpointProvider");
9297
}
9398

9499
return b.build();
@@ -112,6 +117,10 @@ private TypeSpec builderImplSpec() {
112117
.addModifiers(Modifier.PRIVATE, Modifier.STATIC, Modifier.FINAL)
113118
.addSuperinterface(authSchemeSpecUtils.parametersInterfaceBuilderInterfaceName());
114119

120+
if (authSchemeSpecUtils.useEndpointBasedAuthProvider()) {
121+
b.addSuperinterface(authSchemeSpecUtils.parametersEndpointAwareDefaultImplName().nestedClass("Builder"));
122+
}
123+
115124
addBuilderConstructors(b);
116125
addBuilderFieldsAndSetter(b);
117126

@@ -142,6 +151,7 @@ private void addBuilderConstructors(TypeSpec.Builder b) {
142151
builderFromInstance.addStatement("this.$1N = params.$1N", endpointRulesSpecUtils.variableName(name));
143152
}
144153
});
154+
builderFromInstance.addStatement("this.endpointProvider = params.endpointProvider");
145155
}
146156
b.addMethod(builderFromInstance.build());
147157
}
@@ -181,6 +191,17 @@ private void addFieldsAndAccessors(TypeSpec.Builder b) {
181191
.build());
182192
}
183193
});
194+
ClassName endpointProvider = endpointRulesSpecUtils.providerInterfaceName();
195+
b.addField(FieldSpec.builder(endpointProvider, "endpointProvider")
196+
.addModifiers(Modifier.PRIVATE, Modifier.FINAL)
197+
.build());
198+
b.addMethod(MethodSpec.methodBuilder("endpointProvider")
199+
.addModifiers(Modifier.PUBLIC)
200+
.addAnnotation(Override.class)
201+
.returns(endpointProvider)
202+
.addStatement("return endpointProvider")
203+
.build());
204+
184205
}
185206
}
186207

@@ -213,6 +234,10 @@ private void addBuilderFieldsAndSetter(TypeSpec.Builder b) {
213234
b.addMethod(endpointRulesSpecUtils.parameterBuilderSetterMethod(className(), name, model));
214235
}
215236
});
237+
b.addField(FieldSpec.builder(endpointRulesSpecUtils.providerInterfaceName(), "endpointProvider")
238+
.addModifiers(Modifier.PRIVATE)
239+
.build());
240+
b.addMethod(builderSetterMethod("endpointProvider", endpointRulesSpecUtils.providerInterfaceName()));
216241
}
217242
}
218243

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

Lines changed: 75 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,75 @@
1+
/*
2+
* Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License").
5+
* You may not use this file except in compliance with the License.
6+
* A copy of the License is located at
7+
*
8+
* http://aws.amazon.com/apache2.0
9+
*
10+
* or in the "license" file accompanying this file. This file is distributed
11+
* on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
12+
* express or implied. See the License for the specific language governing
13+
* permissions and limitations under the License.
14+
*/
15+
16+
package software.amazon.awssdk.codegen.poet.auth.scheme;
17+
18+
import com.squareup.javapoet.ClassName;
19+
import com.squareup.javapoet.MethodSpec;
20+
import com.squareup.javapoet.TypeSpec;
21+
import javax.lang.model.element.Modifier;
22+
import software.amazon.awssdk.annotations.SdkInternalApi;
23+
import software.amazon.awssdk.codegen.model.intermediate.IntermediateModel;
24+
import software.amazon.awssdk.codegen.poet.ClassSpec;
25+
import software.amazon.awssdk.codegen.poet.PoetUtils;
26+
import software.amazon.awssdk.codegen.poet.rules.EndpointRulesSpecUtils;
27+
28+
public class EndpointAwareAuthSchemeParamsSpec implements ClassSpec {
29+
30+
private final AuthSchemeSpecUtils authSchemeSpecUtils;
31+
private final EndpointRulesSpecUtils endpointRulesSpecUtils;
32+
33+
public EndpointAwareAuthSchemeParamsSpec(IntermediateModel intermediateModel) {
34+
this.authSchemeSpecUtils = new AuthSchemeSpecUtils(intermediateModel);
35+
this.endpointRulesSpecUtils = new EndpointRulesSpecUtils(intermediateModel);
36+
}
37+
38+
@Override
39+
public ClassName className() {
40+
return authSchemeSpecUtils.parametersEndpointAwareDefaultImplName();
41+
}
42+
43+
@Override
44+
public TypeSpec poetSpec() {
45+
TypeSpec.Builder b = PoetUtils.createInterfaceBuilder(className())
46+
.addAnnotation(SdkInternalApi.class)
47+
.addMethod(endpointProviderMethod())
48+
.addType(builderSpec());
49+
return b.build();
50+
}
51+
52+
private ClassName builderClassName() {
53+
return className().nestedClass("Builder");
54+
}
55+
56+
private MethodSpec endpointProviderMethod() {
57+
return MethodSpec.methodBuilder("endpointProvider")
58+
.addModifiers(Modifier.ABSTRACT, Modifier.PUBLIC)
59+
.returns(endpointRulesSpecUtils.providerInterfaceName())
60+
.build();
61+
}
62+
63+
private TypeSpec builderSpec() {
64+
ClassName builderClassName = builderClassName();
65+
TypeSpec.Builder b = TypeSpec.interfaceBuilder(builderClassName)
66+
.addModifiers(Modifier.PUBLIC, Modifier.STATIC);
67+
b.addMethod(MethodSpec.methodBuilder("endpointProvider")
68+
.addModifiers(Modifier.ABSTRACT, Modifier.PUBLIC)
69+
.addParameter(endpointRulesSpecUtils.providerInterfaceName(), "endpointProvider")
70+
.returns(builderClassName)
71+
.build());
72+
73+
return b.build();
74+
}
75+
}

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

Lines changed: 21 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -72,6 +72,7 @@ public TypeSpec poetSpec() {
7272
.addField(endpointDelegateInstance())
7373
.addMethod(createMethod())
7474
.addMethod(resolveAuthSchemeMethod())
75+
.addMethod(endpointProvider())
7576
.build();
7677
}
7778

@@ -93,6 +94,25 @@ private FieldSpec endpointDelegateInstance() {
9394
.build();
9495
}
9596

97+
private MethodSpec endpointProvider() {
98+
ClassName endpointProviderClass = endpointRulesSpecUtils.providerInterfaceName();
99+
MethodSpec.Builder builder = MethodSpec.methodBuilder("endpointProvider")
100+
.addModifiers(Modifier.PRIVATE)
101+
.returns(endpointProviderClass)
102+
.addParameter(authSchemeSpecUtils.parametersInterfaceName(), "params");
103+
104+
ClassName endpointAwareParams = authSchemeSpecUtils.parametersEndpointAwareDefaultImplName();
105+
builder.beginControlFlow("if (params instanceof $T)", endpointAwareParams);
106+
builder.addStatement("$1T endpointAwareParams = ($1T) params", endpointAwareParams);
107+
builder.addStatement("$T endpointProvider = endpointAwareParams.endpointProvider()", endpointProviderClass);
108+
builder.beginControlFlow("if (endpointProvider != null)");
109+
builder.addStatement("return endpointProvider");
110+
builder.endControlFlow();
111+
builder.endControlFlow();
112+
builder.addStatement("return DELEGATE");
113+
return builder.build();
114+
}
115+
96116
private FieldSpec modeledResolverInstance() {
97117
return FieldSpec.builder(authSchemeSpecUtils.providerInterfaceName(), "MODELED_RESOLVER")
98118
.addModifiers(Modifier.PRIVATE, Modifier.STATIC, Modifier.FINAL)
@@ -124,7 +144,7 @@ private MethodSpec resolveAuthSchemeMethod() {
124144
}
125145
});
126146
spec.addStatement(".build()");
127-
spec.addStatement("$T endpoint = $T.joinLikeSync(DELEGATE.resolveEndpoint(endpointParameters))",
147+
spec.addStatement("$T endpoint = $T.joinLikeSync(endpointProvider(params).resolveEndpoint(endpointParameters))",
128148
Endpoint.class, CompletableFutureUtils.class);
129149
spec.addStatement("$T authSchemes = endpoint.attribute($T.AUTH_SCHEMES)",
130150
ParameterizedTypeName.get(List.class, EndpointAuthScheme.class), AwsEndpointAttribute.class);

‎codegen/src/test/resources/software/amazon/awssdk/codegen/poet/auth/scheme/query-endpoint-auth-params-auth-scheme-default-params-with-allowlist.java

Lines changed: 20 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -19,11 +19,12 @@
1919
import software.amazon.awssdk.annotations.SdkInternalApi;
2020
import software.amazon.awssdk.regions.Region;
2121
import software.amazon.awssdk.services.query.auth.scheme.QueryAuthSchemeParams;
22+
import software.amazon.awssdk.services.query.endpoints.QueryEndpointProvider;
2223
import software.amazon.awssdk.utils.Validate;
2324

2425
@Generated("software.amazon.awssdk:codegen")
2526
@SdkInternalApi
26-
public final class DefaultQueryAuthSchemeParams implements QueryAuthSchemeParams {
27+
public final class DefaultQueryAuthSchemeParams implements QueryAuthSchemeParams, QueryEndpointResolverAware {
2728
private final String operation;
2829

2930
private final Region region;
@@ -40,6 +41,8 @@ public final class DefaultQueryAuthSchemeParams implements QueryAuthSchemeParams
4041

4142
private final String operationContextParam;
4243

44+
private final QueryEndpointProvider endpointProvider;
45+
4346
private DefaultQueryAuthSchemeParams(Builder builder) {
4447
this.operation = Validate.paramNotNull(builder.operation, "operation");
4548
this.region = builder.region;
@@ -49,6 +52,7 @@ private DefaultQueryAuthSchemeParams(Builder builder) {
4952
this.booleanContextParam = builder.booleanContextParam;
5053
this.stringContextParam = builder.stringContextParam;
5154
this.operationContextParam = builder.operationContextParam;
55+
this.endpointProvider = builder.endpointProvider;
5256
}
5357

5458
public static QueryAuthSchemeParams.Builder builder() {
@@ -96,12 +100,17 @@ public String operationContextParam() {
96100
return operationContextParam;
97101
}
98102

103+
@Override
104+
public QueryEndpointProvider endpointProvider() {
105+
return endpointProvider;
106+
}
107+
99108
@Override
100109
public QueryAuthSchemeParams.Builder toBuilder() {
101110
return new Builder(this);
102111
}
103112

104-
private static final class Builder implements QueryAuthSchemeParams.Builder {
113+
private static final class Builder implements QueryAuthSchemeParams.Builder, QueryEndpointResolverAware.Builder {
105114
private String operation;
106115

107116
private Region region;
@@ -118,6 +127,8 @@ private static final class Builder implements QueryAuthSchemeParams.Builder {
118127

119128
private String operationContextParam;
120129

130+
private QueryEndpointProvider endpointProvider;
131+
121132
Builder() {
122133
}
123134

@@ -130,6 +141,7 @@ private static final class Builder implements QueryAuthSchemeParams.Builder {
130141
this.booleanContextParam = params.booleanContextParam;
131142
this.stringContextParam = params.stringContextParam;
132143
this.operationContextParam = params.operationContextParam;
144+
this.endpointProvider = params.endpointProvider;
133145
}
134146

135147
@Override
@@ -187,6 +199,12 @@ public Builder operationContextParam(String operationContextParam) {
187199
return this;
188200
}
189201

202+
@Override
203+
public Builder endpointProvider(QueryEndpointProvider endpointProvider) {
204+
this.endpointProvider = endpointProvider;
205+
return this;
206+
}
207+
190208
@Override
191209
public QueryAuthSchemeParams build() {
192210
return new DefaultQueryAuthSchemeParams(this);

0 commit comments

Comments
 (0)