Merge pull request #37 from rabbitmq/rabbitmq-perf-test-35

Support TLS options on the command line

Author: michaelklishin

src/main/java/com/rabbitmq/perf/PerfTest.java

@@ -15,11 +15,12 @@
 
 package com.rabbitmq.perf;
 
+import java.security.NoSuchAlgorithmException;
 import java.text.SimpleDateFormat;
-import java.util.Arrays;
 import java.util.Calendar;
 import java.util.Collections;
 import java.util.List;
+import java.util.Properties;
 
 import org.apache.commons.cli.CommandLine;
 import org.apache.commons.cli.CommandLineParser;
@@ -30,11 +31,17 @@
 import org.apache.commons.cli.ParseException;
 
 import com.rabbitmq.client.ConnectionFactory;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import javax.net.ssl.SSLContext;
 
 import static java.util.Arrays.asList;
 
 public class PerfTest {
 
+    private static final Logger LOGGER = LoggerFactory.getLogger(PerfTest.class);
+
     public static void main(String[] args) {
         Options options = getOptions();
         CommandLineParser parser = new GnuParser();
@@ -100,7 +107,12 @@ public static void main(String[] args) {
                     flags.contains("immediate")),
                 confirm != -1);
 
+            SSLContext sslContext = getSslContextIfNecessary(cmd, System.getProperties());
+
             ConnectionFactory factory = new ConnectionFactory();
+            if (sslContext != null) {
+                factory.useSslProtocol(sslContext);
+            }
             factory.setShutdownTimeout(0); // So we still shut down even with slow consumers
             factory.setUri(uris.get(0));
             factory.setRequestedFrameMax(frameMax);
@@ -150,6 +162,26 @@ public static void main(String[] args) {
         }
     }
 
+    private static SSLContext getSslContextIfNecessary(CommandLine cmd, Properties systemProperties) throws NoSuchAlgorithmException {
+        SSLContext sslContext = null;
+        if (cmd.hasOption("useDefaultSslContext")) {
+            LOGGER.info("Using default SSL context as per command line option");
+            sslContext = SSLContext.getDefault();
+        }
+        for (String propertyName : systemProperties.stringPropertyNames()) {
+            if (propertyName != null && isPropertyTlsRelated(propertyName)) {
+                LOGGER.info("TLS related system properties detected, using default SSL context");
+                sslContext = SSLContext.getDefault();
+                break;
+            }
+        }
+        return sslContext;
+    }
+
+    private static boolean isPropertyTlsRelated(String propertyName) {
+        return propertyName.startsWith("javax.net.ssl") || propertyName.startsWith("jdk.tls");
+    }
+
     private static void usage(Options options) {
         HelpFormatter formatter = new HelpFormatter();
         formatter.printHelp("<program>", options);
@@ -192,7 +224,7 @@ private static Options getOptions() {
         options.addOption(new Option("p", "predeclared",            false,"allow use of predeclared objects"));
         options.addOption(new Option("B", "body",                   true, "comma-separated list of files to use in message bodies"));
         options.addOption(new Option("T", "bodyContenType",         true, "body content-type"));
-
+        options.addOption(new Option("useDefaultSslContext", "useDefaultSslContext",       false,"use JVM default SSL context"));
         return options;
     }