SERVER-36606: Allow construction of large BSON objects

7 files changed, 150 insertions, 13 deletions

diff --git a/src/mongo/bson/bson_obj_test.cpp b/src/mongo/bson/bson_obj_test.cpp
index 82f95be072a..461e97eb636 100644
--- a/src/mongo/bson/bson_obj_test.cpp
+++ b/src/mongo/bson/bson_obj_test.cpp
@@ -675,4 +675,46 @@ TEST(BSONObj, addField) {
     ASSERT_BSONOBJ_EQ(obj, BSON("a" << 1 << "b" << 2));
 }
 
+TEST(BSONObj, sizeChecks) {
+    auto generateBuffer = [](std::int32_t size) {
+        std::vector<char> buffer(size);
+        DataRange bufferRange(&buffer.front(), &buffer.back());
+        ASSERT_OK(bufferRange.write(LittleEndian<int32_t>(size)));
+
+        return buffer;
+    };
+
+    {
+        // Implicitly assert that BSONObj constructor does not throw
+        // with standard size buffers.
+        auto normalBuffer = generateBuffer(15 * 1024 * 1024);
+        BSONObj obj(normalBuffer.data());
+    }
+
+    // Large buffers cause an exception to be thrown.
+    ASSERT_THROWS_CODE(
+        [&] {
+            auto largeBuffer = generateBuffer(17 * 1024 * 1024);
+            BSONObj obj(largeBuffer.data());
+        }(),
+        DBException,
+        ErrorCodes::BSONObjectTooLarge);
+
+
+    // Assert that the max size can be increased by passing BSONObj a tag type.
+    {
+        auto largeBuffer = generateBuffer(17 * 1024 * 1024);
+        BSONObj obj(largeBuffer.data(), BSONObj::LargeSizeTrait{});
+    }
+
+    // But a size is in fact being enforced.
+    ASSERT_THROWS_CODE(
+        [&]() {
+            auto hugeBuffer = generateBuffer(70 * 1024 * 1024);
+            BSONObj obj(hugeBuffer.data(), BSONObj::LargeSizeTrait{});
+        }(),
+        DBException,
+        ErrorCodes::BSONObjectTooLarge);
+}
+
 }  // unnamed namespace
diff --git a/src/mongo/bson/bsonelement.cpp b/src/mongo/bson/bsonelement.cpp
index 42151c23b50..96f2c21410c 100644
--- a/src/mongo/bson/bsonelement.cpp
+++ b/src/mongo/bson/bsonelement.cpp
@@ -568,7 +568,7 @@ bool BSONElement::binaryEqualValues(const BSONElement& rhs) const {
 
 BSONObj BSONElement::embeddedObjectUserCheck() const {
     if (MONGO_likely(isABSONObj()))
-        return BSONObj(value());
+        return BSONObj(value(), BSONObj::LargeSizeTrait{});
     std::stringstream ss;
     ss << "invalid parameter: expected an object (" << fieldName() << ")";
     uasserted(10065, ss.str());
@@ -577,7 +577,7 @@ BSONObj BSONElement::embeddedObjectUserCheck() const {
 
 BSONObj BSONElement::embeddedObject() const {
     verify(isABSONObj());
-    return BSONObj(value());
+    return BSONObj(value(), BSONObj::LargeSizeTrait{});
 }
 
 BSONObj BSONElement::codeWScopeObject() const {
diff --git a/src/mongo/bson/bsonelement_test.cpp b/src/mongo/bson/bsonelement_test.cpp
index 9b4038249a5..a0ca66f5061 100644
--- a/src/mongo/bson/bsonelement_test.cpp
+++ b/src/mongo/bson/bsonelement_test.cpp
@@ -108,5 +108,23 @@ TEST(BSONElement, TimestampToString) {
     ASSERT_EQ(obj["ts3"].toString(false, false), "Timestamp(4294967295, 4294967295)");
 }
 
+TEST(BSONElement, ExtractLargeSubObject) {
+    std::int32_t size = 17 * 1024 * 1024;
+    std::vector<char> buffer(size);
+    DataRange bufferRange(&buffer.front(), &buffer.back());
+    ASSERT_OK(bufferRange.write(LittleEndian<int32_t>(size)));
+
+    BSONObj obj(buffer.data(), BSONObj::LargeSizeTrait{});
+
+    BSONObjBuilder bigObjectBuilder;
+    bigObjectBuilder.append("a", obj);
+    BSONObj bigObj = bigObjectBuilder.obj<BSONObj::LargeSizeTrait>();
+
+    BSONElement element = bigObj["a"];
+    ASSERT_EQ(BSONType::Object, element.type());
+
+    BSONObj subObj = element.Obj();
+}
+
 }  // namespace
 }  // namespace mongo
diff --git a/src/mongo/bson/bsonobj.cpp b/src/mongo/bson/bsonobj.cpp
index 2fd01c2d8e2..0cbb7b91330 100644
--- a/src/mongo/bson/bsonobj.cpp
+++ b/src/mongo/bson/bsonobj.cpp
@@ -87,12 +87,12 @@ using namespace std;
 
 /* BSONObj ------------------------------------------------------------*/
 
-void BSONObj::_assertInvalid() const {
+void BSONObj::_assertInvalid(int maxSize) const {
     StringBuilder ss;
     int os = objsize();
     ss << "BSONObj size: " << os << " (0x" << integerToHex(os) << ") is invalid. "
        << "Size must be between 0 and " << BSONObjMaxInternalSize << "("
-       << (BSONObjMaxInternalSize / (1024 * 1024)) << "MB)";
+       << (maxSize / (1024 * 1024)) << "MB)";
     try {
         BSONElement e = firstElement();
         ss << " First element: " << e.toString();
diff --git a/src/mongo/bson/bsonobj.h b/src/mongo/bson/bsonobj.h
index 2c9f41ceb60..69a06b2e76f 100644
--- a/src/mongo/bson/bsonobj.h
+++ b/src/mongo/bson/bsonobj.h
@@ -95,6 +95,13 @@ class BSONObjStlIterator;
  */
 class BSONObj {
 public:
+    struct DefaultSizeTrait {
+        constexpr static int MaxSize = BSONObjMaxInternalSize;
+    };
+    struct LargeSizeTrait {
+        constexpr static int MaxSize = BufferMaxSize;
+    };
+
     // Declared in bsonobj_comparator_interface.h.
     class ComparatorInterface;
 
@@ -125,8 +132,9 @@ public:
     /** Construct a BSONObj from data in the proper format.
      *  Use this constructor when something else owns bsonData's buffer
     */
-    explicit BSONObj(const char* bsonData) {
-        init(bsonData);
+    template <typename Traits = DefaultSizeTrait>
+    explicit BSONObj(const char* bsonData, Traits t = Traits{}) {
+        init<Traits>(bsonData);
     }
 
     explicit BSONObj(ConstSharedBuffer ownedBuffer)
@@ -379,9 +387,12 @@ public:
     }
 
     /** performs a cursory check on the object's size only. */
+    template <typename Traits = DefaultSizeTrait>
     bool isValid() const {
+        static_assert(Traits::MaxSize > 0 && Traits::MaxSize <= std::numeric_limits<int>::max(),
+                      "BSONObj maximum size must be within possible limits");
         int x = objsize();
-        return x > 0 && x <= BSONObjMaxInternalSize;
+        return x > 0 && x <= Traits::MaxSize;
     }
 
     /**
@@ -564,12 +575,13 @@ public:
     }
 
 private:
-    void _assertInvalid() const;
+    void _assertInvalid(int maxSize) const;
 
+    template <typename Traits = DefaultSizeTrait>
     void init(const char* data) {
         _objdata = data;
-        if (!isValid())
-            _assertInvalid();
+        if (!isValid<Traits>())
+            _assertInvalid(Traits::MaxSize);
     }
 
     const char* _objdata;
diff --git a/src/mongo/bson/bsonobjbuilder.h b/src/mongo/bson/bsonobjbuilder.h
index ef92065d9a0..35402a3a3ed 100644
--- a/src/mongo/bson/bsonobjbuilder.h
+++ b/src/mongo/bson/bsonobjbuilder.h
@@ -666,9 +666,10 @@ public:
      * The returned BSONObj will free the buffer when it is finished.
      * @return owned BSONObj
     */
+    template <typename BSONTraits = BSONObj::DefaultSizeTrait>
     BSONObj obj() {
         massert(10335, "builder does not own memory", owned());
-        auto out = done();
+        auto out = done<BSONTraits>();
         out.shareOwnershipWith(_b.release());
         return out;
     }
@@ -678,8 +679,9 @@ public:
         scope -- very important to keep in mind.  Use obj() if you
         would like the BSONObj to last longer than the builder.
     */
+    template <typename BSONTraits = BSONObj::DefaultSizeTrait>
     BSONObj done() {
-        return BSONObj(_done());
+        return BSONObj(_done(), BSONTraits{});
     }
 
     // Like 'done' above, but does not construct a BSONObj to return to the caller.
@@ -692,7 +694,7 @@ public:
         Intended use case: append a field if not already there.
     */
     BSONObj asTempObj() {
-        BSONObj temp(_done());
+        BSONObj temp(_done(), BSONObj::LargeSizeTrait{});
         _b.setlen(_b.len() - 1);  // next append should overwrite the EOO
         _b.reserveBytes(1);       // Rereserve room for the real EOO
         _doneCalled = false;
diff --git a/src/mongo/bson/bsonobjbuilder_test.cpp b/src/mongo/bson/bsonobjbuilder_test.cpp
index 6debaa2fa3a..39ce08cc588 100644
--- a/src/mongo/bson/bsonobjbuilder_test.cpp
+++ b/src/mongo/bson/bsonobjbuilder_test.cpp
@@ -446,5 +446,68 @@ TEST(BSONObjBuilderTest, SeedingBSONObjBuilderWithNonrootedUnownedBsonWorks) {
     ASSERT_NE(static_cast<const void*>(obj.objdata()), static_cast<const void*>(origObjPtr));
 }
 
+TEST(BSONObjBuilderTest, SizeChecks) {
+    auto generateBuffer = [](std::int32_t size) {
+        std::vector<char> buffer(size);
+        DataRange bufferRange(&buffer.front(), &buffer.back());
+        ASSERT_OK(bufferRange.write(LittleEndian<int32_t>(size)));
+
+        return buffer;
+    };
+
+    {
+        // Implicitly assert that BSONObjBuilder does not throw
+        // with standard size buffers.
+        auto normalBuffer = generateBuffer(15 * 1024 * 1024);
+        BSONObj obj(normalBuffer.data());
+
+        BSONObjBuilder builder;
+        builder.append("a", obj);
+        BSONObj finalObj = builder.obj();
+    }
+
+    // Large buffers cause an exception to be thrown.
+    ASSERT_THROWS_CODE(
+        [&] {
+            auto largeBuffer = generateBuffer(17 * 1024 * 1024);
+            BSONObj obj(largeBuffer.data(), BSONObj::LargeSizeTrait{});
+
+            BSONObjBuilder builder;
+            builder.append("a", obj);
+            BSONObj finalObj = builder.obj();
+        }(),
+        DBException,
+        ErrorCodes::BSONObjectTooLarge);
+
+
+    // Assert that the max size can be increased by passing BSONObj a tag type.
+    {
+        auto largeBuffer = generateBuffer(17 * 1024 * 1024);
+        BSONObj obj(largeBuffer.data(), BSONObj::LargeSizeTrait{});
+
+        BSONObjBuilder builder;
+        builder.append("a", obj);
+        BSONObj finalObj = builder.obj<BSONObj::LargeSizeTrait>();
+    }
+
+    // But a size is in fact being enforced.
+    {
+        auto largeBuffer = generateBuffer(40 * 1024 * 1024);
+        BSONObj obj(largeBuffer.data(), BSONObj::LargeSizeTrait{});
+        BSONObjBuilder builder;
+        ASSERT_THROWS(
+            [&]() {
+
+                for (StringData character : {"a", "b", "c"}) {
+                    builder.append(character, obj);
+                }
+                BSONObj finalObj = builder.obj<BSONObj::LargeSizeTrait>();
+
+            }(),
+            DBException);
+    }
+}
+
+
 }  // namespace
 }  // namespace mongo