Skip to content

Commit 75341f8

Browse files
wwbmmmwwbmmm
authored
Add redis allocation size limit (#3035)
* Add redis allocation size limit * Fix compiler error * Format code
1 parent bea48d7 commit 75341f8

3 files changed

Lines changed: 119 additions & 10 deletions

File tree

‎src/brpc/redis_command.cpp‎

Lines changed: 12 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@
2020
#include "butil/logging.h"
2121
#include "brpc/log.h"
2222
#include "brpc/redis_command.h"
23+
#include "gflags/gflags.h"
2324

2425
namespace {
2526

@@ -29,6 +30,8 @@ const size_t CTX_WIDTH = 5;
2930

3031
namespace brpc {
3132

33+
DECLARE_int32(redis_max_allocation_size);
34+
3235
// Much faster than snprintf(..., "%lu", d);
3336
inline size_t AppendDecimal(char* outbuf, unsigned long d) {
3437
char buf[24]; // enough for decimal 64-bit integers
@@ -403,6 +406,12 @@ ParseError RedisCommandParser::Consume(butil::IOBuf& buf,
403406
return PARSE_ERROR_ABSOLUTELY_WRONG;
404407
}
405408
if (!_parsing_array) {
409+
if (value > (int64_t)(FLAGS_redis_max_allocation_size / sizeof(butil::StringPiece))) {
410+
LOG(ERROR) << "command array size exceeds limit! max="
411+
<< (FLAGS_redis_max_allocation_size / sizeof(butil::StringPiece))
412+
<< ", actually=" << value;
413+
return PARSE_ERROR_ABSOLUTELY_WRONG;
414+
}
406415
buf.pop_front(crlf_pos + 2/*CRLF*/);
407416
_parsing_array = true;
408417
_length = value;
@@ -417,9 +426,9 @@ ParseError RedisCommandParser::Consume(butil::IOBuf& buf,
417426
LOG(ERROR) << "string in command is nil!";
418427
return PARSE_ERROR_ABSOLUTELY_WRONG;
419428
}
420-
if (len > (int64_t)std::numeric_limits<uint32_t>::max()) {
421-
LOG(ERROR) << "string in command is too long! max length=2^32-1,"
422-
" actually=" << len;
429+
if (len > FLAGS_redis_max_allocation_size) {
430+
LOG(ERROR) << "command string exceeds max allocation size! max="
431+
<< FLAGS_redis_max_allocation_size << ", actually=" << len;
423432
return PARSE_ERROR_ABSOLUTELY_WRONG;
424433
}
425434
if (buf.size() < crlf_pos + 2 + (size_t)len + 2/*CRLF*/) {

‎src/brpc/redis_reply.cpp‎

Lines changed: 12 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -20,9 +20,13 @@
2020
#include "butil/logging.h"
2121
#include "butil/string_printf.h"
2222
#include "brpc/redis_reply.h"
23+
#include "gflags/gflags.h"
2324

2425
namespace brpc {
2526

27+
DEFINE_int32(redis_max_allocation_size, 64 * 1024 * 1024,
28+
"Maximum memory allocation size in bytes for a single redis request or reply (64MB by default)");
29+
2630
//BAIDU_CASSERT(sizeof(RedisReply) == 24, size_match);
2731
const int RedisReply::npos = -1;
2832

@@ -175,9 +179,9 @@ ParseError RedisReply::ConsumePartialIOBuf(butil::IOBuf& buf) {
175179
_data.integer = 0;
176180
return PARSE_OK;
177181
}
178-
if (len > (int64_t)std::numeric_limits<uint32_t>::max()) {
179-
LOG(ERROR) << "bulk string is too long! max length=2^32-1,"
180-
" actually=" << len;
182+
if (len > FLAGS_redis_max_allocation_size) {
183+
LOG(ERROR) << "bulk string exceeds max allocation size! max="
184+
<< FLAGS_redis_max_allocation_size << ", actually=" << len;
181185
return PARSE_ERROR_ABSOLUTELY_WRONG;
182186
}
183187
// We provide c_str(), thus even if bulk string is started with
@@ -229,13 +233,14 @@ ParseError RedisReply::ConsumePartialIOBuf(butil::IOBuf& buf) {
229233
_data.array.replies = NULL;
230234
return PARSE_OK;
231235
}
232-
if (count > (int64_t)std::numeric_limits<uint32_t>::max()) {
233-
LOG(ERROR) << "Too many sub replies! max count=2^32-1,"
234-
" actually=" << count;
236+
int64_t array_size = sizeof(RedisReply) * count;
237+
if (array_size > FLAGS_redis_max_allocation_size) {
238+
LOG(ERROR) << "array allocation exceeds max allocation size! max="
239+
<< FLAGS_redis_max_allocation_size << ", actually=" << array_size;
235240
return PARSE_ERROR_ABSOLUTELY_WRONG;
236241
}
237242
// FIXME(gejun): Call allocate_aligned instead.
238-
RedisReply* subs = (RedisReply*)_arena->allocate(sizeof(RedisReply) * count);
243+
RedisReply* subs = (RedisReply*)_arena->allocate(array_size);
239244
if (subs == NULL) {
240245
LOG(FATAL) << "Fail to allocate RedisReply[" << count << "]";
241246
return PARSE_ERROR_ABSOLUTELY_WRONG;

‎test/brpc_redis_unittest.cpp‎

Lines changed: 95 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,9 +26,11 @@
2626
#include <brpc/server.h>
2727
#include <brpc/redis_command.h>
2828
#include <gtest/gtest.h>
29+
#include <gflags/gflags.h>
2930

3031
namespace brpc {
3132
DECLARE_int32(idle_timeout_second);
33+
DECLARE_int32(redis_max_allocation_size);
3234
}
3335

3436
int main(int argc, char* argv[]) {
@@ -1329,4 +1331,97 @@ TEST_F(RedisTest, server_handle_pipeline) {
13291331
ASSERT_STREQ(response.reply(7).c_str(), "world");
13301332
}
13311333

1334+
TEST_F(RedisTest, memory_allocation_limits) {
1335+
int32_t original_limit = brpc::FLAGS_redis_max_allocation_size;
1336+
brpc::FLAGS_redis_max_allocation_size = 1024;
1337+
1338+
butil::Arena arena;
1339+
1340+
// Test redis_reply.cpp limits
1341+
{
1342+
// Test bulk string exceeding limit
1343+
butil::IOBuf buf;
1344+
std::string large_string = "*1\r\n$2000\r\n";
1345+
large_string.append(2000, 'a');
1346+
large_string.append("\r\n");
1347+
buf.append(large_string);
1348+
1349+
brpc::RedisReply reply(&arena);
1350+
brpc::ParseError err = reply.ConsumePartialIOBuf(buf);
1351+
ASSERT_EQ(brpc::PARSE_ERROR_ABSOLUTELY_WRONG, err);
1352+
}
1353+
1354+
{
1355+
// Test array allocation exceeding limit
1356+
butil::IOBuf buf;
1357+
int32_t large_count = brpc::FLAGS_redis_max_allocation_size / sizeof(brpc::RedisReply) + 1;
1358+
std::string large_array = "*" + std::to_string(large_count) + "\r\n";
1359+
buf.append(large_array);
1360+
1361+
brpc::RedisReply reply(&arena);
1362+
brpc::ParseError err = reply.ConsumePartialIOBuf(buf);
1363+
ASSERT_EQ(brpc::PARSE_ERROR_ABSOLUTELY_WRONG, err);
1364+
}
1365+
1366+
// Test redis_command.cpp limits
1367+
{
1368+
// Test command string exceeding limit
1369+
brpc::RedisCommandParser parser;
1370+
butil::IOBuf buf;
1371+
std::string large_cmd = "*2\r\n$3\r\nget\r\n$2000\r\n";
1372+
large_cmd.append(2000, 'b');
1373+
large_cmd.append("\r\n");
1374+
buf.append(large_cmd);
1375+
1376+
std::vector<butil::StringPiece> args;
1377+
brpc::ParseError err = parser.Consume(buf, &args, &arena);
1378+
ASSERT_EQ(brpc::PARSE_ERROR_ABSOLUTELY_WRONG, err);
1379+
}
1380+
1381+
{
1382+
// Test command array size exceeding limit
1383+
brpc::RedisCommandParser parser;
1384+
butil::IOBuf buf;
1385+
int32_t large_array_size = brpc::FLAGS_redis_max_allocation_size / sizeof(butil::StringPiece) + 1;
1386+
std::string large_array_cmd = "*" + std::to_string(large_array_size) + "\r\n";
1387+
buf.append(large_array_cmd);
1388+
1389+
std::vector<butil::StringPiece> args;
1390+
brpc::ParseError err = parser.Consume(buf, &args, &arena);
1391+
ASSERT_EQ(brpc::PARSE_ERROR_ABSOLUTELY_WRONG, err);
1392+
}
1393+
1394+
// Test valid cases within limits
1395+
{
1396+
// Test small bulk string should work
1397+
butil::IOBuf buf;
1398+
std::string small_string = "*1\r\n$10\r\nhelloworld\r\n";
1399+
buf.append(small_string);
1400+
1401+
brpc::RedisReply reply(&arena);
1402+
brpc::ParseError err = reply.ConsumePartialIOBuf(buf);
1403+
ASSERT_EQ(brpc::PARSE_OK, err);
1404+
ASSERT_TRUE(reply.is_array());
1405+
ASSERT_EQ(1, (int)reply.size());
1406+
ASSERT_STREQ("helloworld", reply[0].c_str());
1407+
}
1408+
1409+
{
1410+
// Test small command should work
1411+
brpc::RedisCommandParser parser;
1412+
butil::IOBuf buf;
1413+
std::string small_cmd = "*2\r\n$3\r\nget\r\n$5\r\nmykey\r\n";
1414+
buf.append(small_cmd);
1415+
1416+
std::vector<butil::StringPiece> args;
1417+
brpc::ParseError err = parser.Consume(buf, &args, &arena);
1418+
ASSERT_EQ(brpc::PARSE_OK, err);
1419+
ASSERT_EQ(2, (int)args.size());
1420+
ASSERT_EQ("get", args[0].as_string());
1421+
ASSERT_EQ("mykey", args[1].as_string());
1422+
}
1423+
1424+
brpc::FLAGS_redis_max_allocation_size = original_limit;
1425+
}
1426+
13321427
} //namespace

0 commit comments

Comments
 (0)