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