SERVER-70885 Fix getElement/getField memory bug for stack owned value

2 files changed, 80 insertions, 0 deletions

diff --git a/src/mongo/db/exec/sbe/expressions/sbe_extract_sub_array_builtin_test.cpp b/src/mongo/db/exec/sbe/expressions/sbe_extract_sub_array_builtin_test.cpp
index df3c0d45d50..ca6885a98bf 100644
--- a/src/mongo/db/exec/sbe/expressions/sbe_extract_sub_array_builtin_test.cpp
+++ b/src/mongo/db/exec/sbe/expressions/sbe_extract_sub_array_builtin_test.cpp
@@ -210,4 +210,60 @@ TEST_F(SBEBuiltinExtractSubArrayTest, NotArray) {
         ASSERT_EQ(value, value::bitcastFrom<int64_t>(0));
     }
 }
+
+TEST_F(SBEBuiltinExtractSubArrayTest, MemoryManagement) {
+    {
+        auto array = makeArray(BSON_ARRAY("Item#1"
+                                          << "Item#2"
+                                          << "Item#3"
+                                          << "Item#4"));
+
+        // Use 'extractSubArray' to create a stack owned array and extract object from it, then test
+        // if 'getElement' can return the value with correct memory management.
+        auto extractFromSubArrayExpr = makeE<EFunction>(
+            "getElement",
+            makeEs(makeE<EFunction>("extractSubArray",
+                                    makeEs(makeC(array),
+                                           makeC(value::TypeTags::NumberInt32, 1),
+                                           makeC(value::TypeTags::NumberInt32, 2))),
+                   makeC(value::TypeTags::NumberInt32, 0)));
+
+        auto compiledExpr = compileExpression(*extractFromSubArrayExpr);
+
+        auto [tag, value] = runCompiledExpression(compiledExpr.get());
+        value::ValueGuard guard(tag, value);
+        ASSERT_TRUE(value::isString(tag));
+        ASSERT_EQ("Item#3", value::getStringView(tag, value));
+    }
+    {
+        const auto [objTag, objVal] = value::makeNewObject();
+        auto obj = value::getObjectView(objVal);
+
+        const auto [fieldTag, fieldVal] = value::makeNewString("not so small string"_sd);
+        ASSERT_EQ(value::TypeTags::StringBig, fieldTag);
+        obj->push_back("field"_sd, fieldTag, fieldVal);
+        const auto [arrTag, arrVal] = value::makeNewArray();
+        auto arr = value::getArrayView(arrVal);
+        arr->push_back(objTag, objVal);
+
+        // Use 'extractSubArray' to create a stack owned array and extract object from it, then test
+        // if 'getField' can return the value with correct memory management.
+        auto extractFromSubArrayExpr = makeE<EFunction>(
+            "getField",
+            makeEs(makeE<EFunction>(
+                       "getElement",
+                       makeEs(makeE<EFunction>("extractSubArray",
+                                               makeEs(makeC(arrTag, arrVal),
+                                                      makeC(value::TypeTags::NumberInt32, 1))),
+                              makeC(value::TypeTags::NumberInt32, 0))),
+                   makeC(value::makeNewString("field"_sd))));
+
+        auto compiledExpr = compileExpression(*extractFromSubArrayExpr);
+
+        auto [tag, value] = runCompiledExpression(compiledExpr.get());
+        value::ValueGuard guard(tag, value);
+        ASSERT_TRUE(value::isString(tag));
+        ASSERT_EQ("not so small string"_sd, value::getStringView(tag, value));
+    }
+}
 }  // namespace mongo::sbe
diff --git a/src/mongo/db/exec/sbe/vm/vm.cpp b/src/mongo/db/exec/sbe/vm/vm.cpp
index 15fcc3a2823..872cc884430 100644
--- a/src/mongo/db/exec/sbe/vm/vm.cpp
+++ b/src/mongo/db/exec/sbe/vm/vm.cpp
@@ -5627,6 +5627,12 @@ void ByteCode::runInternal(const CodeFragment* code, int64_t position) {
 
                 auto [owned, tag, val] = getField(lhsTag, lhsVal, rhsTag, rhsVal);
 
+                // Copy value only if needed
+                if (lhsOwned && !owned) {
+                    owned = true;
+                    std::tie(tag, val) = value::copyValue(tag, val);
+                }
+
                 pushStack(owned, tag, val);
 
                 if (rhsOwned && popRhs) {
@@ -5648,6 +5654,12 @@ void ByteCode::runInternal(const CodeFragment* code, int64_t position) {
 
                 auto [owned, tag, val] = getField(lhsTag, lhsVal, fieldName);
 
+                // Copy value only if needed
+                if (lhsOwned && !owned) {
+                    owned = true;
+                    std::tie(tag, val) = value::copyValue(tag, val);
+                }
+
                 pushStack(owned, tag, val);
 
                 if (lhsOwned && popLhs) {
@@ -5664,6 +5676,12 @@ void ByteCode::runInternal(const CodeFragment* code, int64_t position) {
 
                 auto [owned, tag, val] = getElement(lhsTag, lhsVal, rhsTag, rhsVal);
 
+                // Copy value only if needed
+                if (lhsOwned && !owned) {
+                    owned = true;
+                    std::tie(tag, val) = value::copyValue(tag, val);
+                }
+
                 pushStack(owned, tag, val);
 
                 if (rhsOwned && popRhs) {
@@ -5728,6 +5746,12 @@ void ByteCode::runInternal(const CodeFragment* code, int64_t position) {
 
                 auto [owned, tag, val] = getFieldOrElement(lhsTag, lhsVal, rhsTag, rhsVal);
 
+                // Copy value only if needed
+                if (lhsOwned && !owned) {
+                    owned = true;
+                    std::tie(tag, val) = value::copyValue(tag, val);
+                }
+
                 pushStack(owned, tag, val);
 
                 if (rhsOwned && popRhs) {