Revert "Revert "Bratseth/lockdown zk""

3 files changed, 67 insertions, 2 deletions

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..ae7df9ac7cf
--- /dev/null
+++ b/zkfacade/src/main/java/com/yahoo/vespa/zookeeper/RestrictedServerCnxnFactory.java
@@ -0,0 +1,53 @@
+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());
+    private final Set<String> zooKeeperServerHostnames;
+    
+    public RestrictedServerCnxnFactory() throws IOException {
+        super();
+        zooKeeperServerHostnames = toHostnameSet(System.getProperty(ZooKeeperServer.ZOOKEEPER_VESPA_SERVERS_PROPERTY));
+    }
+    
+    private Set<String> toHostnameSet(String commaSeparatedString) {
+        if (commaSeparatedString == null || commaSeparatedString.isEmpty())
+            throw new IllegalArgumentException("We have not received the list of ZooKeeper servers in this system");
+        
+        Set<String> hostnames = new HashSet<>();
+        for (String hostname : commaSeparatedString.split(","))
+            hostnames.add(hostname.trim());
+        return hostnames;
+    }
+
+    @Override
+    protected NIOServerCnxn createConnection(SocketChannel socket, SelectionKey selection) throws IOException {
+        String remoteHost = ((InetSocketAddress)socket.getRemoteAddress()).getHostName();
+        if ( ! remoteHost.equals("localhost") && ! zooKeeperServerHostnames.contains(remoteHost)) {
+            String errorMessage = "Rejecting connection to ZooKeeper from " + remoteHost +
+                                  ": This cluster only allow connection among its own hosts. " +
+                                  "Hosts in this cluster: " + zooKeeperServerHostnames;
+            log.warning(errorMessage);
+            throw new IllegalArgumentException(errorMessage);
+        }
+        return super.createConnection(socket, selection);
+    }
+
+}
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..d4670e97ed8 100644
--- a/zkfacade/src/main/java/com/yahoo/vespa/zookeeper/ZooKeeperServer.java
+++ b/zkfacade/src/main/java/com/yahoo/vespa/zookeeper/ZooKeeperServer.java
@@ -22,6 +22,7 @@ public class ZooKeeperServer extends AbstractComponent implements Runnable {
     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";
+    static final String ZOOKEEPER_VESPA_SERVERS_PROPERTY = "zookeeper.vespa.servers";
     private final Thread zkServerThread;
     private final ZookeeperServerConfig config;
 
@@ -29,6 +30,10 @@ 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_VESPA_SERVERS_PROPERTY, toHostnameString(config.server()));
+        System.setProperty("zookeeper.serverCnxnFactory", "com.yahoo.vespa.zookeeper.RestrictedServerCnxnFactory");
+
         writeConfigToDisk(config);
         zkServerThread = new Thread(this, "zookeeper server");
         if (startServer) {
@@ -40,6 +45,14 @@ public class ZooKeeperServer extends AbstractComponent implements Runnable {
     public ZooKeeperServer(ZookeeperServerConfig config) {
         this(config, true);
     }
+    
+    private String toHostnameString(List<ZookeeperServerConfig.Server> servers) {
+        StringBuilder b = new StringBuilder();
+        for (ZookeeperServerConfig.Server server : servers)
+            b.append(server.hostname()).append(", ");
+        b.setLength(b.length()-1); // remove the last ", "
+        return b.toString();
+    }
 
     private void writeConfigToDisk(ZookeeperServerConfig config) {
        String cfg = transformConfigToString(config);
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());