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).

README.md

@@ -4,7 +4,7 @@
 
 ##### KeyDB is now a part of Snap Inc! Check out the announcement [here](https://docs.keydb.dev/news/2022/05/12/keydb-joins-snap) 
 
-##### [Release v6.3.0](https://github.com/EQ-Alpha/KeyDB/releases/tag/v6.3.0) is here with major improvements as we consolidate our Open Source and Enterprise offerings into a single BSD-3 licensed project. See our [roadmap](https://docs.keydb.dev/docs/coming-soon) for details. 
+##### [Release v6.3.0](https://github.com/EQ-Alpha/KeyDB/releases/tag/v6.3.0) is here with major improvements as we consolodate our Open Source and Enterprise offerings into a single BSD-3 licensed project. See our [roadmap](https://docs.keydb.dev/docs/coming-soon) for details. 
 
 ##### Want to extend KeyDB with Javascript?  Try [ModJS](https://github.com/JohnSully/ModJS)
 
@@ -37,7 +37,7 @@ Because of this difference of opinion features which are right for KeyDB may not
 Project Support
 -------------------
 
-The KeyDB team maintains this project as part of Snap Inc. KeyDB is used by Snap as part of its caching infrastructure and is fully open sourced. There is no separate commercial product and no paid support options available. We really value collaborating with the open source community and welcome PRs, bug reports, and open discussion. For community support or to get involved further with the project check out our community support options [here](https://docs.keydb.dev/docs/support) (slack, forum, meetup, github issues). Our team monitors these channels regularly.
+The KeyDB team maintains this project as part of Snap Inc. KeyDB is used by Snap as part of its caching infrastructure and is fully open sourced. There is no separate commercial product and no paid support options available. We really value collaborating with the open source community and welcome PRs, bug reports, and open discussion. For community support or to get involved further with the project check out our community support options [here](https://docs.keydb.dev/docs/support) (slack, forum, meetup, github issues). Our team monitors these channlels regularly.
 
 
 Additional Resources

src/sort.cpp

@@ -194,8 +194,8 @@ void sortCommand(client *c) {
     list *operations;
     unsigned int outputlen = 0;
     int desc = 0, alpha = 0;
-    long limit_start = 0, limit_count = -1, start, end, vectorlen;
-    int j, dontsort = 0;
+    long limit_start = 0, limit_count = -1, start, end;
+    int j, dontsort = 0, vectorlen;
     int getop = 0; /* GET operation counter */
     int int_conversion_error = 0;
     int syntax_error = 0;
@@ -323,8 +323,8 @@ void sortCommand(client *c) {
 
     /* Perform LIMIT start,count sanity checking.
      * And avoid integer overflow by limiting inputs to object sizes. */
-    start = std::min(std::max(limit_start, (long)0), vectorlen);
-    limit_count = std::min(std::max(limit_count, (long)-1), vectorlen);
+    start = min(max(limit_start, 0), vectorlen);
+    limit_count = min(max(limit_count, -1), vectorlen);
     end = (limit_count < 0) ? vectorlen-1 : start+limit_count-1;
     if (start >= vectorlen) {
         start = vectorlen-1;

src/storage/rocksdb.cpp

@@ -100,14 +100,14 @@ bool RocksDBStorageProvider::erase(const char *key, size_t cchKey)
 {
     rocksdb::Status status;
     std::unique_lock<fastlock> l(m_lock);
-    if (!FKeyExists(key, cchKey))
-        return false;
     if (m_spbatch != nullptr)
     {
         status = m_spbatch->Delete(m_spcolfamily.get(), rocksdb::Slice(key, cchKey));
     }
     else
     {
+        if (!FKeyExists(key, cchKey))
+            return false;
         status = m_spdb->Delete(WriteOptions(), m_spcolfamily.get(), rocksdb::Slice(key, cchKey));
     }
     if (status.ok())
@@ -202,12 +202,12 @@ rocksdb::WriteOptions RocksDBStorageProvider::WriteOptions() const
 void RocksDBStorageProvider::beginWriteBatch()
 {
     m_lock.lock();
-    m_spbatch = std::make_unique<rocksdb::WriteBatchWithIndex>();
+    m_spbatch = std::make_unique<rocksdb::WriteBatch>();
 }
 
 void RocksDBStorageProvider::endWriteBatch()
 {
-    m_spdb->Write(WriteOptions(), m_spbatch.get()->GetWriteBatch());
+    m_spdb->Write(WriteOptions(), m_spbatch.get());
     m_spbatch = nullptr;
     m_lock.unlock();
 }
@@ -230,7 +230,5 @@ void RocksDBStorageProvider::flush()
 bool RocksDBStorageProvider::FKeyExists(const char *key, size_t cch) const
 {
     rocksdb::PinnableSlice slice;
-    if (m_spbatch)
-        return m_spbatch->GetFromBatchAndDB(m_spdb.get(), ReadOptions(), m_spcolfamily.get(), rocksdb::Slice(key, cch), &slice).ok();
     return m_spdb->Get(ReadOptions(), m_spcolfamily.get(), rocksdb::Slice(key, cch), &slice).ok();
 }

src/storage/rocksdb.h

@@ -3,7 +3,6 @@
 #include <memory>
 #include "../IStorage.h"
 #include <rocksdb/db.h>
-#include <rocksdb/utilities/write_batch_with_index.h>
 #include "../fastlock.h"
 
 #define INTERNAL_KEY_PREFIX "\x00\x04\x03\x00\x05\x02\x04"
@@ -16,7 +15,7 @@ class RocksDBStorageProvider : public IStorage
 {
     RocksDBStorageFactory *m_pfactory;
     std::shared_ptr<rocksdb::DB> m_spdb;    // Note: This must be first so it is deleted last
-    std::unique_ptr<rocksdb::WriteBatchWithIndex> m_spbatch;
+    std::unique_ptr<rocksdb::WriteBatch> m_spbatch;
     const rocksdb::Snapshot *m_psnapshot = nullptr;
     std::shared_ptr<rocksdb::ColumnFamilyHandle> m_spcolfamily;
     rocksdb::ReadOptions m_readOptionsTemplate;

src/t_hash.cpp

@@ -1062,6 +1062,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_ZIPLIST) {
             ziplistEntry *keys, *vals = NULL;
@@ -1075,6 +1077,8 @@ void hrandfieldWithCountCommand(client *c, long l, int withvalues) {
                 count -= sample_count;
                 ziplistRandomPairs((unsigned char*)ptrFromObj(hash), sample_count, keys, vals);
                 harndfieldReplyWithZiplist(c, sample_count, keys, vals);
+                if (c->flags & CLIENT_CLOSE_ASAP)
+                    break;
             }
             zfree(keys);
             zfree(vals);
@@ -1225,13 +1229,8 @@ void hrandfieldCommand(client *c) {
         if (c->argc > 4 || (c->argc == 4 && strcasecmp(szFromObj(c->argv[3]),"withvalues"))) {
             addReplyErrorObject(c,shared.syntaxerr);
             return;
-        } else if (c->argc == 4) {
+        } else if (c->argc == 4)
             withvalues = 1;
-            if (l < LONG_MIN/2 || l > LONG_MAX/2) {
-                addReplyError(c,"value is out of range");
-                return;
-            }
-        }
         hrandfieldWithCountCommand(c, l, withvalues);
         return;
     }

src/t_set.cpp

@@ -706,6 +706,8 @@ void srandmemberWithCountCommand(client *c) {
             } else {
                 addReplyBulkCBuffer(c,ele,sdslen(ele));
             }
+            if (c->flags & CLIENT_CLOSE_ASAP)
+                break;
         }
         return;
     }

src/t_string.cpp

@@ -45,7 +45,7 @@ static int checkStringLength(client *c, long long size, long long append) {
     long long total = (uint64_t)size + append;
     /* Test configured max-bulk-len represending a limit of the biggest string object,
      * and also test for overflow. */
-    if (total > g_pserver->proto_max_bulk_len || total < size || total < append) {
+    if (g_pserver->proto_max_bulk_len || total < size || total < append) {
         addReplyError(c,"string exceeds maximum allowed size (proto-max-bulk-len)");
         return C_ERR;
     }

src/t_zset.cpp

@@ -4054,6 +4054,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_ZIPLIST) {
             ziplistEntry *keys, *vals = NULL;
@@ -4067,6 +4069,8 @@ void zrandmemberWithCountCommand(client *c, long l, int withscores) {
                 count -= sample_count;
                 ziplistRandomPairs((unsigned char*)ptrFromObj(zsetobj), sample_count, keys, vals);
                 zarndmemberReplyWithZiplist(c, sample_count, keys, vals);
+                if (c->flags & CLIENT_CLOSE_ASAP)
+                    break;
             }
             zfree(keys);
             zfree(vals);
@@ -4214,13 +4218,8 @@ void zrandmemberCommand(client *c) {
         if (c->argc > 4 || (c->argc == 4 && strcasecmp(szFromObj(c->argv[3]),"withscores"))) {
             addReplyErrorObject(c,shared.syntaxerr);
             return;
-        } else if (c->argc == 4) {
+        } else if (c->argc == 4)
             withscores = 1;
-            if (l < LONG_MIN/2 || l > LONG_MAX/2) {
-                addReplyError(c,"value is out of range");
-                return;
-            }
-        }
         zrandmemberWithCountCommand(c, l, withscores);
         return;
     }

tests/support/test.tcl

@@ -40,12 +40,6 @@ proc assert_failed {expected_err detail} {
      error "assertion:$expected_err $detail"
 }
 
-proc assert_not_equal {value expected {detail ""}} {
-    if {!($expected ne $value)} {
-        assert_failed "Expected '$value' not equal to '$expected'" $detail
-    }
-}
-
 proc assert_equal {value expected {detail ""}} {
     if {$expected ne $value} {
         assert_failed "Expected '$value' to be equal to '$expected'" $detail

tests/unit/obuf-limits.tcl

@@ -182,4 +182,20 @@ start_server {tags {"obuf-limits"} overrides { server-threads 1 }} {
         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
+    }
 }

tests/unit/type/hash.tcl

@@ -68,11 +68,6 @@ start_server {tags {"hash"}} {
         r hrandfield myhash 0
     } {}
 
-    test "HRANDFIELD count overflow" {
-        r hmset myhash a 1
-        assert_error {*value is out of range*} {r hrandfield myhash -9223372036854770000 withvalues}
-    } {}
-
     test "HRANDFIELD with <count> against non existing key" {
         r hrandfield nonexisting_key 100
     } {}

tests/unit/type/zset.tcl

@@ -1714,11 +1714,6 @@ start_server {tags {"zset"}} {
         r zrandmember nonexisting_key 100
     } {}
 
-    test "ZRANDMEMBER count overflow" {
-        r zadd myzset 0 a
-        assert_error {*value is out of range*} {r zrandmember myzset -9223372036854770000 withscores}
-    } {}
-
     # Make sure we can distinguish between an empty array and a null response
     r readraw 1