From d66a06e8183818c035bb78706f46fd62645db07e Mon Sep 17 00:00:00 2001
From: Binbin <binloveplay1314@qq.com>
Date: Tue, 27 Aug 2024 12:04:27 +0800
Subject: [PATCH] Make KEYS to be an exact match if there is no pattern (#792)

Although KEYS is a dangerous command and we recommend people
to avoid using it, some people who are not familiar with it
still using it, and even use KEYS with no pattern at all.

Once KEYS is using with no pattern, we can convert it to an
exact match to avoid iterating over all data.

Signed-off-by: Binbin <binloveplay1314@qq.com>
---
 src/db.c                | 35 ++++++++++++++++++++++++++
 tests/unit/keyspace.tcl | 55 ++++++++++++++++++++++++++++++++++++++++-
 2 files changed, 89 insertions(+), 1 deletion(-)

diff --git a/src/db.c b/src/db.c
index 4a2c0a495e..1a57aa3c44 100644
--- a/src/db.c
+++ b/src/db.c
@@ -808,6 +808,21 @@ void randomkeyCommand(client *c) {
     decrRefCount(key);
 }
 
+/* Returns 1 if the pattern can be an exact match in KEYS context. */
+int patternExactMatch(const char *pattern, int length) {
+    for (int i = 0; i < length; i++) {
+        if (pattern[i] == '*' || pattern[i] == '?' || pattern[i] == '[') {
+            /* Wildcard or character class found. Keys can be in anywhere. */
+            return 0;
+        } else if (pattern[i] == '\\') {
+            /* Escaped character. Computing the key name in this case is not
+             * implemented. We would need a temp buffer. */
+            return 0;
+        }
+    }
+    return 1;
+}
+
 void keysCommand(client *c) {
     dictEntry *de;
     sds pattern = c->argv[1]->ptr;
@@ -817,7 +832,27 @@ void keysCommand(client *c) {
     allkeys = (pattern[0] == '*' && plen == 1);
     if (server.cluster_enabled && !allkeys) {
         pslot = patternHashSlot(pattern, plen);
+    } else if (!server.cluster_enabled) {
+        pslot = 0;
     }
+
+    /* Once the pattern can do an exact match, we can convert
+     * it to a kvstoreDictFind to avoid iterating over all data. */
+    if (patternExactMatch(pattern, plen) && pslot != -1) {
+        de = kvstoreDictFind(c->db->keys, pslot, pattern);
+        if (de) {
+            robj keyobj;
+            sds key = dictGetKey(de);
+            initStaticStringObject(keyobj, key);
+            if (!keyIsExpired(c->db, &keyobj)) {
+                addReplyBulkCBuffer(c, key, sdslen(key));
+                numkeys++;
+            }
+        }
+        setDeferredArrayLen(c, replylen, numkeys);
+        return;
+    }
+
     kvstoreDictIterator *kvs_di = NULL;
     kvstoreIterator *kvs_it = NULL;
     if (pslot != -1) {
diff --git a/tests/unit/keyspace.tcl b/tests/unit/keyspace.tcl
index c8bdcac984..fd2a622c8e 100644
--- a/tests/unit/keyspace.tcl
+++ b/tests/unit/keyspace.tcl
@@ -40,7 +40,25 @@ start_server {tags {"keyspace"}} {
         }
         assert_equal [lsort [r keys "{a}*"]] [list "{a}x" "{a}y" "{a}z"]
         assert_equal [lsort [r keys "*{b}*"]] [list "{b}a" "{b}b" "{b}c"]
-    } 
+    }
+
+    test {KEYS with no pattern} {
+        r debug populate 1000
+        r set foo bar
+        r set foo{t} bar
+        r set foo\\ bar
+        r del non-exist
+        assert_equal [r keys foo] {foo}
+        assert_equal [r keys foo{t}] {foo{t}}
+        assert_equal [r keys foo\\] {foo\\}
+        assert_equal [r keys non-exist] {}
+
+        # Make sure the return value is consistent with the * one
+        assert_equal [r keys foo] [r keys *foo]
+        assert_equal [r keys foo{t}] [r keys *foo{t}]
+        assert_equal [r keys foo\\] [r keys *foo\\]
+        assert_equal [r keys non-exist] [r keys *non-exist]
+    } {} {needs:debug}
 
     test {DEL all keys} {
         foreach key [r keys *] {r del $key}
@@ -548,3 +566,38 @@ foreach {type large} [array get largevalue] {
        r flushall
     } {OK} {singledb:skip}
 }
+
+start_cluster 1 0 {tags {"keyspace external:skip cluster"}} {
+    # SET keys for random slots, for random noise.
+    set num_keys 0
+    while {$num_keys < 1000} {
+        set random_key [randomInt 163840]
+        r SET $random_key VALUE
+        incr num_keys 1
+    }
+
+    test {KEYS with hashtag in cluster mode} {
+        foreach key {"{a}x" "{a}y" "{a}z" "{b}a" "{b}b" "{b}c"} {
+            r set $key hello
+        }
+        assert_equal [lsort [r keys "{a}*"]] [list "{a}x" "{a}y" "{a}z"]
+        assert_equal [lsort [r keys "*{b}*"]] [list "{b}a" "{b}b" "{b}c"]
+    }
+
+    test {KEYS with no pattern in cluster mode} {
+        r set foo bar
+        r set foo{t} bar
+        r set foo\\ bar
+        r del non-exist
+        assert_equal [r keys foo] {foo}
+        assert_equal [r keys foo{t}] {foo{t}}
+        assert_equal [r keys foo\\] {foo\\}
+        assert_equal [r keys non-exist] {}
+
+        # Make sure the return value is consistent with the * one
+        assert_equal [r keys foo] [r keys *foo]
+        assert_equal [r keys foo{t}] [r keys *foo{t}]
+        assert_equal [r keys foo\\] [r keys *foo\\]
+        assert_equal [r keys non-exist] [r keys *non-exist]
+    }
+}