Code changes to support assuming a role which could be a cross account role.

Author: akshayar

README.md

@@ -119,7 +119,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 +138,24 @@ The following is an example of this.
         }
     }
 ```
+
+Dollowing is an example of this config with 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:aws:iam::ACCOUNT_ID:role/ROLE_NAME"
+            }
+            ssl-engine-factory {
+                class = DefaultSslEngineFactory
+            }
+        }
+    }
+```

pom.xml

@@ -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,11 @@
 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.StsClientBuilder;
+import software.amazon.awssdk.services.sts.auth.StsAssumeRoleCredentialsProvider;
+import software.amazon.awssdk.services.sts.auth.StsGetSessionTokenCredentialsProvider;
+import software.amazon.awssdk.services.sts.model.AssumeRoleRequest;
 
 import static software.amazon.awssdk.auth.credentials.DefaultCredentialsProvider.create;
 
@@ -106,6 +112,12 @@ public String getPath() {
             }
         };
 
+    private final static DriverOption ROLE_OPTION = new DriverOption() {
+        public String getPath() {
+            return "advanced.auth-provider.aws-role";
+        }
+    };
+
     /**
      * This constructor is provided so that the driver can create
      * instances of this class based on configuration. For example:
@@ -130,7 +142,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));
     }
 
     /**
@@ -139,8 +152,8 @@ public SigV4AuthProvider(DriverContext driverContext) {
      * null value indicates to use the AWS_REGION environment
      * variable, or the "aws.region" system property to configure it.
      */
-    public SigV4AuthProvider(final String region) {
-        this(create(), region);
+    public SigV4AuthProvider(final String region,final String roleArn) {
+        this(Optional.ofNullable(roleArn).map(r->(AwsCredentialsProvider)createSTSRoleCredentialProvider(r,"keyspaces-session",region)).orElse(create()), region);
     }
 
     /**
@@ -373,4 +386,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 sessionName The name of the session
+     * @param stsRegion The region of the STS endpoint
+     * @return
+     */
+    private static StsAssumeRoleCredentialsProvider createSTSRoleCredentialProvider(String roleArn,
+                                                                     String sessionName, String stsRegion) {
+        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
+     */
+    private static String getDefaultRegion() {
+        DefaultAwsRegionProviderChain chain = new DefaultAwsRegionProviderChain();
+        return Optional.ofNullable(chain.getRegion()).orElse(Region.US_EAST_1).toString();
+    }
 }

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

@@ -55,7 +55,7 @@ public static void main(String[] args) throws Exception {
         //By default the reference.conf is loaded by the driver which contains all defaults.
         //You can override this by providing reference.conf on the classpath
         //to isolate test you can load conf with a custom name
-        URL url = TestSigV4Config.class.getClassLoader().getResource("keyspaces-reference.conf");
+        URL url = TestSigV4Config.class.getClassLoader().getResource("keyspaces-reference-norole.conf");
 
         File file = new File(url.toURI());
         // The CqlSession object is the main entry point of the driver.
@@ -64,18 +64,16 @@ public static void main(String[] args) throws Exception {
         // it throughout your application.
         try (CqlSession session = CqlSession.builder()
                 .withConfigLoader(DriverConfigLoader.fromFile(file))
-                .addContactPoints(contactPoints)
-                .withLocalDatacenter("us-west-2")
              .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 release_version from system.local");
+            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("release_version");
+            String releaseVersion = row.getString("category");
             System.out.printf("Cassandra version is: %s%n", releaseVersion);
         }
     }

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: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
+}