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

tez-api/findbugs-exclude.xml

@@ -142,4 +142,11 @@
     <Method name="get" params="java.lang.Class, org.apache.hadoop.conf.Configuration, java.lang.Class" returns="org.apache.tez.frameworkplugins.FrameworkService" />
     <Bug pattern="REFLC_REFLECTION_MAY_INCREASE_ACCESSIBILITY_OF_CLASS" />
   </Match>
+
+  <!-- TEZ-4749 -->
+  <Match>
+    <Class name="org.apache.tez.client.registry.zookeeper.ZkConfig" />
+    <Method name="isSslEnabled" />
+    <Bug pattern="NP_BOOLEAN_RETURN_NULL" />
+  </Match>
 </FindBugsFilter>

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

@@ -0,0 +1,76 @@
+/*
+ * 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.tez.client.registry.zookeeper;
+
+import org.apache.commons.lang3.StringUtils;
+import org.apache.curator.utils.ZookeeperFactory;
+import org.apache.zookeeper.Watcher;
+import org.apache.zookeeper.ZooKeeper;
+import org.apache.zookeeper.client.ZKClientConfig;
+import org.apache.zookeeper.common.ClientX509Util;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Factory to create Zookeeper clients with the zookeeper.client.secure enabled,
+ * allowing SSL communication with the Zookeeper server.
+ */
+public class SSLZookeeperFactory implements ZookeeperFactory {
+
+  private static final Logger LOG = LoggerFactory.getLogger(SSLZookeeperFactory.class);
+
+  private boolean sslEnabled;
+  private String keyStoreLocation;
+  private String keyStorePassword;
+  private String trustStoreLocation;
+  private String trustStorePassword;
+
+  public SSLZookeeperFactory(boolean sslEnabled, String keyStoreLocation, String keyStorePassword,
+                             String trustStoreLocation, String trustStorePassword) {
+
+    this.sslEnabled = sslEnabled;
+    this.keyStoreLocation = keyStoreLocation;
+    this.keyStorePassword = keyStorePassword;
+    this.trustStoreLocation = trustStoreLocation;
+    this.trustStorePassword = trustStorePassword;
+    if (sslEnabled) {
+      if (StringUtils.isEmpty(keyStoreLocation)) {
+        LOG.warn("Missing keystoreLocation parameter");
+      }
+      if (StringUtils.isEmpty(trustStoreLocation)) {
+        LOG.warn("Missing trustStoreLocation parameter");
+      }
+    }
+  }
+
+  @Override
+  public ZooKeeper newZooKeeper(String connectString, int sessionTimeout, Watcher watcher,
+                                boolean canBeReadOnly) throws Exception {
+    ZKClientConfig clientConfig = new ZKClientConfig();
+    clientConfig.setProperty(ZKClientConfig.SECURE_CLIENT, Boolean.toString(sslEnabled));
+    clientConfig.setProperty(ZKClientConfig.ZOOKEEPER_CLIENT_CNXN_SOCKET, "org.apache.zookeeper.ClientCnxnSocketNetty");
+    ClientX509Util x509Util = new ClientX509Util();
+    clientConfig.setProperty(x509Util.getSslKeystoreLocationProperty(), this.keyStoreLocation);
+    clientConfig.setProperty(x509Util.getSslKeystorePasswdProperty(), this.keyStorePassword);
+    clientConfig.setProperty(x509Util.getSslTruststoreLocationProperty(), this.trustStoreLocation);
+    clientConfig.setProperty(x509Util.getSslTruststorePasswdProperty(), this.trustStorePassword);
+    return new ZooKeeper(connectString, sessionTimeout, watcher, canBeReadOnly, clientConfig);
+  }
+}

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

@@ -49,6 +49,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 +89,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 +125,67 @@ 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 boolean value that indicates whether zookeeper client uses a secure
+   * zookeeper connection. A null value 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 Boolean isSslEnabled() {
+    if (this.sslEnabled == null || this.sslEnabled.isEmpty()) {
+      return null;
+    }
+    return Boolean.parseBoolean(sslEnabled);
+  }
+
   public RetryPolicy getRetryPolicy() {
     return new ExponentialBackoffRetry(getCuratorBackoffSleepMs(), getCuratorMaxRetries());
   }
 
   public CuratorFramework createCuratorFramework() {
-    return CuratorFrameworkFactory.newClient(
-        getZkQuorum(),
-        getSessionTimeoutMs(),
-        getConnectionTimeoutMs(),
-        getRetryPolicy()
-    );
+    if (isSslEnabled() == null) {
+      return CuratorFrameworkFactory.newClient(
+              getZkQuorum(),
+              getSessionTimeoutMs(),
+              getConnectionTimeoutMs(),
+              getRetryPolicy()
+      );
+    }
+
+    return CuratorFrameworkFactory.builder()
+            .connectString(getZkQuorum())
+            .sessionTimeoutMs(getSessionTimeoutMs())
+            .connectionTimeoutMs(getConnectionTimeoutMs())
+            .retryPolicy(getRetryPolicy())
+            .zookeeperFactory(
+                    new SSLZookeeperFactory(isSslEnabled(), getZookeeperKeyStoreLocation(),
+                            getZookeeperKeyStorePassword(), getZookeeperTrustStoreLocation(),
+                            getZookeeperTrustStorePassword()))
+            .build();
+  }
+
+  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

@@ -20,6 +20,7 @@
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
 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 java.util.concurrent.TimeUnit;
@@ -231,4 +232,71 @@ 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);
+
+    assertNull(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);
+
+    assertNull(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);
+
+    assertEquals(zkConf.isSslEnabled(), Boolean.TRUE);
+    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);
+
+    assertEquals(zkConf.isSslEnabled(), Boolean.FALSE);
+    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,12 @@
       <groupId>org.junit.jupiter</groupId>
       <artifactId>junit-jupiter</artifactId>
     </dependency>
+    <dependency>
+      <groupId>org.apache.curator</groupId>
+      <artifactId>curator-test</artifactId>
+      <version>${curator.version}</version>
+      <scope>test</scope>
+    </dependency>
   </dependencies>
 
   <build>