CASSANDRA-17559: FQLTool Replay should support full set of login options

tools/fqltool/src/org/apache/cassandra/fqltool/ConnectionOptions.java

@@ -0,0 +1,203 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you 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.
+ */
+
+package org.apache.cassandra.fqltool;
+
+import javax.net.ssl.SSLContext;
+
+import com.datastax.driver.core.AuthProvider;
+import com.datastax.driver.core.RemoteEndpointAwareJdkSSLOptions;
+import com.datastax.driver.core.SSLOptions;
+
+import org.apache.cassandra.config.EncryptionOptions;
+import org.apache.cassandra.security.SSLFactory;
+
+/**
+ * Holds the SSL and authentication settings used to connect to target hosts during fqltool replay.
+ * Note that providing any SSL-related configuration option implicitly enables SSL.
+ */
+public class ConnectionOptions
+{
+    private final boolean ssl;
+    private final SSLOptions sslOptions;
+    private final String authProviderClass;
+
+    private ConnectionOptions(boolean ssl, SSLOptions sslOptions, String authProviderClass)
+    {
+        this.ssl = ssl;
+        this.sslOptions = sslOptions;
+        this.authProviderClass = authProviderClass;
+    }
+
+    public boolean ssl()
+    {
+        return ssl;
+    }
+
+    public SSLOptions sslOptions()
+    {
+        return sslOptions;
+    }
+
+    public String authProviderClass()
+    {
+        return authProviderClass;
+    }
+
+    /** 
+     * Builds the configured AuthProvider: a (String,String) constructor when credentials are present, otherwise a no-arg constructor.
+     */
+    @SuppressWarnings("unchecked")
+    public AuthProvider instantiateAuthProvider(String user, String password)
+    {
+        try
+        {
+            Class<? extends AuthProvider> clazz = (Class<? extends AuthProvider>) Class.forName(authProviderClass);
+
+            if (user != null && password != null)
+                return clazz.getConstructor(String.class, String.class).newInstance(user, password);
+
+            return clazz.getDeclaredConstructor().newInstance();
+        }
+        catch (NoSuchMethodException e)
+        {
+            throw new RuntimeException("Auth provider " + authProviderClass + " does not support plain text credentials", e);
+        }
+        catch (Exception e)
+        {
+            throw new RuntimeException("Could not instantiate auth provider: " + authProviderClass, e);
+        }
+    }
+
+    public static Builder builder()
+    {
+        return new Builder();
+    }
+
+    public static class Builder
+    {
+        private boolean ssl;
+        private String truststorePath;
+        private String truststorePassword;
+        private String keystorePath;
+        private String keystorePassword;
+        private String authProviderClass;
+
+        public Builder withSsl(boolean ssl)
+        {
+            this.ssl = ssl;
+            return this;
+        }
+
+        public Builder withTruststore(String truststorePath)
+        {
+            this.truststorePath = truststorePath;
+            return this;
+        }
+
+        public Builder withTruststorePassword(String truststorePassword)
+        {
+            this.truststorePassword = truststorePassword;
+            return this;
+        }
+
+        public Builder withKeystore(String keystorePath)
+        {
+            this.keystorePath = keystorePath;
+            return this;
+        }
+
+        public Builder withKeystorePassword(String keystorePassword)
+        {
+            this.keystorePassword = keystorePassword;
+            return this;
+        }
+
+        public Builder withAuthProviderClass(String authProviderClass)
+        {
+            this.authProviderClass = authProviderClass;
+            return this;
+        }
+
+        public ConnectionOptions build()
+        {
+            if (truststorePassword != null && truststorePath == null)
+                throw new IllegalArgumentException("--ssl-truststore-password requires --ssl-truststore to be set");
+            if (keystorePassword != null && keystorePath == null)
+                throw new IllegalArgumentException("--ssl-keystore-password requires --ssl-keystore to be set");
+
+            // any SSL-related option implicitly enables SSL
+            boolean effectiveSsl = ssl || truststorePath != null || keystorePath != null;
+
+            if (authProviderClass != null)
+                validateAuthProviderClass();
+
+            SSLOptions sslOptions = effectiveSsl ? buildSSLOptions() : null;
+
+            return new ConnectionOptions(effectiveSsl, sslOptions, authProviderClass);
+        }
+
+        private void validateAuthProviderClass()
+        {
+            try
+            {
+                Class<?> clazz = Class.forName(authProviderClass);
+                if (!AuthProvider.class.isAssignableFrom(clazz))
+                    throw new IllegalArgumentException(authProviderClass + " does not implement " + AuthProvider.class.getName());
+            }
+            catch (ClassNotFoundException e)
+            {
+                throw new RuntimeException("Could not find auth provider class: " + authProviderClass, e);
+            }
+        }
+
+        private SSLOptions buildSSLOptions()
+        {
+            try
+            {
+                EncryptionOptions.ClientEncryptionOptions.Builder encBuilder = new EncryptionOptions.ClientEncryptionOptions.Builder();
+                encBuilder.withEnabled(true);
+
+                if (truststorePath != null)
+                    encBuilder.withTrustStore(truststorePath);
+                if (truststorePassword != null)
+                    encBuilder.withTrustStorePassword(truststorePassword);
+
+                EncryptionOptions.ClientEncryptionOptions.ClientAuth clientAuth = EncryptionOptions.ClientEncryptionOptions.ClientAuth.NOT_REQUIRED;
+                if (keystorePath != null)
+                {
+                    encBuilder.withKeyStore(keystorePath);
+                    clientAuth = EncryptionOptions.ClientEncryptionOptions.ClientAuth.REQUIRED;
+                }
+                if (keystorePassword != null)
+                    encBuilder.withKeyStorePassword(keystorePassword);
+
+                EncryptionOptions.ClientEncryptionOptions clientEncOptions = encBuilder.build();
+                SSLContext sslContext = SSLFactory.createSSLContext(clientEncOptions, clientAuth);
+
+                return RemoteEndpointAwareJdkSSLOptions.builder()
+                                                        .withSSLContext(sslContext)
+                                                        .build();
+            }
+            catch (Exception e)
+            {
+                throw new RuntimeException("Could not configure SSL for fqltool replay", e);
+            }
+        }
+    }
+}

tools/fqltool/src/org/apache/cassandra/fqltool/QueryReplayer.java

@@ -76,6 +76,20 @@ public QueryReplayer(Iterator<List<FQLQuery>> queryIterator,
         this(queryIterator, targetHosts, resultPaths, filters, queryFilePathString, new DefaultSessionProvider(), null);
     }
 
+    /**
+     * Constructor that takes SSL and auth provider settings via ConnectionOptions.
+     */
+    public QueryReplayer(Iterator<List<FQLQuery>> queryIterator,
+                         List<String> targetHosts,
+                         List<File> resultPaths,
+                         List<Predicate<FQLQuery>> filters,
+                         String queryFilePathString,
+                         ConnectionOptions connectionOptions)
+    {
+        this(queryIterator, targetHosts, resultPaths, filters, queryFilePathString,
+             new DefaultSessionProvider(connectionOptions), null);
+    }
+
     /**
      * Constructor public to allow external users to build their own session provider
      *
@@ -186,7 +200,7 @@ public void close() throws IOException
         resultHandler.close();
     }
 
-    static class ParsedTargetHost
+    public static class ParsedTargetHost
     {
         final int port;
         final String user;
@@ -201,26 +215,42 @@ static class ParsedTargetHost
             this.password = password;
         }
 
-        static ParsedTargetHost fromString(String s)
+        /**
+         * Masks the password in a target host string so it's never logged or written to result paths.
+         */
+        public static String maskPassword(String target)
         {
-            String [] userInfoHostPort = s.split("@");
+            int at = target.lastIndexOf('@');
+            if (at < 0)
+                return target;
+            String userInfo = target.substring(0, at);
+            int colon = userInfo.indexOf(':');
+            if (colon < 0)
+                return target;
+            return userInfo.substring(0, colon) + ":*****@" + target.substring(at + 1);
+        }
 
-            String hostPort = null;
+        static ParsedTargetHost fromString(String s)
+        {
+            int at = s.lastIndexOf('@');
+            String hostPort;
             String user = null;
             String password = null;
-            if (userInfoHostPort.length == 2)
+
+            if (at >= 0)
             {
-                String [] userPassword = userInfoHostPort[0].split(":");
-                if (userPassword.length != 2)
+                String userInfo = s.substring(0, at);
+                hostPort = s.substring(at + 1);
+                int colon = userInfo.indexOf(':');
+                if (colon < 0)
                     throw new RuntimeException("Username provided but no password");
-                hostPort = userInfoHostPort[1];
-                user = userPassword[0];
-                password = userPassword[1];
+                user = userInfo.substring(0, colon);
+                password = userInfo.substring(colon + 1);
             }
-            else if (userInfoHostPort.length == 1)
-                hostPort = userInfoHostPort[0];
             else
-                throw new RuntimeException("Malformed target host: "+s);
+            {
+                hostPort = s;
+            }
 
             String[] splitHostPort = hostPort.split(":");
             int port = 9042;
@@ -241,6 +271,18 @@ private static final class DefaultSessionProvider implements SessionProvider
     {
         private final static Map<String, Session> sessionCache = new HashMap<>();
 
+        private final ConnectionOptions connectionOptions;
+
+        DefaultSessionProvider()
+        {
+            this(ConnectionOptions.builder().build());
+        }
+
+        DefaultSessionProvider(ConnectionOptions connectionOptions)
+        {
+            this.connectionOptions = connectionOptions;
+        }
+
         public synchronized Session connect(String connectionString)
         {
             if (sessionCache.containsKey(connectionString))
@@ -249,8 +291,15 @@ public synchronized Session connect(String connectionString)
             ParsedTargetHost pth = ParsedTargetHost.fromString(connectionString);
             builder.addContactPoint(pth.host);
             builder.withPort(pth.port);
-            if (pth.user != null)
+
+            if (connectionOptions.sslOptions() != null)
+                builder.withSSL(connectionOptions.sslOptions());
+
+            if (connectionOptions.authProviderClass() != null)
+                builder.withAuthProvider(connectionOptions.instantiateAuthProvider(pth.user, pth.password));
+            else if (pth.user != null)
                 builder.withCredentials(pth.user, pth.password);
+
             Cluster c = builder.build();
             sessionCache.put(connectionString, c.connect());
             return sessionCache.get(connectionString);

tools/fqltool/src/org/apache/cassandra/fqltool/ResultHandler.java

@@ -47,7 +47,7 @@ public ResultHandler(List<String> targetHosts, List<File> resultPaths, File quer
 
     public ResultHandler(List<String> targetHosts, List<File> resultPaths, File queryFilePath, MismatchListener mismatchListener)
     {
-        this.targetHosts = targetHosts;
+        this.targetHosts = targetHosts.stream().map(QueryReplayer.ParsedTargetHost::maskPassword).collect(Collectors.toList());
         resultStore = resultPaths != null ? new ResultStore(resultPaths, queryFilePath) : null;
         resultComparator = new ResultComparator(mismatchListener);
     }