SERVER-43328 Enforce BSON size limit in OpMsgBuilder

5 files changed, 64 insertions, 7 deletions

diff --git a/src/mongo/rpc/op_msg.cpp b/src/mongo/rpc/op_msg.cpp
index d29b5ac2578..ca699693170 100644
--- a/src/mongo/rpc/op_msg.cpp
+++ b/src/mongo/rpc/op_msg.cpp
@@ -229,18 +229,32 @@ OpMsg OpMsg::parse(const Message& message) try {
     throw;
 }
 
-Message OpMsg::serialize() const {
-    OpMsgBuilder builder;
+namespace {
+void serializeHelper(const std::vector<OpMsg::DocumentSequence>& sequences,
+                     const BSONObj& body,
+                     OpMsgBuilder* output) {
     for (auto&& seq : sequences) {
-        auto docSeq = builder.beginDocSequence(seq.name);
+        auto docSeq = output->beginDocSequence(seq.name);
         for (auto&& obj : seq.objs) {
             docSeq.append(obj);
         }
     }
-    builder.beginBody().appendElements(body);
+    output->beginBody().appendElements(body);
+}
+}  // namespace
+
+Message OpMsg::serialize() const {
+    OpMsgBuilder builder;
+    serializeHelper(sequences, body, &builder);
     return builder.finish();
 }
 
+Message OpMsg::serializeWithoutSizeChecking() const {
+    OpMsgBuilder builder;
+    serializeHelper(sequences, body, &builder);
+    return builder.finishWithoutSizeChecking();
+}
+
 void OpMsg::shareOwnershipWith(const ConstSharedBuffer& buffer) {
     if (!body.isOwned()) {
         body.shareOwnershipWith(buffer);
@@ -293,6 +307,17 @@ BSONObjBuilder OpMsgBuilder::resumeBody() {
 AtomicWord<bool> OpMsgBuilder::disableDupeFieldCheck_forTest{false};
 
 Message OpMsgBuilder::finish() {
+    const auto size = _buf.len();
+    uassert(ErrorCodes::BSONObjectTooLarge,
+            str::stream() << "BSON size limit hit while building Message. Size: " << size << " (0x"
+                          << integerToHex(size) << "); maxSize: " << BSONObjMaxInternalSize << "("
+                          << (BSONObjMaxInternalSize / (1024 * 1024)) << "MB)",
+            size <= BSONObjMaxInternalSize);
+
+    return finishWithoutSizeChecking();
+}
+
+Message OpMsgBuilder::finishWithoutSizeChecking() {
     if (kDebugBuild && !disableDupeFieldCheck_forTest.load()) {
         std::set<StringData> seenFields;
         for (auto elem : resumeBody().asTempObj()) {
diff --git a/src/mongo/rpc/op_msg.h b/src/mongo/rpc/op_msg.h
index 3b13b586bff..9b677891375 100644
--- a/src/mongo/rpc/op_msg.h
+++ b/src/mongo/rpc/op_msg.h
@@ -121,6 +121,12 @@ struct OpMsg {
     Message serialize() const;
 
     /**
+     * Like serialize() but doesn't enforce max BSON size limits. This is almost never what you
+     * want. Prefer serialize() unless there's a good reason to skip the size check.
+     */
+    Message serializeWithoutSizeChecking() const;
+
+    /**
      * Makes all BSONObjs in this object share ownership with buffer.
      */
     void shareOwnershipWith(const ConstSharedBuffer& buffer);
@@ -235,10 +241,18 @@ public:
     /**
      * Finish building and return a Message ready to give to the networking layer for transmission.
      * It is illegal to call any methods on this object after calling this.
+     * Can throw BSONObjectTooLarge if the internal buffer has grown too large to be converted
+     * to a Message within the BSON size limit.
      */
     Message finish();
 
     /**
+     * Like finish() but doesn't enforce max BSON size limits. This is almost never what you want.
+     * Prefer finish() unless there's a good reason to skip the size check.
+     */
+    Message finishWithoutSizeChecking();
+
+    /**
      * Reset this object to its initial empty state. All previously appended data is lost.
      */
     void reset() {
diff --git a/src/mongo/rpc/op_msg_integration_test.cpp b/src/mongo/rpc/op_msg_integration_test.cpp
index 1f751a0528b..2222d3b6a1d 100644
--- a/src/mongo/rpc/op_msg_integration_test.cpp
+++ b/src/mongo/rpc/op_msg_integration_test.cpp
@@ -148,7 +148,7 @@ TEST(OpMsg, DocumentSequenceLargeDocumentMultiInsertWorks) {
         $db: "test"
     })"));
 
-    Message request = msgBuilder.finish();
+    Message request = msgBuilder.finishWithoutSizeChecking();
     Message reply;
     ASSERT_TRUE(conn->call(request, reply, false));
 
@@ -184,7 +184,7 @@ TEST(OpMsg, DocumentSequenceMaxWriteBatchWorks) {
 
     msgBuilder.setBody(std::move(body));
 
-    Message request = msgBuilder.finish();
+    Message request = msgBuilder.finishWithoutSizeChecking();
     Message reply;
     ASSERT_TRUE(conn->call(request, reply, false));
 
@@ -1064,7 +1064,7 @@ TEST(OpMsg, ServerHandlesReallyLargeMessagesGracefully) {
     OpMsgRequest request;
     request.body = bob.obj<BSONObj::LargeSizeTrait>();
     ASSERT_GT(request.body.objsize(), BSONObjMaxInternalSize);
-    auto requestMsg = request.serialize();
+    auto requestMsg = request.serializeWithoutSizeChecking();
 
     Message replyMsg;
     ASSERT(conn->call(requestMsg, replyMsg));
diff --git a/src/mongo/rpc/reply_builder_interface.h b/src/mongo/rpc/reply_builder_interface.h
index c0f93503c3c..860f52a5d02 100644
--- a/src/mongo/rpc/reply_builder_interface.h
+++ b/src/mongo/rpc/reply_builder_interface.h
@@ -113,6 +113,8 @@ public:
     /**
      * Writes data then transfers ownership of the message to the caller. The behavior of
      * calling any methods on the builder is subsequently undefined.
+     * Can throw BSONObjectTooLarge if the internal buffer has grown too large to be converted
+     * to a Message within the BSON size limit.
      */
     virtual Message done() = 0;
 
diff --git a/src/mongo/rpc/reply_builder_test.cpp b/src/mongo/rpc/reply_builder_test.cpp
index fcdb036c650..3b472bab4f0 100644
--- a/src/mongo/rpc/reply_builder_test.cpp
+++ b/src/mongo/rpc/reply_builder_test.cpp
@@ -166,6 +166,22 @@ TEST(OpMsgReplyBuilder, CommandError) {
     ASSERT_BSONOBJ_EQ(parsed.getCommandReply(), body);
 }
 
+TEST(OpMsgReplyBuilder, MessageOverBSONSizeLimit) {
+    rpc::OpMsgReplyBuilder r;
+    std::string bigStr(1024 * 1024 * 16, 'a');
+
+    {
+        // 'builder' is an unowned BSONObjBuilder and thus does none of its own size checking,
+        // allowing us to grow the OpMsgReplyBuilder past the bson object size limit.
+        auto builder = r.getBodyBuilder();
+        for (auto i = 0; i < 2; i++) {
+            builder.append("field" + std::to_string(i), bigStr);
+        }
+    }
+
+    ASSERT_THROWS_CODE(r.done(), DBException, ErrorCodes::BSONObjectTooLarge);
+}
+
 template <typename T>
 void testRoundTrip(rpc::ReplyBuilderInterface& replyBuilder, bool unifiedBodyAndMetadata) {
     auto metadata = buildMetadata();