ZOOKEEPER-3969: Add whoami API and Cli command

Author: Mohammad Arshad <arshad@apache.org>

Reviewers: Enrico Olivelli <eolivelli@apache.org>, Damien Diederen <dd@crosstwine.com>, Mate Szalay-Beko <symat@apache.org>

Closes #1504 from arshadmohammad/whoami-master

zookeeper-jute/src/main/resources/zookeeper.jute

@@ -51,6 +51,11 @@ module org.apache.zookeeper.data {
         long ephemeralOwner; // owner id if ephemeral, 0 otw
         long pzxid;      // last modified children
     }
+
+    class ClientInfo {
+        ustring authScheme; // Authentication scheme
+        ustring user;       // user name or any other id(for example ip)
+    }
 }
 
 module org.apache.zookeeper.proto {
@@ -248,6 +253,11 @@ module org.apache.zookeeper.proto {
     class GetEphemeralsResponse {
         vector<ustring> ephemerals;
     }
+
+    class WhoAmIResponse {
+        vector<org.apache.zookeeper.data.ClientInfo> clientInfo;
+    }
+
 }
 
 module org.apache.zookeeper.server.quorum {

zookeeper-server/src/main/java/org/apache/zookeeper/ZooDefs.java

@@ -93,6 +93,8 @@ public class ZooDefs {
 
         int addWatch = 106;
 
+        int whoAmI = 107;
+
         int createSession = -10;
 
         int closeSession = -11;

zookeeper-server/src/main/java/org/apache/zookeeper/ZooKeeper.java

@@ -49,6 +49,7 @@ import org.apache.zookeeper.client.ZKClientConfig;
 import org.apache.zookeeper.client.ZooKeeperSaslClient;
 import org.apache.zookeeper.common.PathUtils;
 import org.apache.zookeeper.data.ACL;
+import org.apache.zookeeper.data.ClientInfo;
 import org.apache.zookeeper.data.Stat;
 import org.apache.zookeeper.proto.AddWatchRequest;
 import org.apache.zookeeper.proto.CheckWatchesRequest;
@@ -80,6 +81,7 @@ import org.apache.zookeeper.proto.SetDataRequest;
 import org.apache.zookeeper.proto.SetDataResponse;
 import org.apache.zookeeper.proto.SyncRequest;
 import org.apache.zookeeper.proto.SyncResponse;
+import org.apache.zookeeper.proto.WhoAmIResponse;
 import org.apache.zookeeper.server.DataTree;
 import org.apache.zookeeper.server.EphemeralType;
 import org.slf4j.Logger;
@@ -3070,4 +3072,18 @@ public class ZooKeeper implements AutoCloseable {
         }
     }
 
+    /**
+     * Gives all authentication information added into the current session.
+     *
+     * @return list of authentication info
+     * @throws InterruptedException when interrupted
+     */
+    public synchronized List<ClientInfo> whoAmI() throws InterruptedException {
+        RequestHeader h = new RequestHeader();
+        h.setType(ZooDefs.OpCode.whoAmI);
+        WhoAmIResponse response = new WhoAmIResponse();
+        cnxn.submitRequest(h, null, response, null);
+        return response.getClientInfo();
+    }
+
 }

zookeeper-server/src/main/java/org/apache/zookeeper/cli/CommandFactory.java

@@ -50,7 +50,8 @@ public class CommandFactory {
         GET_EPHEMERALS(GetEphemeralsCommand::new),
         GET_ALL_CHILDREN_NUMBER(GetAllChildrenNumberCommand::new),
         VERSION(VersionCommand::new),
-        ADD_WATCH(AddWatchCommand::new);
+        ADD_WATCH(AddWatchCommand::new),
+        WHO_AM_I(WhoAmICommand::new);
 
         private Supplier<? extends CliCommand> instantiator;
 

zookeeper-server/src/main/java/org/apache/zookeeper/cli/WhoAmICommand.java

@@ -0,0 +1,55 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.zookeeper.cli;
+
+import java.util.List;
+import org.apache.zookeeper.data.ClientInfo;
+
+/**
+ * WhoAmI command for cli
+ */
+public class WhoAmICommand extends CliCommand {
+
+    public WhoAmICommand() {
+        super("whoami", "");
+    }
+
+    @Override
+    public CliCommand parse(String[] cmdArgs) throws CliParseException {
+        return this;
+    }
+
+    @Override
+    public boolean exec() throws CliException {
+        try {
+            List<ClientInfo> clientInfos = zk.whoAmI();
+            out.println("Auth scheme: User");
+            if (clientInfos != null) {
+                // clientInfos will never be null, added null check to pass static checks
+                clientInfos.forEach(clientInfo -> {
+                    out.println(clientInfo.getAuthScheme() + ": " + clientInfo.getUser());
+                });
+            }
+        } catch (Exception ex) {
+            throw new CliWrapperException(ex);
+        }
+        return false;
+    }
+
+}

zookeeper-server/src/main/java/org/apache/zookeeper/server/FinalRequestProcessor.java

@@ -76,8 +76,10 @@ import org.apache.zookeeper.proto.SetWatches;
 import org.apache.zookeeper.proto.SetWatches2;
 import org.apache.zookeeper.proto.SyncRequest;
 import org.apache.zookeeper.proto.SyncResponse;
+import org.apache.zookeeper.proto.WhoAmIResponse;
 import org.apache.zookeeper.server.DataTree.ProcessTxnResult;
 import org.apache.zookeeper.server.quorum.QuorumZooKeeperServer;
+import org.apache.zookeeper.server.util.AuthUtil;
 import org.apache.zookeeper.server.util.RequestPathMetricsCollector;
 import org.apache.zookeeper.txn.ErrorTxn;
 import org.slf4j.Logger;
@@ -546,6 +548,11 @@ public class FinalRequestProcessor implements RequestProcessor {
                 requestPathMetricsCollector.registerRequest(request.type, removeWatches.getPath());
                 break;
             }
+            case OpCode.whoAmI: {
+                lastOp = "HOMI";
+                rsp = new WhoAmIResponse(AuthUtil.getClientInfos(request.authInfo));
+                break;
+             }
             case OpCode.getEphemerals: {
                 lastOp = "GETE";
                 GetEphemeralsRequest getEphemerals = new GetEphemeralsRequest();

zookeeper-server/src/main/java/org/apache/zookeeper/server/PrepRequestProcessor.java

@@ -915,6 +915,7 @@ public class PrepRequestProcessor extends ZooKeeperCriticalThread implements Req
             case OpCode.getEphemerals:
             case OpCode.multiRead:
             case OpCode.addWatch:
+            case OpCode.whoAmI:
                 zks.sessionTracker.checkSession(request.sessionId, request.getOwner());
                 break;
             default:

zookeeper-server/src/main/java/org/apache/zookeeper/server/Request.java

@@ -271,6 +271,7 @@ public class Request {
         case OpCode.checkWatches:
         case OpCode.removeWatches:
         case OpCode.addWatch:
+        case OpCode.whoAmI:
             return true;
         default:
             return false;
@@ -287,6 +288,7 @@ public class Request {
         case OpCode.getData:
         case OpCode.getEphemerals:
         case OpCode.multiRead:
+        case OpCode.whoAmI:
             return false;
         case OpCode.create:
         case OpCode.create2:
@@ -373,6 +375,8 @@ public class Request {
                 return "closeSession";
             case OpCode.error:
                 return "error";
+            case OpCode.whoAmI:
+                return "whoAmI";
             default:
                 return "unknown " + op;
         }

zookeeper-server/src/main/java/org/apache/zookeeper/server/util/AuthUtil.java

@@ -17,6 +17,9 @@
  */
 package org.apache.zookeeper.server.util;
 
+import java.util.ArrayList;
+import java.util.List;
+import org.apache.zookeeper.data.ClientInfo;
 import org.apache.zookeeper.data.Id;
 import org.apache.zookeeper.server.auth.AuthenticationProvider;
 import org.apache.zookeeper.server.auth.ProviderRegistry;
@@ -36,4 +39,19 @@ public final class AuthUtil {
         AuthenticationProvider provider = ProviderRegistry.getProvider(id.getScheme());
         return provider == null ? null : provider.getUserName(id.getId());
     }
+
+    /**
+     * Gets user from id to prepare ClientInfo.
+     *
+     * @param authInfo List of id objects. id contains scheme and authentication info
+     * @return list of client authentication info
+     */
+    public static List<ClientInfo> getClientInfos(List<Id> authInfo) {
+        List<ClientInfo> clientAuthInfo = new ArrayList<>(authInfo.size());
+        authInfo.forEach(id -> {
+            String user = AuthUtil.getUser(id);
+            clientAuthInfo.add(new ClientInfo(id.getScheme(), user == null ? "" : user));
+        });
+        return clientAuthInfo;
+    }
 }

zookeeper-server/src/test/java/org/apache/zookeeper/ZooKeeperTest.java

@@ -39,12 +39,14 @@ import org.apache.zookeeper.cli.LsCommand;
 import org.apache.zookeeper.cli.MalformedCommandException;
 import org.apache.zookeeper.cli.MalformedPathException;
 import org.apache.zookeeper.cli.SyncCommand;
+import org.apache.zookeeper.cli.WhoAmICommand;
 import org.apache.zookeeper.client.ConnectStringParser;
 import org.apache.zookeeper.client.HostProvider;
 import org.apache.zookeeper.client.StaticHostProvider;
 import org.apache.zookeeper.client.ZKClientConfig;
 import org.apache.zookeeper.common.StringUtils;
 import org.apache.zookeeper.data.ACL;
+import org.apache.zookeeper.data.ClientInfo;
 import org.apache.zookeeper.data.Id;
 import org.apache.zookeeper.data.Stat;
 import org.apache.zookeeper.test.ClientBase;
@@ -517,14 +519,18 @@ public class ZooKeeperTest extends ClientBase {
     }
 
     private static void runCommandExpect(CliCommand command, List<String> expectedResults) throws Exception {
+        String result = runCommandExpect(command);
+        assertTrue(result.contains(StringUtils.joinStrings(expectedResults, LINE_SEPARATOR)), result);
+    }
+
+    private static String runCommandExpect(CliCommand command) throws CliException {
         // call command and put result in byteStream
         ByteArrayOutputStream byteStream = new ByteArrayOutputStream();
         PrintStream out = new PrintStream(byteStream);
         command.setOut(out);
         command.exec();
 
-        String result = byteStream.toString();
-        assertTrue(result.contains(StringUtils.joinStrings(expectedResults, LINE_SEPARATOR)), result);
+        return byteStream.toString();
     }
 
     @Test
@@ -698,4 +704,74 @@ public class ZooKeeperTest extends ClientBase {
         assertEquals("Insufficient permission : " + zNodeToBeCreated, errorMessage);
     }
 
+    @Test
+    public void testWhoAmIAPI() throws Exception {
+        final ZooKeeper zk = createClient();
+
+        // Check who ami without authentication/without any user into the session
+        List<ClientInfo> clientInfos = zk.whoAmI();
+        // By default server adds ip as the authentication info
+        assertEquals(1, clientInfos.size());
+        assertEquals("ip", clientInfos.get(0).getAuthScheme());
+
+        // Add one user into the session
+        zk.addAuthInfo("digest", "user1:abcXYZ".getBytes());
+        clientInfos = zk.whoAmI();
+        assertEquals(2, clientInfos.size());
+        ClientInfo user1 = getClientInfos(clientInfos, "user1");
+        assertEquals("digest", user1.getAuthScheme());
+
+        // Add one more user into the session
+        zk.addAuthInfo("digest", "user2:xyzABC".getBytes());
+        clientInfos = zk.whoAmI();
+        assertEquals(3, clientInfos.size());
+        user1 = getClientInfos(clientInfos, "user1");
+        assertEquals("digest", user1.getAuthScheme());
+        ClientInfo user2 = getClientInfos(clientInfos, "user2");
+        assertEquals("digest", user2.getAuthScheme());
+    }
+
+    private ClientInfo getClientInfos(List<ClientInfo> clientInfos, String user) {
+        for (ClientInfo clientInfo : clientInfos) {
+            if (clientInfo.getUser().equals(user)) {
+                return clientInfo;
+            }
+        }
+        throw new AssertionError("User +" + user + " not found");
+    }
+
+    @Test
+    public void testWhoAmICLICommand() throws Exception {
+        final ZooKeeper zk = createClient();
+        WhoAmICommand cmd = new WhoAmICommand();
+        cmd.setZk(zk);
+        List<String> expectedResults = new ArrayList<>();
+        expectedResults.add("Auth scheme: User");
+        expectedResults.add("ip: 127.0.0.1");
+
+        // Check who ami without authentication/without any user into the session
+        cmd.parse(new String[] { "whoami" });
+        String actualResult = runCommandExpect(cmd);
+        assertClientAuthInfo(expectedResults, actualResult);
+
+        // Add one user into the session
+        zk.addAuthInfo("digest", "user1:abcXYZ".getBytes());
+        expectedResults.add("digest: user1");
+        actualResult = runCommandExpect(cmd);
+        assertClientAuthInfo(expectedResults, actualResult);
+
+        // Add one more user into the session
+        zk.addAuthInfo("digest", "user2:xyzABC".getBytes());
+        expectedResults.add("digest: user2");
+        actualResult = runCommandExpect(cmd);
+        assertClientAuthInfo(expectedResults, actualResult);
+    }
+
+    private void assertClientAuthInfo(List<String> expected, String actual) {
+        expected.forEach(s -> {
+            assertTrue(actual.contains(s),
+                "Expected result part '" + s + "' not present in actual result '" + actual + "' ");
+        });
+    }
+
 }