This is an automated email from the ASF dual-hosted git repository.
guangmingchen 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 75341f89 Add redis allocation size limit (#3035)
75341f89 is described below
commit 75341f896d75f3a6724c8df21a211be331c2b6ad
Author: Weibing Wang <[email protected]>
AuthorDate: Wed Jul 23 00:23:36 2025 +0800
Add redis allocation size limit (#3035)
* Add redis allocation size limit
* Fix compiler error
* Format code
---
src/brpc/redis_command.cpp | 15 +++++--
src/brpc/redis_reply.cpp | 19 +++++----
test/brpc_redis_unittest.cpp | 95 ++++++++++++++++++++++++++++++++++++++++++++
3 files changed, 119 insertions(+), 10 deletions(-)
diff --git a/src/brpc/redis_command.cpp b/src/brpc/redis_command.cpp
index 81b734be..6eabd59b 100644
--- a/src/brpc/redis_command.cpp
+++ b/src/brpc/redis_command.cpp
@@ -20,6 +20,7 @@
#include "butil/logging.h"
#include "brpc/log.h"
#include "brpc/redis_command.h"
+#include "gflags/gflags.h"
namespace {
@@ -29,6 +30,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
@@ -403,6 +406,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;
@@ -417,9 +426,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..26fa2449 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,13 +233,14 @@ 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 array_size = sizeof(RedisReply) * count;
+ if (array_size > FLAGS_redis_max_allocation_size) {
+ LOG(ERROR) << "array allocation exceeds max allocation size!
max="
+ << FLAGS_redis_max_allocation_size << ", actually="
<< array_size;
return PARSE_ERROR_ABSOLUTELY_WRONG;
}
// FIXME(gejun): Call allocate_aligned instead.
- RedisReply* subs =
(RedisReply*)_arena->allocate(sizeof(RedisReply) * count);
+ RedisReply* subs = (RedisReply*)_arena->allocate(array_size);
if (subs == NULL) {
LOG(FATAL) << "Fail to allocate RedisReply[" << count << "]";
return PARSE_ERROR_ABSOLUTELY_WRONG;
diff --git a/test/brpc_redis_unittest.cpp b/test/brpc_redis_unittest.cpp
index 7a48d19a..c1e8d059 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[]) {
@@ -1329,4 +1331,97 @@ 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 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: [email protected]
For additional commands, e-mail: [email protected]
