Skip to content

Limit redis alloc size #5

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

Open
thweetkomputer wants to merge 3 commits into
base: master
Choose a base branch
from
Open

Limit redis alloc size #5

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
5d7ed60
fix
thweetkomputer Aug 12, 2025
0129644
make default 256
thweetkomputer Aug 12, 2025
e95de77
format
thweetkomputer Aug 13, 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
16 changes: 12 additions & 4 deletions src/brpc/redis_command.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,10 +20,12 @@
#include "butil/logging.h"
#include "brpc/log.h"
#include "brpc/redis_command.h"
#include "gflags/gflags.h"

namespace brpc {
constexpr size_t CTX_WIDTH = 5;

const size_t CTX_WIDTH = 5;
DECLARE_int32(redis_max_allocation_size);

// Much faster than snprintf(..., "%lu", d);
inline size_t AppendDecimal(char* outbuf, unsigned long d) {
Expand Down Expand Up @@ -451,6 +453,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;
Expand All @@ -465,9 +473,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*/) {
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, 256 * 1024 * 1024,
"Maximum memory allocation size in bytes for a single redis request or reply (256MB 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;
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 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.
Expand Down
107 changes: 107 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 @@ -27,9 +27,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 @@ -1349,4 +1351,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
Loading