Change the way ZkClientUriDomainMappingHelper add watch (#111)

ZkClientUriDomainMappingHelper instance is registered as a watcher during the instantiation process, so it will get the notification every time there's a change in the client URI domain mapping, and it can update each session that are connected to this zookeeper server for the updated access control metatdata.
However, in ZooKeeper server code, there is an implicit assumption that the objects registered as a watcher are ServerCnxn objects, while ZkClientUriDomainMappingHelper is not an instance of ServerCnxn. This incorrect cast of ZkClientUriDomainMappingHelper to ServerCnxn can lead to code failure.
This commit leveraged a dummy server cnxn class `DumbWatcher` which actually functions as a watcher, to prevent this error.

Author: mgao0

zookeeper-server/src/main/java/org/apache/zookeeper/server/auth/znode/groupacl/ZkClientUriDomainMappingHelper.java

@@ -27,8 +27,8 @@
 import java.util.Set;
 import org.apache.zookeeper.KeeperException;
 import org.apache.zookeeper.WatchedEvent;
-import org.apache.zookeeper.Watcher;
 import org.apache.zookeeper.ZooDefs;
+import org.apache.zookeeper.server.DumbWatcher;
 import org.apache.zookeeper.server.ServerCnxn;
 import org.apache.zookeeper.server.ServerCnxnFactory;
 import org.apache.zookeeper.server.ZooKeeperServer;
@@ -58,7 +58,7 @@
  * Note: It is not expected that there would be too many distinct client URIs so as to overwhelm
  * heap usage.
  */
-public class ZkClientUriDomainMappingHelper implements Watcher, ClientUriDomainMappingHelper {
+public class ZkClientUriDomainMappingHelper implements ClientUriDomainMappingHelper {
 
   private static final Logger LOG = LoggerFactory.getLogger(ZkClientUriDomainMappingHelper.class);
 
@@ -106,9 +106,13 @@ boolean setDomainAuthUpdater(ConnectionAuthInfoUpdater updater) {
 
   /**
    * Install a persistent recursive watch on the root path.
+   * The watcher has to be added here instead of in {@link org.apache.zookeeper.server.auth.znode.groupacl.X509ZNodeGroupAclProvider}
+   * because this class is one layer above connections, and is a central place to handle add watch,
+   * process event, etc logic for all the connections in this server.
+   * Therefore, only one watch needs to be added per server.
    */
   private void addWatches() {
-    zks.getZKDatabase().addWatch(rootPath, this, ZooDefs.AddWatchModes.persistentRecursive);
+    zks.getZKDatabase().addWatch(rootPath, new MappingRootWatcher(), ZooDefs.AddWatchModes.persistentRecursive);
   }
 
   /**
@@ -141,33 +145,6 @@ private void parseZNodeMapping() {
     clientUriToDomainNames = newClientUriToDomainNames;
   }
 
-  @Override
-  public void process(WatchedEvent event) {
-    LOG.info("Processing watched event: {}", event.toString());
-    parseZNodeMapping();
-    // Update AuthInfo for all the known connections.
-    // Note : It is not ideal to iterate over all plaintext connections which are connected over non-TLS but right now
-    // there is no way to find out if connection on unified port is using SSLHandler or nonSSLHandler. Anyways, we
-    // should not ideally have any nonSSLHandler connections on unified port after complete rollout.
-
-    // TODO Change to read SecureServerCnxnFactory only. The current logic is to support unit test who is not creating
-    // a secured server cnxn factory. It won't cause any problem but is not technically correct.
-
-    // Since port unification is supported, TLS requests could be made on unified as well as secure port. Hence iterate
-    // over all connections to update auth info.
-    ServerCnxnFactory factory = zks.getServerCnxnFactory();
-    LOG.info("Updating auth info for connections");
-    // TODO Evaluate performance impact and potentially use thread pool to parallelize the AuthInfo update.
-    if (factory != null) {
-      factory.getConnections().forEach(cnxn -> updateDomainBasedAuthInfo(cnxn));
-    }
-    ServerCnxnFactory secureFactory = zks.getSecureServerCnxnFactory();
-    LOG.info("Updating auth info for TLS connections");
-    if (secureFactory != null) {
-      secureFactory.getConnections().forEach(cnxn -> updateDomainBasedAuthInfo(cnxn));
-    }
-  }
-
   @Override
   public Set<String> getDomains(String clientUri) {
     return clientUriToDomainNames.getOrDefault(clientUri, Collections.emptySet());
@@ -183,4 +160,38 @@ public void updateDomainBasedAuthInfo(ServerCnxn cnxn) {
       }
     }
   }
+
+  /**
+   * The watcher used to listen on client uri - domain mapping root path for mapping update
+   * Extends DumbWatcher instead of ServerCnxn because there are some package-private methods in
+   * ServerCnxn which cannot be overridden here, such as {@code abstract void setSessionId(long sessionId);}.
+   */
+  public class MappingRootWatcher extends DumbWatcher {
+    @Override
+    public void process(WatchedEvent event) {
+      LOG.info("Processing watched event: {}", event.toString());
+      parseZNodeMapping();
+      // Update AuthInfo for all the known connections.
+      // Note : It is not ideal to iterate over all plaintext connections which are connected over non-TLS but right now
+      // there is no way to find out if connection on unified port is using SSLHandler or nonSSLHandler. Anyways, we
+      // should not ideally have any nonSSLHandler connections on unified port after complete rollout.
+
+      // TODO Change to read SecureServerCnxnFactory only. The current logic is to support unit test who is not creating
+      // a secured server cnxn factory. It won't cause any problem but is not technically correct.
+
+      // Since port unification is supported, TLS requests could be made on unified as well as secure port. Hence iterate
+      // over all connections to update auth info.
+      ServerCnxnFactory factory = zks.getServerCnxnFactory();
+      LOG.info("Updating auth info for connections");
+      // TODO Evaluate performance impact and potentially use thread pool to parallelize the AuthInfo update.
+      if (factory != null) {
+        factory.getConnections().forEach(cnxn -> updateDomainBasedAuthInfo(cnxn));
+      }
+      ServerCnxnFactory secureFactory = zks.getSecureServerCnxnFactory();
+      LOG.info("Updating auth info for TLS connections");
+      if (secureFactory != null) {
+        secureFactory.getConnections().forEach(cnxn -> updateDomainBasedAuthInfo(cnxn));
+      }
+    }
+  }
 }

zookeeper-server/src/test/java/org/apache/zookeeper/server/auth/znode/groupacl/ZkClientUriDomainMappingHelperTest.java

@@ -32,6 +32,7 @@
 import org.apache.zookeeper.server.ServerCnxn;
 import org.apache.zookeeper.server.ServerCnxnFactory;
 import org.apache.zookeeper.server.ZooKeeperServer;
+import org.apache.zookeeper.server.watch.WatchesReport;
 import org.apache.zookeeper.test.ClientBase;
 import org.junit.After;
 import org.junit.Assert;
@@ -122,7 +123,7 @@ public void cleanUp() throws InterruptedException, IOException, KeeperException
    *         └── bar1 (client URI)
    */
   @Test
-  public void testB_ZkClientUriDomainMappingHelper() throws KeeperException, InterruptedException {
+  public void testA_ZkClientUriDomainMappingHelper() throws KeeperException, InterruptedException {
     for (String path : MAPPING_PATHS) {
       zookeeperClientConnection
           .create(path, null, ZooDefs.Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT);
@@ -137,21 +138,33 @@ public void testB_ZkClientUriDomainMappingHelper() throws KeeperException, Inter
     Assert.assertEquals(new HashSet<>(Arrays.asList("bar", "foo")), helper.getDomains("bar1"));
 
     // Add a new application domain and add bar1 to it
-    zookeeperClientConnection
-        .create(CLIENT_URI_DOMAIN_MAPPING_ROOT_PATH + "/new", null, ZooDefs.Ids.OPEN_ACL_UNSAFE,
-            CreateMode.PERSISTENT);
-    zookeeperClientConnection.create(CLIENT_URI_DOMAIN_MAPPING_ROOT_PATH + "/new/bar1", null,
-        ZooDefs.Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT);
-
-    // For bar1, we should get bar, foo, and new
-    Assert
-        .assertEquals(new HashSet<>(Arrays.asList("bar", "foo", "new")), helper.getDomains("bar1"));
-
-    // Remove the application domain and bar1
-    zookeeperClientConnection.delete(CLIENT_URI_DOMAIN_MAPPING_ROOT_PATH + "/new/bar1", -1);
-    zookeeperClientConnection.delete(CLIENT_URI_DOMAIN_MAPPING_ROOT_PATH + "/new", -1);
+    try {
+      zookeeperClientConnection
+          .create(CLIENT_URI_DOMAIN_MAPPING_ROOT_PATH + "/new", null, ZooDefs.Ids.OPEN_ACL_UNSAFE,
+              CreateMode.PERSISTENT);
+      zookeeperClientConnection.create(CLIENT_URI_DOMAIN_MAPPING_ROOT_PATH + "/new/bar1", null,
+          ZooDefs.Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT);
+
+      // For bar1, we should get bar, foo, and new
+      Assert.assertEquals(new HashSet<>(Arrays.asList("bar", "foo", "new")), helper.getDomains("bar1"));
+    } finally {
+      // Remove the application domain and bar1
+      zookeeperClientConnection.delete(CLIENT_URI_DOMAIN_MAPPING_ROOT_PATH + "/new/bar1", -1);
+      zookeeperClientConnection.delete(CLIENT_URI_DOMAIN_MAPPING_ROOT_PATH + "/new", -1);
+    }
 
     // For bar1, we should get bar and foo
     Assert.assertEquals(new HashSet<>(Arrays.asList("bar", "foo")), helper.getDomains("bar1"));
   }
+
+  @Test
+  /**
+   * Make sure the watcher installed while instantiate ZkClientUriDomainMappingHelper does not break
+   * the functionality of getting watches
+   */
+  public void testB_GetWatches() {
+    ClientUriDomainMappingHelper helper = new ZkClientUriDomainMappingHelper(zookeeperServer);
+    WatchesReport report = zookeeperServer.getZKDatabase().getDataTree().getWatches();
+    Assert.assertEquals(1, report.getPaths(0).size());
+  }
 }