SERVER-101342 Improve handling of invalid input in key_string::getKeySize() (#39397)

GitOrigin-RevId: 02dd2fbfea6ce9bcb044463a15c263a94866e17d

src/mongo/db/storage/key_string.cpp

@@ -2688,21 +2688,20 @@ uint8_t TypeBits::ExplainReader::readDecimalExponent() {
 size_t getKeySize(const char* buffer, size_t len, Ordering ord, Version version) {
     invariant(len > 0);
     BufReader reader(buffer, len);
-    unsigned remainingBytes;
-    for (int i = 0; (remainingBytes = reader.remaining()); i++) {
+    for (int i = 0;; i++) {
+        // We reached the end of the buffer without reading a valid complete key
+        if (reader.remaining() == 0)
+            return 0;
+
         const bool invert = (ord.get(i) == -1);
         uint8_t ctype = readType<uint8_t>(&reader, invert);
-        // We have already read the Key.
+        // We have reached the end of the Key. The Key size is the number of bytes we used.
         if (ctype == kEnd)
-            break;
+            return len - reader.remaining();
 
         // Read the Key that comes after the first byte in KeyString.
         filterKeyFromKeyString(ctype, &reader, invert, version);
     }
-
-    invariant(len > remainingBytes);
-    // Key size = buffer len - number of bytes comprising the RecordId
-    return (len - (remainingBytes - 1));
 }
 
 // This discriminator byte only exists in KeyStrings for queries, not in KeyStrings stored in an

src/mongo/db/storage/key_string_test.cpp

@@ -270,8 +270,8 @@ TEST_F(KeyStringBuilderTest, TooManyElementsInCompoundKey) {
     // Construct a KeyString with more than the limit of 32 elements in a compound index key. Encode
     // 33 kBoolTrue ('o') values.
     // Note that this KeyString encoding is legal, but it may not be legally stored in an index.
-    const char* data = "ooooooooooooooooooooooooooooooooo";
-    const size_t size = 33;
+    const char* data = "ooooooooooooooooooooooooooooooooo\x4";
+    const size_t size = 34;
 
     key_string::Builder ks(key_string::Version::V1);
     ks.resetFromBuffer(data, size);
@@ -279,7 +279,8 @@ TEST_F(KeyStringBuilderTest, TooManyElementsInCompoundKey) {
     // No exceptions should be thrown.
     key_string::toBsonSafe(data, size, ALL_ASCENDING, ks.getTypeBits());
     key_string::decodeDiscriminator(ks.getBuffer(), ks.getSize(), ALL_ASCENDING, ks.getTypeBits());
-    key_string::getKeySize(ks.getBuffer(), ks.getSize(), ALL_ASCENDING, ks.version);
+    ASSERT_EQ(size,
+              key_string::getKeySize(ks.getBuffer(), ks.getSize(), ALL_ASCENDING, ks.version));
 }
 
 TEST_F(KeyStringBuilderTest, MaxElementsInCompoundKey) {
@@ -294,7 +295,38 @@ TEST_F(KeyStringBuilderTest, MaxElementsInCompoundKey) {
     // No exceptions should be thrown.
     key_string::toBsonSafe(data, size, ALL_ASCENDING, ks.getTypeBits());
     key_string::decodeDiscriminator(ks.getBuffer(), ks.getSize(), ALL_ASCENDING, ks.getTypeBits());
-    key_string::getKeySize(ks.getBuffer(), ks.getSize(), ALL_ASCENDING, ks.version);
+    ASSERT_EQ(size,
+              key_string::getKeySize(ks.getBuffer(), ks.getSize(), ALL_ASCENDING, ks.version));
+}
+
+TEST_F(KeyStringBuilderTest, SizeOfIncompleteKey) {
+    // The key portion of a keystring is terminated with kEnd, so missing that means the key size is
+    // zero
+    const char* data = "oooo\x4";
+    const size_t dataSize = 5;
+    ASSERT_EQ(dataSize,
+              key_string::getKeySize(data, dataSize, ALL_ASCENDING, key_string::Version::V1));
+    ASSERT_EQ(0,
+              key_string::getKeySize(data, dataSize - 1, ALL_ASCENDING, key_string::Version::V1));
+}
+
+TEST_F(KeyStringBuilderTest, SizeWithTrailingDataInBuffer) {
+    // Verify that we actually stop counting key bytes when we reach kEnd
+    const char data[] = {'o', 'o', 4, 'a', 'b', 'c'};
+    size_t len = 6;
+    ASSERT_EQ(3, key_string::getKeySize(data, len, ALL_ASCENDING, key_string::Version::V1));
+}
+
+TEST_F(KeyStringBuilderTest, EmbeddedkEnd) {
+    // Construct a KeyString which contains kEnd inside a string key and verify that getKeySize()
+    // does not report that the size ends at that spot
+    key_string::Builder ks(version, ALL_ASCENDING);
+    ks.appendString("_\0_\4_"_sd);
+    ks.appendString("abc"_sd);
+    StringData buffer = ks.finishAndGetBuffer();
+    ASSERT_EQ(buffer.size(), 14);
+    ASSERT_EQ(key_string::getKeySize(buffer.data(), buffer.size(), ALL_ASCENDING, version),
+              buffer.size());
 }
 
 TEST_F(KeyStringBuilderTest, EmbeddedNullString) {