From 6f6d88dbd6fc2aa543c61c8da207683a2799d43e Mon Sep 17 00:00:00 2001
From: Vaibhav Khurana <vaibhav.khurana@razorpay.com>
Date: Fri, 10 Apr 2020 04:38:03 +0530
Subject: [PATCH 1/2] Adding ssl support on ndbench cassandra

---
 ndbench-cass-plugins/README.md                |  2 ++
 .../plugin/cass/CJavaDriverBasePlugin.java    |  6 ++++-
 .../plugin/cass/CassJavaDriverManager.java    |  2 +-
 .../cass/CassJavaDriverManagerImpl.java       | 24 +++++++++++++++++--
 .../configs/CassandraConfigurationBase.java   |  6 +++++
 5 files changed, 36 insertions(+), 4 deletions(-)

diff --git a/ndbench-cass-plugins/README.md b/ndbench-cass-plugins/README.md
index 433c9411..73acd85b 100644
--- a/ndbench-cass-plugins/README.md
+++ b/ndbench-cass-plugins/README.md
@@ -26,6 +26,8 @@ Every config has to be prefixed with configuration prefix mentioned above
 |cass.cfname| 	emp | Destination CF Name|
 |cass.readConsistencyLevel| 	LOCAL_ONE | Read CL|
 |cass.writeConsistencyLevel| 	LOCAL_ONE | Write CL|
+|cass.truststorePath| null | Absolute path of truststore
+|cass.truststorePass| null | Password of truststore
 
 
 #### Cassandra *CassJavaDriverGeneric* Plugin configurations
diff --git a/ndbench-cass-plugins/src/main/java/com/netflix/ndbench/plugin/cass/CJavaDriverBasePlugin.java b/ndbench-cass-plugins/src/main/java/com/netflix/ndbench/plugin/cass/CJavaDriverBasePlugin.java
index 3f4c0d14..c82cf588 100644
--- a/ndbench-cass-plugins/src/main/java/com/netflix/ndbench/plugin/cass/CJavaDriverBasePlugin.java
+++ b/ndbench-cass-plugins/src/main/java/com/netflix/ndbench/plugin/cass/CJavaDriverBasePlugin.java
@@ -46,6 +46,8 @@ public abstract class CJavaDriverBasePlugin<C extends CassandraConfigurationBase
     protected volatile Session session;
     protected volatile PreparedStatement readPstmt;
     protected volatile PreparedStatement writePstmt;
+    protected volatile String truststorePath;
+    protected volatile String truststorePass;
 
     /**
      * Creates an instance of the abstract CJavaDriverBasePlugin class. Subclasses calling this method should use
@@ -70,6 +72,8 @@ public void init(DataGenerator dataGenerator) throws Exception {
         this.connections = config.getConnections();
         this.username = config.getUsername();
         this.password = config.getPassword();
+        this.truststorePath = config.getTruststorePath();
+        this.truststorePass = config.getTruststorePass();
 
         // we do not set ReadConsistencyLevel and WriteConsistencyLevel and MaxColCount here because the
         // enum classes corresponding to the consistency levels differ among the concrete subclasses and because
@@ -100,7 +104,7 @@ private void initDriver() {
         logger.info("Cassandra  Cluster: " + clusterName);
 
         this.cluster = cassJavaDriverManager.registerCluster(clusterName, clusterContactPoint, connections, port,
-                                                             username, password);
+                                                             username, password, truststorePath, truststorePass);
         this.session = cassJavaDriverManager.getSession(cluster);
         if(config.getCreateSchema())
         {
diff --git a/ndbench-cass-plugins/src/main/java/com/netflix/ndbench/plugin/cass/CassJavaDriverManager.java b/ndbench-cass-plugins/src/main/java/com/netflix/ndbench/plugin/cass/CassJavaDriverManager.java
index cd2ba11b..620760eb 100644
--- a/ndbench-cass-plugins/src/main/java/com/netflix/ndbench/plugin/cass/CassJavaDriverManager.java
+++ b/ndbench-cass-plugins/src/main/java/com/netflix/ndbench/plugin/cass/CassJavaDriverManager.java
@@ -12,7 +12,7 @@
  */
 @ImplementedBy(CassJavaDriverManagerImpl.class)
 public interface CassJavaDriverManager {
-    Cluster registerCluster(String clName, String contactPoint, int connections, int port, String username, String password);
+    Cluster registerCluster(String clName, String contactPoint, int connections, int port, String username, String password, String truststorePath, String truststorePass);
     Cluster registerCluster(String clName, String contactPoint, int connections, int port);
     Session getSession(Cluster cluster);
 
diff --git a/ndbench-cass-plugins/src/main/java/com/netflix/ndbench/plugin/cass/CassJavaDriverManagerImpl.java b/ndbench-cass-plugins/src/main/java/com/netflix/ndbench/plugin/cass/CassJavaDriverManagerImpl.java
index 967496c1..c1e1f4eb 100644
--- a/ndbench-cass-plugins/src/main/java/com/netflix/ndbench/plugin/cass/CassJavaDriverManagerImpl.java
+++ b/ndbench-cass-plugins/src/main/java/com/netflix/ndbench/plugin/cass/CassJavaDriverManagerImpl.java
@@ -8,6 +8,12 @@
 import com.datastax.driver.core.PoolingOptions;
 import com.datastax.driver.core.HostDistance;
 import com.datastax.driver.core.policies.*;
+import com.datastax.driver.core.NettySSLOptions;
+import io.netty.handler.ssl.SslContextBuilder;
+import io.netty.handler.ssl.SslContext;
+import javax.net.ssl.TrustManagerFactory;
+import java.security.KeyStore;
+import com.datastax.driver.core.RemoteEndpointAwareJdkSSLOptions;
 
 
 /**
@@ -20,15 +26,25 @@ public class CassJavaDriverManagerImpl implements CassJavaDriverManager {
 
     @Override
     public Cluster registerCluster(String clName, String contactPoint, int connections, int port) {
-        return registerCluster(clName,contactPoint,connections,port,null,null);
+        return registerCluster(clName,contactPoint,connections,port,null,null,null,null);
     }
         @Override
-    public Cluster registerCluster(String clName, String contactPoint, int connections, int port, String username, String password) {
+    public Cluster registerCluster(String clName, String contactPoint, int connections, int port, String username, String password, String truststorePath, String truststorePass) {
     
         PoolingOptions poolingOpts = new PoolingOptions()
                                      .setConnectionsPerHost(HostDistance.LOCAL, connections, connections)
                                      .setMaxRequestsPerConnection(HostDistance.LOCAL, 32768);
 
+        KeyStore ks = KeyStore.getInstance("JKS");
+        InputStream trustStore = new java.io.FileInputStream(truststorePath);
+        ks.load(trustStore, truststorePass.toCharArray());
+        TrustManagerFactory tmf = TrustManagerFactory.getInstance(TrustManagerFactory.getDefaultAlgorithm());
+        tmf.init(ks);
+
+        SslContext sslContext = javax.net.ssl.SSLContext.getInstance("TLS");
+        sslContext.init(null,tmf.getTrustManagers,null);
+
+        SSLOptions sslOptions = RemoteEndpointAwareJdkSSLOptions.builder().withSSLContext(sslContext).build();
 
         Cluster.Builder clusterBuilder = Cluster.builder()
                 .withClusterName(clName)
@@ -40,6 +56,10 @@ public Cluster registerCluster(String clName, String contactPoint, int connectio
             clusterBuilder = clusterBuilder.withCredentials(username, password);
         }
 
+        if ((truststorePath !=null)) {
+            clusterBuilder = clusterBuilder.withSSL(sslOptions);
+        }
+
         cluster = clusterBuilder.build();
         return cluster;
     }
diff --git a/ndbench-cass-plugins/src/main/java/com/netflix/ndbench/plugin/configs/CassandraConfigurationBase.java b/ndbench-cass-plugins/src/main/java/com/netflix/ndbench/plugin/configs/CassandraConfigurationBase.java
index 909ba465..7b6d3511 100644
--- a/ndbench-cass-plugins/src/main/java/com/netflix/ndbench/plugin/configs/CassandraConfigurationBase.java
+++ b/ndbench-cass-plugins/src/main/java/com/netflix/ndbench/plugin/configs/CassandraConfigurationBase.java
@@ -75,4 +75,10 @@ public interface CassandraConfigurationBase {
 
     @DefaultValue("true")
     Boolean getCreateSchema();
+
+    @PropertyName(name = "truststorePath")
+    String getTruststorePath();
+
+    @PropertyName(name = "truststorePass")
+    String getTruststorePass();
 }

From 7e965475fe41e6f47572f00b65f9813386260d4a Mon Sep 17 00:00:00 2001
From: Vaibhav Khurana <vaibhav.khurana@razorpay.com>
Date: Tue, 28 Apr 2020 00:26:27 +0530
Subject: [PATCH 2/2] Adding try catch block

---
 .../cass/CassJavaDriverManagerImpl.java       | 33 +++++++++++--------
 1 file changed, 19 insertions(+), 14 deletions(-)

diff --git a/ndbench-cass-plugins/src/main/java/com/netflix/ndbench/plugin/cass/CassJavaDriverManagerImpl.java b/ndbench-cass-plugins/src/main/java/com/netflix/ndbench/plugin/cass/CassJavaDriverManagerImpl.java
index c1e1f4eb..1de172f4 100644
--- a/ndbench-cass-plugins/src/main/java/com/netflix/ndbench/plugin/cass/CassJavaDriverManagerImpl.java
+++ b/ndbench-cass-plugins/src/main/java/com/netflix/ndbench/plugin/cass/CassJavaDriverManagerImpl.java
@@ -3,17 +3,16 @@
  */
 package com.netflix.ndbench.plugin.cass;
 
-import com.datastax.driver.core.Cluster;
-import com.datastax.driver.core.Session;
-import com.datastax.driver.core.PoolingOptions;
-import com.datastax.driver.core.HostDistance;
+import com.datastax.driver.core.*;
 import com.datastax.driver.core.policies.*;
-import com.datastax.driver.core.NettySSLOptions;
 import io.netty.handler.ssl.SslContextBuilder;
 import io.netty.handler.ssl.SslContext;
+
+import javax.net.ssl.SSLContext;
 import javax.net.ssl.TrustManagerFactory;
+import java.io.InputStream;
 import java.security.KeyStore;
-import com.datastax.driver.core.RemoteEndpointAwareJdkSSLOptions;
+import java.security.KeyStoreException;
 
 
 /**
@@ -34,16 +33,22 @@ public Cluster registerCluster(String clName, String contactPoint, int connectio
         PoolingOptions poolingOpts = new PoolingOptions()
                                      .setConnectionsPerHost(HostDistance.LOCAL, connections, connections)
                                      .setMaxRequestsPerConnection(HostDistance.LOCAL, 32768);
+        KeyStore ks = null;
+        SSLContext sslContext = null ;
 
-        KeyStore ks = KeyStore.getInstance("JKS");
-        InputStream trustStore = new java.io.FileInputStream(truststorePath);
-        ks.load(trustStore, truststorePass.toCharArray());
-        TrustManagerFactory tmf = TrustManagerFactory.getInstance(TrustManagerFactory.getDefaultAlgorithm());
-        tmf.init(ks);
-
-        SslContext sslContext = javax.net.ssl.SSLContext.getInstance("TLS");
-        sslContext.init(null,tmf.getTrustManagers,null);
+        try {
+            ks = KeyStore.getInstance("JKS");
+            InputStream trustStore = new java.io.FileInputStream(truststorePath);
+            ks.load(trustStore, truststorePass.toCharArray());
+            TrustManagerFactory tmf = TrustManagerFactory.getInstance(TrustManagerFactory.getDefaultAlgorithm());
+            tmf.init(ks);
 
+            sslContext = SSLContext.getInstance("TLS");
+            sslContext.init(null,tmf.getTrustManagers(),null);
+        } catch (Exception e) {
+            e.printStackTrace();
+        }
+        
         SSLOptions sslOptions = RemoteEndpointAwareJdkSSLOptions.builder().withSSLContext(sslContext).build();
 
         Cluster.Builder clusterBuilder = Cluster.builder()