Merge pull request #34 from akshayar/master

Support for assuming role before accessing KeySpaces, helps in cross account access

Author: pranayw333

README.md

@@ -8,6 +8,7 @@ This package implements an authentication plugin for the open-source Datastax Ja
 
 The plugin depends on the AWS SDK for Java. It uses `AWSCredentialsProvider` to obtain credentials. Because the IAuthenticator interface operates at the level of `InetSocketAddress`, you must specify the service endpoint to use for the connection.
 You can provide the Region in the constructor programmatically, via the `AWS_REGION` environment variable, or via the `aws.region` system property.
+You can also provide an IAM role to assume for access to KeySpaces, programmatically or via the configuration file.
 
 The full documentation for the plugin is available at
 https://docs.aws.amazon.com/keyspaces/latest/devguide/programmatic.credentials.html#programmatic.credentials.SigV4_KEYSPACES.
@@ -35,28 +36,40 @@ You can specify the Region using one of the following four methods:
 * Constructor
 * Configuration
 
-## Environment Variable
+### Environment Variable
 
 You can use the `AWS_REGION` environment variable to match the endpoint that you are communicating with by setting it as part of your application start-up, as follows.
 
 ``` shell
 $ export AWS_Region=us-east-1
 ```
-## System Property
+### System Property
 
 You can use the `aws.region` Java system property by specifying it on the command line, as follows.
 
 ``` shell
 $ java -Daws.region=us=east-1 ...
 ```
 
-## Constructor
+### Constructor
 
 One of the constructors for `software.aws.mcs.auth.SigV4AuthProvider` takes a `String` representing the Region that will be used for that instance.
 
-## Configuration
+### Configuration
 
-Set the Region explicitly in your `advanced.auth-provider.class` configuration (see example below), by specifying the `advanced.auth-provider.aws-region` property.
+Set the Region explicitly in your `advanced.auth-provider` configuration (see example below), by specifying the `advanced.auth-provider.aws-region` property.
+
+## Assume IAM Role Configuration
+
+You can specify an IAM role to assume for access to KeySpaces using either the constructor or the driver configuration file
+
+### Constructor
+
+One of the constructors for `software.aws.mcs.auth.SigV4AuthProvider` takes two Strings , the first representing the region and the second representing the ARN of the IAM role to assume. 
+
+### Configuration
+
+Set the IAM Role explicitly in your `advanced.auth-provider` configuration (see example below), by specifying the `advanced.auth-provider.aws-role-arn` property.
 
 ## Add the Authentication Plugin to the Application
 
@@ -119,7 +132,7 @@ To use the configuration file, set the `advanced.auth-provider.class` to `softwa
 1. Set the `advanced.auth-provider.class` to `software.aws.mcs.auth.SigV4AuthProvider`.
 1. Set `basic.load-balancing-policy.local-datacenter` to the region name. In this case, use `us-east-2`.
 
-The following is an example of this.
+The following is an example of this config without explicit role to be assumed. 
 
 ``` text
     datastax-java-driver {
@@ -138,3 +151,24 @@ The following is an example of this.
         }
     }
 ```
+
+The following is an example of this config with an explicit role to be assumed.
+
+``` text
+    datastax-java-driver {
+        basic.load-balancing-policy {
+            class = DefaultLoadBalancingPolicy
+            local-datacenter = us-east-2
+        }
+        advanced {
+            auth-provider = {
+                class = software.aws.mcs.auth.SigV4AuthProvider
+                aws-region = us-east-2
+                aws-role-arn = "arn:aws:iam::ACCOUNT_ID:role/ROLE_NAME"
+            }
+            ssl-engine-factory {
+                class = DefaultSslEngineFactory
+            }
+        }
+    }
+```

pom.xml

@@ -3,7 +3,7 @@
     <modelVersion>4.0.0</modelVersion>
     <groupId>software.aws.mcs</groupId>
     <artifactId>aws-sigv4-auth-cassandra-java-driver-plugin</artifactId>
-    <version>4.0.9</version>
+    <version>4.0.10</version>
     <name>AWS SigV4 Auth Java Driver 4.x Plugin</name>
     <description>A Plugin to allow SigV4 authentication for Java Cassandra drivers with Amazon MCS</description>
     <url>https://github.com/aws/aws-sigv4-auth-cassandra-java-driver-plugin</url>
@@ -44,6 +44,7 @@
         <project.build.sourceEncoding>UTF-8</project.build.sourceEncoding>
         <jackson.version>2.13.4</jackson.version>
         <jackson-databind.version>2.13.4.2</jackson-databind.version>
+        <aws-sdk-version>2.15.66</aws-sdk-version>
     </properties>
 
     <dependencies>
@@ -88,7 +89,13 @@
         <dependency>
             <groupId>software.amazon.awssdk</groupId>
             <artifactId>auth</artifactId>
-            <version>2.15.66</version>
+            <version>${aws-sdk-version}</version>
+        </dependency>
+        <!-- https://mvnrepository.com/artifact/software.amazon.awssdk/sts -->
+        <dependency>
+            <groupId>software.amazon.awssdk</groupId>
+            <artifactId>sts</artifactId>
+            <version>${aws-sdk-version}</version>
         </dependency>
     </dependencies>
 

src/main/java/software/aws/mcs/auth/SigV4AuthProvider.java

@@ -32,6 +32,7 @@
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.List;
+import java.util.Optional;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CompletionStage;
 import javax.crypto.Mac;
@@ -53,6 +54,9 @@
 import software.amazon.awssdk.auth.signer.internal.SignerConstant;
 import software.amazon.awssdk.regions.Region;
 import software.amazon.awssdk.regions.providers.DefaultAwsRegionProviderChain;
+import software.amazon.awssdk.services.sts.StsClient;
+import software.amazon.awssdk.services.sts.auth.StsAssumeRoleCredentialsProvider;
+import software.amazon.awssdk.services.sts.model.AssumeRoleRequest;
 
 import static software.amazon.awssdk.auth.credentials.DefaultCredentialsProvider.create;
 
@@ -100,11 +104,9 @@ public SigV4AuthProvider() {
         this(create(), null);
     }
 
-    private final static DriverOption REGION_OPTION = new DriverOption() {
-            public String getPath() {
-                return "advanced.auth-provider.aws-region";
-            }
-        };
+    private final static DriverOption REGION_OPTION = () -> "advanced.auth-provider.aws-region";
+
+    private final static DriverOption ROLE_OPTION = () -> "advanced.auth-provider.aws-role-arn";
 
     /**
      * This constructor is provided so that the driver can create
@@ -130,7 +132,8 @@ public String getPath() {
      * Unused for this plugin.
      */
     public SigV4AuthProvider(DriverContext driverContext) {
-        this(driverContext.getConfig().getDefaultProfile().getString(REGION_OPTION, null));
+        this(driverContext.getConfig().getDefaultProfile().getString(REGION_OPTION, getDefaultRegion()),
+             driverContext.getConfig().getDefaultProfile().getString(ROLE_OPTION, null));
     }
 
     /**
@@ -143,6 +146,17 @@ public SigV4AuthProvider(final String region) {
         this(create(), region);
     }
 
+    /**
+     * Create a new Provider, using the specified region and IAM role to assume.
+     * @param region the region (e.g. us-east-1) to use for signing. A
+     * null value indicates to use the AWS_REGION environment
+     * variable, or the "aws.region" system property to configure it.
+     * @param roleArn The IAM Role ARN which the connecting client should assume before connecting with Amazon Keyspaces.
+     */
+    public SigV4AuthProvider(final String region,final String roleArn) {
+        this(Optional.ofNullable(roleArn).map(r->(AwsCredentialsProvider)createSTSRoleCredentialProvider(r,region)).orElse(create()), region);
+    }
+
     /**
      * Create a new Provider, using the specified AWSCredentialsProvider and region.
      * @param credentialsProvider the credentials provider used to obtain signature material
@@ -154,9 +168,7 @@ public SigV4AuthProvider(@NotNull AwsCredentialsProvider credentialsProvider, fi
         this.credentialsProvider = credentialsProvider;
 
         if (region == null) {
-            DefaultAwsRegionProviderChain chain = new DefaultAwsRegionProviderChain();
-            Region defaultRegion = chain.getRegion();
-            this.signingRegion = defaultRegion.toString().toLowerCase();
+            this.signingRegion = getDefaultRegion();
         } else {
             this.signingRegion = region.toLowerCase();
         }
@@ -373,4 +385,36 @@ static int indexOf(byte[] target, byte[] pattern) {
         // Loop exhaustion means we did not find it
         return -1;
     }
+
+
+    /**
+     * Creates a STS role credential provider
+     * @param roleArn The ARN of the role to assume
+     * @param stsRegion The region of the STS endpoint
+     * @return The STS role credential provider
+     */
+    private static StsAssumeRoleCredentialsProvider createSTSRoleCredentialProvider(@NotNull String roleArn,
+                                                                            @NotNull String stsRegion) {
+        final String sessionName="keyspaces-session-"+System.currentTimeMillis();
+        StsClient stsClient = StsClient.builder()
+                .region(Region.of(stsRegion))
+                .build();
+        AssumeRoleRequest assumeRoleRequest=AssumeRoleRequest.builder()
+                .roleArn(roleArn)
+                .roleSessionName(sessionName)
+                .build();
+        return StsAssumeRoleCredentialsProvider.builder()
+                .stsClient(stsClient)
+                .refreshRequest(assumeRoleRequest)
+                .build();
+    }
+
+    /**
+     * Gets the default region for SigV4 if region is not provided.
+     * @return Default region
+     */
+    private static String getDefaultRegion() {
+        DefaultAwsRegionProviderChain chain = new DefaultAwsRegionProviderChain();
+        return chain.getRegion().toString().toLowerCase();
+    }
 }

src/test/java/software/aws/mcs/auth/TestSigV4AssumeRoleConfig.java

@@ -0,0 +1,73 @@
+package software.aws.mcs.auth;
+
+/*-
+ * #%L
+ * AWS SigV4 Auth Java Driver 4.x Plugin
+ * %%
+ * Copyright (C) 2020-2021 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.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License 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.
+ * #L%
+ */
+
+import com.datastax.oss.driver.api.core.CqlSession;
+import com.datastax.oss.driver.api.core.config.DriverConfigLoader;
+import com.datastax.oss.driver.api.core.cql.ResultSet;
+import com.datastax.oss.driver.api.core.cql.Row;
+import org.junit.platform.commons.util.StringUtils;
+
+import java.io.File;
+import java.net.InetSocketAddress;
+import java.net.URL;
+import java.util.ArrayList;
+import java.util.Optional;
+
+public class TestSigV4AssumeRoleConfig {
+    static String KEYSPACES_DEFAULT_CONF="keyspaces-reference-norole.conf";
+    /**
+     * Before executing this test, ensure that KeySpaces tables are created.
+     * Refer ddl.cql and dml.cql to create and populate tables.
+     * @param args
+     * @throws Exception
+     */
+    public static void main(String[] args) throws Exception {
+        String keySpacesConf=KEYSPACES_DEFAULT_CONF;
+        if(args.length>1){
+            keySpacesConf=args[0];
+            System.out.println("Using key spaces config file: "+keySpacesConf);
+            keySpacesConf=Optional.of(keySpacesConf).filter(StringUtils::isBlank).orElse(KEYSPACES_DEFAULT_CONF);
+        }
+
+        URL url = TestSigV4AssumeRoleConfig.class.getClassLoader().getResource(keySpacesConf);
+
+        File file = new File(url.toURI());
+        // The CqlSession object is the main entry point of the driver.
+        // It holds the known state of the actual Cassandra cluster (notably the Metadata).
+        // This class is thread-safe, you should create a single instance (per target Cassandra cluster), and share
+        // it throughout your application.
+        try (CqlSession session = CqlSession.builder()
+                .withConfigLoader(DriverConfigLoader.fromFile(file))
+             .build()) {
+
+            // We use execute to send a query to Cassandra. This returns a ResultSet, which is essentially a collection
+            // of Row objects.
+            ResultSet rs = session.execute("select * from testkeyspace.testconf");
+            //  Extract the first row (which is the only one in this case).
+            Row row = rs.one();
+
+            // Extract the value of the first (and only) column from the row.
+            String releaseVersion = row.getString("category");
+            System.out.printf("Cassandra version is: %s%n", releaseVersion);
+        }
+    }
+}

src/test/java/software/aws/mcs/auth/TestSigV4Config.java

@@ -66,7 +66,7 @@ public static void main(String[] args) throws Exception {
                 .withConfigLoader(DriverConfigLoader.fromFile(file))
                 .addContactPoints(contactPoints)
                 .withLocalDatacenter("us-west-2")
-             .build()) {
+                .build()) {
 
             // We use execute to send a query to Cassandra. This returns a ResultSet, which is essentially a collection
             // of Row objects.

src/test/resources/ddl.cql

@@ -0,0 +1,18 @@
+CREATE KEYSPACE "testkeyspace" WITH REPLICATION = {'class': 'SingleRegionStrategy'};
+
+CREATE TABLE "testkeyspace"."testconf"(
+  "id" ascii,
+  "category" ascii,
+  PRIMARY KEY("id")
+)
+WITH CUSTOM_PROPERTIES = {
+  'capacity_mode':{
+  'throughput_mode':'PAY_PER_REQUEST'
+  },
+  'point_in_time_recovery':{
+  'status':'enabled'
+  },
+  'encryption_specification':{
+  'encryption_type':'AWS_OWNED_KMS_KEY'
+  }
+} ;

src/test/resources/dml.cql

@@ -0,0 +1 @@
+INSERT into testkeyspace.testconf(id,category) values('first','first Category');

src/test/resources/keyspaces-reference-norole.conf

@@ -0,0 +1,24 @@
+datastax-java-driver {
+        basic.contact-points = ["cassandra.ap-south-1.amazonaws.com:9142"]
+        basic.load-balancing-policy {
+            class = DefaultLoadBalancingPolicy
+            local-datacenter = ap-south-1
+            slow-replica-avoidance = false
+        }
+        basic.request {
+              consistency = LOCAL_QUORUM
+        }
+        advanced {
+                auth-provider = {
+                   class = software.aws.mcs.auth.SigV4AuthProvider
+                   aws-region = ap-south-1
+                 }
+            ssl-engine-factory {
+                class = DefaultSslEngineFactory
+                truststore-path = "<path>/cassandra_truststore.jks"
+                truststore-password = "<password>"
+                hostname-validation=false
+            }
+        }
+        advanced.connection.pool.local.size = 3
+}

src/test/resources/keyspaces-reference-withrole.conf

@@ -0,0 +1,25 @@
+datastax-java-driver {
+        basic.contact-points = ["cassandra.ap-south-1.amazonaws.com:9142"]
+        basic.load-balancing-policy {
+            class = DefaultLoadBalancingPolicy
+            local-datacenter = ap-south-1
+            slow-replica-avoidance = false
+        }
+        basic.request {
+              consistency = LOCAL_QUORUM
+        }
+        advanced {
+                auth-provider = {
+                   class = software.aws.mcs.auth.SigV4AuthProvider
+                   aws-region = ap-south-1
+                   aws-role-arn = "arn:aws:iam::ACCOUNT_ID:role/keyspaces-act2-role"
+                 }
+            ssl-engine-factory {
+                class = DefaultSslEngineFactory
+                truststore-path = "<path>/cassandra_truststore.jks"
+                truststore-password = "<password>"
+                hostname-validation=false
+            }
+        }
+        advanced.connection.pool.local.size = 3
+}