SERVER-43328 Enforce BSON size limit in OpMsgBuilder

1 files changed, 29 insertions, 4 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()) {