Enforce body size limits in protocol parsers

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.

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'.

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);
     }

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);
         }

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.

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);
         }

test/brpc_couchbase_unittest.cpp

@@ -16,11 +16,17 @@
 // under the License.
 
 #include <brpc/couchbase.h>
+#include <brpc/policy/couchbase_protocol.h>
+#include <brpc/socket.h>
+#include <butil/sys_byteorder.h>
+#include <butil/iobuf.h>
 #include <butil/logging.h>
+#include <gflags/gflags.h>
 #include <gtest/gtest.h>
 
 namespace brpc {
 DECLARE_int32(idle_timeout_second);
+DECLARE_uint64(max_body_size);
 }
 
 int main(int argc, char* argv[]) {
@@ -34,6 +40,27 @@ namespace {
 // Unit Tests - No Server Required
 class CouchbaseUnitTest : public testing::Test {};
 
+TEST_F(CouchbaseUnitTest, RejectOversizedResponseBeforeBufferingBody) {
+  GFLAGS_NAMESPACE::FlagSaver flag_saver;
+  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());
+
+}
+
 TEST_F(CouchbaseUnitTest, RequestBuilders) {
   brpc::CouchbaseOperations::CouchbaseRequest req;
   req.Clear();
@@ -82,4 +109,4 @@ TEST_F(CouchbaseUnitTest, EdgeCases) {
   EXPECT_TRUE(true);
 }
 
-}  // namespace
+}  // namespace

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;
 }
 
@@ -344,6 +345,43 @@ class HttpTest : public ::testing::Test{
     MyAuthenticator _auth;
 };
 
+TEST_F(HttpTest, reject_oversized_http_body) {
+    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");
+
+    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) {
+    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");
+
+    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);

test/brpc_memcache_unittest.cpp

@@ -20,10 +20,17 @@
 #include "butil/logging.h"
 #include <brpc/memcache.h>
 #include <brpc/channel.h>
+#include <brpc/policy/memcache_binary_header.h>
+#include <brpc/policy/memcache_binary_protocol.h>
+#include <brpc/socket.h>
+#include <butil/iobuf.h>
+#include <butil/sys_byteorder.h>
+#include <gflags/gflags.h>
 #include <gtest/gtest.h>
 
 namespace brpc {
 DECLARE_int32(idle_timeout_second);
+DECLARE_uint64(max_body_size);
 } 
 
 int main(int argc, char* argv[]) {
@@ -33,6 +40,28 @@ int main(int argc, char* argv[]) {
 }
 
 namespace {
+
+TEST(MemcacheParserTest, RejectOversizedResponseBeforeBufferingBody) {
+    GFLAGS_NAMESPACE::FlagSaver flag_saver;
+    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());
+
+}
+
 static pthread_once_t download_memcached_once = PTHREAD_ONCE_INIT;
 static pid_t g_mc_pid = -1;