TEZ-4749: Support SSL/TLS connections to ZooKeeper in ZkAMRegistry and ZkAMRegistryClient

tez-api/src/main/java/org/apache/tez/client/registry/zookeeper/ZkConfig.java

@@ -18,15 +18,19 @@
  */
 package org.apache.tez.client.registry.zookeeper;
 
+import java.util.Optional;
 import java.util.concurrent.TimeUnit;
 
+import org.apache.commons.lang3.StringUtils;
 import org.apache.curator.RetryPolicy;
 import org.apache.curator.framework.CuratorFramework;
 import org.apache.curator.framework.CuratorFrameworkFactory;
 import org.apache.curator.retry.ExponentialBackoffRetry;
 import org.apache.hadoop.conf.Configuration;
 import org.apache.hadoop.fs.Path;
 import org.apache.tez.dag.api.TezConfiguration;
+import org.apache.zookeeper.client.ZKClientConfig;
+import org.apache.zookeeper.common.ClientX509Util;
 
 import com.google.common.base.Preconditions;
 import com.google.common.base.Strings;
@@ -49,6 +53,11 @@ public class ZkConfig {
   private final int curatorMaxRetries;
   private final int sessionTimeoutMs;
   private final int connectionTimeoutMs;
+  private final String sslEnabled;
+  private final String sslKeystoreLocation;
+  private final String sslKeystorePassword;
+  private final String sslTruststoreLocation;
+  private final String sslTruststorePassword;
 
   public ZkConfig(Configuration conf) {
     zkQuorum = conf.get(TezConfiguration.TEZ_AM_ZOOKEEPER_QUORUM);
@@ -84,6 +93,16 @@ public ZkConfig(Configuration conf) {
         TezConfiguration.TEZ_AM_CURATOR_SESSION_TIMEOUT_DEFAULT, TimeUnit.MILLISECONDS));
     connectionTimeoutMs = Math.toIntExact(conf.getTimeDuration(TezConfiguration.TEZ_AM_CURATOR_CONNECTION_TIMEOUT,
         TezConfiguration.TEZ_AM_CURATOR_CONNECTION_TIMEOUT_DEFAULT, TimeUnit.MILLISECONDS));
+    sslEnabled = conf.get(TezConfiguration.TEZ_AM_ZOOKEEPER_SSL_ENABLE);
+    Preconditions.checkArgument(
+        isValidSslEnabledValue(sslEnabled),
+        "If the optional %s setting is set, then the value should be a boolean value instead of '%s'",
+        TezConfiguration.TEZ_AM_ZOOKEEPER_SSL_ENABLE,
+        sslEnabled);
+    sslKeystoreLocation = conf.get(TezConfiguration.TEZ_AM_ZOOKEEPER_SSL_KEYSTORE_LOCATION);
+    sslKeystorePassword = conf.get(TezConfiguration.TEZ_AM_ZOOKEEPER_SSL_KEYSTORE_PASSWORD);
+    sslTruststoreLocation = conf.get(TezConfiguration.TEZ_AM_ZOOKEEPER_SSL_TRUSTSTORE_LOCATION);
+    sslTruststorePassword = conf.get(TezConfiguration.TEZ_AM_ZOOKEEPER_SSL_TRUSTSTORE_PASSWORD);
   }
 
   public String getZkQuorum() {
@@ -110,17 +129,91 @@ public int getConnectionTimeoutMs() {
     return connectionTimeoutMs;
   }
 
+  public String getZookeeperTrustStorePassword() {
+    return sslTruststorePassword;
+  }
+
+  public String getZookeeperTrustStoreLocation() {
+    return sslTruststoreLocation;
+  }
+
+  public String getZookeeperKeyStorePassword() {
+    return sslKeystorePassword;
+  }
+
+  public String getZookeeperKeyStoreLocation() {
+    return sslKeystoreLocation;
+  }
+
+  /**
+   * Returns whether the zookeeper connection will be secure or insecure.
+   * @return An Optional containing the boolean value that indicates whether zookeeper client
+   * uses a secure zookeeper connection. An empty Optional indicates that it is not specified,
+   * and in this case the default settings of zookeeper are used, which can be controlled by
+   * specific JVM properties.
+   * @see TezConfiguration#TEZ_AM_ZOOKEEPER_SSL_ENABLE
+   */
+  public Optional<Boolean> isSslEnabled() {
+    if (this.sslEnabled == null || this.sslEnabled.isEmpty()) {
+      return Optional.empty();
+    }
+    return Optional.of(Boolean.parseBoolean(sslEnabled));
+  }
+
   public RetryPolicy getRetryPolicy() {
     return new ExponentialBackoffRetry(getCuratorBackoffSleepMs(), getCuratorMaxRetries());
   }
 
   public CuratorFramework createCuratorFramework() {
-    return CuratorFrameworkFactory.newClient(
-        getZkQuorum(),
-        getSessionTimeoutMs(),
-        getConnectionTimeoutMs(),
-        getRetryPolicy()
-    );
+    if (isSslEnabled().isEmpty()) {
+      return CuratorFrameworkFactory.newClient(
+              getZkQuorum(),
+              getSessionTimeoutMs(),
+              getConnectionTimeoutMs(),
+              getRetryPolicy()
+      );
+    }
+
+    ZKClientConfig zkClientConfig = new ZKClientConfig();
+    zkClientConfig.setProperty(ZKClientConfig.SECURE_CLIENT, Boolean.toString(isSslEnabled().get()));
+    zkClientConfig.setProperty(ZKClientConfig.ZOOKEEPER_CLIENT_CNXN_SOCKET,
+            "org.apache.zookeeper.ClientCnxnSocketNetty");
+    if (isSslEnabled().get()) {
+      try (ClientX509Util x509Util = new ClientX509Util()) {
+        setStoreConfig(zkClientConfig, x509Util.getSslKeystoreLocationProperty(), getZookeeperKeyStoreLocation(),
+            x509Util.getSslKeystorePasswdProperty(), getZookeeperKeyStorePassword(), "keystore");
+
+        setStoreConfig(zkClientConfig, x509Util.getSslTruststoreLocationProperty(), getZookeeperTrustStoreLocation(),
+            x509Util.getSslTruststorePasswdProperty(), getZookeeperTrustStorePassword(), "truststore");
+      }
+    }
+
+    return CuratorFrameworkFactory.builder()
+            .connectString(getZkQuorum())
+            .sessionTimeoutMs(getSessionTimeoutMs())
+            .connectionTimeoutMs(getConnectionTimeoutMs())
+            .retryPolicy(getRetryPolicy())
+            .zkClientConfig(zkClientConfig)
+            .build();
+  }
+
+  private void setStoreConfig(ZKClientConfig config, String locationProp, String locationVal, String passwordProp,
+                              String passwordVal, String storeName) {
+    if (StringUtils.isEmpty(locationVal)) {
+      LOG.info("No {} location configured, using ZooKeeper client defaults", storeName);
+      return;
+    }
+
+    config.setProperty(locationProp, locationVal);
+    if (StringUtils.isNotEmpty(passwordVal)) {
+      config.setProperty(passwordProp, passwordVal);
+    }
+  }
+
+  private boolean isValidSslEnabledValue(String value) {
+    return value == null || value.isEmpty()
+        || value.trim().equalsIgnoreCase("true")
+        || value.trim().equalsIgnoreCase("false");
   }
 
   /**

tez-api/src/main/java/org/apache/tez/dag/api/TezConfiguration.java

@@ -2268,6 +2268,63 @@ static Set<String> getPropertySet() {
   public static final String TEZ_SHARED_EXECUTOR_MAX_THREADS = "tez.shared-executor.max-threads";
   public static final int TEZ_SHARED_EXECUTOR_MAX_THREADS_DEFAULT = -1;
 
+  /**
+   * Optional boolean value represented by string type. A value of "true" enables secure
+   * Zookeeper connection in ZkAMRegistry and ZkAMRegistryClient classes, while a value
+   * of "false" disables secure Zookeeper connection.
+   * If not specified or empty string, then zookeeper enables/disables the secure Zookeeper
+   * connection based on JVM properties.
+   * Default: Empty
+   */
+  @ConfigurationScope(Scope.AM)
+  @ConfigurationProperty
+  public static final String TEZ_AM_ZOOKEEPER_SSL_ENABLE = TEZ_AM_PREFIX
+          + "zookeeper.ssl.client.enable";
+
+  /**
+   * String value
+   * An optional setting that specifies the path to the keystore used for the secure
+   * zookeeper connection.
+   * Default: Empty
+   */
+  @ConfigurationScope(Scope.AM)
+  @ConfigurationProperty
+  public static final String TEZ_AM_ZOOKEEPER_SSL_KEYSTORE_LOCATION = TEZ_AM_PREFIX
+          + "zookeeper.ssl.keystore.location";
+
+  /**
+   * String value
+   * An optional setting that specifies the password of the keystore used for the secure
+   * zookeeper connection.
+   * Default: Empty
+   */
+  @ConfigurationScope(Scope.AM)
+  @ConfigurationProperty
+  public static final String TEZ_AM_ZOOKEEPER_SSL_KEYSTORE_PASSWORD = TEZ_AM_PREFIX
+          + "zookeeper.ssl.keystore.password";
+
+  /**
+   * String value
+   * An optional setting that specifies the path to the truststore used for the secure
+   * zookeeper connection.
+   * Default: Empty
+   */
+  @ConfigurationScope(Scope.AM)
+  @ConfigurationProperty
+  public static final String TEZ_AM_ZOOKEEPER_SSL_TRUSTSTORE_LOCATION = TEZ_AM_PREFIX
+          + "zookeeper.ssl.truststore.location";
+
+  /**
+   * String value
+   * An optional setting that specifies the password of the truststore used for the secure
+   * zookeeper connection.
+   * Default: Empty
+   */
+  @ConfigurationScope(Scope.AM)
+  @ConfigurationProperty
+  public static final String TEZ_AM_ZOOKEEPER_SSL_TRUSTSTORE_PASSWORD = TEZ_AM_PREFIX
+          + "zookeeper.ssl.truststore.password";
+
   /**
    *  Acquire all FileSystems info. e.g., all namenodes info of HDFS federation cluster.
    */

tez-api/src/test/java/org/apache/tez/client/registry/zookeeper/TestZkConfig.java

@@ -19,9 +19,13 @@
 package org.apache.tez.client.registry.zookeeper;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
 import static org.junit.jupiter.api.Assertions.assertNotNull;
+import static org.junit.jupiter.api.Assertions.assertNull;
 import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
 
+import java.util.Optional;
 import java.util.concurrent.TimeUnit;
 
 import org.apache.curator.RetryPolicy;
@@ -231,4 +235,73 @@ public void testDefaultNamespace() {
     assertEquals("/tez-external-sessions" + TezConfiguration.TEZ_AM_REGISTRY_NAMESPACE_DEFAULT,
         zkConfig.getZkNamespace());
   }
+
+  @Test
+  public void testZkConfigTezAmZookeeperSslEnableNotSpecified() {
+    Configuration conf = new Configuration();
+    conf.set(TezConfiguration.TEZ_AM_ZOOKEEPER_QUORUM, "dummyZkQuorum");
+    ZkConfig zkConf = new ZkConfig(conf);
+
+    assertEquals(Optional.empty(), zkConf.isSslEnabled());
+    assertNull(zkConf.getZookeeperKeyStoreLocation());
+    assertNull(zkConf.getZookeeperKeyStorePassword());
+    assertNull(zkConf.getZookeeperTrustStoreLocation());
+    assertNull(zkConf.getZookeeperTrustStorePassword());
+  }
+
+  @Test
+  public void testZkConfigTezAmZookeeperSslEnableEmpty() {
+    Configuration conf = new Configuration();
+    conf.set(TezConfiguration.TEZ_AM_ZOOKEEPER_QUORUM, "dummyZkQuorum");
+    conf.set(TezConfiguration.TEZ_AM_ZOOKEEPER_SSL_ENABLE, ""); // empty means not set
+    ZkConfig zkConf = new ZkConfig(conf);
+
+    assertEquals(Optional.empty(), zkConf.isSslEnabled());
+    assertNull(zkConf.getZookeeperKeyStoreLocation());
+    assertNull(zkConf.getZookeeperKeyStorePassword());
+    assertNull(zkConf.getZookeeperTrustStoreLocation());
+    assertNull(zkConf.getZookeeperTrustStorePassword());
+  }
+
+  @Test
+  public void testZkConfigSslEnabled() {
+    Configuration conf = new Configuration();
+    conf.set(TezConfiguration.TEZ_AM_ZOOKEEPER_QUORUM, "dummyZkQuorum");
+    conf.set(TezConfiguration.TEZ_AM_ZOOKEEPER_SSL_ENABLE, "true");
+    conf.set(TezConfiguration.TEZ_AM_ZOOKEEPER_SSL_KEYSTORE_LOCATION, "/keystore.jks");
+    conf.set(TezConfiguration.TEZ_AM_ZOOKEEPER_SSL_KEYSTORE_PASSWORD, "secret");
+    conf.set(TezConfiguration.TEZ_AM_ZOOKEEPER_SSL_TRUSTSTORE_LOCATION, "/truststore.jks");
+    conf.set(TezConfiguration.TEZ_AM_ZOOKEEPER_SSL_TRUSTSTORE_PASSWORD, "changeit");
+    ZkConfig zkConf = new ZkConfig(conf);
+
+    assertTrue(zkConf.isSslEnabled().isPresent());
+    assertTrue(zkConf.isSslEnabled().get());
+    assertEquals(zkConf.getZookeeperKeyStoreLocation(), "/keystore.jks");
+    assertEquals(zkConf.getZookeeperKeyStorePassword(), "secret");
+    assertEquals(zkConf.getZookeeperTrustStoreLocation(), "/truststore.jks");
+    assertEquals(zkConf.getZookeeperTrustStorePassword(), "changeit");
+  }
+
+  @Test
+  public void testZkConfigSslDisabled() {
+    Configuration conf = new Configuration();
+    conf.set(TezConfiguration.TEZ_AM_ZOOKEEPER_QUORUM, "dummyZkQuorum");
+    conf.set(TezConfiguration.TEZ_AM_ZOOKEEPER_SSL_ENABLE, "False");
+    ZkConfig zkConf = new ZkConfig(conf);
+
+    assertTrue(zkConf.isSslEnabled().isPresent());
+    assertFalse(zkConf.isSslEnabled().get());
+    assertNull(zkConf.getZookeeperKeyStoreLocation());
+    assertNull(zkConf.getZookeeperKeyStorePassword());
+    assertNull(zkConf.getZookeeperTrustStoreLocation());
+    assertNull(zkConf.getZookeeperTrustStorePassword());
+  }
+
+  @Test
+  public void testZkConfigAmZookeeperSslEnableInvalid() {
+    Configuration conf = new Configuration();
+    conf.set(TezConfiguration.TEZ_AM_ZOOKEEPER_QUORUM, "dummyZkQuorum");
+    conf.set(TezConfiguration.TEZ_AM_ZOOKEEPER_SSL_ENABLE, "invalidValue");
+    assertThrows(IllegalArgumentException.class, () -> new ZkConfig(conf));
+  }
 }

tez-tests/pom.xml

@@ -133,6 +133,11 @@
       <groupId>org.junit.jupiter</groupId>
       <artifactId>junit-jupiter</artifactId>
     </dependency>
+    <dependency>
+      <groupId>org.apache.curator</groupId>
+      <artifactId>curator-test</artifactId>
+      <scope>test</scope>
+    </dependency>
   </dependencies>
 
   <build>