Skip to content

Commit a53fa62

Browse files
danielfrankcomdanielfrankcom
committed
add plugin for DSQL auth tokens
1 parent 3b32ac2 commit a53fa62

File tree

6 files changed

+383
-0
lines changed

6 files changed

+383
-0
lines changed

‎wrapper/build.gradle.kts

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,7 @@ dependencies {
3737
compileOnly("software.amazon.awssdk:rds:2.31.78")
3838
compileOnly("software.amazon.awssdk:auth:2.31.45") // Required for IAM (light implementation)
3939
compileOnly("software.amazon.awssdk:http-client-spi:2.31.60") // Required for IAM (light implementation)
40+
compileOnly("software.amazon.awssdk:dsql:2.31.78")
4041
compileOnly("software.amazon.awssdk:sts:2.31.78")
4142
compileOnly("com.zaxxer:HikariCP:4.0.3") // Version 4.+ is compatible with Java 8
4243
compileOnly("com.mchange:c3p0:0.11.0")
@@ -73,6 +74,7 @@ dependencies {
7374
testImplementation("software.amazon.awssdk:rds:2.31.78")
7475
testImplementation("software.amazon.awssdk:auth:2.31.45") // Required for IAM (light implementation)
7576
testImplementation("software.amazon.awssdk:http-client-spi:2.31.60") // Required for IAM (light implementation)
77+
testImplementation("software.amazon.awssdk:dsql:2.31.78")
7678
testImplementation("software.amazon.awssdk:ec2:2.31.78")
7779
testImplementation("software.amazon.awssdk:secretsmanager:2.31.12")
7880
testImplementation("software.amazon.awssdk:sts:2.31.78")

‎wrapper/src/main/java/software/amazon/jdbc/ConnectionPluginChainBuilder.java

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,7 @@
4545
import software.amazon.jdbc.plugin.failover.FailoverConnectionPluginFactory;
4646
import software.amazon.jdbc.plugin.federatedauth.FederatedAuthPluginFactory;
4747
import software.amazon.jdbc.plugin.federatedauth.OktaAuthPluginFactory;
48+
import software.amazon.jdbc.plugin.iam.DsqlIamConnectionPluginFactory;
4849
import software.amazon.jdbc.plugin.iam.IamAuthConnectionPluginFactory;
4950
import software.amazon.jdbc.plugin.limitless.LimitlessConnectionPluginFactory;
5051
import software.amazon.jdbc.plugin.readwritesplitting.ReadWriteSplittingPluginFactory;
@@ -75,6 +76,7 @@ public class ConnectionPluginChainBuilder {
7576
put("failover", new FailoverConnectionPluginFactory());
7677
put("failover2", new software.amazon.jdbc.plugin.failover2.FailoverConnectionPluginFactory());
7778
put("iam", new IamAuthConnectionPluginFactory());
79+
put("dsql", new DsqlIamConnectionPluginFactory());
7880
put("awsSecretsManager", new AwsSecretsManagerConnectionPluginFactory());
7981
put("federatedAuth", new FederatedAuthPluginFactory());
8082
put("okta", new OktaAuthPluginFactory());
@@ -114,6 +116,7 @@ public class ConnectionPluginChainBuilder {
114116
put(FastestResponseStrategyPluginFactory.class, 900);
115117
put(LimitlessConnectionPluginFactory.class, 950);
116118
put(IamAuthConnectionPluginFactory.class, 1000);
119+
put(DsqlIamConnectionPluginFactory.class, 1001);
117120
put(AwsSecretsManagerConnectionPluginFactory.class, 1100);
118121
put(FederatedAuthPluginFactory.class, 1200);
119122
put(LogQueryConnectionPluginFactory.class, 1300);

‎wrapper/src/main/java/software/amazon/jdbc/plugin/iam/DsqlIamConnectionPluginFactory.java

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,33 @@
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+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package software.amazon.jdbc.plugin.iam;
18+
19+
import java.util.Properties;
20+
import org.checkerframework.checker.nullness.qual.NonNull;
21+
import software.amazon.jdbc.ConnectionPlugin;
22+
import software.amazon.jdbc.ConnectionPluginFactory;
23+
import software.amazon.jdbc.PluginService;
24+
25+
/**
26+
* Provides {@link ConnectionPlugin} instances which can be used to connect to Amazon Aurora DSQL.
27+
*/
28+
public class DsqlIamConnectionPluginFactory implements ConnectionPluginFactory {
29+
@Override
30+
public ConnectionPlugin getInstance(@NonNull final PluginService pluginService, @NonNull final Properties props) {
31+
return new IamAuthConnectionPlugin(pluginService, new DsqlTokenUtility());
32+
}
33+
}

‎wrapper/src/main/java/software/amazon/jdbc/plugin/iam/DsqlTokenUtility.java

Lines changed: 61 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,61 @@
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+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package software.amazon.jdbc.plugin.iam;
18+
19+
import org.checkerframework.checker.nullness.qual.NonNull;
20+
import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider;
21+
import software.amazon.awssdk.regions.Region;
22+
import software.amazon.awssdk.services.dsql.DsqlUtilities;
23+
24+
/**
25+
* Represents an {@link IamTokenUtility} which provides auth tokens for connecting to Amazon Aurora DSQL.
26+
*/
27+
public class DsqlTokenUtility implements IamTokenUtility {
28+
29+
private DsqlUtilities utilities = null;
30+
31+
public DsqlTokenUtility() { }
32+
33+
// For testing only
34+
DsqlTokenUtility(@NonNull final DsqlUtilities utilities) {
35+
this.utilities = utilities;
36+
}
37+
38+
@Override
39+
public String generateAuthenticationToken(
40+
@NonNull final AwsCredentialsProvider credentialsProvider,
41+
@NonNull final Region region,
42+
@NonNull final String hostname,
43+
final int port,
44+
@NonNull final String username) {
45+
if (this.utilities == null) {
46+
this.utilities = DsqlUtilities.builder()
47+
.credentialsProvider(credentialsProvider)
48+
.region(region)
49+
.build();
50+
}
51+
if (username.equals("admin")) {
52+
return this.utilities.generateDbConnectAdminAuthToken((builder) ->
53+
builder.hostname(hostname).region(region)
54+
);
55+
} else {
56+
return this.utilities.generateDbConnectAuthToken((builder) ->
57+
builder.hostname(hostname).region(region)
58+
);
59+
}
60+
}
61+
}

‎wrapper/src/test/java/software/amazon/jdbc/plugin/iam/DsqlIamConnectionPluginTest.java

Lines changed: 142 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,142 @@
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+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package software.amazon.jdbc.plugin.iam;
18+
19+
import static org.junit.jupiter.api.Assertions.*;
20+
import static org.mockito.ArgumentMatchers.anyString;
21+
import static org.mockito.ArgumentMatchers.eq;
22+
import static org.mockito.Mockito.*;
23+
import static software.amazon.jdbc.plugin.iam.DsqlTokenUtilityTest.*;
24+
25+
import java.sql.Connection;
26+
import java.sql.SQLException;
27+
import java.util.List;
28+
import java.util.Properties;
29+
import org.junit.jupiter.api.AfterEach;
30+
import org.junit.jupiter.api.BeforeEach;
31+
import org.junit.jupiter.api.Test;
32+
import org.junit.jupiter.params.ParameterizedTest;
33+
import org.junit.jupiter.params.provider.ValueSource;
34+
import org.mockito.Mock;
35+
import org.mockito.Mockito;
36+
import org.mockito.MockitoAnnotations;
37+
import software.amazon.awssdk.regions.Region;
38+
import software.amazon.jdbc.*;
39+
import software.amazon.jdbc.dialect.Dialect;
40+
import software.amazon.jdbc.hostavailability.SimpleHostAvailabilityStrategy;
41+
import software.amazon.jdbc.plugin.TokenInfo;
42+
import software.amazon.jdbc.util.IamAuthUtils;
43+
import software.amazon.jdbc.util.telemetry.TelemetryContext;
44+
import software.amazon.jdbc.util.telemetry.TelemetryFactory;
45+
import software.amazon.jdbc.util.telemetry.TelemetryTraceLevel;
46+
47+
class DsqlIamConnectionPluginTest {
48+
49+
private static final Region TEST_REGION = Region.US_EAST_1;
50+
private static final String TEST_HOSTNAME = String.format("foo0bar1baz2quux3quuux4.dsql.%s.on.aws", TEST_REGION);
51+
private static final int TEST_PORT = 5432;
52+
53+
private static final String DRIVER_PROTOCOL = "jdbc:postgresql:";
54+
private static final HostSpec HOST_SPEC = new HostSpecBuilder(new SimpleHostAvailabilityStrategy())
55+
.host(TEST_HOSTNAME).port(TEST_PORT).build();
56+
57+
private static final String DEFAULT_USERNAME = "admin";
58+
59+
private final Properties props = new Properties();
60+
61+
private AutoCloseable cleanMocksCallback;
62+
@Mock private Connection mockConnection;
63+
@Mock private PluginService mockPluginService;
64+
@Mock private Dialect mockDialect;
65+
@Mock private TelemetryFactory mockTelemetryFactory;
66+
@Mock private TelemetryContext mockTelemetryContext;
67+
@Mock private JdbcCallable<Connection, SQLException> mockLambda;
68+
@Mock private ConnectionProvider mockConnectionProvider;
69+
@Mock private PluginManagerService mockPluginManagerService;
70+
71+
@BeforeEach
72+
public void init() {
73+
cleanMocksCallback = MockitoAnnotations.openMocks(this);
74+
75+
IamAuthConnectionPlugin.clearCache();
76+
77+
props.setProperty(PropertyDefinition.USER.name, DEFAULT_USERNAME);
78+
props.setProperty("iamRegion", Region.US_EAST_1.toString());
79+
props.setProperty(PropertyDefinition.PLUGINS.name, "dsql");
80+
81+
doReturn(mockDialect).when(mockPluginService).getDialect();
82+
doReturn(TEST_PORT).when(mockDialect).getDefaultPort();
83+
doReturn(mockTelemetryFactory).when(mockPluginService).getTelemetryFactory();
84+
doReturn(mockTelemetryContext).when(mockTelemetryFactory)
85+
.openTelemetryContext(anyString(), eq(TelemetryTraceLevel.NESTED));
86+
}
87+
88+
@AfterEach
89+
public void cleanup() throws Exception {
90+
cleanMocksCallback.close();
91+
}
92+
93+
@SuppressWarnings("resource") // Prevent Mockito warning when mocking closeable return type.
94+
private void assertPluginProvidesDsqlTokens(final ConnectionPlugin plugin, final String username) throws SQLException {
95+
Mockito.doReturn(mockConnection).when(mockLambda).call();
96+
97+
plugin
98+
.connect(DRIVER_PROTOCOL, HOST_SPEC, props, true, mockLambda)
99+
.close();
100+
101+
final String cacheKey = IamAuthUtils.getCacheKey(
102+
username,
103+
TEST_HOSTNAME,
104+
TEST_PORT,
105+
TEST_REGION);
106+
107+
final TokenInfo info = IamAuthCacheHolder.tokenCache.get(cacheKey);
108+
final String token = info.getToken();
109+
110+
assertTokenContainsProperties(token, TEST_HOSTNAME, username);
111+
}
112+
113+
@Test
114+
public void testDsqlPluginRegistration() throws SQLException {
115+
ConnectionPluginChainBuilder builder = new ConnectionPluginChainBuilder();
116+
117+
final List<ConnectionPlugin> result = builder.getPlugins(
118+
mockPluginService,
119+
mockConnectionProvider,
120+
null,
121+
mockPluginManagerService,
122+
props,
123+
null);
124+
125+
// 2 because default plugin is always included.
126+
assertEquals(2, result.size());
127+
final ConnectionPlugin plugin = result.get(0);
128+
129+
assertInstanceOf(IamAuthConnectionPlugin.class, plugin);
130+
assertPluginProvidesDsqlTokens(plugin, DEFAULT_USERNAME);
131+
}
132+
133+
@ParameterizedTest
134+
@ValueSource(strings = {REGULAR_USER, ADMIN_USER})
135+
public void testDsqlTokenGeneratedBasedOnUser(final String username) throws SQLException {
136+
props.setProperty(PropertyDefinition.USER.name, username);
137+
138+
final DsqlIamConnectionPluginFactory factory = new DsqlIamConnectionPluginFactory();
139+
final ConnectionPlugin plugin = factory.getInstance(mockPluginService, props);
140+
assertPluginProvidesDsqlTokens(plugin, username);
141+
}
142+
}

0 commit comments

Comments
 (0)