Merge pull request #3504 from aws/merge-dsql-customization

Merge dsql customization

.changes/next-release/feature-AxdbFrontend-d6ce1e6.json

@@ -0,0 +1,6 @@
+{
+    "type": "feature",
+    "category": "DSQL",
+    "contributor": "APandher",
+    "description": "Add IAM Token Generation Utility for DSQL"
+}

feature.metadata

@@ -0,0 +1 @@
+{"trebuchetFeatureArn":"arn:aws:trebuchet:::feature:v2:2a236723-7c96-48cd-8055-918b35a8308a","c2jModelsRevision":25,"messageId":1,"serviceId":"DSQL","serviceModule":"dsql"}

services/dsql/pom.xml

@@ -41,6 +41,11 @@
         </plugins>
     </build>
     <dependencies>
+        <dependency>
+            <groupId>nl.jqno.equalsverifier</groupId>
+            <artifactId>equalsverifier</artifactId>
+            <scope>test</scope>
+        </dependency>
         <dependency>
             <groupId>software.amazon.awssdk</groupId>
             <artifactId>protocol-core</artifactId>

services/dsql/src/main/java/software/amazon/awssdk/services/dsql/DefaultDsqlUtilities.java

@@ -0,0 +1,166 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License").
+ * You may not use this file except in compliance with the License.
+ * A copy of the License is located at
+ *
+ *  http://aws.amazon.com/apache2.0
+ *
+ * or in the "license" file accompanying this file. This file is distributed
+ * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ * express or implied. See the License for the specific language governing
+ * permissions and limitations under the License.
+ */
+
+package software.amazon.awssdk.services.dsql;
+
+import java.time.Clock;
+import java.time.Instant;
+import software.amazon.awssdk.annotations.Immutable;
+import software.amazon.awssdk.annotations.SdkInternalApi;
+import software.amazon.awssdk.annotations.SdkTestInternalApi;
+import software.amazon.awssdk.auth.credentials.AwsCredentials;
+import software.amazon.awssdk.auth.credentials.CredentialUtils;
+import software.amazon.awssdk.auth.signer.Aws4Signer;
+import software.amazon.awssdk.auth.signer.params.Aws4PresignerParams;
+import software.amazon.awssdk.awscore.client.config.AwsClientOption;
+import software.amazon.awssdk.core.client.config.SdkClientConfiguration;
+import software.amazon.awssdk.http.SdkHttpFullRequest;
+import software.amazon.awssdk.http.SdkHttpMethod;
+import software.amazon.awssdk.identity.spi.AwsCredentialsIdentity;
+import software.amazon.awssdk.identity.spi.IdentityProvider;
+import software.amazon.awssdk.regions.Region;
+import software.amazon.awssdk.services.dsql.model.GenerateAuthTokenRequest;
+import software.amazon.awssdk.utils.CompletableFutureUtils;
+import software.amazon.awssdk.utils.Logger;
+import software.amazon.awssdk.utils.StringUtils;
+
+@Immutable
+@SdkInternalApi
+public final class DefaultDsqlUtilities implements DsqlUtilities {
+    private static final Logger log = Logger.loggerFor(DsqlUtilities.class);
+    private final Aws4Signer signer = Aws4Signer.create();
+    private final Region region;
+    private final IdentityProvider<? extends AwsCredentialsIdentity> credentialsProvider;
+    private final Clock clock;
+
+    public DefaultDsqlUtilities(DefaultBuilder builder) {
+        this(builder, Clock.systemUTC());
+    }
+
+    /**
+     * For testing purposes only
+     */
+    @SdkTestInternalApi
+    public DefaultDsqlUtilities(DefaultBuilder builder, Clock clock) {
+        this.credentialsProvider = builder.credentialsProvider;
+        this.region = builder.region;
+        this.clock = clock;
+    }
+
+    /**
+     * Used by DSQL low-level client's utilities() method
+     */
+    @SdkInternalApi
+    static DsqlUtilities create(SdkClientConfiguration clientConfiguration) {
+        return new DefaultBuilder().clientConfiguration(clientConfiguration).build();
+    }
+
+    @Override
+    public String generateDbConnectAuthToken(GenerateAuthTokenRequest request) {
+        return generateAuthToken(request, false);
+    }
+
+    @Override
+    public String generateDbConnectAdminAuthToken(GenerateAuthTokenRequest request) {
+        return generateAuthToken(request, true);
+    }
+
+    private String generateAuthToken(GenerateAuthTokenRequest request, boolean isAdmin) {
+        String action = isAdmin ? "DbConnectAdmin" : "DbConnect";
+
+        SdkHttpFullRequest httpRequest = SdkHttpFullRequest.builder()
+                                                           .method(SdkHttpMethod.GET)
+                                                           .protocol("https")
+                                                           .host(request.hostname())
+                                                           .encodedPath("/")
+                                                           .putRawQueryParameter("Action", action)
+                                                           .build();
+
+        Instant expirationTime = Instant.now(clock).plus(request.expiresIn());
+
+        Aws4PresignerParams presignRequest = Aws4PresignerParams.builder()
+                                                                .signingClockOverride(clock)
+                                                                .expirationTime(expirationTime)
+                                                                .awsCredentials(resolveCredentials(request))
+                                                                .signingName("dsql")
+                                                                .signingRegion(resolveRegion(request))
+                                                                .build();
+
+        SdkHttpFullRequest fullRequest = signer.presign(httpRequest, presignRequest);
+        String signedUrl = fullRequest.getUri().toString();
+
+        log.debug(() -> "Generated DSQL authentication token with expiration of " + expirationTime);
+        return StringUtils.replacePrefixIgnoreCase(signedUrl, "https://", "");
+    }
+
+    private Region resolveRegion(GenerateAuthTokenRequest request) {
+        if (request.region() != null) {
+            return request.region();
+        }
+
+        if (this.region != null) {
+            return this.region;
+        }
+
+        throw new IllegalArgumentException("Region must be provided in GenerateAuthTokenRequest or DsqlUtilities");
+    }
+
+    // TODO: update this to use AwsCredentialsIdentity when we migrate Signers to accept the new type.
+    private AwsCredentials resolveCredentials(GenerateAuthTokenRequest request) {
+        if (request.credentialsIdentityProvider() != null) {
+            return CredentialUtils.toCredentials(
+                CompletableFutureUtils.joinLikeSync(request.credentialsIdentityProvider().resolveIdentity()));
+        }
+
+        if (this.credentialsProvider != null) {
+            return CredentialUtils.toCredentials(CompletableFutureUtils.joinLikeSync(this.credentialsProvider.resolveIdentity()));
+        }
+
+        throw new IllegalArgumentException("CredentialsProvider must be provided in GenerateAuthTokenRequest or DsqlUtilities");
+    }
+
+    @SdkInternalApi
+    public static final class DefaultBuilder implements DsqlUtilities.Builder {
+        private Region region;
+        private IdentityProvider<? extends AwsCredentialsIdentity> credentialsProvider;
+
+        Builder clientConfiguration(SdkClientConfiguration clientConfiguration) {
+            this.credentialsProvider = clientConfiguration.option(AwsClientOption.CREDENTIALS_IDENTITY_PROVIDER);
+            this.region = clientConfiguration.option(AwsClientOption.AWS_REGION);
+
+            return this;
+        }
+
+        @Override
+        public Builder region(Region region) {
+            this.region = region;
+            return this;
+        }
+
+        @Override
+        public Builder credentialsProvider(IdentityProvider<? extends AwsCredentialsIdentity> credentialsProvider) {
+            this.credentialsProvider = credentialsProvider;
+            return this;
+        }
+
+        /**
+         * Construct a {@link DsqlUtilities} object.
+         */
+        @Override
+        public DsqlUtilities build() {
+            return new DefaultDsqlUtilities(this);
+        }
+    }
+}

services/dsql/src/main/java/software/amazon/awssdk/services/dsql/DsqlUtilities.java

@@ -0,0 +1,143 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License").
+ * You may not use this file except in compliance with the License.
+ * A copy of the License is located at
+ *
+ *  http://aws.amazon.com/apache2.0
+ *
+ * or in the "license" file accompanying this file. This file is distributed
+ * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ * express or implied. See the License for the specific language governing
+ * permissions and limitations under the License.
+ */
+
+package software.amazon.awssdk.services.dsql;
+
+import java.util.function.Consumer;
+import software.amazon.awssdk.annotations.SdkPublicApi;
+import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider;
+import software.amazon.awssdk.identity.spi.AwsCredentialsIdentity;
+import software.amazon.awssdk.identity.spi.IdentityProvider;
+import software.amazon.awssdk.regions.Region;
+import software.amazon.awssdk.services.dsql.model.GenerateAuthTokenRequest;
+
+/**
+ * Utilities for working with DSQL. An instance of this class can be created by:
+ * <p>
+ * 1) Using the low-level client {@link DsqlClient#utilities()} (or {@link DsqlAsyncClient#utilities()}} method.
+ * This is recommended as SDK will use the same configuration from the {@link DsqlClient} object to create the
+ * {@link DsqlUtilities} object.
+ *
+ * @snippet :
+ * {@code
+ * DsqlClient dsqlClient = DsqlClient.create();
+ * DsqlUtilities utilities = DsqlClient.utilities();
+ * }
+ *
+ * <p>
+ * 2) Directly using the {@link #builder()} method.
+ *
+ * @snippet :
+ * {@code
+ * DsqlUtilities utilities = DsqlUtilities.builder()
+ *  .credentialsProvider(DefaultCredentialsProvider.create())
+ *  .region(Region.US_WEST_2)
+ *  .build()
+ * }
+ *
+ * Note: This class does not make network calls.
+ */
+@SdkPublicApi
+public interface DsqlUtilities {
+    /**
+     * Create a builder that can be used to configure and create a {@link DsqlUtilities}.
+     */
+    static Builder builder() {
+        return new DefaultDsqlUtilities.DefaultBuilder();
+    }
+
+    /**
+     * Generates an authentication token for IAM authentication to an DSQL database.
+     *
+     * @param request The request used to generate the authentication token
+     * @return String to use as the DSQL authentication token
+     * @throws IllegalArgumentException if the required parameters are not valid
+     */
+    default String generateDbConnectAuthToken(Consumer<GenerateAuthTokenRequest.Builder> request) {
+        return generateDbConnectAuthToken(GenerateAuthTokenRequest.builder().applyMutation(request).build());
+    }
+
+    /**
+     * Generates an authentication token for IAM authentication to an DSQL database.
+     *
+     * @param request The request used to generate the authentication token
+     * @return String to use as the DSQL authentication token
+     * @throws IllegalArgumentException if the required parameters are not valid
+     */
+    default String generateDbConnectAuthToken(GenerateAuthTokenRequest request) {
+        throw new UnsupportedOperationException();
+    }
+
+    /**
+     * Generates an admin authentication token for IAM authentication to an DSQL database.
+     *
+     * @param request The request used to generate the admin authentication token
+     * @return String to use as the DSQL authentication token
+     * @throws IllegalArgumentException if the required parameters are not valid
+     */
+    default String generateDbConnectAdminAuthToken(Consumer<GenerateAuthTokenRequest.Builder> request) {
+        return generateDbConnectAdminAuthToken(GenerateAuthTokenRequest.builder().applyMutation(request).build());
+    }
+
+    /**
+     * Generates an admin authentication token for IAM authentication to an DSQL database.
+     *
+     * @param request The request used to generate the admin authentication token
+     * @return String to use as the DSQL authentication token
+     * @throws IllegalArgumentException if the required parameters are not valid
+     */
+    default String generateDbConnectAdminAuthToken(GenerateAuthTokenRequest request) {
+        throw new UnsupportedOperationException();
+    }
+
+
+    /**
+     * Builder for creating an instance of {@link DsqlUtilities}. It can be configured using
+     * {@link DsqlUtilities#builder()}.
+     * Once configured, the {@link DsqlUtilities} can created using {@link #build()}.
+     */
+    @SdkPublicApi
+    interface Builder {
+        /**
+         * The default region to use when working with the methods in {@link DsqlUtilities} class.
+         *
+         * @return This object for method chaining
+         */
+        Builder region(Region region);
+
+        /**
+         * The default credentials provider to use when working with the methods in {@link DsqlUtilities} class.
+         *
+         * @return This object for method chaining
+         */
+        default Builder credentialsProvider(AwsCredentialsProvider credentialsProvider) {
+            return credentialsProvider((IdentityProvider<? extends AwsCredentialsIdentity>) credentialsProvider);
+        }
+
+        /**
+         * The default credentials provider to use when working with the methods in {@link DsqlUtilities} class.
+         *
+         * @return This object for method chaining
+         */
+        default Builder credentialsProvider(IdentityProvider<? extends AwsCredentialsIdentity> credentialsProvider) {
+            throw new UnsupportedOperationException();
+        }
+
+        /**
+         * Create a {@link DsqlUtilities}
+         */
+        DsqlUtilities build();
+    }
+}