Skip to content

Commit e8ca8aa

Browse files
cenedhryncenedhryn
authored
Adding tests for overriding credentials or credentials providers (#5534)
* Adding tests for overriding credentials or credentials providers through request override and execution interceptors * Update tests
1 parent ef181f4 commit e8ca8aa

File tree

1 file changed

+211
-0
lines changed

1 file changed

+211
-0
lines changed

‎test/codegen-generated-classes-test/src/test/java/software/amazon/awssdk/services/IdentityResolutionOverrideTest.java

Lines changed: 211 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,211 @@
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.services;
17+
18+
import static org.assertj.core.api.Assertions.assertThat;
19+
import static org.assertj.core.api.Assertions.assertThatThrownBy;
20+
import static org.assertj.core.api.Assertions.fail;
21+
22+
import java.util.ArrayList;
23+
import java.util.List;
24+
import java.util.Optional;
25+
import java.util.concurrent.CompletableFuture;
26+
import java.util.function.Consumer;
27+
import org.junit.jupiter.api.BeforeEach;
28+
import org.junit.jupiter.api.Test;
29+
import software.amazon.awssdk.auth.credentials.AwsBasicCredentials;
30+
import software.amazon.awssdk.auth.credentials.AwsCredentials;
31+
import software.amazon.awssdk.auth.credentials.StaticCredentialsProvider;
32+
import software.amazon.awssdk.auth.signer.AwsSignerExecutionAttribute;
33+
import software.amazon.awssdk.awscore.AwsRequest;
34+
import software.amazon.awssdk.awscore.AwsRequestOverrideConfiguration;
35+
import software.amazon.awssdk.core.RequestOverrideConfiguration;
36+
import software.amazon.awssdk.core.SdkRequest;
37+
import software.amazon.awssdk.core.SelectedAuthScheme;
38+
import software.amazon.awssdk.core.interceptor.Context;
39+
import software.amazon.awssdk.core.interceptor.ExecutionAttributes;
40+
import software.amazon.awssdk.core.interceptor.ExecutionInterceptor;
41+
import software.amazon.awssdk.core.interceptor.SdkInternalExecutionAttribute;
42+
import software.amazon.awssdk.identity.spi.Identity;
43+
import software.amazon.awssdk.regions.Region;
44+
import software.amazon.awssdk.services.restjsonendpointproviders.RestJsonEndpointProvidersClient;
45+
import software.amazon.awssdk.services.restjsonendpointproviders.RestJsonEndpointProvidersClientBuilder;
46+
import software.amazon.awssdk.utils.CompletableFutureUtils;
47+
48+
class IdentityResolutionOverrideTest {
49+
50+
private static final AwsCredentials CLIENT_CREDENTIALS =
51+
AwsBasicCredentials.create("akid", "skid");
52+
53+
private static final AwsCredentials OVERRIDE_CREDENTIALS =
54+
AwsBasicCredentials.create("akidOverride", "skidOverride");
55+
56+
private CapturingInterceptor capturingInterceptor;
57+
58+
@BeforeEach
59+
public void setup() {
60+
this.capturingInterceptor = new CapturingInterceptor();
61+
}
62+
63+
@Test
64+
void when_apiCall_setsCredentialsProviderInRequestOverride_overrideCredentialsAreUsed() {
65+
RestJsonEndpointProvidersClient syncClient = syncClientBuilder().build();
66+
67+
assertThatThrownBy(() -> syncClient.allTypes(r -> r.overrideConfiguration(
68+
c -> c.credentialsProvider(StaticCredentialsProvider.create(OVERRIDE_CREDENTIALS)))))
69+
.hasMessageContaining("stop");
70+
71+
assertSelectedAuthSchemeBeforeTransmissionContains(OVERRIDE_CREDENTIALS);
72+
}
73+
74+
// Changing the credentials provider in modifyRequest does not work in SRA identity resolution
75+
// Identity is resolved in beforeExecution (and happens before user applied interceptors) and cannot
76+
// be affected by execution interceptors.
77+
@Test
78+
void when_executionInterceptorModifyRequest_setsCredentialProviderInRequestOverride_clientCredentialsAreUsed() {
79+
ExecutionInterceptor overridingInterceptor =
80+
new OverrideInterceptor(b -> b.credentialsProvider(StaticCredentialsProvider.create(OVERRIDE_CREDENTIALS)));
81+
82+
RestJsonEndpointProvidersClient syncClient = syncClientBuilder(overridingInterceptor).build();
83+
84+
assertThatThrownBy(() -> syncClient.allTypes(r -> {})).hasMessageContaining("stop");
85+
86+
assertSelectedAuthSchemeBeforeTransmissionContains(CLIENT_CREDENTIALS);
87+
}
88+
89+
@Test
90+
void when_apiCall_setsCredentialsExecutionAttributeInRequestOverride_clientCredentialsAreUsed() {
91+
RestJsonEndpointProvidersClient syncClient = syncClientBuilder().build();
92+
93+
assertThatThrownBy(() -> syncClient.allTypes(r -> r.overrideConfiguration(
94+
c -> c.putExecutionAttribute(AwsSignerExecutionAttribute.AWS_CREDENTIALS, OVERRIDE_CREDENTIALS))))
95+
.hasMessageContaining("stop");
96+
97+
assertSelectedAuthSchemeBeforeTransmissionContains(CLIENT_CREDENTIALS);
98+
}
99+
100+
// Updating an execution attribute inside the request override, in an interceptor doesn't work
101+
// Execution attributes are merged before interceptor hooks
102+
@Test
103+
void when_executionInterceptorModifyRequest_setsCredentialsExecutionAttributeInRequestOverride_clientCredentialsAreUsed() {
104+
ExecutionInterceptor overridingInterceptor =
105+
new OverrideInterceptor(b -> b.putExecutionAttribute(AwsSignerExecutionAttribute.AWS_CREDENTIALS, OVERRIDE_CREDENTIALS));
106+
107+
RestJsonEndpointProvidersClient syncClient = syncClientBuilder(overridingInterceptor).build();
108+
109+
assertThatThrownBy(() -> syncClient.allTypes(r -> {})).hasMessageContaining("stop");
110+
111+
assertSelectedAuthSchemeBeforeTransmissionContains(CLIENT_CREDENTIALS);
112+
}
113+
114+
// Updating the AWS_CREDENTIALS pre-SRA credentials execution attribute in an interceptor works
115+
// Resolved credentials are synced both ways with AWS_CREDENTIALS.
116+
@Test
117+
void when_executionInterceptorModifyRequest_overridesCredentialsExecutionAttribute_overrideCredentialsAreUsed() {
118+
ExecutionInterceptor overridingInterceptor = new ExecutionInterceptor() {
119+
@Override
120+
public SdkRequest modifyRequest(Context.ModifyRequest context, ExecutionAttributes executionAttributes) {
121+
executionAttributes.putAttribute(AwsSignerExecutionAttribute.AWS_CREDENTIALS, OVERRIDE_CREDENTIALS);
122+
return ExecutionInterceptor.super.modifyRequest(context, executionAttributes);
123+
}
124+
};
125+
126+
RestJsonEndpointProvidersClient syncClient = syncClientBuilder(overridingInterceptor).build();
127+
128+
assertThatThrownBy(() -> syncClient.allTypes(r -> {})).hasMessageContaining("stop");
129+
130+
assertSelectedAuthSchemeBeforeTransmissionContains(OVERRIDE_CREDENTIALS);
131+
}
132+
133+
private void assertSelectedAuthSchemeBeforeTransmissionContains(AwsCredentials overriddenCredentials) {
134+
SelectedAuthScheme<?> auth = capturingInterceptor.executionAttributes().getAttribute(SdkInternalExecutionAttribute.SELECTED_AUTH_SCHEME);
135+
Optional<AwsCredentials> requestCredentials = Optional.empty();
136+
try {
137+
requestCredentials = providerNameFromIdentity(auth);
138+
} catch (Exception e) {
139+
fail("Failed to resolve identity", e);
140+
}
141+
assertThat(requestCredentials).isPresent().contains(overriddenCredentials);
142+
}
143+
144+
private static <T extends Identity> Optional<AwsCredentials> providerNameFromIdentity(SelectedAuthScheme<T> selectedAuthScheme) {
145+
CompletableFuture<? extends T> identityFuture = selectedAuthScheme.identity();
146+
T identity = CompletableFutureUtils.joinLikeSync(identityFuture);
147+
if (identity instanceof AwsBasicCredentials) {
148+
return Optional.of((AwsBasicCredentials) identity);
149+
}
150+
return Optional.empty();
151+
}
152+
153+
private RestJsonEndpointProvidersClientBuilder syncClientBuilder() {
154+
return syncClientBuilder(null);
155+
}
156+
private RestJsonEndpointProvidersClientBuilder syncClientBuilder(ExecutionInterceptor additionalInterceptor) {
157+
List<ExecutionInterceptor> interceptors = new ArrayList<>();
158+
interceptors.add(capturingInterceptor);
159+
if (additionalInterceptor != null) {
160+
interceptors.add(additionalInterceptor);
161+
}
162+
return RestJsonEndpointProvidersClient.builder()
163+
.region(Region.US_WEST_2)
164+
.credentialsProvider(StaticCredentialsProvider.create(CLIENT_CREDENTIALS))
165+
.overrideConfiguration(c -> c.executionInterceptors(interceptors));
166+
}
167+
168+
public static class OverrideInterceptor implements ExecutionInterceptor {
169+
170+
private final Consumer<AwsRequestOverrideConfiguration.Builder> modifier;
171+
172+
public OverrideInterceptor(Consumer<AwsRequestOverrideConfiguration.Builder> m) {
173+
this.modifier = m;
174+
}
175+
176+
@Override
177+
public SdkRequest modifyRequest(Context.ModifyRequest context, ExecutionAttributes executionAttributes) {
178+
SdkRequest request = context.request();
179+
AwsRequest awsRequest = (AwsRequest) request;
180+
181+
AwsRequestOverrideConfiguration.Builder overrideConfigBuilder = getOrCreateOverrideConfig(request);
182+
modifier.accept(overrideConfigBuilder);
183+
184+
return awsRequest.toBuilder().overrideConfiguration(overrideConfigBuilder.build()).build();
185+
}
186+
187+
private AwsRequestOverrideConfiguration.Builder getOrCreateOverrideConfig(SdkRequest request) {
188+
Optional<? extends RequestOverrideConfiguration> requestOverrideConfiguration = request.overrideConfiguration();
189+
if (requestOverrideConfiguration.isPresent()) {
190+
return ((AwsRequestOverrideConfiguration) requestOverrideConfiguration.get()).toBuilder();
191+
}
192+
return AwsRequestOverrideConfiguration.builder();
193+
}
194+
}
195+
196+
public static class CapturingInterceptor implements ExecutionInterceptor {
197+
private Context.BeforeTransmission context;
198+
private ExecutionAttributes executionAttributes;
199+
200+
@Override
201+
public void beforeTransmission(Context.BeforeTransmission context, ExecutionAttributes executionAttributes) {
202+
this.context = context;
203+
this.executionAttributes = executionAttributes;
204+
throw new RuntimeException("stop");
205+
}
206+
207+
public ExecutionAttributes executionAttributes() {
208+
return executionAttributes;
209+
}
210+
}
211+
}

0 commit comments

Comments
 (0)