Skip to content

Commit b9b54f5

Browse files
author
Luke Sikina
committed
[ALS-7539] - Get rid of super complex self refreshing client
1 parent 0b7e349 commit b9b54f5

File tree

8 files changed

+269
-177
lines changed

8 files changed

+269
-177
lines changed

‎uploader/src/main/java/edu/harvard/dbmi/avillach/dataupload/aws/AWSClientBuilder.java

+91
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,91 @@
1+
package edu.harvard.dbmi.avillach.dataupload.aws;
2+
3+
import org.slf4j.Logger;
4+
import org.slf4j.LoggerFactory;
5+
import org.springframework.beans.factory.annotation.Autowired;
6+
import org.springframework.stereotype.Service;
7+
import software.amazon.awssdk.auth.credentials.AwsSessionCredentials;
8+
import software.amazon.awssdk.auth.credentials.StaticCredentialsProvider;
9+
import software.amazon.awssdk.http.SdkHttpClient;
10+
import software.amazon.awssdk.services.s3.S3Client;
11+
import software.amazon.awssdk.services.s3.S3ClientBuilder;
12+
import software.amazon.awssdk.services.sts.model.AssumeRoleRequest;
13+
import software.amazon.awssdk.services.sts.model.AssumeRoleResponse;
14+
import software.amazon.awssdk.services.sts.model.Credentials;
15+
16+
import java.util.Map;
17+
import java.util.Optional;
18+
19+
@Service
20+
public class AWSClientBuilder {
21+
22+
private static final Logger log = LoggerFactory.getLogger(AWSClientBuilder.class);
23+
24+
private final Map<String, SiteAWSInfo> sites;
25+
private final StsClientProvider stsClientProvider;
26+
private final S3ClientBuilder s3ClientBuilder;
27+
private final SdkHttpClient sdkHttpClient;
28+
29+
@Autowired
30+
public AWSClientBuilder(
31+
Map<String, SiteAWSInfo> sites,
32+
StsClientProvider stsClientProvider,
33+
S3ClientBuilder s3ClientBuilder,
34+
@Autowired(required = false) SdkHttpClient sdkHttpClient
35+
) {
36+
this.sites = sites;
37+
this.stsClientProvider = stsClientProvider;
38+
this.s3ClientBuilder = s3ClientBuilder;
39+
this.sdkHttpClient = sdkHttpClient;
40+
}
41+
42+
public Optional<S3Client> buildClientForSite(String siteName) {
43+
log.info("Building client for site {}", siteName);
44+
if (!sites.containsKey(siteName)) {
45+
log.warn("Could not find site {}", siteName);
46+
return Optional.empty();
47+
}
48+
49+
log.info("Found site, making assume role request");
50+
SiteAWSInfo site = sites.get(siteName);
51+
AssumeRoleRequest roleRequest = AssumeRoleRequest.builder()
52+
.roleArn(site.roleARN())
53+
.roleSessionName("test_session" + System.nanoTime())
54+
.externalId(site.externalId())
55+
.durationSeconds(60*60) // 1 hour
56+
.build();
57+
Optional<Credentials> assumeRoleResponse = stsClientProvider.createClient()
58+
.map(c -> c.assumeRole(roleRequest))
59+
.map(AssumeRoleResponse::credentials);
60+
if (assumeRoleResponse.isEmpty() ) {
61+
log.error("Error assuming role {} , no credentials returned", site.roleARN());
62+
return Optional.empty();
63+
}
64+
log.info("Successfully assumed role {} for site {}", site.roleARN(), site.siteName());
65+
66+
log.info("Building S3 client for site {}", site.siteName());
67+
// Use the credentials from the role to create the S3 client
68+
Credentials credentials = assumeRoleResponse.get();
69+
AwsSessionCredentials sessionCredentials = AwsSessionCredentials.builder()
70+
.accessKeyId(credentials.accessKeyId())
71+
.secretAccessKey(credentials.secretAccessKey())
72+
.sessionToken(credentials.sessionToken())
73+
.expirationTime(credentials.expiration())
74+
.build();
75+
StaticCredentialsProvider provider = StaticCredentialsProvider.create(sessionCredentials);
76+
return Optional.of(buildFromProvider(provider));
77+
}
78+
79+
private S3Client buildFromProvider(StaticCredentialsProvider provider) {
80+
if (sdkHttpClient == null) {
81+
return s3ClientBuilder.credentialsProvider(provider).build();
82+
}
83+
log.info("Http proxy detected and added to S3 client");
84+
return s3ClientBuilder
85+
.credentialsProvider(provider)
86+
.httpClient(sdkHttpClient)
87+
.build();
88+
89+
}
90+
91+
}

‎uploader/src/main/java/edu/harvard/dbmi/avillach/dataupload/aws/AWSConfiguration.java

+14
Original file line numberDiff line numberDiff line change
@@ -11,9 +11,12 @@
1111
import org.springframework.context.annotation.Configuration;
1212
import org.springframework.scheduling.annotation.Scheduled;
1313
import org.springframework.util.StringUtils;
14+
import org.springframework.web.context.annotation.RequestScope;
1415
import software.amazon.awssdk.auth.credentials.*;
1516
import software.amazon.awssdk.http.SdkHttpClient;
1617
import software.amazon.awssdk.regions.Region;
18+
import software.amazon.awssdk.services.s3.S3Client;
19+
import software.amazon.awssdk.services.s3.S3ClientBuilder;
1720
import software.amazon.awssdk.services.sts.StsClient;
1821
import software.amazon.awssdk.services.sts.StsClientBuilder;
1922
import software.amazon.encryption.s3.S3EncryptionClient;
@@ -82,4 +85,15 @@ StsClientBuilder stsClientBuilder() {
8285
// This is a bean for mocking purposes
8386
return StsClient.builder();
8487
}
88+
89+
@Bean
90+
S3ClientBuilder s3ClientBuilder() {
91+
return S3Client.builder();
92+
}
93+
94+
@Bean
95+
@RequestScope
96+
StsClient getStsClient() {
97+
return StsClient.builder().region(Region.US_EAST_1).build();
98+
}
8599
}

‎uploader/src/main/java/edu/harvard/dbmi/avillach/dataupload/aws/S3StateVerifier.java

+10-7
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,7 @@ public class S3StateVerifier {
2828
private Map<String, SiteAWSInfo> sites;
2929

3030
@Autowired
31-
private SelfRefreshingS3Client client;
31+
private AWSClientBuilder clientBuilder;
3232

3333
@PostConstruct
3434
private void verifyS3Status() {
@@ -39,7 +39,7 @@ private void verifyS3Status() {
3939
private void asyncVerify(SiteAWSInfo institution) {
4040
LOG.info("Checking S3 connection to {} ...", institution.siteName());
4141
createTempFileWithText(institution)
42-
.map(p -> uploadFileFromPath(p, institution))
42+
.flatMap(p -> uploadFileFromPath(p, institution))
4343
.map(this::waitABit)
4444
.flatMap(s1 -> deleteFileFromBucket(s1, institution))
4545
.orElseThrow();
@@ -49,8 +49,10 @@ private void asyncVerify(SiteAWSInfo institution) {
4949
private Optional<String> deleteFileFromBucket(String s, SiteAWSInfo info) {
5050
LOG.info("Verifying delete capabilities");
5151
DeleteObjectRequest request = DeleteObjectRequest.builder().bucket(info.bucket()).key(s).build();
52-
DeleteObjectResponse deleteObjectResponse = client.getS3Client(info.siteName()).deleteObject(request);
53-
return deleteObjectResponse.deleteMarker() ? Optional.of(s) : Optional.empty();
52+
return clientBuilder.buildClientForSite(info.siteName())
53+
.map(c -> c.deleteObject(request))
54+
.map(DeleteObjectResponse::deleteMarker)
55+
.map((ignored) -> s);
5456
}
5557

5658
private String waitABit(String s) {
@@ -62,7 +64,7 @@ private String waitABit(String s) {
6264
return s;
6365
}
6466

65-
private String uploadFileFromPath(Path p, SiteAWSInfo info) {
67+
private Optional<String> uploadFileFromPath(Path p, SiteAWSInfo info) {
6668
LOG.info("Verifying upload capabilities");
6769
RequestBody body = RequestBody.fromFile(p.toFile());
6870
PutObjectRequest request = PutObjectRequest.builder()
@@ -71,8 +73,9 @@ private String uploadFileFromPath(Path p, SiteAWSInfo info) {
7173
.ssekmsKeyId(info.kmsKeyID())
7274
.key(p.getFileName().toString())
7375
.build();
74-
client.getS3Client(info.siteName()).putObject(request, body);
75-
return p.getFileName().toString();
76+
return clientBuilder.buildClientForSite(info.siteName())
77+
.map(client -> client.putObject(request, body))
78+
.map(resp -> p.getFileName().toString());
7679
}
7780

7881
private Optional<Path> createTempFileWithText(SiteAWSInfo info) {

‎uploader/src/main/java/edu/harvard/dbmi/avillach/dataupload/aws/SelfRefreshingS3Client.java

-159
This file was deleted.

‎uploader/src/main/java/edu/harvard/dbmi/avillach/dataupload/aws/StsClientProvider.java

+20
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,20 @@
1+
package edu.harvard.dbmi.avillach.dataupload.aws;
2+
3+
import org.slf4j.Logger;
4+
import org.slf4j.LoggerFactory;
5+
import org.springframework.stereotype.Service;
6+
import software.amazon.awssdk.regions.Region;
7+
import software.amazon.awssdk.services.sts.StsClient;
8+
9+
import java.util.Optional;
10+
11+
@Service
12+
public class StsClientProvider {
13+
14+
private static final Logger log = LoggerFactory.getLogger(StsClientProvider.class);
15+
16+
public Optional<StsClient> createClient() {
17+
StsClient client = StsClient.builder().region(Region.US_EAST_1).build();
18+
return Optional.of(client);
19+
}
20+
}

0 commit comments

Comments
 (0)