SERVER-119319 fix ExprCtx use-after-free for $lookup/$graphLookup clone (#49956)

Co-authored-by: Zixuan <leozzx@users.noreply.github.com>
GitOrigin-RevId: 61bea16b08ea8b5ff20ba7282dc6d8ab5bdcfb3b

src/mongo/db/pipeline/document_source_graph_lookup.cpp

@@ -850,7 +850,6 @@ DocumentSourceGraphLookUp::DocumentSourceGraphLookUp(
       _as(original._as),
       _connectFromField(original._connectFromField),
       _connectToField(original._connectToField),
-      _startWith(original._startWith),
       _additionalFilter(original._additionalFilter),
       _depthField(original._depthField),
       _maxDepth(original._maxDepth),
@@ -868,6 +867,10 @@ DocumentSourceGraphLookUp::DocumentSourceGraphLookUp(
       _cache(pExpCtx->getValueComparator()),
       _variables(original._variables),
       _variablesParseState(original._variablesParseState.copyWith(_variables.useIdGenerator())) {
+    if (original._startWith) {
+        // re-create startWith expression using newExpCtx.
+        _startWith = original._startWith->cloneUsingNewExpCtx(newExpCtx.get());
+    }
     if (original._unwind) {
         _unwind =
             static_cast<DocumentSourceUnwind*>(original._unwind.value()->clone(pExpCtx).get());

src/mongo/db/pipeline/document_source_graph_lookup_test.cpp

@@ -1063,5 +1063,43 @@ TEST_F(DocumentSourceGraphLookupServerlessTest,
     ASSERT_EQ(1ul, involvedNssSet.count(graphLookupNs));
 }
 
+TEST_F(DocumentSourceGraphLookUpTest, StartWithCloneRebindsExpressionContext) {
+    NamespaceString nss =
+        NamespaceString::createNamespaceString_forTest(boost::none, "test", "coll");
+    auto resolvedNss = ResolvedNamespaceMap{{nss, {nss, std::vector<BSONObj>()}}};
+
+    auto opCtx = getOpCtx();
+    auto expCtx = make_intrusive<ExpressionContextForTest>(opCtx, nss);
+    expCtx->setResolvedNamespaces(resolvedNss);
+
+    auto spec = fromjson(R"({
+        "$graphLookup": {
+            "from": "coll",
+            "startWith": { "$cond": [ { "$eq": ["$f", "v"] }, "$f", "$$REMOVE" ] },
+            "connectFromField": "f",
+            "connectToField": "f",
+            "as": "j"
+        }
+    })");
+
+    auto ds = DocumentSourceGraphLookUp::createFromBson(spec.firstElement(), expCtx);
+    DocumentSourceGraphLookUp* docSource = static_cast<DocumentSourceGraphLookUp*>(ds.get());
+
+    // docSource points to the ExpressionContext
+    ASSERT_EQ(docSource->getStartWithField()->getExpressionContext(), expCtx);
+
+    // Clone with a new top-level ExpressionContext
+    auto newExpCtx = make_intrusive<ExpressionContextForTest>(getOpCtx(), nss);
+    newExpCtx->setResolvedNamespaces(resolvedNss);
+    auto dsClone = docSource->clone(newExpCtx);
+    DocumentSourceGraphLookUp* docSourceClone =
+        static_cast<DocumentSourceGraphLookUp*>(dsClone.get());
+
+    // docSource still points to the original ExpressionContext
+    ASSERT_EQ(docSource->getStartWithField()->getExpressionContext(), expCtx);
+    // clonedDocSource points to the new ExpressionContext
+    ASSERT_EQ(docSourceClone->getStartWithField()->getExpressionContext(), newExpCtx);
+}
+
 }  // namespace
 }  // namespace mongo

src/mongo/db/pipeline/document_source_lookup.cpp

@@ -422,8 +422,7 @@ DocumentSourceLookUp::DocumentSourceLookUp(const DocumentSourceLookUp& original,
                                                  original._fromExpCtx->getView())),
       _resolvedPipeline(original._resolvedPipeline),
       _userPipeline(original._userPipeline),
-      _resolvedIntrospectionPipeline(original._resolvedIntrospectionPipeline->clone(_fromExpCtx)),
-      _letVariables(original._letVariables) {
+      _resolvedIntrospectionPipeline(original._resolvedIntrospectionPipeline->clone(_fromExpCtx)) {
     if (!_localField && !_foreignField) {
         _cache.emplace(internalDocumentSourceCursorBatchSizeBytes.load());
     }
@@ -433,6 +432,8 @@ DocumentSourceLookUp::DocumentSourceLookUp(const DocumentSourceLookUp& original,
     if (original._unwindSrc) {
         _unwindSrc = static_cast<DocumentSourceUnwind*>(original._unwindSrc->clone(pExpCtx).get());
     }
+    // clone let variables with new expCtx in case the original expCtx is deleted.
+    copyLetVariablesWithNewExpCtx(original._letVariables, newExpCtx.get());
 }
 
 boost::intrusive_ptr<DocumentSource> DocumentSourceLookUp::clone(
@@ -440,6 +441,16 @@ boost::intrusive_ptr<DocumentSource> DocumentSourceLookUp::clone(
     return make_intrusive<DocumentSourceLookUp>(*this, newExpCtx);
 }
 
+void DocumentSourceLookUp::copyLetVariablesWithNewExpCtx(const std::vector<LetVariable>& src,
+                                                         ExpressionContext* newExpCtx) {
+    _letVariables.clear();
+    _letVariables.reserve(src.size());
+
+    for (const auto& var : src) {
+        _letVariables.emplace_back(var.cloneUsingNewExpCtx(newExpCtx));
+    }
+}
+
 void validateLookupCollectionlessPipeline(const std::vector<BSONObj>& pipeline) {
     uassert(ErrorCodes::FailedToParse,
             "$lookup stage without explicit collection must have a pipeline with $documents as "

src/mongo/db/pipeline/document_source_lookup.h

@@ -403,6 +403,12 @@ private:
      */
     void resolveLetVariables(const Document& localDoc, Variables* variables);
 
+    /**
+     * Clones the given vector of LetVariable objects using the newExpCtx.
+     */
+    void copyLetVariablesWithNewExpCtx(const std::vector<LetVariable>& src,
+                                       ExpressionContext* newExpCtx);
+
     /**
      * Builds a parsed pipeline for introspection (e.g. constraints, dependencies). Any sub-$lookup
      * pipelines will be built recursively.

src/mongo/db/pipeline/document_source_lookup_test.cpp

@@ -1839,6 +1839,42 @@ TEST_F(DocumentSourceLookUpTest, LookupParseSerializedStageWithAbsorbedUnwind) {
     ASSERT(dynamic_cast<DocumentSourceLookUp*>(lookup.get())->hasUnwindSrc());
 }
 
+static boost::intrusive_ptr<DocumentSourceLookUp> makeLookupWithLet(
+    const boost::intrusive_ptr<ExpressionContext>& expCtx, NamespaceString fromNs) {
+    expCtx->setResolvedNamespaces(ResolvedNamespaceMap{{fromNs, {fromNs, std::vector<BSONObj>()}}});
+
+    auto spec =
+        BSON("$lookup" << BSON("from" << fromNs.coll() << "let" << BSON("v" << "$x") << "pipeline"
+                                      << BSON_ARRAY(BSON("$match" << BSON("y" << "$$v"))) << "as"
+                                      << "out"));
+    auto ds = DocumentSourceLookUp::createFromBson(spec.firstElement(), expCtx);
+    return boost::static_pointer_cast<DocumentSourceLookUp>(ds);
+}
+
+TEST_F(DocumentSourceLookUpTest, LetVariablesCloneRebindsExpressionContext) {
+    NamespaceString nss =
+        NamespaceString::createNamespaceString_forTest(boost::none, "test", "coll");
+    auto opCtx = getOpCtx();
+    auto expCtx = make_intrusive<ExpressionContextForTest>(opCtx, nss);
+
+    // Build an original $lookup with a let expression
+    auto lookup = makeLookupWithLet(expCtx, nss);
+
+    // Sanity: expressions in _letVariables use original expCtx
+    for (auto& var : lookup->getLetVariables()) {
+        ASSERT_EQ(var.expression->getExpressionContext(), expCtx);
+    }
+
+    // Clone with a new top-level ExpressionContext
+    auto newExpCtx = make_intrusive<ExpressionContextForTest>(opCtx, nss);
+    auto lookupClone = static_pointer_cast<DocumentSourceLookUp>(lookup->clone(newExpCtx));
+
+    // Check that every let expression in the clone now points to the new context
+    for (auto& var : lookupClone->getLetVariables()) {
+        ASSERT_EQ(var.expression->getExpressionContext(), newExpCtx);
+    }
+}
+
 using DocumentSourceLookUpServerlessTest = ServerlessAggregationContextFixture;
 
 TEST_F(DocumentSourceLookUpServerlessTest,

src/mongo/db/pipeline/expression.cpp

@@ -89,6 +89,23 @@ using std::pair;
 using std::string;
 using std::vector;
 
+// Clone by serializing and reparsing.
+boost::intrusive_ptr<Expression> Expression::cloneUsingNewExpCtx(
+    ExpressionContext* newExpCtx) const {
+    // Serialize this expression to a generic Value.
+    SerializationOptions opts{.serializeForCloning = true};
+    auto val = serialize(opts);
+
+    // Wrap it into a BSONObj as the value of a dummy field.
+    BSONObjBuilder bob;
+    val.addToBsonObj(&bob, "");
+    BSONObj obj = bob.obj();
+    BSONElement elem = obj.firstElement();
+
+    // Re-parse as an operand in the new ExpressionContext.
+    return Expression::parseOperand(newExpCtx, elem, newExpCtx->variablesParseState);
+}
+
 Value ExpressionConstant::serializeConstant(const SerializationOptions& opts,
                                             Value val,
                                             bool wrapRepresentativeValue) {

src/mongo/db/pipeline/expression.h

@@ -251,6 +251,14 @@ public:
         return this;
     }
 
+    /**
+     * Every node in the expression tree maintains an unowned pointer to the query's
+     * ExpressionContext. This variant of clone() creates a deep copy of the expression tree where
+     * every node in the tree points to newExprCtx (instead of whatever expression context it used
+     * to point to).
+     */
+    boost::intrusive_ptr<Expression> cloneUsingNewExpCtx(ExpressionContext* newExpCtx) const;
+
     /**
      * Serialize the Expression tree recursively.
      *

src/mongo/db/pipeline/variables.cpp

@@ -525,4 +525,9 @@ std::pair<LegacyRuntimeConstants, BSONObj> VariablesParseState::transitionalComp
 
     return {vars.transitionalExtractRuntimeConstants(), bob.obj()};
 }
+
+LetVariable LetVariable::cloneUsingNewExpCtx(ExpressionContext* newExpCtx) const {
+    auto clonedExpr = expression ? expression->cloneUsingNewExpCtx(newExpCtx) : nullptr;
+    return {name, std::move(clonedExpr), id};
+}
 }  // namespace mongo

src/mongo/db/pipeline/variables.h

@@ -297,6 +297,7 @@ private:
  */
 struct LetVariable {
     LetVariable(std::string name, boost::intrusive_ptr<Expression> expression, Variables::Id id);
+    LetVariable cloneUsingNewExpCtx(ExpressionContext* newExpCtx) const;
 
     std::string name;
     boost::intrusive_ptr<Expression> expression;