SERVER-42060 Limit maximum recursion depth for KeyString toBsonValue

2 files changed, 67 insertions, 12 deletions

diff --git a/src/mongo/db/storage/key_string.cpp b/src/mongo/db/storage/key_string.cpp
index 6dfe799ba8d..a1f23adba73 100644
--- a/src/mongo/db/storage/key_string.cpp
+++ b/src/mongo/db/storage/key_string.cpp
@@ -38,6 +38,7 @@
 
 #include "mongo/base/data_cursor.h"
 #include "mongo/base/data_view.h"
+#include "mongo/bson/bson_depth.h"
 #include "mongo/platform/bits.h"
 #include "mongo/platform/strnlen.h"
 #include "mongo/util/decimal_counter.h"
@@ -1145,24 +1146,36 @@ void toBsonValue(uint8_t ctype,
                  TypeBits::Reader* typeBits,
                  bool inverted,
                  Version version,
-                 BSONObjBuilderValueStream* stream);
+                 BSONObjBuilderValueStream* stream,
+                 uint32_t depth);
 
 void toBson(BufReader* reader,
             TypeBits::Reader* typeBits,
             bool inverted,
             Version version,
-            BSONObjBuilder* builder) {
+            BSONObjBuilder* builder,
+            uint32_t depth) {
     while (readType<uint8_t>(reader, inverted) != 0) {
         if (inverted) {
             std::string name = readInvertedCString(reader);
             BSONObjBuilderValueStream& stream = *builder << name;
-            toBsonValue(
-                readType<uint8_t>(reader, inverted), reader, typeBits, inverted, version, &stream);
+            toBsonValue(readType<uint8_t>(reader, inverted),
+                        reader,
+                        typeBits,
+                        inverted,
+                        version,
+                        &stream,
+                        depth);
         } else {
             StringData name = readCString(reader);
             BSONObjBuilderValueStream& stream = *builder << name;
-            toBsonValue(
-                readType<uint8_t>(reader, inverted), reader, typeBits, inverted, version, &stream);
+            toBsonValue(readType<uint8_t>(reader, inverted),
+                        reader,
+                        typeBits,
+                        inverted,
+                        version,
+                        &stream,
+                        depth);
         }
     }
 }
@@ -1196,7 +1209,12 @@ void toBsonValue(uint8_t ctype,
                  TypeBits::Reader* typeBits,
                  bool inverted,
                  Version version,
-                 BSONObjBuilderValueStream* stream) {
+                 BSONObjBuilderValueStream* stream,
+                 uint32_t depth) {
+    uassert(ErrorCodes::Overflow,
+            "KeyString encoding exceeded maximum allowable BSON nesting depth",
+            depth <= BSONDepth::getMaxAllowableDepth());
+
     // This is only used by the kNumeric.*ByteInt types, but needs to be declared up here
     // since it is used across a fallthrough.
     bool isNegative = false;
@@ -1288,7 +1306,7 @@ void toBsonValue(uint8_t ctype,
             }
             // Not going to optimize CodeWScope.
             BSONObjBuilder scope;
-            toBson(reader, typeBits, inverted, version, &scope);
+            toBson(reader, typeBits, inverted, version, &scope, depth + 1);
             *stream << BSONCodeWScope(code, scope.done());
             break;
         }
@@ -1343,7 +1361,7 @@ void toBsonValue(uint8_t ctype,
 
         case CType::kObject: {
             BSONObjBuilder subObj(stream->subobjStart());
-            toBson(reader, typeBits, inverted, version, &subObj);
+            toBson(reader, typeBits, inverted, version, &subObj, depth + 1);
             break;
         }
 
@@ -1352,8 +1370,13 @@ void toBsonValue(uint8_t ctype,
             DecimalCounter<unsigned> index;
             uint8_t elemType;
             while ((elemType = readType<uint8_t>(reader, inverted)) != 0) {
-                toBsonValue(
-                    elemType, reader, typeBits, inverted, version, &(subArr << StringData{index}));
+                toBsonValue(elemType,
+                            reader,
+                            typeBits,
+                            inverted,
+                            version,
+                            &(subArr << StringData{index}),
+                            depth + 1);
                 ++index;
             }
             break;
@@ -2347,7 +2370,7 @@ BSONObj toBsonSafe(const char* buffer, size_t len, Ordering ord, const TypeBits&
 
         if (ctype == kEnd)
             break;
-        toBsonValue(ctype, &reader, &typeBitsReader, invert, typeBits.version, &(builder << ""));
+        toBsonValue(ctype, &reader, &typeBitsReader, invert, typeBits.version, &(builder << ""), 1);
     }
     return builder.obj();
 }
diff --git a/src/mongo/db/storage/key_string_test.cpp b/src/mongo/db/storage/key_string_test.cpp
index 826b3130684..aeb38a76d2d 100644
--- a/src/mongo/db/storage/key_string_test.cpp
+++ b/src/mongo/db/storage/key_string_test.cpp
@@ -42,6 +42,8 @@
 
 #include "mongo/base/owned_pointer_vector.h"
 #include "mongo/base/simple_string_data_comparator.h"
+#include "mongo/bson/bson_depth.h"
+#include "mongo/bson/bson_validate.h"
 #include "mongo/bson/bsonobj_comparator.h"
 #include "mongo/bson/simple_bsonobj_comparator.h"
 #include "mongo/config.h"
@@ -226,6 +228,36 @@ TEST_F(KeyStringBuilderTest, TooManyElementsInCompoundKey) {
                        ErrorCodes::Overflow);
 }
 
+TEST_F(KeyStringBuilderTest, ExceededBSONDepth) {
+    KeyString::Builder ks(KeyString::Version::V1);
+
+    // Construct an illegal KeyString encoding with excessively nested BSON arrays '80' (P).
+    const auto nestedArr = std::string(BSONDepth::getMaxAllowableDepth() + 1, 'P');
+    ks.resetFromBuffer(nestedArr.c_str(), nestedArr.size());
+    ASSERT_THROWS_CODE(
+        KeyString::toBsonSafe(ks.getBuffer(), ks.getSize(), ALL_ASCENDING, ks.getTypeBits()),
+        AssertionException,
+        ErrorCodes::Overflow);
+
+    // Construct an illegal BSON object with excessive nesting.
+    BSONObj nestedObj;
+    for (unsigned i = 0; i < BSONDepth::getMaxAllowableDepth() + 1; i++) {
+        nestedObj = BSON("" << nestedObj);
+    }
+    // This BSON object should not be valid.
+    auto validateStatus =
+        validateBSON(nestedObj.objdata(), nestedObj.objsize(), BSONVersion::kV1_1);
+    ASSERT_EQ(ErrorCodes::Overflow, validateStatus.code());
+
+    // Construct a KeyString from the invalid BSON, and confirm that it fails to convert back to
+    // BSON.
+    ks.resetToKey(nestedObj, ALL_ASCENDING, RecordId());
+    ASSERT_THROWS_CODE(
+        KeyString::toBsonSafe(ks.getBuffer(), ks.getSize(), ALL_ASCENDING, ks.getTypeBits()),
+        AssertionException,
+        ErrorCodes::Overflow);
+}
+
 TEST_F(KeyStringBuilderTest, Simple1) {
     BSONObj a = BSON("" << 5);
     BSONObj b = BSON("" << 6);