From 046166ea8182f5d249f5af1da27aa71375f04c12 Mon Sep 17 00:00:00 2001 From: Alif Merchant Date: Mon, 2 Dec 2024 11:11:08 -0800 Subject: [PATCH 1/3] use aws sdk dsql utilities for password token generation --- pom.xml | 19 +++--- .../java/com/oltpbenchmark/util/IAMUtil.java | 59 +++++++------------ .../com/oltpbenchmark/util/TestIAMUtil.java | 25 -------- 3 files changed, 28 insertions(+), 75 deletions(-) diff --git a/pom.xml b/pom.xml index f001b13..b4fed84 100644 --- a/pom.xml +++ b/pom.xml @@ -378,29 +378,24 @@ - software.amazon.awssdk - http-auth-aws - 2.25.28 - - - software.amazon.awssdk - auth - 2.25.28 + software.amazon.awssdk.services + dsql + 2.29.23 software.amazon.awssdk regions - 2.25.28 + 2.25.60 software.amazon.awssdk sdk-core - 2.25.28 + 2.25.60 software.amazon.awssdk - http-client-spi - 2.25.28 + auth + 2.25.60 diff --git a/src/main/java/com/oltpbenchmark/util/IAMUtil.java b/src/main/java/com/oltpbenchmark/util/IAMUtil.java index 24df45a..a47a8d1 100644 --- a/src/main/java/com/oltpbenchmark/util/IAMUtil.java +++ b/src/main/java/com/oltpbenchmark/util/IAMUtil.java @@ -1,19 +1,12 @@ package com.oltpbenchmark.util; -import java.net.URI; -import java.net.URISyntaxException; -import java.time.Clock; import java.time.Duration; import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider; import software.amazon.awssdk.auth.credentials.DefaultCredentialsProvider; -import software.amazon.awssdk.auth.signer.Aws4Signer; -import software.amazon.awssdk.auth.signer.params.Aws4PresignerParams; import software.amazon.awssdk.core.exception.SdkClientException; -import software.amazon.awssdk.http.SdkHttpFullRequest; -import software.amazon.awssdk.http.SdkHttpMethod; -import software.amazon.awssdk.regions.Region; import software.amazon.awssdk.regions.providers.AwsRegionProviderChain; import software.amazon.awssdk.regions.providers.DefaultAwsRegionProviderChain; +import software.amazon.awssdk.services.dsql.DsqlUtilities; public class IAMUtil { // Default token validity is one hour @@ -21,12 +14,6 @@ public class IAMUtil { private static final String ADMIN_USERNAME = "admin"; - private static final String SIGNING_NAME = "dsql"; - - private static final String DB_CONNECT_ADMIN = "DbConnectAdmin"; - - private static final String DB_CONNECT = "DbConnect"; - public static String generateAuroraDsqlPasswordToken(String url, String username) { return generateAuroraDsqlPasswordToken( url, @@ -40,33 +27,29 @@ public static String generateAuroraDsqlPasswordToken( String username, AwsCredentialsProvider credentialsProvider, AwsRegionProviderChain regionProvider) { + DsqlUtilities utilities = + DsqlUtilities.builder() + .region(regionProvider.getRegion()) + .credentialsProvider(credentialsProvider) + .build(); + try { IAMUtil.validateUrl(url); String host = url.split("//")[1].split(":")[0]; - - Clock now = Clock.systemUTC(); - Region region = regionProvider.getRegion(); - if (region == null) region = Region.US_EAST_1; - - Aws4Signer signer = Aws4Signer.create(); - Aws4PresignerParams presignerParams = - Aws4PresignerParams.builder() - .signingName(SIGNING_NAME) - .signingRegion(region) - .awsCredentials(credentialsProvider.resolveCredentials()) - .signingClockOverride(now) - .expirationTime(now.instant().plus(DEFAULT_VALIDITY)) - .build(); - SdkHttpFullRequest request = - SdkHttpFullRequest.builder() - .uri(new URI("https", host, "/", null)) - .appendRawQueryParameter( - "Action", (username.equals(ADMIN_USERNAME)) ? DB_CONNECT_ADMIN : DB_CONNECT) - .method(SdkHttpMethod.GET) - .build(); - - return signer.presign(request, presignerParams).getUri().toString().replace("https://", ""); - } catch (URISyntaxException | SdkClientException e) { + return username.equals(ADMIN_USERNAME) + ? utilities.generateDbConnectAdminAuthToken( + builder -> + builder + .hostname(host) + .region(regionProvider.getRegion()) + .expiresIn(DEFAULT_VALIDITY)) + : utilities.generateDbConnectAuthToken( + builder -> + builder + .hostname(host) + .region(regionProvider.getRegion()) + .expiresIn(DEFAULT_VALIDITY)); + } catch (SdkClientException e) { throw new RuntimeException(e); } } diff --git a/src/test/java/com/oltpbenchmark/util/TestIAMUtil.java b/src/test/java/com/oltpbenchmark/util/TestIAMUtil.java index e6b647f..b9f58a1 100644 --- a/src/test/java/com/oltpbenchmark/util/TestIAMUtil.java +++ b/src/test/java/com/oltpbenchmark/util/TestIAMUtil.java @@ -1,8 +1,6 @@ package com.oltpbenchmark.util; -import static org.junit.Assert.assertNotNull; import static org.junit.Assert.assertThrows; -import static org.junit.Assert.assertTrue; import org.junit.Before; import org.junit.Test; @@ -42,29 +40,6 @@ public String secretAccessKey() { Mockito.when(regionProvider.getRegion()).thenReturn(Region.US_EAST_2); } - @Test - public void testGenerateAuroraDsqlPasswordToken() { - String token = - IAMUtil.generateAuroraDsqlPasswordToken( - VALID_URL, VALID_ADMIN_USERNAME, credentialsProvider, regionProvider); - assertNotNull(token); - assertTrue(token.contains("localhost/?")); - assertTrue(token.contains("X-Amz-Algorithm=AWS4-HMAC-SHA256")); - assertTrue(token.contains("X-Amz-Expires=3600")); - assertTrue(token.contains("Action=DbConnectAdmin")); - assertTrue(token.contains("X-Amz-Credential=ACCESS_KEY")); - assertTrue(token.contains("X-Amz-Signature")); - } - - @Test - public void testGenerateAuroraDsqlPasswordTokenNonAdminUser() { - String token = - IAMUtil.generateAuroraDsqlPasswordToken( - VALID_URL, "other", credentialsProvider, regionProvider); - assertNotNull(token); - assertTrue(token.contains("Action=DbConnect")); - } - @Test public void testGenerateAuroraDsqlPasswordTokenInvalidUrl() { assertThrows( From 19c3b8f8fa003a1fdacab8eae9afcb9078512193 Mon Sep 17 00:00:00 2001 From: Srudeep Date: Tue, 3 Dec 2024 23:25:13 -0800 Subject: [PATCH 2/3] Update pom.xml --- pom.xml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pom.xml b/pom.xml index b4fed84..6a3e823 100644 --- a/pom.xml +++ b/pom.xml @@ -380,7 +380,7 @@ software.amazon.awssdk.services dsql - 2.29.23 + 2.29.27 software.amazon.awssdk From a4045c6469288cdde1d04bd4fff3d989f9e27cca Mon Sep 17 00:00:00 2001 From: Srudeep Date: Tue, 3 Dec 2024 23:26:54 -0800 Subject: [PATCH 3/3] Update pom.xml --- pom.xml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pom.xml b/pom.xml index 6a3e823..cc4a5d7 100644 --- a/pom.xml +++ b/pom.xml @@ -378,7 +378,7 @@ - software.amazon.awssdk.services + software.amazon.awssdk dsql 2.29.27