Revert "Revert "Revert "Revert "Revert "Revert "Bratseth/lockdown zk""""""

5 files changed, 78 insertions, 8 deletions

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 4c08924f8de..135154bba92 100644
--- a/zkfacade/src/main/java/com/yahoo/vespa/curator/Curator.java
+++ b/zkfacade/src/main/java/com/yahoo/vespa/curator/Curator.java
@@ -25,6 +25,7 @@ import java.time.Duration;
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.List;
+import java.util.Objects;
 import java.util.Optional;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.TimeUnit;
@@ -54,7 +55,7 @@ public class Curator {
     private final String connectionSpec;
     private final int serverCount;
 
-    /** Creates a curator instance from a comma-separated string of ZooKeeper host names */
+    /** Creates a curator instance from a comma-separated string of ZooKeeper host:port strings */
     public static Curator create(String connectionSpec) {
         return new Curator(connectionSpec);
     }
@@ -64,7 +65,7 @@ public class Curator {
     public Curator(ConfigserverConfig configserverConfig, ZooKeeperServer server) {
         this(createConnectionSpec(configserverConfig));
     }
-
+    
     private static String createConnectionSpec(ConfigserverConfig config) {
         StringBuilder sb = new StringBuilder();
         for (int i = 0; i < config.zookeeperserver().size(); i++) {
@@ -80,6 +81,7 @@ public class Curator {
     }
 
     private Curator(String connectionSpec) {
+        Objects.requireNonNull(connectionSpec, "The curator connection spec cannot be null");
         this.connectionSpec = connectionSpec;
         this.serverCount = connectionSpec.split(",").length;
         validateConnectionSpec(connectionSpec);
@@ -103,15 +105,17 @@ public class Curator {
     }
 
     private static void validateConnectionSpec(String connectionSpec) {
-        if (connectionSpec == null || connectionSpec.isEmpty()) {
+        if (connectionSpec == null || connectionSpec.isEmpty())
             throw new IllegalArgumentException(String.format("Connections spec '%s' is not valid", connectionSpec));
-        }
     }
 
     /** Returns the number of zooKeeper servers in this cluster */
     public int serverCount() { return serverCount; }
 
-    /** Returns a comma-separated list of the zookeeper servers in this cluster */
+    /** 
+     * Returns the servers in this cluster as a comma-separated list of host:port strings. 
+     * This may be empty but never null 
+     */
     public String connectionSpec() { return connectionSpec; }
 
     /** For internal use; prefer creating a {@link CuratorCounter} */
diff --git a/zkfacade/src/main/java/com/yahoo/vespa/curator/mock/MockCurator.java b/zkfacade/src/main/java/com/yahoo/vespa/curator/mock/MockCurator.java
index 92db255a67f..4615eb4e34a 100644
--- a/zkfacade/src/main/java/com/yahoo/vespa/curator/mock/MockCurator.java
+++ b/zkfacade/src/main/java/com/yahoo/vespa/curator/mock/MockCurator.java
@@ -101,6 +101,7 @@ public class MockCurator extends Curator {
     private boolean shouldTimeoutOnEnter = false;
     private int monotonicallyIncreasingNumber = 0;
     private final boolean stableOrdering;
+    private String connectionSpec = "";
 
     /** The file system used by this mock to store zookeeper files and directories */
     private final MemoryFileSystem fileSystem = new MemoryFileSystem();
@@ -143,6 +144,12 @@ public class MockCurator extends Curator {
     public Optional<DistributedAtomicLong> counter(String path) {
         return Optional.ofNullable(atomicCounters.get(path));
     }
+    
+    /** Assigns the connection string, which must be on the form host1:port,host2:port ... */
+    public void setConnectionSpec(String connectionSpec) { this.connectionSpec = connectionSpec; }
+
+    @Override
+    public String connectionSpec() { return connectionSpec; }
 
     // ----- Start of adaptor methods from Curator to the mock file system -----
 
diff --git a/zkfacade/src/main/java/com/yahoo/vespa/zookeeper/RestrictedServerCnxnFactory.java b/zkfacade/src/main/java/com/yahoo/vespa/zookeeper/RestrictedServerCnxnFactory.java
new file mode 100644
index 00000000000..d8561c67767
--- /dev/null
+++ b/zkfacade/src/main/java/com/yahoo/vespa/zookeeper/RestrictedServerCnxnFactory.java
@@ -0,0 +1,56 @@
+package com.yahoo.vespa.zookeeper;
+
+import org.apache.zookeeper.server.NIOServerCnxn;
+import org.apache.zookeeper.server.NIOServerCnxnFactory;
+
+import java.io.IOException;
+import java.net.InetSocketAddress;
+import java.nio.channels.SelectionKey;
+import java.nio.channels.SocketChannel;
+import java.util.HashSet;
+import java.util.Set;
+import java.util.logging.Logger;
+
+/**
+ * This class is created by zookeeper by reflection, see the ZooKeeperServer constructor.
+ * 
+ * @author bratseth
+ */
+@SuppressWarnings("unused")
+public class RestrictedServerCnxnFactory extends NIOServerCnxnFactory {
+
+    private static final Logger log = Logger.getLogger(RestrictedServerCnxnFactory.class.getName());
+    
+    public RestrictedServerCnxnFactory() throws IOException {
+        super();
+    }
+    
+    @Override
+    protected NIOServerCnxn createConnection(SocketChannel socket, SelectionKey selection) throws IOException {
+        String remoteHost = ((InetSocketAddress)socket.getRemoteAddress()).getHostName(); // TODO: Move this line down
+
+        String zookeeperClients = System.getProperty(ZooKeeperServer.ZOOKEEPER_VESPA_CLIENTS_PROPERTY);
+        if (zookeeperClients == null || zookeeperClients.isEmpty()) {
+            log.info("Allowing connection to ZooKeeper from " + remoteHost + ", as " + ZooKeeperServer.ZOOKEEPER_VESPA_CLIENTS_PROPERTY + " is not set"); // TODO: Remove this line
+            return super.createConnection(socket, selection); // client checking is not activated
+        }
+
+        Set<String> zooKeeperClients = toHostnameSet(zookeeperClients);
+        if ( ! remoteHost.equals("localhost") && ! zooKeeperClients.contains(remoteHost)) {
+            String errorMessage = "Rejecting connection to ZooKeeper from " + remoteHost +
+                                  ": This cluster only allow connection from hosts in: " + zooKeeperClients;
+            log.warning(errorMessage);
+            throw new IllegalArgumentException(errorMessage);
+        }
+        log.info("Allowing connection to ZooKeeper from " + remoteHost + ", as it is in " + zookeeperClients); // TODO: Remove this line
+        return super.createConnection(socket, selection);
+    }
+
+    private Set<String> toHostnameSet(String commaSeparatedString) {
+        Set<String> hostnames = new HashSet<>();
+        for (String hostname : commaSeparatedString.split(","))
+            hostnames.add(hostname.trim());
+        return hostnames;
+    }
+
+}
diff --git a/zkfacade/src/main/java/com/yahoo/vespa/zookeeper/ZooKeeperServer.java b/zkfacade/src/main/java/com/yahoo/vespa/zookeeper/ZooKeeperServer.java
index ffae797561c..f037e3c9265 100644
--- a/zkfacade/src/main/java/com/yahoo/vespa/zookeeper/ZooKeeperServer.java
+++ b/zkfacade/src/main/java/com/yahoo/vespa/zookeeper/ZooKeeperServer.java
@@ -19,6 +19,8 @@ import java.util.List;
  */
 public class ZooKeeperServer extends AbstractComponent implements Runnable {
 
+    public static final String ZOOKEEPER_VESPA_CLIENTS_PROPERTY = "zookeeper.vespa.clients";
+
     private static final java.util.logging.Logger log = java.util.logging.Logger.getLogger(ZooKeeperServer.class.getName());
     private static final String ZOOKEEPER_JMX_LOG4J_DISABLE = "zookeeper.jmx.log4j.disable";
     static final String ZOOKEEPER_JUTE_MAX_BUFFER = "jute.maxbuffer";
@@ -29,6 +31,8 @@ public class ZooKeeperServer extends AbstractComponent implements Runnable {
         this.config = config;
         System.setProperty("zookeeper.jmx.log4j.disable", "true");
         System.setProperty(ZOOKEEPER_JUTE_MAX_BUFFER, "" + config.juteMaxBuffer());
+        System.setProperty("zookeeper.serverCnxnFactory", "com.yahoo.vespa.zookeeper.RestrictedServerCnxnFactory");
+
         writeConfigToDisk(config);
         zkServerThread = new Thread(this, "zookeeper server");
         if (startServer) {
@@ -40,7 +44,7 @@ public class ZooKeeperServer extends AbstractComponent implements Runnable {
     public ZooKeeperServer(ZookeeperServerConfig config) {
         this(config, true);
     }
-
+    
     private void writeConfigToDisk(ZookeeperServerConfig config) {
        String cfg = transformConfigToString(config);
        try (FileWriter writer = new FileWriter(Defaults.getDefaults().underVespaHome(config.zooKeeperConfigFile()))) {
diff --git a/zkfacade/src/test/java/com/yahoo/vespa/zookeeper/ZooKeeperServerTest.java b/zkfacade/src/test/java/com/yahoo/vespa/zookeeper/ZooKeeperServerTest.java
index 4e231631f28..8c4db5dc8e4 100644
--- a/zkfacade/src/test/java/com/yahoo/vespa/zookeeper/ZooKeeperServerTest.java
+++ b/zkfacade/src/test/java/com/yahoo/vespa/zookeeper/ZooKeeperServerTest.java
@@ -3,7 +3,6 @@ package com.yahoo.vespa.zookeeper;
 
 import com.yahoo.cloud.config.ZookeeperServerConfig;
 import com.yahoo.io.IOUtils;
-import com.yahoo.vespa.curator.Curator;
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.rules.TemporaryFolder;
@@ -12,7 +11,6 @@ import java.io.File;
 import java.io.IOException;
 
 import static org.hamcrest.CoreMatchers.is;
-import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertThat;
 import com.yahoo.vespa.defaults.Defaults;
 
@@ -74,6 +72,7 @@ public class ZooKeeperServerTest {
         File idFile = folder.newFile();
         File cfgFile = folder.newFile();
 
+        builder.server(new ZookeeperServerConfig.Server.Builder().id(0).hostname("testhost"));
         builder.zooKeeperConfigFile(cfgFile.getAbsolutePath());
         builder.myidFile(idFile.getAbsolutePath());