diff options
-rw-r--r-- | src/mongo/rpc/op_msg.cpp | 33 | ||||
-rw-r--r-- | src/mongo/rpc/op_msg.h | 14 | ||||
-rw-r--r-- | src/mongo/rpc/op_msg_integration_test.cpp | 6 | ||||
-rw-r--r-- | src/mongo/rpc/reply_builder_interface.h | 2 | ||||
-rw-r--r-- | src/mongo/rpc/reply_builder_test.cpp | 16 |
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(); |