Skip to content

Commit fec46b0

Browse files
alextwoodsalextwoods
committed
Improve test coverage
1 parent 23c5b58 commit fec46b0

3 files changed

Lines changed: 171 additions & 4 deletions

File tree

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

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

82-
public ClassName getEnvironmentTokenMetricsInterceptorClass() {
83-
return ClassName.get(model.getMetadata().getFullClientInternalPackageName(), "EnvironmentTokenMetricsInterceptor");
84-
}
85-
8682
public ClassName getEnvironmentTokenSystemSettingsClass() {
8783
return ClassName.get(model.getMetadata().getFullClientInternalPackageName(), "EnvironmentTokenSystemSettings");
8884
}

‎codegen/src/test/java/software/amazon/awssdk/codegen/poet/auth/scheme/AuthSchemeSpecTest.java‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -220,6 +220,13 @@ static List<TestCase> parameters() {
220220
.classSpecProvider(AuthSchemeInterceptorSpec::new)
221221
.caseName("ops-auth-sigv4a-value")
222222
.outputFileSuffix("interceptor")
223+
.build(),
224+
// service with environment bearer token enabled
225+
TestCase.builder()
226+
.modelProvider(ClientTestModels::envBearerTokenServiceModels)
227+
.classSpecProvider(AuthSchemeInterceptorSpec::new)
228+
.caseName("env-bearer-token")
229+
.outputFileSuffix("interceptor")
223230
.build()
224231
);
225232
}

‎codegen/src/test/resources/software/amazon/awssdk/codegen/poet/auth/scheme/env-bearer-token-auth-scheme-interceptor.java‎

Lines changed: 164 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,164 @@
1+
package software.amazon.awssdk.services.json.auth.scheme.internal;
2+
3+
import java.time.Duration;
4+
import java.util.ArrayList;
5+
import java.util.List;
6+
import java.util.Map;
7+
import java.util.concurrent.CompletableFuture;
8+
import java.util.function.Supplier;
9+
import java.util.stream.Collectors;
10+
import software.amazon.awssdk.annotations.Generated;
11+
import software.amazon.awssdk.annotations.SdkInternalApi;
12+
import software.amazon.awssdk.core.SdkRequest;
13+
import software.amazon.awssdk.core.SelectedAuthScheme;
14+
import software.amazon.awssdk.core.exception.SdkException;
15+
import software.amazon.awssdk.core.interceptor.Context;
16+
import software.amazon.awssdk.core.interceptor.ExecutionAttributes;
17+
import software.amazon.awssdk.core.interceptor.ExecutionInterceptor;
18+
import software.amazon.awssdk.core.interceptor.SdkExecutionAttribute;
19+
import software.amazon.awssdk.core.interceptor.SdkInternalExecutionAttribute;
20+
import software.amazon.awssdk.core.internal.util.MetricUtils;
21+
import software.amazon.awssdk.core.metrics.CoreMetric;
22+
import software.amazon.awssdk.core.useragent.BusinessMetricFeatureId;
23+
import software.amazon.awssdk.http.auth.scheme.BearerAuthScheme;
24+
import software.amazon.awssdk.http.auth.spi.scheme.AuthScheme;
25+
import software.amazon.awssdk.http.auth.spi.scheme.AuthSchemeOption;
26+
import software.amazon.awssdk.http.auth.spi.signer.HttpSigner;
27+
import software.amazon.awssdk.identity.spi.AwsCredentialsIdentity;
28+
import software.amazon.awssdk.identity.spi.Identity;
29+
import software.amazon.awssdk.identity.spi.IdentityProvider;
30+
import software.amazon.awssdk.identity.spi.IdentityProviders;
31+
import software.amazon.awssdk.identity.spi.ResolveIdentityRequest;
32+
import software.amazon.awssdk.identity.spi.TokenIdentity;
33+
import software.amazon.awssdk.metrics.MetricCollector;
34+
import software.amazon.awssdk.metrics.SdkMetric;
35+
import software.amazon.awssdk.services.json.auth.scheme.JsonAuthSchemeParams;
36+
import software.amazon.awssdk.services.json.auth.scheme.JsonAuthSchemeProvider;
37+
import software.amazon.awssdk.utils.Logger;
38+
import software.amazon.awssdk.utils.Validate;
39+
40+
@Generated("software.amazon.awssdk:codegen")
41+
@SdkInternalApi
42+
public final class JsonAuthSchemeInterceptor implements ExecutionInterceptor {
43+
private static Logger LOG = Logger.loggerFor(JsonAuthSchemeInterceptor.class);
44+
45+
@Override
46+
public void beforeExecution(Context.BeforeExecution context, ExecutionAttributes executionAttributes) {
47+
List<AuthSchemeOption> authOptions = resolveAuthOptions(context, executionAttributes);
48+
SelectedAuthScheme<? extends Identity> selectedAuthScheme = selectAuthScheme(authOptions, executionAttributes);
49+
putSelectedAuthScheme(executionAttributes, selectedAuthScheme);
50+
recordEnvironmentTokenBusinessMetric(selectedAuthScheme, executionAttributes);
51+
}
52+
53+
private List<AuthSchemeOption> resolveAuthOptions(Context.BeforeExecution context, ExecutionAttributes executionAttributes) {
54+
JsonAuthSchemeProvider authSchemeProvider = Validate.isInstanceOf(JsonAuthSchemeProvider.class,
55+
executionAttributes.getAttribute(SdkInternalExecutionAttribute.AUTH_SCHEME_RESOLVER),
56+
"Expected an instance of JsonAuthSchemeProvider");
57+
JsonAuthSchemeParams params = authSchemeParams(context.request(), executionAttributes);
58+
return authSchemeProvider.resolveAuthScheme(params);
59+
}
60+
61+
private SelectedAuthScheme<? extends Identity> selectAuthScheme(List<AuthSchemeOption> authOptions,
62+
ExecutionAttributes executionAttributes) {
63+
MetricCollector metricCollector = executionAttributes.getAttribute(SdkExecutionAttribute.API_CALL_METRIC_COLLECTOR);
64+
Map<String, AuthScheme<?>> authSchemes = executionAttributes.getAttribute(SdkInternalExecutionAttribute.AUTH_SCHEMES);
65+
IdentityProviders identityProviders = executionAttributes.getAttribute(SdkInternalExecutionAttribute.IDENTITY_PROVIDERS);
66+
List<Supplier<String>> discardedReasons = new ArrayList<>();
67+
for (AuthSchemeOption authOption : authOptions) {
68+
AuthScheme<?> authScheme = authSchemes.get(authOption.schemeId());
69+
SelectedAuthScheme<? extends Identity> selectedAuthScheme = trySelectAuthScheme(authOption, authScheme,
70+
identityProviders, discardedReasons, metricCollector, executionAttributes);
71+
if (selectedAuthScheme != null) {
72+
if (!discardedReasons.isEmpty()) {
73+
LOG.debug(() -> String.format("%s auth will be used, discarded: '%s'", authOption.schemeId(),
74+
discardedReasons.stream().map(Supplier::get).collect(Collectors.joining(", "))));
75+
}
76+
return selectedAuthScheme;
77+
}
78+
}
79+
throw SdkException
80+
.builder()
81+
.message(
82+
"Failed to determine how to authenticate the user: "
83+
+ discardedReasons.stream().map(Supplier::get).collect(Collectors.joining(", "))).build();
84+
}
85+
86+
private JsonAuthSchemeParams authSchemeParams(SdkRequest request, ExecutionAttributes executionAttributes) {
87+
String operation = executionAttributes.getAttribute(SdkExecutionAttribute.OPERATION_NAME);
88+
JsonAuthSchemeParams.Builder builder = JsonAuthSchemeParams.builder().operation(operation);
89+
return builder.build();
90+
}
91+
92+
private <T extends Identity> SelectedAuthScheme<T> trySelectAuthScheme(AuthSchemeOption authOption, AuthScheme<T> authScheme,
93+
IdentityProviders identityProviders, List<Supplier<String>> discardedReasons, MetricCollector metricCollector,
94+
ExecutionAttributes executionAttributes) {
95+
if (authScheme == null) {
96+
discardedReasons.add(() -> String.format("'%s' is not enabled for this request.", authOption.schemeId()));
97+
return null;
98+
}
99+
IdentityProvider<T> identityProvider = authScheme.identityProvider(identityProviders);
100+
if (identityProvider == null) {
101+
discardedReasons
102+
.add(() -> String.format("'%s' does not have an identity provider configured.", authOption.schemeId()));
103+
return null;
104+
}
105+
HttpSigner<T> signer;
106+
try {
107+
signer = authScheme.signer();
108+
} catch (RuntimeException e) {
109+
discardedReasons.add(() -> String.format("'%s' signer could not be retrieved: %s", authOption.schemeId(),
110+
e.getMessage()));
111+
return null;
112+
}
113+
ResolveIdentityRequest.Builder identityRequestBuilder = ResolveIdentityRequest.builder();
114+
authOption.forEachIdentityProperty(identityRequestBuilder::putProperty);
115+
CompletableFuture<? extends T> identity;
116+
SdkMetric<Duration> metric = getIdentityMetric(identityProvider);
117+
if (metric == null) {
118+
identity = identityProvider.resolveIdentity(identityRequestBuilder.build());
119+
} else {
120+
identity = MetricUtils.reportDuration(() -> identityProvider.resolveIdentity(identityRequestBuilder.build()),
121+
metricCollector, metric);
122+
}
123+
return new SelectedAuthScheme<>(identity, signer, authOption);
124+
}
125+
126+
private SdkMetric<Duration> getIdentityMetric(IdentityProvider<?> identityProvider) {
127+
Class<?> identityType = identityProvider.identityType();
128+
if (identityType == AwsCredentialsIdentity.class) {
129+
return CoreMetric.CREDENTIALS_FETCH_DURATION;
130+
}
131+
if (identityType == TokenIdentity.class) {
132+
return CoreMetric.TOKEN_FETCH_DURATION;
133+
}
134+
return null;
135+
}
136+
137+
private <T extends Identity> void putSelectedAuthScheme(ExecutionAttributes attributes,
138+
SelectedAuthScheme<T> selectedAuthScheme) {
139+
SelectedAuthScheme<?> existingAuthScheme = attributes.getAttribute(SdkInternalExecutionAttribute.SELECTED_AUTH_SCHEME);
140+
if (existingAuthScheme != null) {
141+
AuthSchemeOption.Builder selectedOption = selectedAuthScheme.authSchemeOption().toBuilder();
142+
existingAuthScheme.authSchemeOption().forEachIdentityProperty(selectedOption::putIdentityPropertyIfAbsent);
143+
existingAuthScheme.authSchemeOption().forEachSignerProperty(selectedOption::putSignerPropertyIfAbsent);
144+
selectedAuthScheme = new SelectedAuthScheme<>(selectedAuthScheme.identity(), selectedAuthScheme.signer(),
145+
selectedOption.build());
146+
}
147+
attributes.putAttribute(SdkInternalExecutionAttribute.SELECTED_AUTH_SCHEME, selectedAuthScheme);
148+
}
149+
150+
private <T extends Identity> void recordEnvironmentTokenBusinessMetric(SelectedAuthScheme<T> selectedAuthScheme,
151+
ExecutionAttributes executionAttributes) {
152+
String tokenFromEnv = executionAttributes.getAttribute(SdkInternalExecutionAttribute.TOKEN_CONFIGURED_FROM_ENV);
153+
if (selectedAuthScheme != null && selectedAuthScheme.authSchemeOption().schemeId().equals(BearerAuthScheme.SCHEME_ID)
154+
&& selectedAuthScheme.identity().isDone()) {
155+
if (selectedAuthScheme.identity().getNow(null) instanceof TokenIdentity) {
156+
TokenIdentity configuredToken = (TokenIdentity) selectedAuthScheme.identity().getNow(null);
157+
if (configuredToken.token().equals(tokenFromEnv)) {
158+
executionAttributes.getAttribute(SdkInternalExecutionAttribute.BUSINESS_METRICS).addMetric(
159+
BusinessMetricFeatureId.BEARER_SERVICE_ENV_VARS.value());
160+
}
161+
}
162+
}
163+
}
164+
}

0 commit comments

Comments
 (0)