From a55685d336e6c4401dc6b09ae2b39cf0e755d11c Mon Sep 17 00:00:00 2001 From: wwbmmm Date: Sun, 12 Jul 2026 20:55:30 +0800 Subject: [PATCH 1/2] Enforce body size limits in protocol parsers --- src/brpc/details/http_message.cpp | 9 ++++ src/brpc/details/http_message.h | 3 ++ src/brpc/policy/couchbase_protocol.cpp | 4 ++ src/brpc/policy/http2_rpc_protocol.cpp | 3 ++ src/brpc/policy/http_rpc_protocol.cpp | 19 ++++++++ src/brpc/policy/memcache_binary_protocol.cpp | 4 ++ test/brpc_couchbase_unittest.cpp | 29 +++++++++++- test/brpc_http_rpc_protocol_unittest.cpp | 48 ++++++++++++++++++++ test/brpc_memcache_unittest.cpp | 29 ++++++++++++ 9 files changed, 147 insertions(+), 1 deletion(-) diff --git a/src/brpc/details/http_message.cpp b/src/brpc/details/http_message.cpp index 0ffe5b1143..003bafa074 100644 --- a/src/brpc/details/http_message.cpp +++ b/src/brpc/details/http_message.cpp @@ -39,6 +39,7 @@ DEFINE_bool(http_verbose, false, DEFINE_int32(http_verbose_max_body_length, 512, "[DEBUG] Max body length printed when -http_verbose is on"); DECLARE_int64(socket_max_unwritten_bytes); +DECLARE_uint64(max_body_size); // Implement callbacks for http parser @@ -254,6 +255,14 @@ int HttpMessage::on_message_complete_cb(http_parser *parser) { } int HttpMessage::OnBody(const char *at, const size_t length) { + if (!_read_body_progressively) { + if (length > FLAGS_max_body_size || + _body_size > FLAGS_max_body_size - length) { + _body_too_large = true; + return -1; + } + _body_size += length; + } if (_vmsgbuilder) { if (_stage != HTTP_ON_BODY) { // only add prefix at first entry. diff --git a/src/brpc/details/http_message.h b/src/brpc/details/http_message.h index 655ba1a9d8..ae4a016dc9 100644 --- a/src/brpc/details/http_message.h +++ b/src/brpc/details/http_message.h @@ -71,6 +71,7 @@ class HttpMessage { HttpHeader& header() { return _header; } const HttpHeader& header() const { return _header; } size_t parsed_length() const { return _parsed_length; } + bool body_too_large() const { return _body_too_large; } // Http parser callback functions static int on_message_begin(http_parser *); @@ -115,6 +116,8 @@ class HttpMessage { // Read body progressively ProgressiveReader* _body_reader{NULL}; butil::IOBuf _body; + size_t _body_size{0}; + bool _body_too_large{false}; // Store the IOBuf information in `ParseFromIOBuf' // for later zero-copy usage in `OnBody'. diff --git a/src/brpc/policy/couchbase_protocol.cpp b/src/brpc/policy/couchbase_protocol.cpp index 0ece53dbfb..0ab79bfb90 100644 --- a/src/brpc/policy/couchbase_protocol.cpp +++ b/src/brpc/policy/couchbase_protocol.cpp @@ -39,6 +39,7 @@ namespace brpc { DECLARE_bool(enable_rpcz); +DECLARE_uint64(max_body_size); namespace policy { @@ -96,6 +97,9 @@ ParseResult ParseCouchbaseMessage(butil::IOBuf* source, Socket* socket, } const CouchbaseResponseHeader* header = (const CouchbaseResponseHeader*)p; uint32_t total_body_length = butil::NetToHost32(header->total_body_length); + if (total_body_length > FLAGS_max_body_size) { + return MakeParseError(PARSE_ERROR_TOO_BIG_DATA); + } if (source->size() < sizeof(*header) + total_body_length) { return MakeParseError(PARSE_ERROR_NOT_ENOUGH_DATA); } diff --git a/src/brpc/policy/http2_rpc_protocol.cpp b/src/brpc/policy/http2_rpc_protocol.cpp index 043f53ebea..b2de496cde 100644 --- a/src/brpc/policy/http2_rpc_protocol.cpp +++ b/src/brpc/policy/http2_rpc_protocol.cpp @@ -739,6 +739,9 @@ H2ParseResult H2StreamContext::OnData( for (size_t i = 0; i < data.backing_block_num(); ++i) { const butil::StringPiece blk = data.backing_block(i); if (OnBody(blk.data(), blk.size()) != 0) { + if (body_too_large()) { + return MakeH2Error(H2_ENHANCE_YOUR_CALM, stream_id()); + } LOG(ERROR) << "Fail to parse data"; return MakeH2Error(H2_PROTOCOL_ERROR); } diff --git a/src/brpc/policy/http_rpc_protocol.cpp b/src/brpc/policy/http_rpc_protocol.cpp index e5e22d924a..8cbe06980f 100644 --- a/src/brpc/policy/http_rpc_protocol.cpp +++ b/src/brpc/policy/http_rpc_protocol.cpp @@ -1217,6 +1217,25 @@ ParseResult ParseHttpMessage(butil::IOBuf *source, Socket *socket, // comments in http_message.h rc = http_imsg->ParseFromIOBuf(*source); } + if (rc < 0 && http_imsg->body_too_large()) { + if (socket->CreatedByConnect()) { + return MakeParseError(PARSE_ERROR_TOO_BIG_DATA); + } + const int release_rc = socket->ReleaseAdditionalReference(); + if (release_rc == 0) { + butil::IOBuf resp; + HttpHeader header; + header.set_status_code(HTTP_STATUS_REQUEST_ENTITY_TOO_LARGE); + header.SetHeader("Connection", "close"); + MakeRawHttpResponse(&resp, &header, NULL); + Socket::WriteOptions wopt; + wopt.ignore_eovercrowded = true; + socket->Write(&resp, &wopt); + } else if (release_rc > 0) { + LOG(ERROR) << "Impossible: Recycled!"; + } + return MakeParseError(PARSE_ERROR_NOT_ENOUGH_DATA); + } if (http_imsg->is_stage2()) { // The header part is already parsed as an intact HTTP message // to the ProcessHttpXXX. Here parses the body part. diff --git a/src/brpc/policy/memcache_binary_protocol.cpp b/src/brpc/policy/memcache_binary_protocol.cpp index 46432c4f7e..e3174be588 100644 --- a/src/brpc/policy/memcache_binary_protocol.cpp +++ b/src/brpc/policy/memcache_binary_protocol.cpp @@ -40,6 +40,7 @@ namespace brpc { DECLARE_bool(enable_rpcz); +DECLARE_uint64(max_body_size); namespace policy { @@ -90,6 +91,9 @@ ParseResult ParseMemcacheMessage(butil::IOBuf* source, } const MemcacheResponseHeader* header = (const MemcacheResponseHeader*)p; uint32_t total_body_length = butil::NetToHost32(header->total_body_length); + if (total_body_length > FLAGS_max_body_size) { + return MakeParseError(PARSE_ERROR_TOO_BIG_DATA); + } if (source->size() < sizeof(*header) + total_body_length) { return MakeParseError(PARSE_ERROR_NOT_ENOUGH_DATA); } diff --git a/test/brpc_couchbase_unittest.cpp b/test/brpc_couchbase_unittest.cpp index dacc9a09ce..a4ca0c92b9 100644 --- a/test/brpc_couchbase_unittest.cpp +++ b/test/brpc_couchbase_unittest.cpp @@ -16,11 +16,16 @@ // under the License. #include +#include +#include +#include +#include #include #include namespace brpc { DECLARE_int32(idle_timeout_second); +DECLARE_uint64(max_body_size); } int main(int argc, char* argv[]) { @@ -34,6 +39,28 @@ namespace { // Unit Tests - No Server Required class CouchbaseUnitTest : public testing::Test {}; +TEST_F(CouchbaseUnitTest, RejectOversizedResponseBeforeBufferingBody) { + const uint64_t saved_max_body_size = brpc::FLAGS_max_body_size; + brpc::FLAGS_max_body_size = 1024; + + brpc::SocketId id; + brpc::SocketOptions options; + ASSERT_EQ(0, brpc::Socket::Create(options, &id)); + brpc::SocketUniquePtr socket; + ASSERT_EQ(0, brpc::Socket::Address(id, &socket)); + + brpc::policy::CouchbaseResponseHeader couchbase_header = {}; + couchbase_header.magic = brpc::policy::CB_MAGIC_RESPONSE; + couchbase_header.total_body_length = butil::HostToNet32(1025); + butil::IOBuf couchbase_buf; + couchbase_buf.append(&couchbase_header, sizeof(couchbase_header)); + EXPECT_EQ(brpc::PARSE_ERROR_TOO_BIG_DATA, + brpc::policy::ParseCouchbaseMessage( + &couchbase_buf, socket.get(), false, NULL).error()); + + brpc::FLAGS_max_body_size = saved_max_body_size; +} + TEST_F(CouchbaseUnitTest, RequestBuilders) { brpc::CouchbaseOperations::CouchbaseRequest req; req.Clear(); @@ -82,4 +109,4 @@ TEST_F(CouchbaseUnitTest, EdgeCases) { EXPECT_TRUE(true); } -} // namespace \ No newline at end of file +} // namespace diff --git a/test/brpc_http_rpc_protocol_unittest.cpp b/test/brpc_http_rpc_protocol_unittest.cpp index 3f3290bdc4..a99c834f21 100644 --- a/test/brpc_http_rpc_protocol_unittest.cpp +++ b/test/brpc_http_rpc_protocol_unittest.cpp @@ -60,6 +60,7 @@ DECLARE_string(rpc_dump_dir); DECLARE_int32(rpc_dump_max_requests_in_one_file); DECLARE_bool(allow_chunked_length); DECLARE_int32(max_connection_pool_size); +DECLARE_uint64(max_body_size); extern bvar::CollectorSpeedLimit g_rpc_dump_sl; } @@ -79,6 +80,18 @@ int main(int argc, char* argv[]) { namespace { +class MaxBodySizeGuard { +public: + explicit MaxBodySizeGuard(uint64_t value) + : _saved(brpc::FLAGS_max_body_size) { + brpc::FLAGS_max_body_size = value; + } + ~MaxBodySizeGuard() { brpc::FLAGS_max_body_size = _saved; } + +private: + uint64_t _saved; +}; + static const std::string EXP_REQUEST = "hello"; static const std::string EXP_RESPONSE = "world"; static const std::string EXP_RESPONSE_CONTENT_LENGTH = "1024"; @@ -344,6 +357,41 @@ class HttpTest : public ::testing::Test{ MyAuthenticator _auth; }; +TEST_F(HttpTest, reject_oversized_http_body) { + MaxBodySizeGuard guard(4); + butil::IOBuf buf; + buf.append("POST / HTTP/1.1\r\nContent-Length: 5\r\n\r\nhello"); + + brpc::ParseResult result = + brpc::policy::ParseHttpMessage(&buf, _socket.get(), false, NULL); + EXPECT_EQ(brpc::PARSE_ERROR_NOT_ENOUGH_DATA, result.error()); + int bytes_in_pipe = 0; + ASSERT_EQ(0, ioctl(_pipe_fds[0], FIONREAD, &bytes_in_pipe)); + ASSERT_GT(bytes_in_pipe, 0); + butil::IOPortal response; + ASSERT_EQ(bytes_in_pipe, + response.append_from_file_descriptor(_pipe_fds[0], bytes_in_pipe)); + EXPECT_NE(std::string::npos, response.to_string().find(" 413 ")); +} + +TEST_F(HttpTest, reject_oversized_chunked_http_body) { + MaxBodySizeGuard guard(4); + butil::IOBuf buf; + buf.append("POST / HTTP/1.1\r\nTransfer-Encoding: chunked\r\n\r\n" + "3\r\nabc\r\n2\r\nde\r\n0\r\n\r\n"); + + brpc::ParseResult result = + brpc::policy::ParseHttpMessage(&buf, _socket.get(), false, NULL); + EXPECT_EQ(brpc::PARSE_ERROR_NOT_ENOUGH_DATA, result.error()); + int bytes_in_pipe = 0; + ASSERT_EQ(0, ioctl(_pipe_fds[0], FIONREAD, &bytes_in_pipe)); + ASSERT_GT(bytes_in_pipe, 0); + butil::IOPortal response; + ASSERT_EQ(bytes_in_pipe, + response.append_from_file_descriptor(_pipe_fds[0], bytes_in_pipe)); + EXPECT_NE(std::string::npos, response.to_string().find(" 413 ")); +} + int AllocateFreePortOrDie() { butil::fd_guard fd(tcp_listen(butil::EndPoint(butil::my_ip(), 0))); EXPECT_GE(fd, 0); diff --git a/test/brpc_memcache_unittest.cpp b/test/brpc_memcache_unittest.cpp index f929767fa5..2903547990 100644 --- a/test/brpc_memcache_unittest.cpp +++ b/test/brpc_memcache_unittest.cpp @@ -20,10 +20,16 @@ #include "butil/logging.h" #include #include +#include +#include +#include +#include +#include #include namespace brpc { DECLARE_int32(idle_timeout_second); +DECLARE_uint64(max_body_size); } int main(int argc, char* argv[]) { @@ -33,6 +39,29 @@ int main(int argc, char* argv[]) { } namespace { + +TEST(MemcacheParserTest, RejectOversizedResponseBeforeBufferingBody) { + const uint64_t saved_max_body_size = brpc::FLAGS_max_body_size; + brpc::FLAGS_max_body_size = 1024; + + brpc::SocketId id; + brpc::SocketOptions options; + ASSERT_EQ(0, brpc::Socket::Create(options, &id)); + brpc::SocketUniquePtr socket; + ASSERT_EQ(0, brpc::Socket::Address(id, &socket)); + + brpc::policy::MemcacheResponseHeader header = {}; + header.magic = brpc::policy::MC_MAGIC_RESPONSE; + header.total_body_length = butil::HostToNet32(1025); + butil::IOBuf buf; + buf.append(&header, sizeof(header)); + EXPECT_EQ(brpc::PARSE_ERROR_TOO_BIG_DATA, + brpc::policy::ParseMemcacheMessage( + &buf, socket.get(), false, NULL).error()); + + brpc::FLAGS_max_body_size = saved_max_body_size; +} + static pthread_once_t download_memcached_once = PTHREAD_ONCE_INIT; static pid_t g_mc_pid = -1; From 67a9f419bd3e5bb994fa32cc7c59ed6423c1921b Mon Sep 17 00:00:00 2001 From: wwbmmm Date: Sun, 12 Jul 2026 21:27:03 +0800 Subject: [PATCH 2/2] Use scoped FlagSaver --- test/brpc_couchbase_unittest.cpp | 4 ++-- test/brpc_http_rpc_protocol_unittest.cpp | 18 ++++-------------- test/brpc_memcache_unittest.cpp | 4 ++-- 3 files changed, 8 insertions(+), 18 deletions(-) diff --git a/test/brpc_couchbase_unittest.cpp b/test/brpc_couchbase_unittest.cpp index a4ca0c92b9..e2f2fc0ad2 100644 --- a/test/brpc_couchbase_unittest.cpp +++ b/test/brpc_couchbase_unittest.cpp @@ -21,6 +21,7 @@ #include #include #include +#include #include namespace brpc { @@ -40,7 +41,7 @@ namespace { class CouchbaseUnitTest : public testing::Test {}; TEST_F(CouchbaseUnitTest, RejectOversizedResponseBeforeBufferingBody) { - const uint64_t saved_max_body_size = brpc::FLAGS_max_body_size; + GFLAGS_NAMESPACE::FlagSaver flag_saver; brpc::FLAGS_max_body_size = 1024; brpc::SocketId id; @@ -58,7 +59,6 @@ TEST_F(CouchbaseUnitTest, RejectOversizedResponseBeforeBufferingBody) { brpc::policy::ParseCouchbaseMessage( &couchbase_buf, socket.get(), false, NULL).error()); - brpc::FLAGS_max_body_size = saved_max_body_size; } TEST_F(CouchbaseUnitTest, RequestBuilders) { diff --git a/test/brpc_http_rpc_protocol_unittest.cpp b/test/brpc_http_rpc_protocol_unittest.cpp index a99c834f21..a75a1225fb 100644 --- a/test/brpc_http_rpc_protocol_unittest.cpp +++ b/test/brpc_http_rpc_protocol_unittest.cpp @@ -80,18 +80,6 @@ int main(int argc, char* argv[]) { namespace { -class MaxBodySizeGuard { -public: - explicit MaxBodySizeGuard(uint64_t value) - : _saved(brpc::FLAGS_max_body_size) { - brpc::FLAGS_max_body_size = value; - } - ~MaxBodySizeGuard() { brpc::FLAGS_max_body_size = _saved; } - -private: - uint64_t _saved; -}; - static const std::string EXP_REQUEST = "hello"; static const std::string EXP_RESPONSE = "world"; static const std::string EXP_RESPONSE_CONTENT_LENGTH = "1024"; @@ -358,7 +346,8 @@ class HttpTest : public ::testing::Test{ }; TEST_F(HttpTest, reject_oversized_http_body) { - MaxBodySizeGuard guard(4); + GFLAGS_NAMESPACE::FlagSaver flag_saver; + brpc::FLAGS_max_body_size = 4; butil::IOBuf buf; buf.append("POST / HTTP/1.1\r\nContent-Length: 5\r\n\r\nhello"); @@ -375,7 +364,8 @@ TEST_F(HttpTest, reject_oversized_http_body) { } TEST_F(HttpTest, reject_oversized_chunked_http_body) { - MaxBodySizeGuard guard(4); + GFLAGS_NAMESPACE::FlagSaver flag_saver; + brpc::FLAGS_max_body_size = 4; butil::IOBuf buf; buf.append("POST / HTTP/1.1\r\nTransfer-Encoding: chunked\r\n\r\n" "3\r\nabc\r\n2\r\nde\r\n0\r\n\r\n"); diff --git a/test/brpc_memcache_unittest.cpp b/test/brpc_memcache_unittest.cpp index 2903547990..b1da958449 100644 --- a/test/brpc_memcache_unittest.cpp +++ b/test/brpc_memcache_unittest.cpp @@ -25,6 +25,7 @@ #include #include #include +#include #include namespace brpc { @@ -41,7 +42,7 @@ int main(int argc, char* argv[]) { namespace { TEST(MemcacheParserTest, RejectOversizedResponseBeforeBufferingBody) { - const uint64_t saved_max_body_size = brpc::FLAGS_max_body_size; + GFLAGS_NAMESPACE::FlagSaver flag_saver; brpc::FLAGS_max_body_size = 1024; brpc::SocketId id; @@ -59,7 +60,6 @@ TEST(MemcacheParserTest, RejectOversizedResponseBeforeBufferingBody) { brpc::policy::ParseMemcacheMessage( &buf, socket.get(), false, NULL).error()); - brpc::FLAGS_max_body_size = saved_max_body_size; } static pthread_once_t download_memcached_once = PTHREAD_ONCE_INIT;