Bound wire-declared sizes and decompressed output during parsing

Several parsers that handle client-controlled input trusted sizes
declared on the wire without bounding them, which could make a
connection consume far more memory than the message itself:

- gzip/zlib/snappy decompression had no output cap: -max_body_size
  is only checked against the compressed bytes, so a small body
  could decompress to tens of GiB. Add -max_decompressed_body_size
  (default 32x -max_body_size, 0 means use the default) and enforce
  it in all three decompressors.
- AMF string readers resized the output buffer to the declared
  length before checking how many bytes were actually available;
  read the string in bounded chunks instead.
- RTMP chunk headers may re-declare a message while a previous
  message is still being assembled on the same chunk stream. Bound
  the declared message length by -max_body_size and drop the stale
  partial body when a new message header arrives.
- mcpack2pb trusted the wire-decla- mcpack2pb trusted the wire-decla- mcpack2pb trusted the wire-dected Reserve() c- mcpack2pb trusted the wire-decla- mcpack2pb trusted tand cap the reserved size.

Author: wwbmmm

src/brpc/amf.cpp

@@ -16,6 +16,7 @@
 // under the License.
 
 
+#include <algorithm>
 #include <google/protobuf/descriptor.h>
 #include "butil/sys_byteorder.h"
 #include "butil/logging.h"
@@ -286,6 +287,31 @@ AMFArray* AMFObject::MutableArray(const std::string& name) {
     return _fields[name].MutableArray();
 }
 
+// Read `len' bytes of string data in bounded chunks. The declared length
+// comes from the (untrusted) stream and may be much larger than the data
+// actually present, so growing the output as bytes arrive keeps a tiny
+// truncated message from forcing an allocation of up to
+// FLAGS_amf_max_string_size bytes before the availability check.
+static const size_t AMF_STRING_READ_CHUNK_SIZE = 64 * 1024;
+
+static bool ReadAMFStringData(std::string* str, AMFInputStream* stream,
+                              uint32_t len) {
+    str->clear();
+    size_t nread = 0;
+    while (nread < len) {
+        const size_t to_read =
+            std::min((size_t)len - nread, AMF_STRING_READ_CHUNK_SIZE);
+        str->resize(nread + to_read);
+        if (stream->cutn(&(*str)[nread], to_read) != to_read) {
+            str->clear();
+            LOG(ERROR) << "stream is not long enough";
+            return false;
+        }
+        nread += to_read;
+    }
+    return true;
+}
+
 static bool ReadAMFShortStringBody(std::string* str, AMFInputStream* stream) {
     uint16_t len = 0;
     if (stream->cut_u16(&len) != 2u) {
@@ -295,13 +321,7 @@ static bool ReadAMFShortStringBody(std::string* str, AMFInputStream* stream) {
     if (!CheckAMFStringSize(len)) {
         return false;
     }
-    str->resize(len);
-    if (len != 0 && stream->cutn(&(*str)[0], len) != len) {
-        str->clear();
-        LOG(ERROR) << "stream is not long enough";
-        return false;
-    }
-    return true;
+    return ReadAMFStringData(str, stream, len);
 }
 
 static bool ReadAMFLongStringBody(std::string* str, AMFInputStream* stream) {
@@ -313,13 +333,7 @@ static bool ReadAMFLongStringBody(std::string* str, AMFInputStream* stream) {
     if (!CheckAMFStringSize(len)) {
         return false;
     }
-    str->resize(len);
-    if (len != 0 && stream->cutn(&(*str)[0], len) != len) {
-        str->clear();
-        LOG(ERROR) << "stream is not long enough";
-        return false;
-    }
-    return true;
+    return ReadAMFStringData(str, stream, len);
 }
 
 bool ReadAMFString(std::string* str, AMFInputStream* stream) {

src/brpc/compress.cpp

@@ -16,6 +16,8 @@
 // under the License.
 
 
+#include <limits>
+#include <gflags/gflags.h>
 #include "butil/logging.h"
 #include "json2pb/json_to_pb.h"
 #include "brpc/compress.h"
@@ -24,6 +26,24 @@
 
 namespace brpc {
 
+DEFINE_uint64(max_decompressed_body_size, 0,
+              "Maximum size (in bytes) that a single compressed message body"
+              " may decompress to, guarding against decompression bombs."
+              " 0 (the default) means 32 times -max_body_size. Raise this"
+              " flag explicitly if larger decompressed messages are expected");
+
+uint64_t MaxDecompressedBodySize() {
+    const uint64_t limit = FLAGS_max_decompressed_body_size;
+    if (limit > 0) {
+        return limit;
+    }
+    const uint64_t base = FLAGS_max_body_size;
+    if (base > std::numeric_limits<uint64_t>::max() / 32) {
+        return std::numeric_limits<uint64_t>::max();
+    }
+    return base * 32;
+}
+
 static const int MAX_HANDLER_SIZE = 1024;
 static CompressHandler s_handler_map[MAX_HANDLER_SIZE] = { { nullptr, nullptr, nullptr } };
 

src/brpc/compress.h

@@ -27,6 +27,15 @@
 
 namespace brpc {
 
+DECLARE_uint64(max_decompressed_body_size);
+
+// Effective limit (in bytes) on the decompressed size of a single message
+// body: FLAGS_max_decompressed_body_size, or 32 x FLAGS_max_body_size when
+// the flag is 0 (the default). Decompressors must fail once their output
+// exceeds this limit, otherwise a small compressed body that passes
+// -max_body_size may expand to tens of GiB (decompression bomb).
+uint64_t MaxDecompressedBodySize();
+
 // Serializer can be used to implement custom serialization
 // before compression with user callback.
 class Serializer : public NonreflectableMessage<Serializer> {

src/brpc/policy/gzip_compress.cpp

@@ -16,7 +16,9 @@
 // under the License.
 
 
+#include <limits>
 #include <google/protobuf/io/gzip_stream.h>    // GzipXXXStream
+#include <google/protobuf/io/zero_copy_stream_impl_lite.h> // LimitingInputStream
 #include <google/protobuf/text_format.h>
 #include "butil/logging.h"
 #include "brpc/policy/gzip_compress.h"
@@ -73,11 +75,26 @@ static bool Decompress(const butil::IOBuf& data, google::protobuf::Message* msg,
                        google::protobuf::io::GzipInputStream::Format format) {
     butil::IOBufAsZeroCopyInputStream wrapper(data);
     google::protobuf::io::GzipInputStream gzip(&wrapper, format);
+    // Cap the decompressed size: zlib expands up to ~1032x, so a body that
+    // passes -max_body_size in compressed form may still decompress to tens
+    // of GiB (decompression bomb). The limiting stream stops feeding the
+    // parser once the cap is hit, bounding the memory materialized here.
+    const uint64_t limit = MaxDecompressedBodySize();
+    const int64_t hard_limit =
+        limit < (uint64_t)std::numeric_limits<int64_t>::max()
+        ? (int64_t)limit + 1 : std::numeric_limits<int64_t>::max();
+    google::protobuf::io::LimitingInputStream limited_in(&gzip, hard_limit);
     bool ok;
     if (msg->GetDescriptor() == Deserializer::descriptor()) {
-        ok = ((Deserializer*)msg)->DeserializeFrom(&gzip);
+        ok = ((Deserializer*)msg)->DeserializeFrom(&limited_in);
     } else {
-        ok = msg->ParseFromZeroCopyStream(&gzip);
+        ok = msg->ParseFromZeroCopyStream(&limited_in);
+    }
+    if (ok && (uint64_t)limited_in.ByteCount() > limit) {
+        LOG(WARNING) << "Decompressed size exceeds"
+                        " -max_decompressed_body_size=" << limit
+                     << ", format=" << Format2CStr(format);
+        return false;
     }
     if (!ok) {
         LOG(WARNING) << "Fail to deserialize input message="
@@ -141,6 +158,11 @@ inline bool GzipDecompressBase(
     butil::IOBufAsZeroCopyInputStream wrapper(data);
     google::protobuf::io::GzipInputStream in(&wrapper, format);
     butil::IOBufAsZeroCopyOutputStream out(msg);
+    // Cap the decompressed size: zlib expands up to ~1032x, so a body that
+    // passes -max_body_size in compressed form may still decompress to tens
+    // of GiB (decompression bomb).
+    const uint64_t limit = MaxDecompressedBodySize();
+    uint64_t total_out = 0;
     const void* data_in = nullptr;
     int size_in = 0;
     void* data_out = nullptr;
@@ -154,6 +176,13 @@ inline bool GzipDecompressBase(
         }
         const int size_cp = std::min(size_in, size_out);
         memcpy(data_out, data_in, size_cp);
+        total_out += size_cp;
+        if (total_out > limit) {
+            LOG(WARNING) << "Decompressed size exceeds"
+                            " -max_decompressed_body_size=" << limit
+                         << ", format=" << Format2CStr(format);
+            return false;
+        }
         size_in -= size_cp;
         data_in = (char*)data_in + size_cp;
         size_out -= size_cp;

src/brpc/policy/rtmp_protocol.cpp

@@ -1460,6 +1460,23 @@ ParseResult RtmpChunkStream::Feed(const RtmpBasicHeader& bh,
         }
         timestamp_delta = mh.timestamp;
         mh.message_length = ReadBigEndian3Bytes(p + 3);
+        if (mh.message_length > FLAGS_max_body_size) {
+            LOG(ERROR) << socket->remote_side() << ": message_length="
+                       << mh.message_length << " in chunk_stream=" << _cs_id
+                       << " is too large";
+            return MakeParseError(PARSE_ERROR_TOO_BIG_DATA);
+        }
+        if (!_r.msg_body.empty()) {
+            // The new message header arrived before the previous message on
+            // this chunk stream completed. Drop the stale partial body,
+            // otherwise it would prefix the new message and, with repeated
+            // mid-message headers, grow `msg_body' without bound.
+            LOG(WARNING) << socket->remote_side() << ": Discard "
+                         << _r.msg_body.size() << " bytes of an incomplete"
+                         " message in chunk_stream=" << _cs_id
+                         << " overridden by a ChunkType0 header";
+            _r.msg_body.clear();
+        }
         _r.left_message_length = mh.message_length;
         cur_chunk_size = std::min(chunk_size_in, _r.left_message_length);
         if (source->size() < header_len + cur_chunk_size) {
@@ -1507,6 +1524,23 @@ ParseResult RtmpChunkStream::Feed(const RtmpBasicHeader& bh,
         }
         mh.timestamp = _r.last_msg_header.timestamp + timestamp_delta;
         mh.message_length = ReadBigEndian3Bytes(p + 3);
+        if (mh.message_length > FLAGS_max_body_size) {
+            LOG(ERROR) << socket->remote_side() << ": message_length="
+                       << mh.message_length << " in chunk_stream=" << _cs_id
+                       << " is too large";
+            return MakeParseError(PARSE_ERROR_TOO_BIG_DATA);
+        }
+        if (!_r.msg_body.empty()) {
+            // The new message header arrived before the previous message on
+            // this chunk stream completed. Drop the stale partial body,
+            // otherwise it would prefix the new message and, with repeated
+            // mid-message headers, grow `msg_body' without bound.
+            LOG(WARNING) << socket->remote_side() << ": Discard "
+                         << _r.msg_body.size() << " bytes of an incomplete"
+                         " message in chunk_stream=" << _cs_id
+                         << " overridden by a ChunkType1 header";
+            _r.msg_body.clear();
+        }
         _r.left_message_length = mh.message_length;
         cur_chunk_size = std::min(chunk_size_in, _r.left_message_length);
         if (source->size() < header_len + cur_chunk_size) {

src/brpc/policy/snappy_compress.cpp

@@ -76,6 +76,23 @@ bool SnappyCompress(const butil::IOBuf& in, butil::IOBuf* out) {
 }
 
 bool SnappyDecompress(const butil::IOBuf& in, butil::IOBuf* out) {
+    {
+        // Reject bodies whose declared uncompressed length exceeds the
+        // decompression cap (decompression bomb): -max_body_size is checked
+        // against the compressed bytes only.
+        butil::IOBufAsSnappySource length_source(in);
+        uint32_t uncompressed_len = 0;
+        if (!butil::snappy::GetUncompressedLength(&length_source,
+                                                  &uncompressed_len)) {
+            return false;
+        }
+        if (uncompressed_len > MaxDecompressedBodySize()) {
+            LOG(WARNING) << "Uncompressed size=" << uncompressed_len
+                         << " exceeds -max_decompressed_body_size="
+                         << MaxDecompressedBodySize();
+            return false;
+        }
+    }
     butil::IOBufAsSnappySource source(in);
     butil::IOBufAsSnappySink sink(*out);
     return butil::snappy::Uncompress(&source, &sink);

src/mcpack2pb/generator.cpp

@@ -227,14 +227,14 @@ bool generate_declarations(const std::set<std::string>& ref_msgs,
     "  $msg$* const msg = static_cast<$msg$*>(msg_base);\n"             \
     "  if (value.type() == ::mcpack2pb::FIELD_ISOARRAY) {\n"               \
     "    ::mcpack2pb::ISOArrayIterator it(value);\n"                       \
-    "    msg->mutable_$lcfield$()->Reserve(it.item_count());\n"         \
+    "    msg->mutable_$lcfield$()->Reserve(::mcpack2pb::capped_reserve_count(it.item_count()));\n"         \
     "    for (; it != NULL; ++it) {\n"                                  \
     "      msg->add_$lcfield$(it.as_"#fntype "());\n"                   \
     "    }\n"                                                           \
     "    return value.stream()->good();\n"                              \
     "  } else if (value.type() == ::mcpack2pb::FIELD_ARRAY) {\n"           \
     "    ::mcpack2pb::ArrayIterator it(value);\n"                          \
-    "    msg->mutable_$lcfield$()->Reserve(it.item_count());\n"         \
+    "    msg->mutable_$lcfield$()->Reserve(::mcpack2pb::capped_reserve_count(it.item_count()));\n"         \
     "    for (; it != NULL; ++it) {\n"                                  \
     "      msg->add_$lcfield$(it->as_"#fntype "(\"$field$\"));\n" \
     "    }\n"                                                           \
@@ -323,14 +323,14 @@ static bool generate_parsing(const google::protobuf::Descriptor* d,
                     "  $msg$* const msg = static_cast<$msg$*>(msg_base);\n"
                     "  if (value.type() == ::mcpack2pb::FIELD_ISOARRAY) {\n"
                     "    ::mcpack2pb::ISOArrayIterator it(value);\n"
-                    "    msg->mutable_$lcfield$()->Reserve(it.item_count());\n" 
+                    "    msg->mutable_$lcfield$()->Reserve(::mcpack2pb::capped_reserve_count(it.item_count()));\n" 
                     "    for (; it != NULL; ++it) {\n"
                     "      msg->add_$lcfield$(($enum$)it.as_int32());\n"
                     "    }\n"
                     "    return value.stream()->good();\n"
                     "  } else if (value.type() == ::mcpack2pb::FIELD_ARRAY) {\n"            
                     "    ::mcpack2pb::ArrayIterator it(value);\n"
-                    "    msg->mutable_$lcfield$()->Reserve(it.item_count());\n"
+                    "    msg->mutable_$lcfield$()->Reserve(::mcpack2pb::capped_reserve_count(it.item_count()));\n"
                     "    for (; it != NULL; ++it) {\n"
                     "      msg->add_$lcfield$(($enum$)it->as_int32(\"$enum$\"));\n"
                     "    }\n"
@@ -361,7 +361,7 @@ static bool generate_parsing(const google::protobuf::Descriptor* d,
                     "  $msg$* const msg = static_cast<$msg$*>(msg_base);\n"
                     "  if (value.type() == ::mcpack2pb::FIELD_ARRAY) {\n"
                     "    ::mcpack2pb::ArrayIterator it(value);\n"
-                    "    msg->mutable_$lcfield$()->Reserve(it.item_count());\n"
+                    "    msg->mutable_$lcfield$()->Reserve(::mcpack2pb::capped_reserve_count(it.item_count()));\n"
                     "    for (; it != NULL; ++it) {\n"
                     "      if (it->type() == ::mcpack2pb::FIELD_STRING) {\n"
                     "        it->as_string(msg->add_$lcfield$(), \"$field$\");\n"
@@ -457,7 +457,7 @@ static bool generate_parsing(const google::protobuf::Descriptor* d,
                     "    return value.stream()->good();\n"
                     "  } else if (value.type() == ::mcpack2pb::FIELD_ARRAY) {\n"
                     "    ::mcpack2pb::ArrayIterator it(value);\n"
-                    "    msg->mutable_$lcfield$()->Reserve(it.item_count());\n"
+                    "    msg->mutable_$lcfield$()->Reserve(::mcpack2pb::capped_reserve_count(it.item_count()));\n"
                     "    for (; it != NULL; ++it) {\n"
                     "      if (it->type() == ::mcpack2pb::FIELD_OBJECT) {\n"
                     "        if (!parse_$vmsg2$_body_internal(msg->add_$lcfield$(), *it)) {\n"

src/mcpack2pb/parser-inl.h

@@ -149,6 +149,16 @@ inline void ObjectIterator::init(InputStream* stream, size_t size) {
         CHECK(false) << "buffer(size=" << size << ") is not enough";
         return set_bad();
     }
+    // `size' covers ItemsHead plus all fields and every field head takes at
+    // least 2 bytes (FieldFixedHead), thus a valid item_count never exceeds
+    // half of the remaining value size. The count is copied verbatim from
+    // the wire, reject inconsistent values instead of trusting them.
+    if (size < sizeof(ItemsHead) ||
+        items_head.item_count > (size - sizeof(ItemsHead)) / 2) {
+        CHECK(false) << "inconsistent item_count(" << items_head.item_count
+                     << ") and value_size(" << size << ")";
+        return set_bad();
+    }
     _field_count = items_head.item_count;
     operator++();
 }

src/mcpack2pb/parser.h

@@ -88,6 +88,17 @@ class ObjectIterator;
 class ArrayIterator;
 class ISOArrayIterator;
 
+// Bound the argument of RepeatedField::Reserve() calls that generated
+// parsing code derives from a wire-declared item count. The declared count
+// is under control of the remote side and is not necessarily backed by
+// actual bytes, so reserving it verbatim lets a tiny message trigger a huge
+// allocation. Repeated fields grow on demand past this bound, thus parsing
+// of genuinely large arrays is unaffected.
+inline int capped_reserve_count(uint32_t item_count) {
+    const uint32_t MAX_RESERVE_COUNT = 1024;
+    return (int)(item_count < MAX_RESERVE_COUNT ? item_count : MAX_RESERVE_COUNT);
+}
+
 // Represent a piece of unparsed(and unread) data of InputStream.
 struct UnparsedValue {
     UnparsedValue()

test/brpc_mcpack2pb_unittest.cpp

@@ -156,4 +156,41 @@ TEST(Mcpack2pbParserTest, ArrayItemCountIsZeroWhenPayloadSmallerThanHeader) {
     EXPECT_EQ(0u, it.item_count());
 }
 
+TEST(Mcpack2pbParserTest, ObjectItemCountIsRejectedWhenInconsistentWithSize) {
+    // An mcpack object whose ItemsHead declares an absurd field count for
+    // the given payload must be rejected instead of being trusted: the
+    // count is copied verbatim from the wire and every field head takes at
+    // least 2 bytes, so a valid item count never exceeds half of the
+    // remaining bytes.
+    const unsigned char data[] = {
+        0xff, 0xff, 0xff, 0x7f,  // item_count = 0x7fffffff
+        0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+        0x00, 0x00, 0x00, 0x00,
+    };
+    butil::IOBuf body;
+    body.append(data, sizeof(data));
+
+    butil::IOBufAsZeroCopyInputStream zc_stream(body);
+    mcpack2pb::InputStream stream(&zc_stream);
+    mcpack2pb::ObjectIterator it(&stream, sizeof(data));
+    EXPECT_TRUE(it == NULL);
+    EXPECT_FALSE(stream.good());
+}
+
+TEST(Mcpack2pbParserTest, EmptyObjectStillParses) {
+    // A consistent empty object head (item_count 0) must still initialize
+    // an empty iterator normally.
+    const unsigned char data[] = {
+        0x00, 0x00, 0x00, 0x00,  // item_count = 0
+    };
+    butil::IOBuf body;
+    body.append(data, sizeof(data));
+
+    butil::IOBufAsZeroCopyInputStream zc_stream(body);
+    mcpack2pb::InputStream stream(&zc_stream);
+    mcpack2pb::ObjectIterator it(&stream, sizeof(data));
+    EXPECT_TRUE(it == NULL);
+    EXPECT_TRUE(stream.good());
+}
+
 }  // namespace