Set TLS config for client based on VESPA_USE_TLS_FOR_ZOOKEEPER_CLIENT

5 files changed, 54 insertions, 23 deletions

diff --git a/zkfacade/abi-spec.json b/zkfacade/abi-spec.json
index efe6fbdaa08..25b652b7312 100644
--- a/zkfacade/abi-spec.json
+++ b/zkfacade/abi-spec.json
@@ -68,7 +68,6 @@
     "methods": [
       "public static com.yahoo.vespa.curator.Curator create(java.lang.String)",
       "public static com.yahoo.vespa.curator.Curator create(java.lang.String, java.util.Optional)",
-      "public void <init>(com.yahoo.cloud.config.ConfigserverConfig)",
       "public void <init>(com.yahoo.cloud.config.ConfigserverConfig, com.yahoo.vespa.zookeeper.VespaZooKeeperServer)",
       "protected void <init>(java.lang.String, java.lang.String, java.util.function.Function)",
       "public java.lang.String connectionSpec()",
diff --git a/zkfacade/src/main/java/com/yahoo/vespa/curator/Curator.java b/zkfacade/src/main/java/com/yahoo/vespa/curator/Curator.java
index b76bad5b97b..9d74306d3d5 100644
--- a/zkfacade/src/main/java/com/yahoo/vespa/curator/Curator.java
+++ b/zkfacade/src/main/java/com/yahoo/vespa/curator/Curator.java
@@ -3,9 +3,12 @@ package com.yahoo.vespa.curator;
 
 import com.google.inject.Inject;
 import com.yahoo.cloud.config.ConfigserverConfig;
+import com.yahoo.io.IOUtils;
 import com.yahoo.net.HostName;
 import com.yahoo.path.Path;
+import com.yahoo.text.Utf8;
 import com.yahoo.vespa.curator.recipes.CuratorCounter;
+import com.yahoo.vespa.defaults.Defaults;
 import com.yahoo.vespa.zookeeper.VespaZooKeeperServer;
 import org.apache.curator.RetryPolicy;
 import org.apache.curator.framework.CuratorFramework;
@@ -20,7 +23,9 @@ import org.apache.curator.framework.recipes.locks.InterProcessLock;
 import org.apache.curator.framework.recipes.locks.InterProcessMutex;
 import org.apache.curator.retry.ExponentialBackoffRetry;
 import org.apache.zookeeper.KeeperException;
+import org.apache.zookeeper.client.ZKClientConfig;
 import org.apache.zookeeper.data.Stat;
+import org.apache.zookeeper.server.quorum.QuorumPeerConfig;
 
 import java.io.File;
 import java.time.Duration;
@@ -63,28 +68,21 @@ public class Curator implements AutoCloseable {
 
     /** Creates a curator instance from a comma-separated string of ZooKeeper host:port strings */
     public static Curator create(String connectionSpec, Optional<File> clientConfigFile) {
-        return new Curator(connectionSpec, connectionSpec);
-    }
-
-    // For testing
-    public Curator(ConfigserverConfig configserverConfig) {
-        this(configserverConfig, createConnectionSpec(configserverConfig));
+        return new Curator(connectionSpec, connectionSpec, clientConfigFile);
     }
 
     // Depend on ZooKeeperServer to make sure it is started first
     // TODO: Move zookeeperserver config out of configserverconfig (requires update of controller services.xml as well)
     @Inject
     public Curator(ConfigserverConfig configserverConfig, VespaZooKeeperServer server) {
-        this(configserverConfig, createConnectionSpec(configserverConfig));
+        this(configserverConfig, Optional.empty());
     }
 
-    private Curator(ConfigserverConfig configserverConfig, String zooKeeperEnsembleConnectionSpec) {
-        this((configserverConfig.zookeeperLocalhostAffinity()) ?
-                createConnectionSpecForLocalhost(configserverConfig) : zooKeeperEnsembleConnectionSpec,
-                zooKeeperEnsembleConnectionSpec);
+    Curator(ConfigserverConfig configserverConfig, Optional<File> clientConfigFile) {
+        this(createConnectionSpec(configserverConfig), createEnsembleConnectionSpec(configserverConfig), clientConfigFile);
     }
 
-    private Curator(String connectionSpec, String zooKeeperEnsembleConnectionSpec) {
+    private Curator(String connectionSpec, String zooKeeperEnsembleConnectionSpec, Optional<File> clientConfigFile) {
         this(connectionSpec,
                 zooKeeperEnsembleConnectionSpec,
                 (retryPolicy) -> CuratorFrameworkFactory
@@ -93,7 +91,7 @@ public class Curator implements AutoCloseable {
                         .sessionTimeoutMs(ZK_SESSION_TIMEOUT)
                         .connectionTimeoutMs(ZK_CONNECTION_TIMEOUT)
                         .connectString(connectionSpec)
-                        .zookeeperFactory(new VespaZooKeeperFactory())
+                        .zookeeperFactory(new VespaZooKeeperFactory(createClientConfig(clientConfigFile)))
                         .dontUseContainerParents() // TODO: Remove when we know ZooKeeper 3.5 works fine, consider waiting until Vespa 8
                         .build());
     }
@@ -123,7 +121,29 @@ public class Curator implements AutoCloseable {
         this.zooKeeperEnsembleCount = zooKeeperEnsembleConnectionSpec.split(",").length;
     }
 
-    private static String createConnectionSpec(ConfigserverConfig config) {
+    private static String createConnectionSpec(ConfigserverConfig configserverConfig) {
+        return configserverConfig.zookeeperLocalhostAffinity()
+                ? createConnectionSpecForLocalhost(configserverConfig)
+                : createEnsembleConnectionSpec(configserverConfig);
+    }
+
+    private static ZKClientConfig createClientConfig(Optional<File> file) {
+        boolean useSecureClient = Boolean.parseBoolean(getEnvironmentVariable("VESPA_USE_TLS_FOR_ZOOKEEPER_CLIENT").orElse("false"));
+        String config = "zookeeper.client.secure=" + useSecureClient + "\n";
+
+        File clientConfigFile =
+                file.orElseGet(() -> new File(Defaults.getDefaults().underVespaHome("conf/zookeeper/zookeeper-client.cfg")));
+        clientConfigFile.getParentFile().mkdirs();
+        IOUtils.writeFile(clientConfigFile, Utf8.toBytes(config));
+
+        try {
+            return new ZKClientConfig(clientConfigFile);
+        } catch (QuorumPeerConfig.ConfigException e) {
+            throw new RuntimeException("Unable to create ZooKeeper client config file " + file);
+        }
+    }
+
+    private static String createEnsembleConnectionSpec(ConfigserverConfig config) {
         StringBuilder connectionSpec = new StringBuilder();
         for (int i = 0; i < config.zookeeperserver().size(); i++) {
             if (connectionSpec.length() > 0) {
@@ -405,4 +425,10 @@ public class Curator implements AutoCloseable {
      * TODO: Move method out of this class.
      */
     public int zooKeeperEnsembleCount() { return zooKeeperEnsembleCount; }
+
+    private static Optional<String> getEnvironmentVariable(String variableName) {
+        return Optional.ofNullable(System.getenv().get(variableName))
+                .filter(var -> !var.isEmpty());
+    }
+
 }
diff --git a/zkfacade/src/main/java/com/yahoo/vespa/curator/VespaZooKeeperFactory.java b/zkfacade/src/main/java/com/yahoo/vespa/curator/VespaZooKeeperFactory.java
index 7c08168c536..84e2cb65a1a 100644
--- a/zkfacade/src/main/java/com/yahoo/vespa/curator/VespaZooKeeperFactory.java
+++ b/zkfacade/src/main/java/com/yahoo/vespa/curator/VespaZooKeeperFactory.java
@@ -4,19 +4,24 @@ package com.yahoo.vespa.curator;
 import org.apache.curator.utils.ZookeeperFactory;
 import org.apache.zookeeper.Watcher;
 import org.apache.zookeeper.ZooKeeper;
+import org.apache.zookeeper.client.ZKClientConfig;
 
 /**
  * A ZooKeeper factory for creating a ZooKeeper client
  *
  * @author hmusum
  */
-// TODO: add constructor that takes feature flag so that we can write ZooKeeper client config and start
-// ZooKeeper client with that config
 class VespaZooKeeperFactory implements ZookeeperFactory {
 
+    private final ZKClientConfig zkClientConfig;
+
+    VespaZooKeeperFactory(ZKClientConfig zkClientConfig) {
+        this.zkClientConfig = zkClientConfig;
+    }
+
     @Override
     public ZooKeeper newZooKeeper(String connectString, int sessionTimeout, Watcher watcher, boolean canBeReadOnly) throws Exception {
-        return new ZooKeeper(connectString, sessionTimeout, watcher);
+        return new ZooKeeper(connectString, sessionTimeout, watcher, zkClientConfig);
     }
 
 }
diff --git a/zkfacade/src/test/java/com/yahoo/vespa/curator/CuratorCounterTest.java b/zkfacade/src/test/java/com/yahoo/vespa/curator/CuratorCounterTest.java
index 9b7e0250f2f..6b85953a1ff 100644
--- a/zkfacade/src/test/java/com/yahoo/vespa/curator/CuratorCounterTest.java
+++ b/zkfacade/src/test/java/com/yahoo/vespa/curator/CuratorCounterTest.java
@@ -8,7 +8,6 @@ import static org.junit.Assert.assertEquals;
 
 /**
  * @author Ulf Lilleengen
- * @date 19.08.13
  */
 public class CuratorCounterTest {
 
diff --git a/zkfacade/src/test/java/com/yahoo/vespa/curator/CuratorTest.java b/zkfacade/src/test/java/com/yahoo/vespa/curator/CuratorTest.java
index a8342dfe5bc..4cd2c708d1a 100644
--- a/zkfacade/src/test/java/com/yahoo/vespa/curator/CuratorTest.java
+++ b/zkfacade/src/test/java/com/yahoo/vespa/curator/CuratorTest.java
@@ -9,6 +9,8 @@ import org.junit.Before;
 import org.junit.Test;
 
 import java.io.IOException;
+import java.nio.file.Files;
+import java.util.Optional;
 
 import static org.hamcrest.core.Is.is;
 import static org.junit.Assert.assertThat;
@@ -52,7 +54,7 @@ public class CuratorTest {
     }
 
     @Test
-    public void require_curator_is_created_from_config() {
+    public void require_curator_is_created_from_config() throws IOException {
         try (Curator curator = createCurator(createTestConfig())) {
             assertThat(curator.zooKeeperEnsembleConnectionSpec(), is(spec1 + "," + spec2));
             assertThat(curator.zooKeeperEnsembleCount(), is(2));
@@ -60,7 +62,7 @@ public class CuratorTest {
     }
 
     @Test
-    public void require_that_server_count_is_correct() {
+    public void require_that_server_count_is_correct() throws IOException {
         ConfigserverConfig.Builder builder = new ConfigserverConfig.Builder();
         builder.zookeeperserver(createZKBuilder(localhost, port1));
         try (Curator curator = createCurator(new ConfigserverConfig(builder))) {
@@ -98,8 +100,8 @@ public class CuratorTest {
         return zkBuilder;
     }
 
-    private Curator createCurator(ConfigserverConfig configserverConfig) {
-        return new Curator(configserverConfig);
+    private Curator createCurator(ConfigserverConfig configserverConfig) throws IOException {
+        return new Curator(configserverConfig, Optional.of(Files.createTempFile("zookeeper-client", "cfg").toFile()));
     }
 
     private static class PortAllocator {