support tls in performance tool (apache#662)

willome · Aug 9, 2017 · de24aaf · de24aaf
1 parent 005aa12
commit de24aaf
Show file tree

Hide file tree

Showing 3 changed files with 59 additions and 2 deletions.
diff --git a/pulsar-testclient/src/main/java/org/apache/pulsar/testclient/PerformanceConsumer.java b/pulsar-testclient/src/main/java/org/apache/pulsar/testclient/PerformanceConsumer.java
@@ -18,6 +18,7 @@
  */
 package org.apache.pulsar.testclient;
 
+import static org.apache.commons.lang3.StringUtils.isBlank;
 import static org.apache.commons.lang3.StringUtils.isNotBlank;
 
 import java.io.FileInputStream;
@@ -100,6 +101,14 @@ static class Arguments {
         @Parameter(names = {
                 "--auth_params" }, description = "Authentication parameters, e.g., \"key1:val1,key2:val2\"")
         public String authParams;
+
+        @Parameter(names = {
+                "--use-tls" }, description = "Use TLS encryption on the connection")
+        public boolean useTls;
+
+        @Parameter(names = {
+                "--trust-cert-file" }, description = "Path for the trusted TLS certificate file")
+        public String tlsTrustCertsFilePath = "";
     }
 
     public static void main(String[] args) throws Exception {
@@ -150,6 +159,14 @@ public static void main(String[] args) throws Exception {
             if (arguments.authParams == null) {
                 arguments.authParams = prop.getProperty("authParams", null);
             }
+
+            if (arguments.useTls == false) {
+                arguments.useTls = Boolean.parseBoolean(prop.getProperty("useTls"));
+            }
+
+            if (isBlank(arguments.tlsTrustCertsFilePath)) {
+                arguments.tlsTrustCertsFilePath = prop.getProperty("tlsTrustCertsFilePath", "");
+            }
         }
 
         // Dump config variables
@@ -181,6 +198,8 @@ public void received(Consumer consumer, Message msg) {
         if (isNotBlank(arguments.authPluginClassName)) {
             clientConf.setAuthentication(arguments.authPluginClassName, arguments.authParams);
         }
+        clientConf.setUseTls(arguments.useTls);
+        clientConf.setTlsTrustCertsFilePath(arguments.tlsTrustCertsFilePath);
         PulsarClient pulsarClient = new PulsarClientImpl(arguments.serviceURL, clientConf);
 
         List<Future<Consumer>> futures = Lists.newArrayList();
@@ -235,4 +254,4 @@ public void received(Consumer consumer, Message msg) {
     }
 
     private static final Logger log = LoggerFactory.getLogger(PerformanceConsumer.class);
-}
+}
diff --git a/pulsar-testclient/src/main/java/org/apache/pulsar/testclient/PerformanceProducer.java b/pulsar-testclient/src/main/java/org/apache/pulsar/testclient/PerformanceProducer.java
@@ -19,6 +19,7 @@
 package org.apache.pulsar.testclient;
 
 import static java.util.concurrent.TimeUnit.NANOSECONDS;
+import static org.apache.commons.lang3.StringUtils.isBlank;
 import static org.apache.commons.lang3.StringUtils.isNotBlank;
 
 import java.io.FileInputStream;
@@ -132,6 +133,14 @@ static class Arguments {
         @Parameter(names = { "-time",
                 "--test-duration" }, description = "Test duration in secs. If 0, it will keep publishing")
         public long testTime = 0;
+
+        @Parameter(names = {
+                "--use-tls" }, description = "Use TLS encryption on the connection")
+        public boolean useTls;
+
+        @Parameter(names = {
+                "--trust-cert-file" }, description = "Path for the trusted TLS certificate file")
+        public String tlsTrustCertsFilePath = "";
     }
 
     public static void main(String[] args) throws Exception {
@@ -182,6 +191,14 @@ public static void main(String[] args) throws Exception {
             if (arguments.authParams == null) {
                 arguments.authParams = prop.getProperty("authParams", null);
             }
+
+            if (arguments.useTls == false) {
+               arguments.useTls = Boolean.parseBoolean(prop.getProperty("useTls"));
+            }
+
+            if (isBlank(arguments.tlsTrustCertsFilePath)) {
+               arguments.tlsTrustCertsFilePath = prop.getProperty("tlsTrustCertsFilePath", "");
+            }
         }
 
         arguments.testTime = TimeUnit.SECONDS.toMillis(arguments.testTime);
@@ -210,6 +227,8 @@ public static void main(String[] args) throws Exception {
         if (isNotBlank(arguments.authPluginClassName)) {
             clientConf.setAuthentication(arguments.authPluginClassName, arguments.authParams);
         }
+        clientConf.setUseTls(arguments.useTls);
+        clientConf.setTlsTrustCertsFilePath(arguments.tlsTrustCertsFilePath);
 
         PulsarClient client = new PulsarClientImpl(arguments.serviceURL, clientConf);
 

diff --git a/pulsar-testclient/src/main/java/org/apache/pulsar/testclient/PerformanceReader.java b/pulsar-testclient/src/main/java/org/apache/pulsar/testclient/PerformanceReader.java
@@ -18,6 +18,7 @@
  */
 package org.apache.pulsar.testclient;
 
+import static org.apache.commons.lang3.StringUtils.isBlank;
 import static org.apache.commons.lang3.StringUtils.isNotBlank;
 
 import java.io.FileInputStream;
@@ -95,6 +96,14 @@ static class Arguments {
         @Parameter(names = {
                 "--auth-params" }, description = "Authentication parameters, e.g., \"key1:val1,key2:val2\"")
         public String authParams;
+
+        @Parameter(names = {
+                "--use-tls" }, description = "Use TLS encryption on the connection")
+        public boolean useTls;
+
+        @Parameter(names = {
+                "--trust-cert-file" }, description = "Path for the trusted TLS certificate file")
+        public String tlsTrustCertsFilePath = "";
     }
 
     public static void main(String[] args) throws Exception {
@@ -145,6 +154,14 @@ public static void main(String[] args) throws Exception {
             if (arguments.authParams == null) {
                 arguments.authParams = prop.getProperty("authParams", null);
             }
+
+            if (arguments.useTls == false) {
+                arguments.useTls = Boolean.parseBoolean(prop.getProperty("useTls"));
+            }
+
+            if (isBlank(arguments.tlsTrustCertsFilePath)) {
+                arguments.tlsTrustCertsFilePath = prop.getProperty("tlsTrustCertsFilePath", "");
+            }
         }
 
         // Dump config variables
@@ -172,6 +189,8 @@ public static void main(String[] args) throws Exception {
         if (isNotBlank(arguments.authPluginClassName)) {
             clientConf.setAuthentication(arguments.authPluginClassName, arguments.authParams);
         }
+        clientConf.setUseTls(arguments.useTls);
+        clientConf.setTlsTrustCertsFilePath(arguments.tlsTrustCertsFilePath);
         PulsarClient pulsarClient = new PulsarClientImpl(arguments.serviceURL, clientConf);
 
         List<CompletableFuture<Reader>> futures = Lists.newArrayList();
@@ -222,4 +241,4 @@ public static void main(String[] args) throws Exception {
     }
 
     private static final Logger log = LoggerFactory.getLogger(PerformanceReader.class);
-}
+}