From b4123663c31aaf1e97f4dbd630cbd7b8d0e91e31 Mon Sep 17 00:00:00 2001
From: Oran Agra <oran@redislabs.com>
Date: Mon, 16 Jan 2023 13:51:18 +0200
Subject: [PATCH] Obuf limit, exit during loop in *RAND* commands and KEYS
 (#11676)

Related to the hang reported in #11671
Currently, redis can disconnect a client due to reaching output buffer limit,
it'll also avoid feeding that output buffer with more data, but it will keep
running the loop in the command (despite the client already being marked for
disconnection)

This PR is an attempt to mitigate the problem, specifically for commands that
are easy to abuse, specifically: KEYS, HRANDFIELD, SRANDMEMBER, ZRANDMEMBER.
The RAND family of commands can take a negative COUNT argument (which is not
bound to the number of elements in the key), so it's enough to create a key
with one field, and then these commands can be used to hang redis.
For KEYS the caller can use the existing keyspace in redis (if big enough).
---
 src/db.c                   |  2 ++
 src/t_hash.c               |  4 ++++
 src/t_set.c                |  4 ++++
 src/t_zset.c               |  4 ++++
 tests/unit/obuf-limits.tcl | 16 ++++++++++++++++
 5 files changed, 30 insertions(+)

diff --git a/src/db.c b/src/db.c
index 8511135c6..e397183c9 100644
--- a/src/db.c
+++ b/src/db.c
@@ -783,6 +783,8 @@ void keysCommand(client *c) {
             }
             decrRefCount(keyobj);
         }
+        if (c->flags & CLIENT_CLOSE_ASAP)
+            break;
     }
     dictReleaseIterator(di);
     setDeferredArrayLen(c,replylen,numkeys);
diff --git a/src/t_hash.c b/src/t_hash.c
index 9ebd99ccc..757ccd4b3 100644
--- a/src/t_hash.c
+++ b/src/t_hash.c
@@ -956,6 +956,8 @@ void hrandfieldWithCountCommand(client *c, long l, int withvalues) {
                 addReplyBulkCBuffer(c, key, sdslen(key));
                 if (withvalues)
                     addReplyBulkCBuffer(c, value, sdslen(value));
+                if (c->flags & CLIENT_CLOSE_ASAP)
+                    break;
             }
         } else if (hash->encoding == OBJ_ENCODING_LISTPACK) {
             listpackEntry *keys, *vals = NULL;
@@ -970,6 +972,8 @@ void hrandfieldWithCountCommand(client *c, long l, int withvalues) {
                 count -= sample_count;
                 lpRandomPairs(hash->ptr, sample_count, keys, vals);
                 hrandfieldReplyWithListpack(c, sample_count, keys, vals);
+                if (c->flags & CLIENT_CLOSE_ASAP)
+                    break;
             }
             zfree(keys);
             zfree(vals);
diff --git a/src/t_set.c b/src/t_set.c
index 6a4840676..6d8023f31 100644
--- a/src/t_set.c
+++ b/src/t_set.c
@@ -1028,6 +1028,8 @@ void srandmemberWithCountCommand(client *c) {
                     else
                         addReplyBulkLongLong(c, entries[i].lval);
                 }
+                if (c->flags & CLIENT_CLOSE_ASAP)
+                    break;
             }
             zfree(entries);
             return;
@@ -1040,6 +1042,8 @@ void srandmemberWithCountCommand(client *c) {
             } else {
                 addReplyBulkCBuffer(c, str, len);
             }
+            if (c->flags & CLIENT_CLOSE_ASAP)
+                break;
         }
         return;
     }
diff --git a/src/t_zset.c b/src/t_zset.c
index e0d8509cf..a5627f031 100644
--- a/src/t_zset.c
+++ b/src/t_zset.c
@@ -4154,6 +4154,8 @@ void zrandmemberWithCountCommand(client *c, long l, int withscores) {
                 addReplyBulkCBuffer(c, key, sdslen(key));
                 if (withscores)
                     addReplyDouble(c, *(double*)dictGetVal(de));
+                if (c->flags & CLIENT_CLOSE_ASAP)
+                    break;
             }
         } else if (zsetobj->encoding == OBJ_ENCODING_LISTPACK) {
             listpackEntry *keys, *vals = NULL;
@@ -4167,6 +4169,8 @@ void zrandmemberWithCountCommand(client *c, long l, int withscores) {
                 count -= sample_count;
                 lpRandomPairs(zsetobj->ptr, sample_count, keys, vals);
                 zrandmemberReplyWithListpack(c, sample_count, keys, vals);
+                if (c->flags & CLIENT_CLOSE_ASAP)
+                    break;
             }
             zfree(keys);
             zfree(vals);
diff --git a/tests/unit/obuf-limits.tcl b/tests/unit/obuf-limits.tcl
index 38048935a..7eb6def58 100644
--- a/tests/unit/obuf-limits.tcl
+++ b/tests/unit/obuf-limits.tcl
@@ -211,4 +211,20 @@ start_server {tags {"obuf-limits external:skip"}} {
         assert_equal "v2" [r get k2]
         assert_equal "v3" [r get k3]
     }
+
+    test "Obuf limit, HRANDFIELD with huge count stopped mid-run" {
+        r config set client-output-buffer-limit {normal 1000000 0 0}
+        r hset myhash a b
+        catch {r hrandfield myhash -999999999} e
+        assert_match "*I/O error*" $e
+        reconnect
+    }
+
+    test "Obuf limit, KEYS stopped mid-run" {
+        r config set client-output-buffer-limit {normal 100000 0 0}
+        populate 1000 "long-key-name-prefix-of-100-chars-------------------------------------------------------------------"
+        catch {r keys *} e
+        assert_match "*I/O error*" $e
+        reconnect
+    }
 }