Skip to content

Add redis allocation size limit #3035

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
chenBright merged 3 commits into from
Jul 22, 2025
Merged

Add redis allocation size limit #3035

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
c560e28
Add redis allocation size limit
wwbmmm Jul 22, 2025
d548ac3
Fix compiler error
wwbmmm Jul 22, 2025
6210af9
Format code
wwbmmm Jul 22, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 12 additions & 3 deletions src/brpc/redis_command.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@
#include "butil/logging.h"
#include "brpc/log.h"
#include "brpc/redis_command.h"
#include "gflags/gflags.h"

namespace {

Expand All @@ -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
Expand Down Expand Up @@ -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;
Copy link
Copy Markdown
Contributor

@chenBright chenBright Jul 22, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Format.

return PARSE_ERROR_ABSOLUTELY_WRONG;
}
buf.pop_front(crlf_pos + 2/*CRLF*/);
_parsing_array = true;
_length = value;
Expand All @@ -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;
Copy link
Copy Markdown
Contributor

@chenBright chenBright Jul 22, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Format

return PARSE_ERROR_ABSOLUTELY_WRONG;
}
if (buf.size() < crlf_pos + 2 + (size_t)len + 2/*CRLF*/) {
Expand Down
17 changes: 11 additions & 6 deletions src/brpc/redis_reply.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand Down Expand Up @@ -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;
Copy link
Copy Markdown
Contributor

@chenBright chenBright Jul 22, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Format

return PARSE_ERROR_ABSOLUTELY_WRONG;
}
// We provide c_str(), thus even if bulk string is started with
Expand Down Expand Up @@ -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 array_size = sizeof(RedisReply) * count;
Copy link
Copy Markdown
Contributor

@chenBright chenBright Jul 22, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Use array_size instead of sizeof(RedisReply) * count in line 243.

Copy link
Copy Markdown
Contributor

@jamesge jamesge Jul 23, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

array_size should better be modified with "const"

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;
Copy link
Copy Markdown
Contributor

@chenBright chenBright Jul 22, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Format

return PARSE_ERROR_ABSOLUTELY_WRONG;
}
// FIXME(gejun): Call allocate_aligned instead.
Expand Down
95 changes: 95 additions & 0 deletions test/brpc_redis_unittest.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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[]) {
Expand Down Expand Up @@ -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
Loading