Bound wire-declared sizes and decompressed output during parsing

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

@@ -20,13 +20,23 @@
 #define BRPC_COMPRESS_H
 
 #include <google/protobuf/message.h>              // Message
+#include <gflags/gflags_declare.h>               // DECLARE_uint64
 #include "butil/iobuf.h"                           // butil::IOBuf
 #include "butil/logging.h"
 #include "brpc/options.pb.h"                     // CompressType
 #include "brpc/nonreflectable_message.h"
 
 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.h>   // ZeroCopyInputStream
 #include <google/protobuf/text_format.h>
 #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<int64_t>::max()
+        ? (int64_t)limit + 1 : std::numeric_limits<int64_t>::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;

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