This is an automated email from the ASF dual-hosted git repository.

wwbmmm pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/brpc.git


The following commit(s) were added to refs/heads/master by this push:
     new 0e5d0220 Limit redis alloc size (#3050)
0e5d0220 is described below

commit 0e5d0220ecf7243a9f25a5234f6d0e567d829cb6
Author: Weibing Wang <wwb...@163.com>
AuthorDate: Sat Aug 2 13:15:12 2025 +0800

    Limit redis alloc size (#3050)
    
    * Reapply "Add redis allocation size limit (#3035)"
    
    This reverts commit d0f8f8fec6a8d97092671578f4ecd9e2c5f6bc34.
    
    * Fix int overflow
    
    * Update log
---
 src/brpc/redis_command.cpp   |  15 ++++--
 src/brpc/redis_reply.cpp     |  17 ++++---
 test/brpc_redis_unittest.cpp | 107 +++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 130 insertions(+), 9 deletions(-)

diff --git a/src/brpc/redis_command.cpp b/src/brpc/redis_command.cpp
index d2620323..5803aaa8 100644
--- a/src/brpc/redis_command.cpp
+++ b/src/brpc/redis_command.cpp
@@ -21,6 +21,7 @@
 #include "butil/logging.h"
 #include "brpc/log.h"
 #include "brpc/redis_command.h"
+#include "gflags/gflags.h"
 
 namespace {
 
@@ -30,6 +31,8 @@ const size_t CTX_WIDTH = 5;
 
 namespace brpc {
 
+DECLARE_int32(redis_max_allocation_size);
+
 // Much faster than snprintf(..., "%lu", d);
 inline size_t AppendDecimal(char* outbuf, unsigned long d) {
     char buf[24];  // enough for decimal 64-bit integers
@@ -456,6 +459,12 @@ ParseError RedisCommandParser::Consume(butil::IOBuf& buf,
         return PARSE_ERROR_ABSOLUTELY_WRONG;
     }
     if (!_parsing_array) {
+        if (value > (int64_t)(FLAGS_redis_max_allocation_size / 
sizeof(butil::StringPiece))) {
+            LOG(ERROR) << "command array size exceeds limit! max=" 
+                       << (FLAGS_redis_max_allocation_size / 
sizeof(butil::StringPiece)) 
+                       << ", actually=" << value;
+            return PARSE_ERROR_ABSOLUTELY_WRONG;
+        }
         buf.pop_front(crlf_pos + 2/*CRLF*/);
         _parsing_array = true;
         _length = value;
@@ -470,9 +479,9 @@ ParseError RedisCommandParser::Consume(butil::IOBuf& buf,
         LOG(ERROR) << "string in command is nil!";
         return PARSE_ERROR_ABSOLUTELY_WRONG;
     }
-    if (len > (int64_t)std::numeric_limits<uint32_t>::max()) {
-        LOG(ERROR) << "string in command is too long! max length=2^32-1,"
-            " actually=" << len;
+    if (len > FLAGS_redis_max_allocation_size) {
+        LOG(ERROR) << "command string exceeds max allocation size! max=" 
+                   << FLAGS_redis_max_allocation_size << ", actually=" << len;
         return PARSE_ERROR_ABSOLUTELY_WRONG;
     }
     if (buf.size() < crlf_pos + 2 + (size_t)len + 2/*CRLF*/) {
diff --git a/src/brpc/redis_reply.cpp b/src/brpc/redis_reply.cpp
index 8710df16..e2053a36 100644
--- a/src/brpc/redis_reply.cpp
+++ b/src/brpc/redis_reply.cpp
@@ -20,9 +20,13 @@
 #include "butil/logging.h"
 #include "butil/string_printf.h"
 #include "brpc/redis_reply.h"
+#include "gflags/gflags.h"
 
 namespace brpc {
 
+DEFINE_int32(redis_max_allocation_size, 64 * 1024 * 1024, 
+             "Maximum memory allocation size in bytes for a single redis 
request or reply (64MB by default)");
+
 //BAIDU_CASSERT(sizeof(RedisReply) == 24, size_match);
 const int RedisReply::npos = -1;
 
@@ -175,9 +179,9 @@ ParseError RedisReply::ConsumePartialIOBuf(butil::IOBuf& 
buf) {
                 _data.integer = 0;
                 return PARSE_OK;
             }
-            if (len > (int64_t)std::numeric_limits<uint32_t>::max()) {
-                LOG(ERROR) << "bulk string is too long! max length=2^32-1,"
-                    " actually=" << len;
+            if (len > FLAGS_redis_max_allocation_size) {
+                LOG(ERROR) << "bulk string exceeds max allocation size! max=" 
+                           << FLAGS_redis_max_allocation_size << ", actually=" 
<< len;
                 return PARSE_ERROR_ABSOLUTELY_WRONG;
             }
             // We provide c_str(), thus even if bulk string is started with
@@ -229,9 +233,10 @@ ParseError RedisReply::ConsumePartialIOBuf(butil::IOBuf& 
buf) {
                 _data.array.replies = NULL;
                 return PARSE_OK;
             }
-            if (count > (int64_t)std::numeric_limits<uint32_t>::max()) {
-                LOG(ERROR) << "Too many sub replies! max count=2^32-1,"
-                    " actually=" << count;
+            int64_t max_count = FLAGS_redis_max_allocation_size / 
sizeof(RedisReply);
+            if (count > max_count) {
+                LOG(ERROR) << "array allocation exceeds max allocation size! 
max=" 
+                           << max_count << ", actually=" << count;
                 return PARSE_ERROR_ABSOLUTELY_WRONG;
             }
             // FIXME(gejun): Call allocate_aligned instead.
diff --git a/test/brpc_redis_unittest.cpp b/test/brpc_redis_unittest.cpp
index ac6872e0..5c8aef12 100644
--- a/test/brpc_redis_unittest.cpp
+++ b/test/brpc_redis_unittest.cpp
@@ -26,9 +26,11 @@
 #include <brpc/server.h>
 #include <brpc/redis_command.h>
 #include <gtest/gtest.h>
+#include <gflags/gflags.h>
 
 namespace brpc {
 DECLARE_int32(idle_timeout_second);
+DECLARE_int32(redis_max_allocation_size);
 }
 
 int main(int argc, char* argv[]) {
@@ -1350,4 +1352,109 @@ TEST_F(RedisTest, server_handle_pipeline) {
     ASSERT_STREQ(response.reply(7).c_str(), "world");
 }
 
+TEST_F(RedisTest, memory_allocation_limits) {
+    int32_t original_limit = brpc::FLAGS_redis_max_allocation_size;
+    brpc::FLAGS_redis_max_allocation_size = 1024;
+    
+    butil::Arena arena;
+    
+    // Test redis_reply.cpp limits
+    {
+        // Test bulk string exceeding limit
+        butil::IOBuf buf;
+        std::string large_string = "*1\r\n$2000\r\n";
+        large_string.append(2000, 'a');
+        large_string.append("\r\n");
+        buf.append(large_string);
+        
+        brpc::RedisReply reply(&arena);
+        brpc::ParseError err = reply.ConsumePartialIOBuf(buf);
+        ASSERT_EQ(brpc::PARSE_ERROR_ABSOLUTELY_WRONG, err);
+    }
+    
+    {
+        // Test array allocation exceeding limit
+        butil::IOBuf buf;
+        int32_t large_count = brpc::FLAGS_redis_max_allocation_size / 
sizeof(brpc::RedisReply) + 1;
+        std::string large_array = "*" + std::to_string(large_count) + "\r\n";
+        buf.append(large_array);
+        
+        brpc::RedisReply reply(&arena);
+        brpc::ParseError err = reply.ConsumePartialIOBuf(buf);
+        ASSERT_EQ(brpc::PARSE_ERROR_ABSOLUTELY_WRONG, err);
+    }
+
+    {
+        // Test large count
+        butil::IOBuf buf;
+        int64_t large_count = 9223372036854775807;
+        std::string large_array = "*" + std::to_string(large_count) + "\r\n";
+        buf.append(large_array);
+        
+        brpc::RedisReply reply(&arena);
+        brpc::ParseError err = reply.ConsumePartialIOBuf(buf);
+        ASSERT_EQ(brpc::PARSE_ERROR_ABSOLUTELY_WRONG, err);
+    }
+    
+    // Test redis_command.cpp limits
+    {
+        // Test command string exceeding limit
+        brpc::RedisCommandParser parser;
+        butil::IOBuf buf;
+        std::string large_cmd = "*2\r\n$3\r\nget\r\n$2000\r\n";
+        large_cmd.append(2000, 'b');
+        large_cmd.append("\r\n");
+        buf.append(large_cmd);
+        
+        std::vector<butil::StringPiece> args;
+        brpc::ParseError err = parser.Consume(buf, &args, &arena);
+        ASSERT_EQ(brpc::PARSE_ERROR_ABSOLUTELY_WRONG, err);
+    }
+    
+    {
+        // Test command array size exceeding limit
+        brpc::RedisCommandParser parser;
+        butil::IOBuf buf;
+        int32_t large_array_size = brpc::FLAGS_redis_max_allocation_size / 
sizeof(butil::StringPiece) + 1;
+        std::string large_array_cmd = "*" + std::to_string(large_array_size) + 
"\r\n";
+        buf.append(large_array_cmd);
+        
+        std::vector<butil::StringPiece> args;
+        brpc::ParseError err = parser.Consume(buf, &args, &arena);
+        ASSERT_EQ(brpc::PARSE_ERROR_ABSOLUTELY_WRONG, err);
+    }
+    
+    // Test valid cases within limits
+    {
+        // Test small bulk string should work
+        butil::IOBuf buf;
+        std::string small_string = "*1\r\n$10\r\nhelloworld\r\n";
+        buf.append(small_string);
+        
+        brpc::RedisReply reply(&arena);
+        brpc::ParseError err = reply.ConsumePartialIOBuf(buf);
+        ASSERT_EQ(brpc::PARSE_OK, err);
+        ASSERT_TRUE(reply.is_array());
+        ASSERT_EQ(1, (int)reply.size());
+        ASSERT_STREQ("helloworld", reply[0].c_str());
+    }
+    
+    {
+        // Test small command should work
+        brpc::RedisCommandParser parser;
+        butil::IOBuf buf;
+        std::string small_cmd = "*2\r\n$3\r\nget\r\n$5\r\nmykey\r\n";
+        buf.append(small_cmd);
+        
+        std::vector<butil::StringPiece> args;
+        brpc::ParseError err = parser.Consume(buf, &args, &arena);
+        ASSERT_EQ(brpc::PARSE_OK, err);
+        ASSERT_EQ(2, (int)args.size());
+        ASSERT_EQ("get", args[0].as_string());
+        ASSERT_EQ("mykey", args[1].as_string());
+    }
+    
+    brpc::FLAGS_redis_max_allocation_size = original_limit;
+}
+
 } //namespace


---------------------------------------------------------------------
To unsubscribe, e-mail: dev-unsubscr...@brpc.apache.org
For additional commands, e-mail: dev-h...@brpc.apache.org

Reply via email to