Skip to content

Commit 513f034

Browse files
gosargosar
authored
Check schemes match for SignerProperties in SRA client (#4505)
1 parent 43b5b43 commit 513f034

File tree

4 files changed

+227
-23
lines changed

4 files changed

+227
-23
lines changed

‎codegen/src/main/java/software/amazon/awssdk/codegen/poet/rules/EndpointResolverInterceptorSpec.java‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -561,6 +561,14 @@ private MethodSpec authSchemeWithEndpointSignerPropertiesMethod() {
561561

562562
method.beginControlFlow("for ($T endpointAuthScheme : endpointAuthSchemes)", EndpointAuthScheme.class);
563563

564+
if (useSraAuth) {
565+
// Don't include signer properties for auth options that don't match our selected auth scheme
566+
method.beginControlFlow("if (!endpointAuthScheme.schemeId()"
567+
+ ".equals(selectedAuthScheme.authSchemeOption().schemeId()))");
568+
method.addStatement("continue");
569+
method.endControlFlow();
570+
}
571+
564572
method.addStatement("$T option = selectedAuthScheme.authSchemeOption().toBuilder()", AuthSchemeOption.Builder.class);
565573

566574
if (dependsOnHttpAuthAws) {

‎codegen/src/test/java/software/amazon/awssdk/codegen/poet/rules/EndpointResolverInterceptorSpecTest.java‎

Lines changed: 15 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,13 +19,27 @@
1919
import static software.amazon.awssdk.codegen.poet.PoetMatchers.generatesTo;
2020

2121
import org.junit.jupiter.api.Test;
22+
import software.amazon.awssdk.codegen.model.intermediate.IntermediateModel;
2223
import software.amazon.awssdk.codegen.poet.ClassSpec;
2324
import software.amazon.awssdk.codegen.poet.ClientTestModels;
2425

2526
public class EndpointResolverInterceptorSpecTest {
2627
@Test
2728
public void endpointResolverInterceptorClass() {
28-
ClassSpec endpointProviderInterceptor = new EndpointResolverInterceptorSpec(ClientTestModels.queryServiceModels());
29+
ClassSpec endpointProviderInterceptor = new EndpointResolverInterceptorSpec(getModel(true));
2930
assertThat(endpointProviderInterceptor, generatesTo("endpoint-resolve-interceptor.java"));
3031
}
32+
33+
// TODO(post-sra-identity-auth): This can be deleted when useSraAuth is removed
34+
@Test
35+
public void endpointResolverInterceptorClass_preSra() {
36+
ClassSpec endpointProviderInterceptor = new EndpointResolverInterceptorSpec(getModel(false));
37+
assertThat(endpointProviderInterceptor, generatesTo("endpoint-resolve-interceptor-preSra.java"));
38+
}
39+
40+
private static IntermediateModel getModel(boolean useSraAuth) {
41+
IntermediateModel model = ClientTestModels.queryServiceModels();
42+
model.getCustomizationConfig().setUseSraAuth(useSraAuth);
43+
return model;
44+
}
3145
}

‎codegen/src/test/resources/software/amazon/awssdk/codegen/poet/rules/endpoint-resolve-interceptor-preSra.java‎

Lines changed: 201 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,201 @@
1+
package software.amazon.awssdk.services.query.endpoints.internal;
2+
3+
import java.util.List;
4+
import java.util.Optional;
5+
import java.util.concurrent.CompletionException;
6+
import java.util.function.Supplier;
7+
import software.amazon.awssdk.annotations.Generated;
8+
import software.amazon.awssdk.annotations.SdkInternalApi;
9+
import software.amazon.awssdk.auth.signer.Aws4Signer;
10+
import software.amazon.awssdk.auth.signer.SignerLoader;
11+
import software.amazon.awssdk.awscore.AwsExecutionAttribute;
12+
import software.amazon.awssdk.awscore.endpoints.AwsEndpointAttribute;
13+
import software.amazon.awssdk.awscore.endpoints.authscheme.EndpointAuthScheme;
14+
import software.amazon.awssdk.awscore.endpoints.authscheme.SigV4AuthScheme;
15+
import software.amazon.awssdk.awscore.endpoints.authscheme.SigV4aAuthScheme;
16+
import software.amazon.awssdk.awscore.util.SignerOverrideUtils;
17+
import software.amazon.awssdk.core.SdkRequest;
18+
import software.amazon.awssdk.core.SelectedAuthScheme;
19+
import software.amazon.awssdk.core.exception.SdkClientException;
20+
import software.amazon.awssdk.core.interceptor.Context;
21+
import software.amazon.awssdk.core.interceptor.ExecutionAttributes;
22+
import software.amazon.awssdk.core.interceptor.ExecutionInterceptor;
23+
import software.amazon.awssdk.core.interceptor.SdkExecutionAttribute;
24+
import software.amazon.awssdk.core.interceptor.SdkInternalExecutionAttribute;
25+
import software.amazon.awssdk.core.signer.Signer;
26+
import software.amazon.awssdk.endpoints.Endpoint;
27+
import software.amazon.awssdk.http.SdkHttpRequest;
28+
import software.amazon.awssdk.http.auth.aws.signer.AwsV4HttpSigner;
29+
import software.amazon.awssdk.http.auth.aws.signer.AwsV4aHttpSigner;
30+
import software.amazon.awssdk.http.auth.aws.signer.RegionSet;
31+
import software.amazon.awssdk.http.auth.spi.scheme.AuthSchemeOption;
32+
import software.amazon.awssdk.identity.spi.Identity;
33+
import software.amazon.awssdk.services.query.endpoints.QueryClientContextParams;
34+
import software.amazon.awssdk.services.query.endpoints.QueryEndpointParams;
35+
import software.amazon.awssdk.services.query.endpoints.QueryEndpointProvider;
36+
import software.amazon.awssdk.services.query.model.OperationWithContextParamRequest;
37+
import software.amazon.awssdk.utils.AttributeMap;
38+
39+
@Generated("software.amazon.awssdk:codegen")
40+
@SdkInternalApi
41+
public final class QueryResolveEndpointInterceptor implements ExecutionInterceptor {
42+
@Override
43+
public SdkRequest modifyRequest(Context.ModifyRequest context, ExecutionAttributes executionAttributes) {
44+
SdkRequest result = context.request();
45+
if (AwsEndpointProviderUtils.endpointIsDiscovered(executionAttributes)) {
46+
return result;
47+
}
48+
QueryEndpointProvider provider = (QueryEndpointProvider) executionAttributes
49+
.getAttribute(SdkInternalExecutionAttribute.ENDPOINT_PROVIDER);
50+
try {
51+
Endpoint endpoint = provider.resolveEndpoint(ruleParams(result, executionAttributes)).join();
52+
if (!AwsEndpointProviderUtils.disableHostPrefixInjection(executionAttributes)) {
53+
Optional<String> hostPrefix = hostPrefix(executionAttributes.getAttribute(SdkExecutionAttribute.OPERATION_NAME),
54+
result);
55+
if (hostPrefix.isPresent()) {
56+
endpoint = AwsEndpointProviderUtils.addHostPrefix(endpoint, hostPrefix.get());
57+
}
58+
}
59+
List<EndpointAuthScheme> endpointAuthSchemes = endpoint.attribute(AwsEndpointAttribute.AUTH_SCHEMES);
60+
SelectedAuthScheme<?> selectedAuthScheme = executionAttributes
61+
.getAttribute(SdkInternalExecutionAttribute.SELECTED_AUTH_SCHEME);
62+
if (endpointAuthSchemes != null && selectedAuthScheme != null) {
63+
selectedAuthScheme = authSchemeWithEndpointSignerProperties(endpointAuthSchemes, selectedAuthScheme);
64+
executionAttributes.putAttribute(SdkInternalExecutionAttribute.SELECTED_AUTH_SCHEME, selectedAuthScheme);
65+
}
66+
if (endpointAuthSchemes != null) {
67+
EndpointAuthScheme chosenAuthScheme = AuthSchemeUtils.chooseAuthScheme(endpointAuthSchemes);
68+
Supplier<Signer> signerProvider = signerProvider(chosenAuthScheme);
69+
result = SignerOverrideUtils.overrideSignerIfNotOverridden(result, executionAttributes, signerProvider);
70+
}
71+
executionAttributes.putAttribute(SdkInternalExecutionAttribute.RESOLVED_ENDPOINT, endpoint);
72+
return result;
73+
} catch (CompletionException e) {
74+
Throwable cause = e.getCause();
75+
if (cause instanceof SdkClientException) {
76+
throw (SdkClientException) cause;
77+
} else {
78+
throw SdkClientException.create("Endpoint resolution failed", cause);
79+
}
80+
}
81+
}
82+
83+
@Override
84+
public SdkHttpRequest modifyHttpRequest(Context.ModifyHttpRequest context, ExecutionAttributes executionAttributes) {
85+
Endpoint resolvedEndpoint = executionAttributes.getAttribute(SdkInternalExecutionAttribute.RESOLVED_ENDPOINT);
86+
if (resolvedEndpoint.headers().isEmpty()) {
87+
return context.httpRequest();
88+
}
89+
SdkHttpRequest.Builder httpRequestBuilder = context.httpRequest().toBuilder();
90+
resolvedEndpoint.headers().forEach((name, values) -> {
91+
values.forEach(v -> httpRequestBuilder.appendHeader(name, v));
92+
});
93+
return httpRequestBuilder.build();
94+
}
95+
96+
public static QueryEndpointParams ruleParams(SdkRequest request, ExecutionAttributes executionAttributes) {
97+
QueryEndpointParams.Builder builder = QueryEndpointParams.builder();
98+
builder.region(AwsEndpointProviderUtils.regionBuiltIn(executionAttributes));
99+
builder.useDualStackEndpoint(AwsEndpointProviderUtils.dualStackEnabledBuiltIn(executionAttributes));
100+
builder.useFipsEndpoint(AwsEndpointProviderUtils.fipsEnabledBuiltIn(executionAttributes));
101+
setClientContextParams(builder, executionAttributes);
102+
setContextParams(builder, executionAttributes.getAttribute(AwsExecutionAttribute.OPERATION_NAME), request);
103+
setStaticContextParams(builder, executionAttributes.getAttribute(AwsExecutionAttribute.OPERATION_NAME));
104+
return builder.build();
105+
}
106+
107+
private static void setContextParams(QueryEndpointParams.Builder params, String operationName, SdkRequest request) {
108+
switch (operationName) {
109+
case "OperationWithContextParam":
110+
setContextParams(params, (OperationWithContextParamRequest) request);
111+
break;
112+
default:
113+
break;
114+
}
115+
}
116+
117+
private static void setContextParams(QueryEndpointParams.Builder params, OperationWithContextParamRequest request) {
118+
params.operationContextParam(request.stringMember());
119+
}
120+
121+
private static void setStaticContextParams(QueryEndpointParams.Builder params, String operationName) {
122+
switch (operationName) {
123+
case "OperationWithStaticContextParams":
124+
operationWithStaticContextParamsStaticContextParams(params);
125+
break;
126+
default:
127+
break;
128+
}
129+
}
130+
131+
private static void operationWithStaticContextParamsStaticContextParams(QueryEndpointParams.Builder params) {
132+
params.staticStringParam("hello");
133+
}
134+
135+
private <T extends Identity> SelectedAuthScheme<T> authSchemeWithEndpointSignerProperties(
136+
List<EndpointAuthScheme> endpointAuthSchemes, SelectedAuthScheme<T> selectedAuthScheme) {
137+
for (EndpointAuthScheme endpointAuthScheme : endpointAuthSchemes) {
138+
AuthSchemeOption.Builder option = selectedAuthScheme.authSchemeOption().toBuilder();
139+
if (endpointAuthScheme instanceof SigV4AuthScheme) {
140+
SigV4AuthScheme v4AuthScheme = (SigV4AuthScheme) endpointAuthScheme;
141+
if (v4AuthScheme.isDisableDoubleEncodingSet()) {
142+
option.putSignerProperty(AwsV4HttpSigner.DOUBLE_URL_ENCODE, !v4AuthScheme.disableDoubleEncoding());
143+
}
144+
if (v4AuthScheme.signingRegion() != null) {
145+
option.putSignerProperty(AwsV4HttpSigner.REGION_NAME, v4AuthScheme.signingRegion());
146+
}
147+
if (v4AuthScheme.signingName() != null) {
148+
option.putSignerProperty(AwsV4HttpSigner.SERVICE_SIGNING_NAME, v4AuthScheme.signingName());
149+
}
150+
return new SelectedAuthScheme<>(selectedAuthScheme.identity(), selectedAuthScheme.signer(), option.build());
151+
}
152+
if (endpointAuthScheme instanceof SigV4aAuthScheme) {
153+
SigV4aAuthScheme v4aAuthScheme = (SigV4aAuthScheme) endpointAuthScheme;
154+
if (v4aAuthScheme.isDisableDoubleEncodingSet()) {
155+
option.putSignerProperty(AwsV4aHttpSigner.DOUBLE_URL_ENCODE, !v4aAuthScheme.disableDoubleEncoding());
156+
}
157+
if (v4aAuthScheme.signingRegionSet() != null) {
158+
RegionSet regionSet = RegionSet.create(v4aAuthScheme.signingRegionSet());
159+
option.putSignerProperty(AwsV4aHttpSigner.REGION_SET, regionSet);
160+
}
161+
if (v4aAuthScheme.signingName() != null) {
162+
option.putSignerProperty(AwsV4aHttpSigner.SERVICE_SIGNING_NAME, v4aAuthScheme.signingName());
163+
}
164+
return new SelectedAuthScheme<>(selectedAuthScheme.identity(), selectedAuthScheme.signer(), option.build());
165+
}
166+
throw new IllegalArgumentException("Endpoint auth scheme '" + endpointAuthScheme.name()
167+
+ "' cannot be mapped to the SDK auth scheme. Was it declared in the service's model?");
168+
}
169+
return selectedAuthScheme;
170+
}
171+
172+
private static void setClientContextParams(QueryEndpointParams.Builder params, ExecutionAttributes executionAttributes) {
173+
AttributeMap clientContextParams = executionAttributes.getAttribute(SdkInternalExecutionAttribute.CLIENT_CONTEXT_PARAMS);
174+
Optional.ofNullable(clientContextParams.get(QueryClientContextParams.BOOLEAN_CONTEXT_PARAM)).ifPresent(
175+
params::booleanContextParam);
176+
Optional.ofNullable(clientContextParams.get(QueryClientContextParams.STRING_CONTEXT_PARAM)).ifPresent(
177+
params::stringContextParam);
178+
}
179+
180+
private static Optional<String> hostPrefix(String operationName, SdkRequest request) {
181+
switch (operationName) {
182+
case "APostOperation": {
183+
return Optional.of("foo-");
184+
}
185+
default:
186+
return Optional.empty();
187+
}
188+
}
189+
190+
private Supplier<Signer> signerProvider(EndpointAuthScheme authScheme) {
191+
switch (authScheme.name()) {
192+
case "sigv4":
193+
return Aws4Signer::create;
194+
case "sigv4a":
195+
return SignerLoader::getSigV4aSigner;
196+
default:
197+
break;
198+
}
199+
throw SdkClientException.create("Don't know how to create signer for auth scheme: " + authScheme.name());
200+
}
201+
}

‎codegen/src/test/resources/software/amazon/awssdk/codegen/poet/rules/endpoint-resolve-interceptor.java‎

Lines changed: 3 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -3,17 +3,13 @@
33
import java.util.List;
44
import java.util.Optional;
55
import java.util.concurrent.CompletionException;
6-
import java.util.function.Supplier;
76
import software.amazon.awssdk.annotations.Generated;
87
import software.amazon.awssdk.annotations.SdkInternalApi;
9-
import software.amazon.awssdk.auth.signer.Aws4Signer;
10-
import software.amazon.awssdk.auth.signer.SignerLoader;
118
import software.amazon.awssdk.awscore.AwsExecutionAttribute;
129
import software.amazon.awssdk.awscore.endpoints.AwsEndpointAttribute;
1310
import software.amazon.awssdk.awscore.endpoints.authscheme.EndpointAuthScheme;
1411
import software.amazon.awssdk.awscore.endpoints.authscheme.SigV4AuthScheme;
1512
import software.amazon.awssdk.awscore.endpoints.authscheme.SigV4aAuthScheme;
16-
import software.amazon.awssdk.awscore.util.SignerOverrideUtils;
1713
import software.amazon.awssdk.core.SdkRequest;
1814
import software.amazon.awssdk.core.SelectedAuthScheme;
1915
import software.amazon.awssdk.core.exception.SdkClientException;
@@ -22,7 +18,6 @@
2218
import software.amazon.awssdk.core.interceptor.ExecutionInterceptor;
2319
import software.amazon.awssdk.core.interceptor.SdkExecutionAttribute;
2420
import software.amazon.awssdk.core.interceptor.SdkInternalExecutionAttribute;
25-
import software.amazon.awssdk.core.signer.Signer;
2621
import software.amazon.awssdk.endpoints.Endpoint;
2722
import software.amazon.awssdk.http.SdkHttpRequest;
2823
import software.amazon.awssdk.http.auth.aws.signer.AwsV4HttpSigner;
@@ -63,11 +58,6 @@ public SdkRequest modifyRequest(Context.ModifyRequest context, ExecutionAttribut
6358
selectedAuthScheme = authSchemeWithEndpointSignerProperties(endpointAuthSchemes, selectedAuthScheme);
6459
executionAttributes.putAttribute(SdkInternalExecutionAttribute.SELECTED_AUTH_SCHEME, selectedAuthScheme);
6560
}
66-
if (endpointAuthSchemes != null) {
67-
EndpointAuthScheme chosenAuthScheme = AuthSchemeUtils.chooseAuthScheme(endpointAuthSchemes);
68-
Supplier<Signer> signerProvider = signerProvider(chosenAuthScheme);
69-
result = SignerOverrideUtils.overrideSignerIfNotOverridden(result, executionAttributes, signerProvider);
70-
}
7161
executionAttributes.putAttribute(SdkInternalExecutionAttribute.RESOLVED_ENDPOINT, endpoint);
7262
return result;
7363
} catch (CompletionException e) {
@@ -135,6 +125,9 @@ private static void operationWithStaticContextParamsStaticContextParams(QueryEnd
135125
private <T extends Identity> SelectedAuthScheme<T> authSchemeWithEndpointSignerProperties(
136126
List<EndpointAuthScheme> endpointAuthSchemes, SelectedAuthScheme<T> selectedAuthScheme) {
137127
for (EndpointAuthScheme endpointAuthScheme : endpointAuthSchemes) {
128+
if (!endpointAuthScheme.schemeId().equals(selectedAuthScheme.authSchemeOption().schemeId())) {
129+
continue;
130+
}
138131
AuthSchemeOption.Builder option = selectedAuthScheme.authSchemeOption().toBuilder();
139132
if (endpointAuthScheme instanceof SigV4AuthScheme) {
140133
SigV4AuthScheme v4AuthScheme = (SigV4AuthScheme) endpointAuthScheme;
@@ -186,16 +179,4 @@ private static Optional<String> hostPrefix(String operationName, SdkRequest requ
186179
return Optional.empty();
187180
}
188181
}
189-
190-
private Supplier<Signer> signerProvider(EndpointAuthScheme authScheme) {
191-
switch (authScheme.name()) {
192-
case "sigv4":
193-
return Aws4Signer::create;
194-
case "sigv4a":
195-
return SignerLoader::getSigV4aSigner;
196-
default:
197-
break;
198-
}
199-
throw SdkClientException.create("Don't know how to create signer for auth scheme: " + authScheme.name());
200-
}
201182
}

0 commit comments

Comments
 (0)