From 2be45154177da4ab98dd70af7c9fc47b6679c106 Mon Sep 17 00:00:00 2001 From: wwbmmm Date: Sat, 22 Aug 2026 19:44:56 +0800 Subject: [PATCH] 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. The output cap is applied via a local ZeroCopyInputStream wrapper so it builds against every protobuf version CI uses (>= 3.5.1); the wrapper clips oversized blocks without double-backing-up the wrapped stream. - 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 pr- RTMP chunk headers may re-declare a message while a pr- RTMP chunk headers may re-deed - RTMP chunk headers may re-declare a message while a pr- RTMP ody- RTMP chunk headers may re-declare a message while a pr- RTwir- RTMP chunk headers may re-declare a message while a pr- cou- RTMP chunk headers may re-declare a message while a pr- RTMP ch r- RTMP chunk headers may re-declare a message while a pr- RTMPvalues instead of CHECK-fatal. --- src/brpc/amf.cpp | 42 +++++--- src/brpc/compress.cpp | 20 ++++ src/brpc/compress.h | 10 ++ src/brpc/policy/gzip_compress.cpp | 112 ++++++++++++++++++++- src/brpc/policy/rtmp_protocol.cpp | 34 +++++++ src/brpc/policy/snappy_compress.cpp | 17 ++++ src/mcpack2pb/generator.cpp | 12 +-- src/mcpack2pb/parser-inl.h | 18 +++- src/mcpack2pb/parser.h | 11 ++ test/brpc_mcpack2pb_unittest.cpp | 37 +++++++ test/brpc_rtmp_unittest.cpp | 76 ++++++++++++++ test/brpc_snappy_compress_unittest.cpp | 32 ++++++ test/brpc_sofa_pbrpc_protocol_unittest.cpp | 14 +++ 13 files changed, 412 insertions(+), 23 deletions(-) diff --git a/src/brpc/amf.cpp b/src/brpc/amf.cpp index 3cd44d3728..290ce64258 100644 --- a/src/brpc/amf.cpp +++ b/src/brpc/amf.cpp @@ -16,6 +16,7 @@ // under the License. +#include #include #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) { diff --git a/src/brpc/compress.cpp b/src/brpc/compress.cpp index 9f9939c9f3..9b026a207a 100644 --- a/src/brpc/compress.cpp +++ b/src/brpc/compress.cpp @@ -16,6 +16,8 @@ // under the License. +#include +#include #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::max() / 32) { + return std::numeric_limits::max(); + } + return base * 32; +} + static const int MAX_HANDLER_SIZE = 1024; static CompressHandler s_handler_map[MAX_HANDLER_SIZE] = { { nullptr, nullptr, nullptr } }; diff --git a/src/brpc/compress.h b/src/brpc/compress.h index a6c61648e3..c9e39f879f 100644 --- a/src/brpc/compress.h +++ b/src/brpc/compress.h @@ -20,6 +20,7 @@ #define BRPC_COMPRESS_H #include // Message +#include // DECLARE_uint64 #include "butil/iobuf.h" // butil::IOBuf #include "butil/logging.h" #include "brpc/options.pb.h" // CompressType @@ -27,6 +28,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 { diff --git a/src/brpc/policy/gzip_compress.cpp b/src/brpc/policy/gzip_compress.cpp index 73a6f02f42..5f1ad72b8f 100644 --- a/src/brpc/policy/gzip_compress.cpp +++ b/src/brpc/policy/gzip_compress.cpp @@ -16,7 +16,9 @@ // under the License. +#include #include // GzipXXXStream +#include // ZeroCopyInputStream #include #include "butil/logging.h" #include "brpc/policy/gzip_compress.h" @@ -26,6 +28,81 @@ namespace brpc { namespace policy { +namespace { + +// A ZeroCopyInputStream wrapper that stops reading from the underlying stream +// once a limit of bytes has been handed out. Different protobuf releases +// disagree on the availability/location of the stock LimitingInputStream (it +// does not exist before ~3.19), so implement the same behaviour locally to +// stay portable across the protobuf versions CI builds against. +class DelegatingLimitingInputStream : public google::protobuf::io::ZeroCopyInputStream { +public: + DelegatingLimitingInputStream(google::protobuf::io::ZeroCopyInputStream* input, + int64_t limit) + : _input(input), _limit(limit), _bytes_read(0), _excess(0) {} + + bool Next(const void** data, int* size) override { + if (_bytes_read >= _limit) { + return false; + } + if (!_input->Next(data, size)) { + return false; + } + const int64_t total = _bytes_read + *size; + if (total > _limit) { + // Clip the tail that does not fit below the limit and record how + // much of the wrapped stream's block is being withheld instead of + // calling BackUp() right now: a consumer backing up the exposed + // prefix would otherwise trigger a second BackUp() on the + // wrapped stream without an intervening Next(), which violates + // the ZeroCopyInputStream contract. + _excess = (int)(total - _limit); + *size -= _excess; + _bytes_read = _limit; + } else { + _excess = 0; + _bytes_read = total; + } + return true; + } + + void BackUp(int count) override { + _bytes_read -= count; + if (_excess > 0) { + // Give back the requested prefix together with the clipped tail + // in the single BackUp() the wrapped stream allows after the last + // Next(), landing both at the right offset. + _input->BackUp(count + _excess); + _excess = 0; + } else { + _input->BackUp(count); + } + } + + bool Skip(int count) override { + // Bound the skip by what still fits below the limit. + const int64_t remaining = _limit - _bytes_read; + if (count > remaining) { + return false; + } + if (_input->Skip(count)) { + _bytes_read += count; + return true; + } + return false; + } + + int64_t ByteCount() const override { return _bytes_read; } + +private: + google::protobuf::io::ZeroCopyInputStream* _input; + int64_t _limit; + int64_t _bytes_read; + int _excess; +}; + +} // namespace + const char* Format2CStr(google::protobuf::io::GzipOutputStream::Format format) { switch (format) { case google::protobuf::io::GzipOutputStream::GZIP: @@ -73,11 +150,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::max() + ? (int64_t)limit + 1 : std::numeric_limits::max(); + DelegatingLimitingInputStream 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 +233,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 +251,17 @@ 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); + // out.Next() already moved the whole output block into `msg'; + // give back the unwritten tail (still uninitialized) before + // leaving, otherwise it stays in the caller's IOBuf. + out.BackUp(size_out); + return false; + } size_in -= size_cp; data_in = (char*)data_in + size_cp; size_out -= size_cp; diff --git a/src/brpc/policy/rtmp_protocol.cpp b/src/brpc/policy/rtmp_protocol.cpp index 2b50c6a3f0..8ad63fb958 100644 --- a/src/brpc/policy/rtmp_protocol.cpp +++ b/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) { diff --git a/src/brpc/policy/snappy_compress.cpp b/src/brpc/policy/snappy_compress.cpp index 8019b97b3c..e78d3c8b85 100644 --- a/src/brpc/policy/snappy_compress.cpp +++ b/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); diff --git a/src/mcpack2pb/generator.cpp b/src/mcpack2pb/generator.cpp index 26aa4b5bd8..f5a46a5c4b 100644 --- a/src/mcpack2pb/generator.cpp +++ b/src/mcpack2pb/generator.cpp @@ -227,14 +227,14 @@ bool generate_declarations(const std::set& 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" diff --git a/src/mcpack2pb/parser-inl.h b/src/mcpack2pb/parser-inl.h index 235bb5405e..bdfb95e5ad 100644 --- a/src/mcpack2pb/parser-inl.h +++ b/src/mcpack2pb/parser-inl.h @@ -144,9 +144,25 @@ inline void ObjectIterator::init(InputStream* stream, size_t size) { _stream = stream; _expected_popped_bytes = _stream->popped_bytes() + sizeof(ItemsHead); _expected_popped_end = _stream->popped_bytes() + size; + // Every field head takes at least 2 bytes (FieldFixedHead), so 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. Guard the size before reading ItemsHead so payloads + // shorter than the header are not read past their declared boundary. + // Note that these are wire-controlled inputs, so reject (set_bad) rather + // than CHECK-fatal, which would terminate the process. + if (size < sizeof(ItemsHead)) { + LOG(ERROR) << "buffer(size=" << size << ") is not enough"; + return set_bad(); + } ItemsHead items_head; if (_stream->cut_packed_pod(&items_head) != sizeof(ItemsHead)) { - CHECK(false) << "buffer(size=" << size << ") is not enough"; + LOG(ERROR) << "buffer(size=" << size << ") is not enough"; + return set_bad(); + } + if (items_head.item_count > (size - sizeof(ItemsHead)) / 2) { + LOG(ERROR) << "inconsistent item_count(" << items_head.item_count + << ") and value_size(" << size << ")"; return set_bad(); } _field_count = items_head.item_count; diff --git a/src/mcpack2pb/parser.h b/src/mcpack2pb/parser.h index d897c5e1b4..d9cc672698 100644 --- a/src/mcpack2pb/parser.h +++ b/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() diff --git a/test/brpc_mcpack2pb_unittest.cpp b/test/brpc_mcpack2pb_unittest.cpp index c0261540d3..ad040a5bf9 100644 --- a/test/brpc_mcpack2pb_unittest.cpp +++ b/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 diff --git a/test/brpc_rtmp_unittest.cpp b/test/brpc_rtmp_unittest.cpp index 9fce391986..07cbc07ea8 100644 --- a/test/brpc_rtmp_unittest.cpp +++ b/test/brpc_rtmp_unittest.cpp @@ -659,6 +659,82 @@ TEST(RtmpTest, amf_rejects_oversized_ecma_array_count) { EXPECT_FALSE(brpc::ReadAMFObject(&obj, &istream)); } +TEST(RtmpTest, amf_truncated_long_string_does_not_allocate_declared_size) { + // Regression: a tiny message declaring a huge (but under-the-cap) + // string length must not cause the declared size to be allocated + // before the bytes are actually available in the stream. + const uint32_t declared_len = 8 * 1024 * 1024; + std::string req_buf; + AppendAMFLongStringHeader(&req_buf, declared_len); + req_buf.append("only-a-few-bytes", 16); + + google::protobuf::io::ArrayInputStream zc_stream(req_buf.data(), req_buf.size()); + brpc::AMFInputStream istream(&zc_stream); + std::string result; + EXPECT_FALSE(brpc::ReadAMFString(&result, &istream)); + EXPECT_TRUE(result.empty()); + // Reading is chunked, so a truncated stream leaves at most one chunk + // of capacity behind instead of the full declared length. + EXPECT_LT(result.capacity(), (size_t)declared_len); +} + +TEST(RtmpTest, amf_reads_long_string_larger_than_one_chunk) { + const std::string big(200 * 1024, 'x'); + std::string req_buf; + { + google::protobuf::io::StringOutputStream zc_stream(&req_buf); + brpc::AMFOutputStream ostream(&zc_stream); + brpc::WriteAMFString(big, &ostream); + ASSERT_TRUE(ostream.good()); + } + google::protobuf::io::ArrayInputStream zc_stream(req_buf.data(), req_buf.size()); + brpc::AMFInputStream istream(&zc_stream); + std::string result; + ASSERT_TRUE(brpc::ReadAMFString(&result, &istream)); + ASSERT_EQ(big, result); +} + +TEST(RtmpTest, chunk_stream_rejects_message_length_over_max_body_size) { + int pipe_fds[2]; + ASSERT_EQ(0, pipe(pipe_fds)); + butil::fd_guard guard0(pipe_fds[0]); // read end, closed by this guard + butil::fd_guard guard1(pipe_fds[1]); // write end, handed over to Socket + + brpc::SocketId id; + brpc::SocketOptions options; + options.fd = guard1.release(); // Socket takes ownership of the fd + ASSERT_EQ(0, brpc::Socket::Create(options, &id)); + brpc::SocketUniquePtr sock; + ASSERT_EQ(0, brpc::Socket::Address(id, &sock)); + + brpc::policy::RtmpContext ctx(nullptr, nullptr); + ctx.SetState(sock->remote_side(), + brpc::policy::RtmpContext::STATE_RECEIVED_C2); + + // The message length declared by a chunk header is remote-controlled and + // was never bounded: with repeated mid-message headers a connection's + // reassembly buffer could grow without limit. A type-0 header declaring + // a length above -max_body_size must be rejected up front. + GFLAGS_NAMESPACE::FlagSaver flag_saver; + brpc::FLAGS_max_body_size = 1024; + + std::string chunk; + chunk.push_back((char)0x02); // basic header: fmt=0, cs_id=2 + chunk.append(3, '\0'); // timestamp = 0 + chunk.push_back('\0'); // message_length (3 bytes) = 4096 + chunk.push_back((char)0x10); + chunk.push_back('\0'); + chunk.push_back((char)0x02); // message_type = Abort + chunk.append(4, '\0'); // stream_id = 0 (little endian) + chunk.append(128, '\0'); // one full chunk of payload + + butil::IOBuf buf; + buf.append(chunk); + ASSERT_EQ(brpc::PARSE_ERROR_TOO_BIG_DATA, + ctx.Feed(&buf, sock.get()).error()); +} + + TEST(RtmpTest, amf_rejects_oversized_strict_array_count) { ScopedAMFLimit scoped_limit(&brpc::FLAGS_amf_max_array_size, 1); diff --git a/test/brpc_snappy_compress_unittest.cpp b/test/brpc_snappy_compress_unittest.cpp index 94b54dfdb9..a71b176656 100644 --- a/test/brpc_snappy_compress_unittest.cpp +++ b/test/brpc_snappy_compress_unittest.cpp @@ -20,12 +20,14 @@ // Date: 2015/01/20 19:01:06 #include +#include #include "gperftools_helper.h" #include "butil/third_party/snappy/snappy.h" #include "butil/macros.h" #include "butil/iobuf.h" #include "butil/time.h" #include "snappy_message.pb.h" +#include "brpc/compress.h" #include "brpc/policy/snappy_compress.h" #include "brpc/policy/gzip_compress.h" @@ -253,3 +255,33 @@ TEST_F(test_compress_method, mass_snappy_iobuf) { ASSERT_TRUE(strcmp(check_str.c_str(), text) == 0); delete [] text; } + +TEST_F(test_compress_method, decompressed_size_capped) { + // Regression test: decompressors used to enforce no output limit, so a + // small compressed body (checked against -max_body_size in compressed + // form only) could decompress to tens of GiB (decompression bomb). + GFLAGS_NAMESPACE::FlagSaver flag_saver; + brpc::FLAGS_max_decompressed_body_size = 1024; + + butil::IOBuf raw; + raw.append(std::string(64 * 1024, '\0')); + + butil::IOBuf gzipped; + ASSERT_TRUE(brpc::policy::GzipCompress(raw, &gzipped, nullptr)); + butil::IOBuf out; + ASSERT_FALSE(brpc::policy::GzipDecompress(gzipped, &out)); + + butil::IOBuf snappied; + ASSERT_TRUE(brpc::policy::SnappyCompress(raw, &snappied)); + out.clear(); + ASSERT_FALSE(brpc::policy::SnappyDecompress(snappied, &out)); + + // Payloads under the cap still decompress fine. + brpc::FLAGS_max_decompressed_body_size = 1024 * 1024; + out.clear(); + ASSERT_TRUE(brpc::policy::GzipDecompress(gzipped, &out)); + ASSERT_EQ(raw.size(), out.size()); + out.clear(); + ASSERT_TRUE(brpc::policy::SnappyDecompress(snappied, &out)); + ASSERT_EQ(raw.size(), out.size()); +} diff --git a/test/brpc_sofa_pbrpc_protocol_unittest.cpp b/test/brpc_sofa_pbrpc_protocol_unittest.cpp index 5a44e89c89..e03575e59c 100644 --- a/test/brpc_sofa_pbrpc_protocol_unittest.cpp +++ b/test/brpc_sofa_pbrpc_protocol_unittest.cpp @@ -282,6 +282,20 @@ TEST_F(SofaTest, process_response_after_eof) { ASSERT_TRUE(_socket->Failed()); } +TEST_F(SofaTest, reject_huge_meta_size) { + // A header declaring a huge meta_size must be rejected up front, + // otherwise the parser keeps buffering meta_size bytes although the + // body itself passes the -max_body_size check. + butil::IOBuf buf; + const uint32_t meta_size = 0xFFFFFFFFu; + const uint64_t body_size = 0; + const uint64_t msg_size = meta_size + body_size; + AppendSofaTestHeader(&buf, meta_size, body_size, msg_size); + brpc::ParseResult pr = + brpc::policy::ParseSofaMessage(&buf, _socket.get(), false, nullptr); + ASSERT_EQ(brpc::PARSE_ERROR_TOO_BIG_DATA, pr.error()); +} + TEST_F(SofaTest, process_response_error_code) { const int ERROR_CODE = 12345; brpc::policy::SofaRpcMeta meta;