Allow all application hosts to talk to ZooKeeper

2 files changed, 17 insertions, 26 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
index ae7df9ac7cf..90c68461699 100644
--- a/zkfacade/src/main/java/com/yahoo/vespa/zookeeper/RestrictedServerCnxnFactory.java
+++ b/zkfacade/src/main/java/com/yahoo/vespa/zookeeper/RestrictedServerCnxnFactory.java
@@ -18,36 +18,35 @@ import java.util.logging.Logger;
  */
 @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 zookeeperClients = System.getProperty(ZooKeeperServer.ZOOKEEPER_VESPA_CLIENTS_PROPERTY);
+        if (zookeeperClients == null || zookeeperClients.isEmpty())
+            return super.createConnection(socket, selection); // client checking is not activated
+
+        Set<String> zooKeeperClients = toHostnameSet(zookeeperClients);
         String remoteHost = ((InetSocketAddress)socket.getRemoteAddress()).getHostName();
-        if ( ! remoteHost.equals("localhost") && ! zooKeeperServerHostnames.contains(remoteHost)) {
+        if ( ! remoteHost.equals("localhost") && ! zooKeeperClients.contains(remoteHost)) {
             String errorMessage = "Rejecting connection to ZooKeeper from " + remoteHost +
-                                  ": This cluster only allow connection among its own hosts. " +
-                                  "Hosts in this cluster: " + zooKeeperServerHostnames;
+                                  ": This cluster only allow connection from hosts in: " + zooKeeperClients;
             log.warning(errorMessage);
             throw new IllegalArgumentException(errorMessage);
         }
         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 d4670e97ed8..65c44981fb2 100644
--- a/zkfacade/src/main/java/com/yahoo/vespa/zookeeper/ZooKeeperServer.java
+++ b/zkfacade/src/main/java/com/yahoo/vespa/zookeeper/ZooKeeperServer.java
@@ -19,10 +19,11 @@ 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";
-    static final String ZOOKEEPER_VESPA_SERVERS_PROPERTY = "zookeeper.vespa.servers";
     private final Thread zkServerThread;
     private final ZookeeperServerConfig config;
 
@@ -31,7 +32,6 @@ public class ZooKeeperServer extends AbstractComponent implements Runnable {
         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);
@@ -46,14 +46,6 @@ public class ZooKeeperServer extends AbstractComponent implements Runnable {
         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);
        try (FileWriter writer = new FileWriter(Defaults.getDefaults().underVespaHome(config.zooKeeperConfigFile()))) {