Add TLS support for Druid

Co-authored-by: Ananthu-Nair <[email protected]>

Author: infvg

presto-druid/src/main/java/com/facebook/presto/druid/DruidConfig.java

@@ -15,6 +15,7 @@
 
 import com.facebook.airlift.configuration.Config;
 import com.facebook.airlift.configuration.ConfigDescription;
+import com.facebook.airlift.configuration.ConfigSecuritySensitive;
 import com.google.common.base.Splitter;
 import com.google.common.base.StandardSystemProperty;
 import com.google.common.collect.ImmutableList;
@@ -45,6 +46,10 @@ public class DruidConfig
     private boolean caseInsensitiveNameMatching;
     private Duration caseInsensitiveNameMatchingCacheTtl = new Duration(1, MINUTES);
 
+    private boolean tlsEnabled;
+    private String trustStorePath;
+    private String truststorePassword;
+
     public enum DruidAuthenticationType
     {
         NONE,
@@ -226,4 +231,41 @@ public DruidConfig setCaseInsensitiveNameMatchingCacheTtl(Duration caseInsensiti
         this.caseInsensitiveNameMatchingCacheTtl = caseInsensitiveNameMatchingCacheTtl;
         return this;
     }
+
+    public boolean isTlsEnabled()
+    {
+        return tlsEnabled;
+    }
+
+    @Config("druid.tls.enabled")
+    public DruidConfig setTlsEnabled(boolean tlsEnabled)
+    {
+        this.tlsEnabled = tlsEnabled;
+        return this;
+    }
+
+    public String getTrustStorePath()
+    {
+        return trustStorePath;
+    }
+
+    @Config("druid.tls.truststore-path")
+    public DruidConfig setTrustStorePath(String path)
+    {
+        this.trustStorePath = path;
+        return this;
+    }
+
+    public String getTrustStorePassword()
+    {
+        return truststorePassword;
+    }
+
+    @Config("druid.tls.truststore-password")
+    @ConfigSecuritySensitive
+    public DruidConfig setTrustStorePassword(String password)
+    {
+        this.truststorePassword = password;
+        return this;
+    }
 }

presto-druid/src/main/java/com/facebook/presto/druid/DruidConnectorFactory.java

@@ -18,6 +18,7 @@
 import com.facebook.presto.common.type.TypeManager;
 import com.facebook.presto.druid.authentication.DruidAuthenticationModule;
 import com.facebook.presto.spi.ConnectorHandleResolver;
+import com.facebook.presto.spi.classloader.ThreadContextClassLoader;
 import com.facebook.presto.spi.connector.Connector;
 import com.facebook.presto.spi.connector.ConnectorContext;
 import com.facebook.presto.spi.connector.ConnectorFactory;
@@ -35,6 +36,13 @@
 public class DruidConnectorFactory
         implements ConnectorFactory
 {
+    private final ClassLoader classLoader;
+
+    public DruidConnectorFactory(ClassLoader classLoader)
+    {
+        this.classLoader = classLoader;
+    }
+
     @Override
     public String getName()
     {
@@ -51,7 +59,7 @@ public ConnectorHandleResolver getHandleResolver()
     public Connector create(String catalogName, Map<String, String> config, ConnectorContext context)
     {
         requireNonNull(config, "config is null");
-        try {
+        try (ThreadContextClassLoader ignore = new ThreadContextClassLoader(classLoader)) {
             Bootstrap app = new Bootstrap(
                     new JsonModule(),
                     new DruidModule(),

presto-druid/src/main/java/com/facebook/presto/druid/DruidPlugin.java

@@ -23,6 +23,15 @@ public class DruidPlugin
     @Override
     public Iterable<ConnectorFactory> getConnectorFactories()
     {
-        return ImmutableList.of(new DruidConnectorFactory());
+        return ImmutableList.of(new DruidConnectorFactory(getClassLoader()));
+    }
+
+    private static ClassLoader getClassLoader()
+    {
+        ClassLoader classLoader = Thread.currentThread().getContextClassLoader();
+        if (classLoader == null) {
+            classLoader = DruidPlugin.class.getClassLoader();
+        }
+        return classLoader;
     }
 }

presto-druid/src/main/java/com/facebook/presto/druid/authentication/DruidAuthenticationModule.java

@@ -33,17 +33,19 @@ public class DruidAuthenticationModule
     @Override
     protected void setup(Binder binder)
     {
+        DruidConfig druidConfig = buildConfigObject(DruidConfig.class);
+
         bindAuthenticationModule(
                 config -> config.getDruidAuthenticationType() == NONE,
                 noneAuthenticationModule());
 
         bindAuthenticationModule(
                 config -> config.getDruidAuthenticationType() == BASIC,
-                basicAuthenticationModule());
+                basicAuthenticationModule(druidConfig));
 
         bindAuthenticationModule(
                 config -> config.getDruidAuthenticationType() == KERBEROS,
-                kerberosbAuthenticationModule());
+                kerberosbAuthenticationModule(druidConfig));
     }
 
     private void bindAuthenticationModule(Predicate<DruidConfig> predicate, Module module)
@@ -56,19 +58,31 @@ private static Module noneAuthenticationModule()
         return binder -> httpClientBinder(binder).bindHttpClient("druid-client", ForDruidClient.class);
     }
 
-    private static Module basicAuthenticationModule()
+    private static Module basicAuthenticationModule(DruidConfig druidConfig)
     {
         return binder -> httpClientBinder(binder).bindHttpClient("druid-client", ForDruidClient.class)
                 .withConfigDefaults(
-                        config -> config.setAuthenticationEnabled(false) //disable Kerberos auth
+                        config -> {
+                            config.setAuthenticationEnabled(false); //disable Kerberos auth
+                            if (druidConfig.isTlsEnabled()) {
+                                config.setTrustStorePath(druidConfig.getTrustStorePath());
+                                config.setTrustStorePassword(druidConfig.getTrustStorePassword());
+                            }
+                        }
                 ).withFilter(
                         DruidBasicAuthHttpRequestFilter.class);
     }
 
-    private static Module kerberosbAuthenticationModule()
+    private static Module kerberosbAuthenticationModule(DruidConfig druidConfig)
     {
         return binder -> httpClientBinder(binder).bindHttpClient("druid-client", ForDruidClient.class)
                 .withConfigDefaults(
-                        config -> config.setAuthenticationEnabled(true));
+                        config -> {
+                            config.setAuthenticationEnabled(true);
+                            if (druidConfig.isTlsEnabled()) {
+                                config.setTrustStorePath(druidConfig.getTrustStorePath());
+                                config.setTrustStorePassword(druidConfig.getTrustStorePassword());
+                            }
+                        });
     }
 }

presto-druid/src/test/java/com/facebook/presto/druid/TestDruidConfig.java

@@ -45,15 +45,18 @@ public void testDefaults()
                 .setBasicAuthenticationPassword(null)
                 .setIngestionStoragePath(StandardSystemProperty.JAVA_IO_TMPDIR.value())
                 .setCaseInsensitiveNameMatching(false)
-                .setCaseInsensitiveNameMatchingCacheTtl(new Duration(1, MINUTES)));
+                .setCaseInsensitiveNameMatchingCacheTtl(new Duration(1, MINUTES))
+                .setTlsEnabled(false)
+                .setTrustStorePath(null)
+                .setTrustStorePassword(null));
     }
 
     @Test
     public void testExplicitPropertyMappings()
     {
         Map<String, String> properties = new ImmutableMap.Builder<String, String>()
-                .put("druid.broker-url", "http://druid.broker:1234")
-                .put("druid.coordinator-url", "http://druid.coordinator:4321")
+                .put("druid.broker-url", "https://druid.broker:1234")
+                .put("druid.coordinator-url", "https://druid.coordinator:4321")
                 .put("druid.schema-name", "test")
                 .put("druid.compute-pushdown-enabled", "true")
                 .put("druid.hadoop.config.resources", "/etc/core-site.xml,/etc/hdfs-site.xml")
@@ -63,6 +66,9 @@ public void testExplicitPropertyMappings()
                 .put("druid.ingestion.storage.path", "hdfs://foo/bar/")
                 .put("druid.case-insensitive-name-matching", "true")
                 .put("druid.case-insensitive-name-matching.cache-ttl", "1s")
+                .put("druid.tls.enabled", "true")
+                .put("druid.tls.truststore-path", "/tmp/truststore")
+                .put("druid.tls.truststore-password", "truststore-password")
                 .build();
 
         DruidConfig expected = new DruidConfig()
@@ -76,7 +82,10 @@ public void testExplicitPropertyMappings()
                 .setBasicAuthenticationPassword("http_basic_password")
                 .setIngestionStoragePath("hdfs://foo/bar/")
                 .setCaseInsensitiveNameMatching(true)
-                .setCaseInsensitiveNameMatchingCacheTtl(new Duration(1, SECONDS));
+                .setCaseInsensitiveNameMatchingCacheTtl(new Duration(1, SECONDS))
+                .setTlsEnabled(true)
+                .setTrustStorePath(("/tmp/truststore"))
+                .setTrustStorePassword("truststore-password");
 
         ConfigAssertions.assertFullMapping(properties, expected);
     }