Skip to content

Commit daf9d4a

Browse files
Luke SikinaLuke-Sikina
Luke Sikina
authored and
Luke-Sikina
committed
[ALS-7539] - Get rid of super complex self refreshing client
1 parent 0b7e349 commit daf9d4a

File tree

10 files changed

+277
-177
lines changed

10 files changed

+277
-177
lines changed

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

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

‎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)