Limit recursion depth of mcpack2pb parser

src/mcpack2pb/parser-inl.h

@@ -128,20 +128,27 @@ struct IsoItemsHead {
 } __attribute__((__packed__));
 
 inline ObjectIterator UnparsedValue::as_object() {
-    return ObjectIterator(_stream, _size);
+    return ObjectIterator(_stream, _size, _depth + 1);
 }
 
 inline ArrayIterator UnparsedValue::as_array() {
-    return ArrayIterator(_stream, _size);
+    return ArrayIterator(_stream, _size, _depth + 1);
 }
 
 inline ISOArrayIterator UnparsedValue::as_iso_array() {
     return ISOArrayIterator(_stream, _size);
 }
 
-inline void ObjectIterator::init(InputStream* stream, size_t size) {
+inline void ObjectIterator::init(InputStream* stream, size_t size, size_t depth) {
+    _depth = depth;
     _field_count = 0;
     _stream = stream;
+    if (_depth > (size_t)MAX_DEPTH) {
+        // The input is nested too deep. Parsing it would recurse until the
+        // stack overflows (CWE-674), so fail like the serializer does when
+        // the nesting level exceeds MAX_DEPTH.
+        return set_bad();
+    }
     _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
@@ -169,9 +176,16 @@ inline void ObjectIterator::init(InputStream* stream, size_t size) {
     operator++();
 }
 
-inline void ArrayIterator::init(InputStream* stream, size_t size) {
+inline void ArrayIterator::init(InputStream* stream, size_t size, size_t depth) {
+    _depth = depth;
     _item_count = 0;
     _stream = stream;
+    if (_depth > (size_t)MAX_DEPTH) {
+        // The input is nested too deep. Parsing it would recurse until the
+        // stack overflows (CWE-674), so fail like the serializer does when
+        // the nesting level exceeds MAX_DEPTH.
+        return set_bad();
+    }
     _expected_popped_bytes = _stream->popped_bytes() + sizeof(ItemsHead);
     _expected_popped_end = _stream->popped_bytes() + size;
     if (size < sizeof(ItemsHead)) {

src/mcpack2pb/parser.cpp

@@ -122,7 +122,7 @@ void ObjectIterator::operator++() {
         if (!_current_field.name.empty()) {
             _current_field.name.remove_suffix(1);
         }
-        _current_field.value.set((FieldType)head.type(), _stream, head.value_size());
+        _current_field.value.set((FieldType)head.type(), _stream, head.value_size(), _depth);
     } else if (first_byte & FIELD_SHORT_MASK) {
         FieldShortHead head;
         if (_stream->cut_packed_pod(&head) != sizeof(FieldShortHead) ||
@@ -143,7 +143,7 @@ void ObjectIterator::operator++() {
         if (!_current_field.name.empty()) {
             _current_field.name.remove_suffix(1);
         }
-        _current_field.value.set(type, _stream, head.value_size());
+        _current_field.value.set(type, _stream, head.value_size(), _depth);
     } else {
         FieldLongHead head;
         if (_stream->cut_packed_pod(&head) != sizeof(FieldLongHead) ||
@@ -162,7 +162,7 @@ void ObjectIterator::operator++() {
         if (!_current_field.name.empty()) {
             _current_field.name.remove_suffix(1);
         }
-        _current_field.value.set((FieldType)head.type(), _stream, head.value_size());
+        _current_field.value.set((FieldType)head.type(), _stream, head.value_size(), _depth);
     }
 }
 
@@ -206,7 +206,7 @@ void ArrayIterator::operator++() {
         if (name_size) {
             _stream->popn(name_size);
         }
-        _current_field.set((FieldType)head.type(), _stream, head.value_size());
+        _current_field.set((FieldType)head.type(), _stream, head.value_size(), _depth);
     } else if (first_byte & FIELD_SHORT_MASK) {
         FieldShortHead head;
         if (_stream->cut_packed_pod(&head) != sizeof(FieldShortHead) ||
@@ -227,7 +227,7 @@ void ArrayIterator::operator++() {
         if (name_size) {
             _stream->popn(name_size);
         }
-        _current_field.set(type, _stream, head.value_size());
+        _current_field.set(type, _stream, head.value_size(), _depth);
     } else {
         FieldLongHead head;
         if (_stream->cut_packed_pod(&head) != sizeof(FieldLongHead) ||
@@ -246,7 +246,7 @@ void ArrayIterator::operator++() {
         if (name_size) {
             _stream->popn(name_size);
         }
-        _current_field.set((FieldType)head.type(), _stream, head.value_size());
+        _current_field.set((FieldType)head.type(), _stream, head.value_size(), _depth);
     }
 }
 

src/mcpack2pb/parser.h

@@ -102,19 +102,34 @@ inline int capped_reserve_count(uint32_t item_count) {
 // Represent a piece of unparsed(and unread) data of InputStream.
 struct UnparsedValue {
     UnparsedValue()
-        : _type(FIELD_UNKNOWN), _stream(NULL), _size(0) {}
+        : _type(FIELD_UNKNOWN), _stream(NULL), _size(0), _depth(0) {}
     UnparsedValue(FieldType type, InputStream* stream, size_t size)
-        : _type(type), _stream(stream), _size(size) {}
+        : _type(type), _stream(stream), _size(size), _depth(0) {}
+    // `depth' is the number of containers this value is nested in
+    // (0 for the top-level object). An iterator unfolded from a value gets
+    // depth+1, i.e. the nesting level of the container itself. Input nested
+    // deeper than MAX_DEPTH is rejected, so that parsing a deeply nested
+    // object fails instead of recursing until the stack overflows.
+    UnparsedValue(FieldType type, InputStream* stream, size_t size, size_t depth)
+        : _type(type), _stream(stream), _size(size), _depth(depth) {}
+    // Sets the value, keeping the depth of the previous value (a reused
+    // value stays at the same nesting level). Internal code populates
+    // nested values with the 4-arg overload below.
     void set(FieldType type, InputStream* stream, size_t size) {
+        set(type, stream, size, _depth);
+    }
+    void set(FieldType type, InputStream* stream, size_t size, size_t depth) {
         _type = type;
         _stream = stream;
         _size = size;
+        _depth = depth;
     }
 
     FieldType type() const { return _type; }
     InputStream* stream() { return _stream; }
     const InputStream* stream() const { return _stream; }
     size_t size() const { return _size; }
+    size_t depth() const { return _depth; }
 
     // Convert to concrete value. These functions can only be called once!
     ObjectIterator as_object();
@@ -138,10 +153,11 @@ struct UnparsedValue {
 friend class ObjectIterator;
 friend class ArrayIterator;
     void set_end() { _type = FIELD_UNKNOWN; }
-    
+
     FieldType _type;
     InputStream* _stream;
     size_t _size;
+    size_t _depth;
 };
 
 std::ostream& operator<<(std::ostream& os, const UnparsedValue& value);
@@ -163,9 +179,15 @@ class ObjectIterator {
     };
 
     // Parse `n' bytes from `stream' as fields of an object.
-    ObjectIterator(InputStream* stream, size_t n) { init(stream, n); }
+    // `depth' is the nesting level of the container being iterated
+    // (1 for the top-level object, since the provided value is already
+    // nested in one container). Input nested deeper than MAX_DEPTH is
+    // rejected to avoid stack overflow on unbounded recursion (CWE-674),
+    // mirroring the serializer's limit.
+    ObjectIterator(InputStream* stream, size_t n, size_t depth = 1)
+    { init(stream, n, depth); }
     explicit ObjectIterator(UnparsedValue& value)
-    { init(value.stream(), value.size()); }
+    { init(value.stream(), value.size(), value.depth() + 1); }
     ~ObjectIterator() {}
 
     Field* operator->() { return &_current_field; }
@@ -177,21 +199,22 @@ class ObjectIterator {
     uint32_t field_count() const { return _field_count; }
 
 private:
-    void init(InputStream* stream, size_t n);
+    void init(InputStream* stream, size_t n, size_t depth);
     void set_bad() {
         set_end();
         _stream->set_bad();
     }
     void set_end() { _current_field.value._type = FIELD_UNKNOWN; }
     size_t left_size() const
     { return _expected_popped_end - _expected_popped_bytes; }
-    
+
     Field _current_field;
     uint32_t _field_count;
     std::string _name_backup_string;
     InputStream* _stream;
     size_t _expected_popped_bytes;
     size_t _expected_popped_end;
+    size_t _depth;
 };
 
 // Iterator all items in a (mcpack) array which should be created like this:
@@ -203,9 +226,10 @@ class ArrayIterator {
 public:
     typedef UnparsedValue Field;
 
-    ArrayIterator(InputStream* stream, size_t size) { init(stream, size); }
+    ArrayIterator(InputStream* stream, size_t size, size_t depth = 1)
+    { init(stream, size, depth); }
     explicit ArrayIterator(UnparsedValue& value)
-    { init(value.stream(), value.size()); }
+    { init(value.stream(), value.size(), value.depth() + 1); }
     ~ArrayIterator() {}
 
     Field* operator->() { return &_current_field; }
@@ -217,7 +241,7 @@ class ArrayIterator {
     uint32_t item_count() const { return _item_count; }
 
 private:
-    void init(InputStream* stream, size_t n);
+    void init(InputStream* stream, size_t n, size_t depth);
     void set_bad() {
         set_end();
         _stream->set_bad();
@@ -231,6 +255,7 @@ class ArrayIterator {
     InputStream* _stream;
     size_t _expected_popped_bytes;
     size_t _expected_popped_end;
+    size_t _depth;
 };
 
 // Iterator all items in an isomorphic array which should be created like this:

test/brpc_mcpack2pb_unittest.cpp

@@ -18,6 +18,7 @@
 // Unit tests for the mcpack2pb parser.
 
 #include <gtest/gtest.h>
+#include <pthread.h>
 #include "butil/iobuf.h"
 #include "mcpack2pb/parser.h"
 
@@ -156,6 +157,146 @@ TEST(Mcpack2pbParserTest, ArrayItemCountIsZeroWhenPayloadSmallerThanHeader) {
     EXPECT_EQ(0u, it.item_count());
 }
 
+// Builds the wire bytes of a recursive message `Node { repeated Node
+// children = 1; }' nesting `depth' levels of { children: [ ... ] }, the way
+// protoc-gen-mcpack serializes such a message.
+static void AppendU32(std::string* out, uint32_t value) {
+    char buf[4];
+    buf[0] = (char)(value & 0xff);
+    buf[1] = (char)((value >> 8) & 0xff);
+    buf[2] = (char)((value >> 16) & 0xff);
+    buf[3] = (char)((value >> 24) & 0xff);
+    out->append(buf, 4);
+}
+
+static std::string BuildRecursivePayload(int depth) {
+    // The innermost level is an empty object: an ItemsHead with no items.
+    std::string body;
+    AppendU32(&body, 0);
+    const std::string name = std::string("children\0", 9);  // trailing '\0'
+    // Reserve enough space for the final payload so that appending at each
+    // nesting level does not reallocate (the size of each wrapping level is
+    // 1 + 1 + 4 + body + 4 + item + 1 + 1 + 4 + name + arr).
+    body.reserve((size_t)depth * 30);
+    for (int i = 1; i < depth; ++i) {
+        std::string item;  // a FIELD_OBJECT item wrapping the inner payload
+        item.push_back(0x10);  // FIELD_OBJECT
+        item.push_back(0x00);  // name_size = 0
+        AppendU32(&item, (uint32_t)body.size());  // value_size
+        item.append(body);
+
+        std::string arr;  // an array holding a single item
+        AppendU32(&arr, 1);
+        arr.append(item);
+
+        std::string child;  // FIELD_ARRAY "children"
+        child.push_back(0x20);  // FIELD_ARRAY
+        child.push_back((char)name.size());  // name_size
+        AppendU32(&child, (uint32_t)arr.size());  // value_size
+        child.append(name);
+        child.append(arr);
+
+        std::string new_body;  // an object holding a single field
+        AppendU32(&new_body, 1);
+        new_body.append(child);
+        body.swap(new_body);
+    }
+    return body;
+}
+
+// Simulates the recursion pattern of the functions generated by
+// protoc-gen-mcpack for a message with a repeated message field
+// (e.g. Node.children): parse_<msg>_body_internal creates an ObjectIterator,
+// set_<msg>_<field> creates an ArrayIterator and calls
+// parse_<msg>_body_internal for each item.
+static bool ParseNodeInternal(mcpack2pb::UnparsedValue& value) {
+    mcpack2pb::ObjectIterator it(value);
+    for (; it != nullptr; ++it) {
+        if (it->name == "children") {
+            if (it->value.type() != mcpack2pb::FIELD_ARRAY) {
+                return false;
+            }
+            mcpack2pb::ArrayIterator it2(it->value);
+            for (; it2 != nullptr; ++it2) {
+                if (it2->type() != mcpack2pb::FIELD_OBJECT ||
+                    !ParseNodeInternal(*it2)) {
+                    return false;
+                }
+            }
+        }
+    }
+    return value.stream()->good();
+}
+
+struct ParseArgs {
+    const std::string* payload;
+    bool parse_ok;
+};
+
+static void* ParseOnSmallStack(void* arg) {
+    ParseArgs* args = static_cast<ParseArgs*>(arg);
+    butil::IOBuf buf;
+    buf.append(args->payload->data(), args->payload->size());
+    butil::IOBufAsZeroCopyInputStream zc_stream(buf);
+    mcpack2pb::InputStream stream(&zc_stream);
+    mcpack2pb::UnparsedValue value(mcpack2pb::FIELD_OBJECT, &stream,
+                                   buf.size());
+    args->parse_ok = ParseNodeInternal(value);
+    return nullptr;
+}
+
+// Parses `payload' on a 1 MB stack thread, the size of a NORMAL bthread
+// stack in brpc, so that the test behaves like a request served by brpc.
+static int ParseRecursivePayloadWith1MBStack(const std::string& payload,
+                                             bool* ok) {
+    ParseArgs args = { &payload, false };
+    pthread_attr_t attr;
+    int rc = pthread_attr_init(&attr);
+    if (rc != 0) {
+        return -1;
+    }
+    rc = pthread_attr_setstacksize(&attr, 1024 * 1024);
+    if (rc != 0) {
+        pthread_attr_destroy(&attr);
+        return -1;
+    }
+    pthread_t tid;
+    rc = pthread_create(&tid, &attr, ParseOnSmallStack, &args);
+    // The attribute is not needed right after create, destroy it here so
+    // that every path below shares one cleanup point.
+    pthread_attr_destroy(&attr);
+    if (rc != 0) {
+        return -1;
+    }
+    rc = pthread_join(tid, nullptr);
+    if (rc != 0) {
+        return -2;
+    }
+    *ok = args.parse_ok;
+    return 0;
+}
+
+TEST(Mcpack2pbParserTest, DeeplyNestedPayloadIsRejectedWithoutStackOverflow) {
+    // A few levels beyond MAX_DEPTH (128) suffice: such input must be
+    // rejected by the depth limit. Before the fix the parse accepted it
+    // (and far deeper input would overflow the 1 MB stack and crash the
+    // process), so this fails without the guard. Keeping the depth small
+    // keeps the test fast even under sanitizers.
+    const std::string payload = BuildRecursivePayload(mcpack2pb::MAX_DEPTH + 2);
+    bool ok = true;
+    ASSERT_EQ(0, ParseRecursivePayloadWith1MBStack(payload, &ok));
+    // The parse must fail cleanly instead of crashing the process.
+    EXPECT_FALSE(ok);
+}
+
+TEST(Mcpack2pbParserTest, ModeratelyNestedPayloadParses) {
+    // 32 levels are well within the depth limit and must parse cleanly.
+    const std::string payload = BuildRecursivePayload(32);
+    bool ok = false;
+    ASSERT_EQ(0, ParseRecursivePayloadWith1MBStack(payload, &ok));
+    EXPECT_TRUE(ok);
+}
+
 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