add user-client mapping.

valkey-io · Apr 30, 2024 · a4b222e · a4b222e
1 parent b0d5a0f
commit a4b222e
Show file tree

Hide file tree

Showing 7 changed files with 127 additions and 5 deletions.
diff --git a/src/acl.c b/src/acl.c
@@ -456,6 +456,8 @@ user *ACLCreateUser(const char *name, size_t namelen) {
     aclSelector *s = ACLCreateSelector(SELECTOR_FLAG_ROOT);
     listAddNodeHead(u->selectors, s);
 
+    u->clients = listCreate();
+
     raxInsert(Users,(unsigned char*)name,namelen,u,NULL);
     return u;
 }
@@ -487,6 +489,7 @@ void ACLFreeUser(user *u) {
     }
     listRelease(u->passwords);
     listRelease(u->selectors);
+    listRelease(u->clients);
     zfree(u);
 }
 
@@ -3241,6 +3244,11 @@ void authCommand(client *c) {
     robj *err = NULL;
     int result = ACLAuthenticateUser(c, username, password, &err);
     if (result == AUTH_OK) {
+        listNode *ln=listSearchKey(c->user->clients, c);
+        if(ln == NULL) {
+            listAddNodeTail(c->user->clients, c);
+            c->user_client_node = listLast(c->user->clients);
+        }
         addReply(c, shared.ok);
     } else if (result == AUTH_ERR) {
         addAuthErrReply(c, err);

diff --git a/src/commands.def b/src/commands.def
@@ -1080,6 +1080,23 @@ struct COMMAND_ARG CLIENT_CACHING_Args[] = {
 {MAKE_ARG("mode",ARG_TYPE_ONEOF,-1,NULL,NULL,NULL,CMD_ARG_NONE,2,NULL),.subargs=CLIENT_CACHING_mode_Subargs},
 };
 
+/********** CLIENT COUNT ********************/
+
+#ifndef SKIP_CMD_HISTORY_TABLE
+/* CLIENT COUNT history */
+#define CLIENT_COUNT_History NULL
+#endif
+
+#ifndef SKIP_CMD_TIPS_TABLE
+/* CLIENT COUNT tips */
+#define CLIENT_COUNT_Tips NULL
+#endif
+
+#ifndef SKIP_CMD_KEY_SPECS_TABLE
+/* CLIENT COUNT key specs */
+#define CLIENT_COUNT_Keyspecs NULL
+#endif
+
 /********** CLIENT GETNAME ********************/
 
 #ifndef SKIP_CMD_HISTORY_TABLE
@@ -1266,6 +1283,7 @@ struct COMMAND_ARG CLIENT_LIST_client_type_Subargs[] = {
 struct COMMAND_ARG CLIENT_LIST_Args[] = {
 {MAKE_ARG("client-type",ARG_TYPE_ONEOF,-1,"TYPE",NULL,"5.0.0",CMD_ARG_OPTIONAL,4,NULL),.subargs=CLIENT_LIST_client_type_Subargs},
 {MAKE_ARG("client-id",ARG_TYPE_INTEGER,-1,"ID",NULL,"6.2.0",CMD_ARG_OPTIONAL|CMD_ARG_MULTIPLE,0,NULL)},
+{MAKE_ARG("user",ARG_TYPE_STRING,-1,NULL,NULL,"7.2.5",CMD_ARG_OPTIONAL,0,NULL)},
 };
 
 /********** CLIENT NO_EVICT ********************/
@@ -1540,13 +1558,14 @@ struct COMMAND_ARG CLIENT_UNBLOCK_Args[] = {
 /* CLIENT command table */
 struct COMMAND_STRUCT CLIENT_Subcommands[] = {
 {MAKE_CMD("caching","Instructs the server whether to track the keys in the next request.","O(1)","6.0.0",CMD_DOC_NONE,NULL,NULL,"connection",COMMAND_GROUP_CONNECTION,CLIENT_CACHING_History,0,CLIENT_CACHING_Tips,0,clientCommand,3,CMD_NOSCRIPT|CMD_LOADING|CMD_STALE|CMD_SENTINEL,ACL_CATEGORY_CONNECTION,CLIENT_CACHING_Keyspecs,0,NULL,1),.args=CLIENT_CACHING_Args},
+{MAKE_CMD("count","Returns client count among the user.","O(1)","7.2.5",CMD_DOC_NONE,NULL,NULL,"connection",COMMAND_GROUP_CONNECTION,CLIENT_COUNT_History,0,CLIENT_COUNT_Tips,0,clientCommand,3,CMD_ADMIN|CMD_NOSCRIPT|CMD_LOADING|CMD_STALE|CMD_SENTINEL,ACL_CATEGORY_CONNECTION,CLIENT_COUNT_Keyspecs,0,NULL,0)},
 {MAKE_CMD("getname","Returns the name of the connection.","O(1)","2.6.9",CMD_DOC_NONE,NULL,NULL,"connection",COMMAND_GROUP_CONNECTION,CLIENT_GETNAME_History,0,CLIENT_GETNAME_Tips,0,clientCommand,2,CMD_NOSCRIPT|CMD_LOADING|CMD_STALE|CMD_SENTINEL,ACL_CATEGORY_CONNECTION,CLIENT_GETNAME_Keyspecs,0,NULL,0)},
 {MAKE_CMD("getredir","Returns the client ID to which the connection's tracking notifications are redirected.","O(1)","6.0.0",CMD_DOC_NONE,NULL,NULL,"connection",COMMAND_GROUP_CONNECTION,CLIENT_GETREDIR_History,0,CLIENT_GETREDIR_Tips,0,clientCommand,2,CMD_NOSCRIPT|CMD_LOADING|CMD_STALE|CMD_SENTINEL,ACL_CATEGORY_CONNECTION,CLIENT_GETREDIR_Keyspecs,0,NULL,0)},
 {MAKE_CMD("help","Returns helpful text about the different subcommands.","O(1)","5.0.0",CMD_DOC_NONE,NULL,NULL,"connection",COMMAND_GROUP_CONNECTION,CLIENT_HELP_History,0,CLIENT_HELP_Tips,0,clientCommand,2,CMD_LOADING|CMD_STALE|CMD_SENTINEL,ACL_CATEGORY_CONNECTION,CLIENT_HELP_Keyspecs,0,NULL,0)},
 {MAKE_CMD("id","Returns the unique client ID of the connection.","O(1)","5.0.0",CMD_DOC_NONE,NULL,NULL,"connection",COMMAND_GROUP_CONNECTION,CLIENT_ID_History,0,CLIENT_ID_Tips,0,clientCommand,2,CMD_NOSCRIPT|CMD_LOADING|CMD_STALE|CMD_SENTINEL,ACL_CATEGORY_CONNECTION,CLIENT_ID_Keyspecs,0,NULL,0)},
 {MAKE_CMD("info","Returns information about the connection.","O(1)","6.2.0",CMD_DOC_NONE,NULL,NULL,"connection",COMMAND_GROUP_CONNECTION,CLIENT_INFO_History,0,CLIENT_INFO_Tips,1,clientCommand,2,CMD_NOSCRIPT|CMD_LOADING|CMD_STALE|CMD_SENTINEL,ACL_CATEGORY_CONNECTION,CLIENT_INFO_Keyspecs,0,NULL,0)},
 {MAKE_CMD("kill","Terminates open connections.","O(N) where N is the number of client connections","2.4.0",CMD_DOC_NONE,NULL,NULL,"connection",COMMAND_GROUP_CONNECTION,CLIENT_KILL_History,6,CLIENT_KILL_Tips,0,clientCommand,-3,CMD_ADMIN|CMD_NOSCRIPT|CMD_LOADING|CMD_STALE|CMD_SENTINEL,ACL_CATEGORY_CONNECTION,CLIENT_KILL_Keyspecs,0,NULL,1),.args=CLIENT_KILL_Args},
-{MAKE_CMD("list","Lists open connections.","O(N) where N is the number of client connections","2.4.0",CMD_DOC_NONE,NULL,NULL,"connection",COMMAND_GROUP_CONNECTION,CLIENT_LIST_History,6,CLIENT_LIST_Tips,1,clientCommand,-2,CMD_ADMIN|CMD_NOSCRIPT|CMD_LOADING|CMD_STALE|CMD_SENTINEL,ACL_CATEGORY_CONNECTION,CLIENT_LIST_Keyspecs,0,NULL,2),.args=CLIENT_LIST_Args},
+{MAKE_CMD("list","Lists open connections.","O(N) where N is the number of client connections","2.4.0",CMD_DOC_NONE,NULL,NULL,"connection",COMMAND_GROUP_CONNECTION,CLIENT_LIST_History,6,CLIENT_LIST_Tips,1,clientCommand,-2,CMD_ADMIN|CMD_NOSCRIPT|CMD_LOADING|CMD_STALE|CMD_SENTINEL,ACL_CATEGORY_CONNECTION,CLIENT_LIST_Keyspecs,0,NULL,3),.args=CLIENT_LIST_Args},
 {MAKE_CMD("no-evict","Sets the client eviction mode of the connection.","O(1)","7.0.0",CMD_DOC_NONE,NULL,NULL,"connection",COMMAND_GROUP_CONNECTION,CLIENT_NO_EVICT_History,0,CLIENT_NO_EVICT_Tips,0,clientCommand,3,CMD_ADMIN|CMD_NOSCRIPT|CMD_LOADING|CMD_STALE|CMD_SENTINEL,ACL_CATEGORY_CONNECTION,CLIENT_NO_EVICT_Keyspecs,0,NULL,1),.args=CLIENT_NO_EVICT_Args},
 {MAKE_CMD("no-touch","Controls whether commands sent by the client affect the LRU/LFU of accessed keys.","O(1)","7.2.0",CMD_DOC_NONE,NULL,NULL,"connection",COMMAND_GROUP_CONNECTION,CLIENT_NO_TOUCH_History,0,CLIENT_NO_TOUCH_Tips,0,clientCommand,3,CMD_NOSCRIPT|CMD_LOADING|CMD_STALE,ACL_CATEGORY_CONNECTION,CLIENT_NO_TOUCH_Keyspecs,0,NULL,1),.args=CLIENT_NO_TOUCH_Args},
 {MAKE_CMD("pause","Suspends commands processing.","O(1)","3.0.0",CMD_DOC_NONE,NULL,NULL,"connection",COMMAND_GROUP_CONNECTION,CLIENT_PAUSE_History,1,CLIENT_PAUSE_Tips,0,clientCommand,-3,CMD_ADMIN|CMD_NOSCRIPT|CMD_LOADING|CMD_STALE|CMD_SENTINEL,ACL_CATEGORY_CONNECTION,CLIENT_PAUSE_Keyspecs,0,NULL,2),.args=CLIENT_PAUSE_Args},

diff --git a/src/commands/client-count.json b/src/commands/client-count.json
@@ -0,0 +1,25 @@
+{
+  "COUNT": {
+    "summary": "Returns client count among the user.",
+    "complexity": "O(1)",
+    "group": "connection",
+    "since": "7.2.5",
+    "arity": 3,
+    "container": "CLIENT",
+    "function": "clientCommand",
+    "command_flags": [
+      "ADMIN",
+      "NOSCRIPT",
+      "LOADING",
+      "STALE",
+      "SENTINEL"
+    ],
+    "acl_categories": [
+      "CONNECTION"
+    ],
+    "reply_schema": {
+      "type": "string",
+      "description": "client count among user"
+    }
+  }
+}
diff --git a/src/commands/client-list.json b/src/commands/client-list.json
@@ -87,6 +87,13 @@
                 "optional": true,
                 "multiple": true,
                 "since": "6.2.0"
+            },
+            {
+                "name": "user",
+                "type": "string",
+                "optional": true,
+                "multiple":false,
+                "since": "7.2.5"
             }
         ]
     }

diff --git a/src/networking.c b/src/networking.c
@@ -198,6 +198,7 @@ client *createClient(connection *conn) {
     c->peerid = NULL;
     c->sockname = NULL;
     c->client_list_node = NULL;
+    c->user_client_node = NULL;
     c->postponed_list_node = NULL;
     c->pending_read_list_node = NULL;
     c->client_tracking_redirection = 0;
@@ -1538,6 +1539,10 @@ void clearClientConnectionState(client *c) {
 #else
     c->resp = 2;
 #endif
+    if(c->user && c->user != DefaultUser && c->user_client_node){
+        listDelNode(c->user->clients,c->user_client_node);
+        c->user_client_node = NULL;
+    }
 
     clientSetDefaultAuth(c);
     moduleNotifyUserChanged(c);
@@ -3012,6 +3017,21 @@ void quitCommand(client *c) {
     c->flags |= CLIENT_CLOSE_AFTER_REPLY;
 }
 
+sds getUserClientsInfoString(user *user) {
+    listNode *ln;
+    listIter li;
+    client *client;
+    sds o = sdsnewlen(SDS_NOINIT, 200 * listLength(server.clients));
+    sdsclear(o);
+    listRewind(user->clients, &li);
+    while ((ln = listNext(&li)) != NULL) {
+        client = listNodeValue(ln);
+        o = catClientInfoString(o, client);
+        o = sdscatlen(o, "\n", 1);
+    }
+    return o;
+}
+
 void clientCommand(client *c) {
     listNode *ln;
     listIter li;
@@ -3050,6 +3070,10 @@ void clientCommand(client *c) {
 "    Return information about client connections. Options:",
 "    * TYPE (NORMAL|MASTER|REPLICA|PUBSUB)",
 "      Return clients of specified type.",
+"    * USER <username>",
+"      Return clients of specified user.",
+"COUNT <username>",
+"    Return the count of clients connected with the specified username.",
 "UNPAUSE",
 "    Stop the current client pause, resuming traffic.",
 "PAUSE <timeout> [WRITE|ALL]",
@@ -3089,6 +3113,7 @@ NULL
         /* CLIENT LIST */
         int type = -1;
         sds o = NULL;
+        user *user = NULL;
         if (c->argc == 4 && !strcasecmp(c->argv[2]->ptr,"type")) {
             type = getClientTypeByName(c->argv[3]->ptr);
             if (type == -1) {
@@ -3112,6 +3137,16 @@ NULL
                     o = sdscatlen(o, "\n", 1);
                 }
             }
+        } else if(c->argc == 4 && !strcasecmp(c->argv[2]->ptr, "user")) {
+            o = sdsempty();
+            user = ACLGetUserByName(c->argv[3]->ptr, sdslen(c->argv[3]->ptr));
+            if (user == NULL) {
+                addReplyErrorFormat(c, "No such user '%s'",
+                                    (char *) c->argv[3]->ptr);
+                sdsfree(o);
+                return;
+            }
+            o = getUserClientsInfoString(user);
         } else if (c->argc != 2) {
             addReplyErrorObject(c,shared.syntaxerr);
             return;
@@ -3264,6 +3299,23 @@ NULL
         /* If this client has to be closed, flag it as CLOSE_AFTER_REPLY
          * only after we queued the reply to its output buffers. */
         if (close_this_client) c->flags |= CLIENT_CLOSE_AFTER_REPLY;
+    } else if (!strcasecmp(c->argv[1]->ptr, "count")) {
+        user *user = NULL;
+        if (c->argc == 3) {
+            user = ACLGetUserByName(c->argv[2]->ptr,
+                                    sdslen(c->argv[2]->ptr));
+            if (user == NULL) {
+                addReplyErrorFormat(c, "No such user '%s'",
+                                    (char *) c->argv[2]->ptr);
+                return;
+            }
+            sds o = sdscatprintf(sdsempty(), "%lu", listLength(user->clients));
+            addReplyVerbatim(c, o, sdslen(o), "txt");
+            sdsfree(o);
+        } else {
+            addReplyErrorObject(c, shared.syntaxerr);
+            return;
+        }
     } else if (!strcasecmp(c->argv[1]->ptr,"unblock") && (c->argc == 3 ||
                                                           c->argc == 4))
     {

diff --git a/src/server.h b/src/server.h
@@ -1102,6 +1102,7 @@ typedef struct {
                         against. This list will always contain at least
                         one selector for backwards compatibility. */
     robj *acl_string; /* cached string represent of ACLs */
+    list *clients;  /* A list of clients associated with this user. */
 } user;
 
 /* With multiplexing we need to take per-client state.
@@ -1230,6 +1231,7 @@ typedef struct client {
     sds peerid;             /* Cached peer ID. */
     sds sockname;           /* Cached connection target address. */
     listNode *client_list_node; /* list node in client list */
+    listNode *user_client_node; /* list node in user->client list */
     listNode *postponed_list_node; /* list node within the postponed list */
     listNode *pending_read_list_node; /* list node in clients pending read list */
     void *module_blocked_client; /* Pointer to the ValkeyModuleBlockedClient associated with this

diff --git a/tests/unit/acl.tcl b/tests/unit/acl.tcl
@@ -188,13 +188,22 @@ start_server {tags {"acl external:skip"}} {
 
         assert_equal {0} [r PUBLISH foo bar]
         catch {r PUBLISH bar game} e
-
-        # Falling back to psuser for the below tests
-        r AUTH psuser pspass
-        r ACL deluser hpuser
         set e
     } {*NOPERM*channel*}
 
+    test {client list users} {
+        set rd [valkey_deferring_client]
+        $rd AUTH psuser pspass
+        $rd read
+        r AUTH hpuser pass
+        $rd client list
+        assert {[llength [$rd read]] >= 2}
+        $rd client list user psuser
+        assert {[llength [$rd read]] >= 1}
+        assert_equal 1 [r client count hpuser]
+        r AUTH psuser pspass
+        r ACL deluser hpuser
+    }
     test {In transaction queue publish/subscribe/psubscribe to unauthorized channel will fail} {
         r ACL setuser psuser +multi +discard
         r MULTI