Skip to content

Commit 5d29684

Browse files
committed
Adding New StsWebIdentityTokenFileCredentialsProvider in sts module that accepts STSClient
1 parent afedc88 commit 5d29684

File tree

11 files changed

+592
-51
lines changed

11 files changed

+592
-51
lines changed
Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,6 @@
1+
{
2+
"type": "feature",
3+
"category": "AWS STS",
4+
"contributor": "",
5+
"description": "Adding New WebIdentityTokenFileCredentialsProvider in sts that accepts STSClient"
6+
}

bom-internal/pom.xml

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -314,6 +314,12 @@
314314
<version>${junit5.version}</version>
315315
<scope>test</scope>
316316
</dependency>
317+
<dependency>
318+
<groupId>org.mockito</groupId>
319+
<artifactId>mockito-junit-jupiter</artifactId>
320+
<version>${mockito.junit5.version}</version>
321+
<scope>test</scope>
322+
</dependency>
317323
<dependency>
318324
<groupId>junit</groupId>
319325
<artifactId>junit</artifactId>

core/auth/src/main/java/software/amazon/awssdk/auth/credentials/WebIdentityTokenFileCredentialsProvider.java

Lines changed: 11 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -29,8 +29,17 @@
2929
/**
3030
* A credential provider that will read web identity token file path, aws role arn
3131
* and aws session name from system properties or environment variables for using
32-
* web identity token credentials with STS. Use of this credentials provider requires
33-
* the 'sts' module to be on the classpath.
32+
* web identity token credentials with STS.
33+
* <p>
34+
* Use of this credentials provider requires the 'sts' module to be on the classpath.
35+
* </p>
36+
* <p>
37+
* StsWebIdentityTokenFileCredentialsProvider in sts package can be used instead of this class if any one of following is required
38+
*<ul>
39+
* <li>Pass a custom StsClient to the provider. </li>
40+
* <li>Periodically update credentials </li>
41+
*</ul>
42+
* @see AwsCredentialsProvider
3443
*/
3544
@SdkPublicApi
3645
public class WebIdentityTokenFileCredentialsProvider implements AwsCredentialsProvider {

pom.xml

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -118,6 +118,7 @@
118118

119119
<!--Test dependencies -->
120120
<junit5.version>5.8.1</junit5.version>
121+
<mockito.junit5.version>4.6.0</mockito.junit5.version>
121122
<junit4.version>4.13.2</junit4.version> <!-- Used to resolve conflicts in transitive dependencies -->
122123
<hamcrest.version>1.3</hamcrest.version>
123124
<mockito.version>4.3.1</mockito.version>

services/sts/pom.xml

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -79,6 +79,11 @@
7979
<artifactId>junit-jupiter</artifactId>
8080
<scope>test</scope>
8181
</dependency>
82+
<dependency>
83+
<groupId>org.mockito</groupId>
84+
<artifactId>mockito-junit-jupiter</artifactId>
85+
<scope>test</scope>
86+
</dependency>
8287
<dependency>
8388
<groupId>org.junit.vintage</groupId>
8489
<artifactId>junit-vintage-engine</artifactId>
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,280 @@
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.sts.auth;
17+
18+
import static software.amazon.awssdk.utils.StringUtils.trim;
19+
import static software.amazon.awssdk.utils.Validate.notNull;
20+
21+
import java.nio.file.Path;
22+
import java.nio.file.Paths;
23+
import java.util.function.Consumer;
24+
import java.util.function.Supplier;
25+
import software.amazon.awssdk.annotations.SdkPublicApi;
26+
import software.amazon.awssdk.auth.credentials.AwsCredentials;
27+
import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider;
28+
import software.amazon.awssdk.auth.credentials.internal.WebIdentityTokenCredentialProperties;
29+
import software.amazon.awssdk.core.SdkSystemSetting;
30+
import software.amazon.awssdk.services.sts.StsClient;
31+
import software.amazon.awssdk.services.sts.internal.AssumeRoleWithWebIdentityRequestSupplier;
32+
import software.amazon.awssdk.services.sts.model.AssumeRoleWithWebIdentityRequest;
33+
import software.amazon.awssdk.services.sts.model.Credentials;
34+
import software.amazon.awssdk.services.sts.model.IdpCommunicationErrorException;
35+
import software.amazon.awssdk.utils.ToString;
36+
37+
/**
38+
* A credential provider that will read web identity token file path, aws role arn, aws session name from system properties or
39+
* environment variables and StsClient for using web identity token credentials with STS
40+
* <p>
41+
* StsWebIdentityTokenFileCredentialsProvider allows passing of a custom StsClient at the time of instantiation of the provider.
42+
* The user needs to make sure that this StsClient handles {@link IdpCommunicationErrorException} in retry policy.
43+
* </p>
44+
* <p>
45+
* This Web Identity Token File Credentials Provider extends {@link StsCredentialsProvider} that supports periodically
46+
* updating session credentials. Refer {@link StsCredentialsProvider} for more details.
47+
* </p>
48+
* <p>STWebIdentityTokenFileCredentialsProvider differs from StsWebIdentityTokenFileCredentialsProvider which is in sts package:
49+
* <ol>
50+
* <li>This Credentials Provider supports custom StsClient</li>
51+
* <li>This Credentials Provider supports periodically updating session credentials. Refer {@link StsCredentialsProvider}</li>
52+
* </ol>
53+
* If asyncCredentialUpdateEnabled is set to true then this provider creates a thread in the background to periodically update
54+
* credentials. If this provider is no longer needed, the background thread can be shut down using {@link #close()}.
55+
* @see StsCredentialsProvider
56+
*/
57+
@SdkPublicApi
58+
public final class StsWebIdentityTokenFileCredentialsProvider extends StsCredentialsProvider {
59+
60+
private final AwsCredentialsProvider credentialsProvider;
61+
private final RuntimeException loadException;
62+
private final Supplier<AssumeRoleWithWebIdentityRequest> assumeRoleWithWebIdentityRequest;
63+
64+
private StsWebIdentityTokenFileCredentialsProvider(Builder builder) {
65+
super(builder, "sts-assume-role-with-web-identity-credentials-provider");
66+
Path webIdentityTokenFile =
67+
builder.webIdentityTokenFile != null ? builder.webIdentityTokenFile
68+
: Paths.get(trim(SdkSystemSetting.AWS_WEB_IDENTITY_TOKEN_FILE
69+
.getStringValueOrThrow()));
70+
71+
String roleArn = builder.roleArn != null ? builder.roleArn
72+
: trim(SdkSystemSetting.AWS_ROLE_ARN.getStringValueOrThrow());
73+
74+
String sessionName = builder.roleSessionName != null ? builder.roleSessionName :
75+
SdkSystemSetting.AWS_ROLE_SESSION_NAME.getStringValue()
76+
.orElse("aws-sdk-java-" + System.currentTimeMillis());
77+
78+
WebIdentityTokenCredentialProperties credentialProperties =
79+
WebIdentityTokenCredentialProperties.builder()
80+
.roleArn(roleArn)
81+
.roleSessionName(builder.roleSessionName)
82+
.webIdentityTokenFile(webIdentityTokenFile)
83+
.build();
84+
85+
this.assumeRoleWithWebIdentityRequest = builder.assumeRoleWithWebIdentityRequestSupplier != null
86+
? builder.assumeRoleWithWebIdentityRequestSupplier
87+
: () -> AssumeRoleWithWebIdentityRequest.builder()
88+
.roleArn(credentialProperties.roleArn())
89+
.roleSessionName(sessionName)
90+
.build();
91+
92+
AwsCredentialsProvider credentialsProviderLocal = null;
93+
RuntimeException loadExceptionLocal = null;
94+
try {
95+
AssumeRoleWithWebIdentityRequestSupplier supplier =
96+
AssumeRoleWithWebIdentityRequestSupplier.builder()
97+
.assumeRoleWithWebIdentityRequest(assumeRoleWithWebIdentityRequest.get())
98+
.webIdentityTokenFile(credentialProperties.webIdentityTokenFile())
99+
.build();
100+
credentialsProviderLocal =
101+
StsAssumeRoleWithWebIdentityCredentialsProvider.builder()
102+
.stsClient(builder.stsClient)
103+
.refreshRequest(supplier)
104+
.build();
105+
} catch (RuntimeException e) {
106+
// If we couldn't load the credentials provider for some reason, save an exception describing why. This exception
107+
// will only be raised on calls to getCredentials. We don't want to raise an exception here because it may be
108+
// expected (eg. in the default credential chain).
109+
loadExceptionLocal = e;
110+
}
111+
this.loadException = loadExceptionLocal;
112+
this.credentialsProvider = credentialsProviderLocal;
113+
}
114+
115+
public static Builder builder() {
116+
return new Builder();
117+
}
118+
119+
@Override
120+
public AwsCredentials resolveCredentials() {
121+
if (loadException != null) {
122+
throw loadException;
123+
}
124+
return credentialsProvider.resolveCredentials();
125+
}
126+
127+
@Override
128+
public String toString() {
129+
return ToString.builder("StsWebIdentityTokenFileCredentialsProvider")
130+
.add("refreshRequest", assumeRoleWithWebIdentityRequest)
131+
.build();
132+
}
133+
134+
@Override
135+
protected Credentials getUpdatedCredentials(StsClient stsClient) {
136+
AssumeRoleWithWebIdentityRequest request = assumeRoleWithWebIdentityRequest.get();
137+
notNull(request, "AssumeRoleWithWebIdentityRequest can't be null");
138+
return stsClient.assumeRoleWithWebIdentity(request).credentials();
139+
}
140+
141+
public static final class Builder extends BaseBuilder<Builder, StsWebIdentityTokenFileCredentialsProvider> {
142+
private String roleArn;
143+
private String roleSessionName;
144+
private Path webIdentityTokenFile;
145+
private Supplier<AssumeRoleWithWebIdentityRequest> assumeRoleWithWebIdentityRequestSupplier;
146+
private StsClient stsClient;
147+
148+
private Builder() {
149+
super(StsWebIdentityTokenFileCredentialsProvider::new);
150+
}
151+
152+
/**
153+
* The Custom {@link StsClient} that will be used to fetch AWS service credentials.
154+
* <ul>
155+
* <li>This SDK client must be closed by the caller when it is ready to be disposed.</li>
156+
* <li>This SDK client's retry policy should handle IdpCommunicationErrorException </li>
157+
* </ul>
158+
* @param stsClient The STS client to use for communication with STS.
159+
* Make sure IdpCommunicationErrorException is retried in the retry policy for this client.
160+
* Make sure the custom STS client is closed when it is ready to be disposed.
161+
* @return Returns a reference to this object so that method calls can be chained together.
162+
*/
163+
@Override
164+
public Builder stsClient(StsClient stsClient) {
165+
this.stsClient = stsClient;
166+
return super.stsClient(stsClient);
167+
}
168+
169+
/**
170+
* <p>
171+
* The Amazon Resource Name (ARN) of the IAM role that is associated with the Sts.
172+
* If not provided this will be read from SdkSystemSetting.AWS_ROLE_ARN.
173+
* </p>
174+
*
175+
* @param roleArn The Amazon Resource Name (ARN) of the IAM role that is associated with the Sts cluster.
176+
* @return Returns a reference to this object so that method calls can be chained together.
177+
*/
178+
public Builder roleArn(String roleArn) {
179+
this.roleArn = roleArn;
180+
return this;
181+
}
182+
183+
/**
184+
* <p>
185+
* Sets Amazon Resource Name (ARN) of the IAM role that is associated with the Sts.
186+
* By default this will be read from SdkSystemSetting.AWS_ROLE_ARN.
187+
* </p>
188+
*
189+
* @param roleArn The Amazon Resource Name (ARN) of the IAM role that is associated with the Sts cluster.
190+
*/
191+
public void setRoleArn(String roleArn) {
192+
roleArn(roleArn);
193+
}
194+
195+
/**
196+
* <p>
197+
* Sets the role session name that should be used by this credentials provider.
198+
* By default this is read from SdkSystemSetting.AWS_ROLE_SESSION_NAME
199+
* </p>
200+
*
201+
* @param roleSessionName role session name that should be used by this credentials provider
202+
* @return Returns a reference to this object so that method calls can be chained together.
203+
*/
204+
public Builder roleSessionName(String roleSessionName) {
205+
this.roleSessionName = roleSessionName;
206+
return this;
207+
}
208+
209+
/**
210+
* <p>
211+
* Sets the role session name that should be used by this credentials provider.
212+
* By default this is read from SdkSystemSetting.AWS_ROLE_SESSION_NAME
213+
* </p>
214+
*
215+
* @param roleSessionName role session name that should be used by this credentials provider.
216+
*/
217+
public void setRoleSessionName(String roleSessionName) {
218+
roleSessionName(roleSessionName);
219+
}
220+
221+
/**
222+
* <p>
223+
* Sets the absolute path to the web identity token file that should be used by this credentials provider.
224+
* By default this will be read from SdkSystemSetting.AWS_WEB_IDENTITY_TOKEN_FILE.
225+
* </p>
226+
*
227+
* @param webIdentityTokenFile absolute path to the web identity token file that should be used by this credentials
228+
* provider.
229+
* @return Returns a reference to this object so that method calls can be chained together.
230+
*/
231+
public Builder webIdentityTokenFile(Path webIdentityTokenFile) {
232+
this.webIdentityTokenFile = webIdentityTokenFile;
233+
return this;
234+
}
235+
236+
public void setWebIdentityTokenFile(Path webIdentityTokenFile) {
237+
webIdentityTokenFile(webIdentityTokenFile);
238+
}
239+
240+
241+
/**
242+
* Configure the {@link AssumeRoleWithWebIdentityRequest} that should be periodically sent to the STS service to update
243+
* the session token when it gets close to expiring.
244+
*
245+
* @param assumeRoleWithWebIdentityRequest The request to send to STS whenever the assumed session expires.
246+
* @return This object for chained calls.
247+
*/
248+
public Builder refreshRequest(AssumeRoleWithWebIdentityRequest assumeRoleWithWebIdentityRequest) {
249+
return refreshRequest(() -> assumeRoleWithWebIdentityRequest);
250+
}
251+
252+
/**
253+
* Similar to {@link #refreshRequest(AssumeRoleWithWebIdentityRequest)}, but takes a {@link Supplier} to supply the
254+
* request to STS.
255+
*
256+
* @param assumeRoleWithWebIdentityRequestSupplier A supplier
257+
* @return This object for chained calls.
258+
*/
259+
public Builder refreshRequest(Supplier<AssumeRoleWithWebIdentityRequest> assumeRoleWithWebIdentityRequestSupplier) {
260+
this.assumeRoleWithWebIdentityRequestSupplier = assumeRoleWithWebIdentityRequestSupplier;
261+
return this;
262+
}
263+
264+
/**
265+
* Similar to {@link #refreshRequest(AssumeRoleWithWebIdentityRequest)}, but takes a lambda to configure a new {@link
266+
* AssumeRoleWithWebIdentityRequest.Builder}. This removes the need to call {@link
267+
* AssumeRoleWithWebIdentityRequest#builder()} and {@link AssumeRoleWithWebIdentityRequest.Builder#build()}.
268+
*/
269+
public Builder refreshRequest(Consumer<AssumeRoleWithWebIdentityRequest.Builder> assumeRoleWithWebIdentityRequest) {
270+
return refreshRequest(AssumeRoleWithWebIdentityRequest.builder().applyMutation(assumeRoleWithWebIdentityRequest)
271+
.build());
272+
}
273+
274+
@Override
275+
public StsWebIdentityTokenFileCredentialsProvider build() {
276+
return new StsWebIdentityTokenFileCredentialsProvider(this);
277+
}
278+
279+
}
280+
}

0 commit comments

Comments
 (0)