SERVER-58581 Add SEARCH_META variable

(cherry-picked from commit 0f554217d6476d804f9ed1051722447fec017fb4)
(cherry-picked from commit 9fdcc4cc39cfde80bce740555632eb34f65102aa)

SERVER-58886 Ban SEARCH_META on sharded collections on variable access

(cherry-picked from commit 4ce48b4608b0b3e0528b4faba0d1dc0075f71c32)
(cherry-picked from commit 04c2c34467babc19d3b175fd0c98cab112c49f16)

12 files changed, 169 insertions, 16 deletions

diff --git a/src/mongo/db/pipeline/expression.cpp b/src/mongo/db/pipeline/expression.cpp
index 2cc52c45bdc..0f5650b8c2d 100644
--- a/src/mongo/db/pipeline/expression.cpp
+++ b/src/mongo/db/pipeline/expression.cpp
@@ -2090,6 +2090,14 @@ intrusive_ptr<ExpressionFieldPath> ExpressionFieldPath::parse(
         const StringData varName = fieldPath.substr(0, fieldPath.find('.'));
         Variables::uassertValidNameForUserRead(varName);
         auto varId = vps.getVariable(varName);
+        if (varName.compare(Variables::getBuiltinVariableName(Variables::kSearchMetaId)) == 0) {
+            return new ExpressionFieldPathNonSharded(
+                expCtx.get(),
+                fieldPath.toString(),
+                varId,
+                std::string("Search queries accessing $$SEARCH_META are not supported in sharded "
+                            "pipelines"));
+        }
         return new ExpressionFieldPath(expCtx, fieldPath.toString(), varId);
     } else {
         return new ExpressionFieldPath(expCtx,
@@ -2207,6 +2215,12 @@ Value ExpressionFieldPath::serialize(bool explain) const {
     }
 }
 
+Value ExpressionFieldPathNonSharded::evaluate(const Document& root, Variables* variables) const {
+    uassert(
+        5858100, _errMsg, !getExpressionContext()->needsMerge && !getExpressionContext()->inMongos);
+    return ExpressionFieldPath::evaluate(root, variables);
+}
+
 Expression::ComputedPaths ExpressionFieldPath::getComputedPaths(const std::string& exprFieldPath,
                                                                 Variables::Id renamingVar) const {
     // An expression field path is either considered a rename or a computed path. We need to find
diff --git a/src/mongo/db/pipeline/expression.h b/src/mongo/db/pipeline/expression.h
index 3eb35d8d98e..eab38a48fcf 100644
--- a/src/mongo/db/pipeline/expression.h
+++ b/src/mongo/db/pipeline/expression.h
@@ -1337,14 +1337,14 @@ public:
 };
 
 
-class ExpressionFieldPath final : public Expression {
+class ExpressionFieldPath : public Expression {
 public:
     bool isRootFieldPath() const {
         return _variable == Variables::kRootId;
     }
 
     boost::intrusive_ptr<Expression> optimize() final;
-    Value evaluate(const Document& root, Variables* variables) const final;
+    Value evaluate(const Document& root, Variables* variables) const;
     Value serialize(bool explain) const final;
 
     /*
@@ -1393,11 +1393,12 @@ public:
 protected:
     void _doAddDependencies(DepsTracker* deps) const final;
 
-private:
     ExpressionFieldPath(const boost::intrusive_ptr<ExpressionContext>& expCtx,
                         const std::string& fieldPath,
                         Variables::Id variable);
 
+
+private:
     /*
       Internal implementation of evaluate(), used recursively.
 
@@ -1420,6 +1421,22 @@ private:
     const Variables::Id _variable;
 };
 
+/**
+ * A version of ExpressionFieldPath that will throw if evaluated in a sharded pipeline.
+ */
+class ExpressionFieldPathNonSharded : public ExpressionFieldPath {
+public:
+    Value evaluate(const Document& root, Variables* variables) const final;
+
+    ExpressionFieldPathNonSharded(ExpressionContext* const expCtx,
+                                  const std::string& fieldPath,
+                                  Variables::Id variable,
+                                  std::string errMsg)
+        : ExpressionFieldPath(expCtx, fieldPath, variable), _errMsg(std::move(errMsg)) {}
+
+private:
+    std::string _errMsg;
+};
 
 class ExpressionFilter final : public Expression {
 public:
diff --git a/src/mongo/db/pipeline/variables.cpp b/src/mongo/db/pipeline/variables.cpp
index 6a75d022883..1a5aa243702 100644
--- a/src/mongo/db/pipeline/variables.cpp
+++ b/src/mongo/db/pipeline/variables.cpp
@@ -30,6 +30,7 @@
 #include "mongo/db/pipeline/variables.h"
 #include "mongo/db/client.h"
 #include "mongo/db/logical_clock.h"
+#include "mongo/db/query/query_knobs_gen.h"
 #include "mongo/platform/basic.h"
 #include "mongo/platform/random.h"
 #include "mongo/util/str.h"
@@ -40,8 +41,11 @@ namespace mongo {
 constexpr Variables::Id Variables::kRootId;
 constexpr Variables::Id Variables::kRemoveId;
 
-const StringMap<Variables::Id> Variables::kBuiltinVarNameToId = {
-    {"ROOT", kRootId}, {"REMOVE", kRemoveId}, {"NOW", kNowId}, {"CLUSTER_TIME", kClusterTimeId}};
+const StringMap<Variables::Id> Variables::kBuiltinVarNameToId = {{"ROOT", kRootId},
+                                                                 {"REMOVE", kRemoveId},
+                                                                 {"NOW", kNowId},
+                                                                 {"CLUSTER_TIME", kClusterTimeId},
+                                                                 {"SEARCH_META", kSearchMetaId}};
 
 void Variables::uassertValidNameForUserWrite(StringData varName) {
     // System variables users allowed to write to (currently just one)
@@ -113,6 +117,24 @@ void Variables::setValue(Id id, const Value& value, bool isConstant) {
     _valueList[idAsSizeT] = ValueAndState(value, isConstant);
 }
 
+void Variables::setReservedValue(Id id, const Value& value, bool isConstant) {
+    // If a value has already been set for 'id', and that value was marked as constant, then it
+    // is illegal to modify.
+    switch (id) {
+        case Variables::kSearchMetaId:
+            uassert(5858101,
+                    "Can't set SEARCH_META multiple times per query",
+                    _runtimeConstantsMap.find(id) == _runtimeConstantsMap.end());
+            _runtimeConstantsMap[id] = value;
+            break;
+        default:
+            // Currently it is only allowed to manually set the SEARCH_META builtin variable.
+            uasserted(5858102,
+                      str::stream() << "Attempted to set '$$" << getBuiltinVariableName(id)
+                                    << "' which is not permitted");
+    }
+}
+
 void Variables::setValue(Variables::Id id, const Value& value) {
     const bool isConstant = false;
     setValue(id, value, isConstant);
@@ -154,6 +176,18 @@ Value Variables::getValue(Id id, const Document& root) const {
                 uasserted(4631100, "Use of undefined variable '$$JS_SCOPE'.");
             case Variables::kIsMapReduceId:
                 uasserted(4631101, "Use of undefined variable '$$IS_MR'.");
+            case Variables::kSearchMetaId: {
+                uassert(5858104,
+                        str::stream() << "Must be fully upgraded to 4.4 to access '$$SEARCH_META",
+                        serverGlobalParams.featureCompatibility.getVersion() >=
+                            ServerGlobalParams::FeatureCompatibility::Version::kFullyUpgradedTo44);
+                uassert(5858105,
+                        str::stream() << "Must enable 'enableSearchMeta' server parameter to "
+                                      << "use $searchMeta or $$SEARCH_META variable",
+                        enableSearchMeta.load());
+                auto metaIt = _runtimeConstantsMap.find(id);
+                return metaIt == _runtimeConstantsMap.end() ? Value() : metaIt->second;
+            }
             default:
                 MONGO_UNREACHABLE;
         }
diff --git a/src/mongo/db/pipeline/variables.h b/src/mongo/db/pipeline/variables.h
index fd5f19f2309..c4b785d4b9d 100644
--- a/src/mongo/db/pipeline/variables.h
+++ b/src/mongo/db/pipeline/variables.h
@@ -85,6 +85,7 @@ public:
     static constexpr Variables::Id kClusterTimeId = Id(-4);
     static constexpr Variables::Id kJsScopeId = Id(-5);
     static constexpr Variables::Id kIsMapReduceId = Id(-6);
+    static constexpr Variables::Id kSearchMetaId = Id(-7);
 
     // Map from builtin var name to reserved id number.
     static const StringMap<Id> kBuiltinVarNameToId;
@@ -102,6 +103,12 @@ public:
     void setConstantValue(Variables::Id id, const Value& value);
 
     /**
+     * Same as 'setValue' but is only allowed on reserved, builtin, variables. Should not be used
+     * when setting from user input.
+     */
+    void setReservedValue(Variables::Id id, const Value& value, bool isConstant);
+
+    /**
      * Gets the value of a user-defined or system variable. If the 'id' provided represents the
      * special ROOT variable, then we return 'root' in Value form.
      */
@@ -148,6 +155,15 @@ public:
      */
     void setDefaultRuntimeConstants(OperationContext* opCtx);
 
+    static auto getBuiltinVariableName(Variables::Id variable) {
+        for (auto& [name, id] : kBuiltinVarNameToId) {
+            if (variable == id) {
+                return name;
+            }
+        }
+        return std::string();
+    }
+
 private:
     struct ValueAndState {
         ValueAndState() = default;
@@ -160,15 +176,6 @@ private:
 
     void setValue(Id id, const Value& value, bool isConstant);
 
-    static auto getBuiltinVariableName(Variables::Id variable) {
-        for (auto& [name, id] : kBuiltinVarNameToId) {
-            if (variable == id) {
-                return name;
-            }
-        }
-        return std::string();
-    }
-
     IdGenerator _idGenerator;
     std::vector<ValueAndState> _valueList;
     stdx::unordered_map<Id, Value> _runtimeConstantsMap;
diff --git a/src/mongo/db/query/cursor_response.cpp b/src/mongo/db/query/cursor_response.cpp
index a6b271253bc..f259b185c5a 100644
--- a/src/mongo/db/query/cursor_response.cpp
+++ b/src/mongo/db/query/cursor_response.cpp
@@ -44,6 +44,7 @@ const char kCursorsField[] = "cursors";
 const char kCursorField[] = "cursor";
 const char kIdField[] = "id";
 const char kNsField[] = "ns";
+const char kVarsField[] = "vars";
 const char kBatchField[] = "nextBatch";
 const char kBatchFieldInitial[] = "firstBatch";
 const char kBatchDocSequenceField[] = "cursor.nextBatch";
@@ -128,6 +129,7 @@ CursorResponse::CursorResponse(NamespaceString nss,
                                boost::optional<long long> numReturnedSoFar,
                                boost::optional<BSONObj> postBatchResumeToken,
                                boost::optional<BSONObj> writeConcernError,
+                               boost::optional<BSONObj> varsField,
                                bool partialResultsReturned)
     : _nss(std::move(nss)),
       _cursorId(cursorId),
@@ -135,8 +137,8 @@ CursorResponse::CursorResponse(NamespaceString nss,
       _numReturnedSoFar(numReturnedSoFar),
       _postBatchResumeToken(std::move(postBatchResumeToken)),
       _writeConcernError(std::move(writeConcernError)),
+      _varsField(std::move(varsField)),
       _partialResultsReturned(partialResultsReturned) {}
-
 std::vector<StatusWith<CursorResponse>> CursorResponse::parseFromBSONMany(
     const BSONObj& cmdResponse) {
     std::vector<StatusWith<CursorResponse>> cursors;
@@ -196,6 +198,13 @@ StatusWith<CursorResponse> CursorResponse::parseFromBSON(const BSONObj& cmdRespo
     }
     fullns = nsElt.String();
 
+    BSONElement varsElt = cmdResponse[kVarsField];
+    if (!varsElt.eoo() && varsElt.type() != BSONType::Object) {
+        return {ErrorCodes::TypeMismatch,
+                str::stream() << "Field '" << kVarsField
+                              << "' must be of type object in: " << cmdResponse};
+    }
+
     BSONElement batchElt = cursorObj[kBatchField];
     if (batchElt.eoo()) {
         batchElt = cursorObj[kBatchFieldInitial];
@@ -257,6 +266,7 @@ StatusWith<CursorResponse> CursorResponse::parseFromBSON(const BSONObj& cmdRespo
              postBatchResumeTokenElem ? postBatchResumeTokenElem.Obj().getOwned()
                                       : boost::optional<BSONObj>{},
              writeConcernError ? writeConcernError.Obj().getOwned() : boost::optional<BSONObj>{},
+             varsElt ? varsElt.Obj().getOwned() : boost::optional<BSONObj>{},
              partialResultsReturned.trueValue()}};
 }
 
diff --git a/src/mongo/db/query/cursor_response.h b/src/mongo/db/query/cursor_response.h
index 509600103e4..ddea9c8d6e2 100644
--- a/src/mongo/db/query/cursor_response.h
+++ b/src/mongo/db/query/cursor_response.h
@@ -194,6 +194,7 @@ public:
                    boost::optional<long long> numReturnedSoFar = boost::none,
                    boost::optional<BSONObj> postBatchResumeToken = boost::none,
                    boost::optional<BSONObj> writeConcernError = boost::none,
+                   boost::optional<BSONObj> varsField = boost::none,
                    bool partialResultsReturned = false);
 
     CursorResponse(CursorResponse&& other) = default;
@@ -231,6 +232,10 @@ public:
         return _writeConcernError;
     }
 
+    boost::optional<BSONObj> getVarsField() const {
+        return _varsField;
+    }
+
     bool getPartialResultsReturned() const {
         return _partialResultsReturned;
     }
@@ -251,6 +256,7 @@ private:
     boost::optional<long long> _numReturnedSoFar;
     boost::optional<BSONObj> _postBatchResumeToken;
     boost::optional<BSONObj> _writeConcernError;
+    boost::optional<BSONObj> _varsField;
     bool _partialResultsReturned = false;
 };
 
diff --git a/src/mongo/db/query/cursor_response_test.cpp b/src/mongo/db/query/cursor_response_test.cpp
index db1a70b6a74..ffff9aea99f 100644
--- a/src/mongo/db/query/cursor_response_test.cpp
+++ b/src/mongo/db/query/cursor_response_test.cpp
@@ -211,6 +211,53 @@ TEST(CursorResponseTest, parseFromBSONPartialResultsReturnedFieldWrongType) {
     ASSERT_NOT_OK(result.getStatus());
 }
 
+TEST(CursorResponseTest, parseFromBSONVarsFieldCorrect) {
+    BSONObj varsContents = BSON("randomVar" << 7);
+    StatusWith<CursorResponse> result = CursorResponse::parseFromBSON(BSON(
+        "cursor" << BSON("id" << CursorId(123) << "ns"
+                              << "db.coll"
+                              << "firstBatch" << BSON_ARRAY(BSON("_id" << 1) << BSON("_id" << 2)))
+                 << "vars" << varsContents << "ok" << 1));
+    ASSERT_OK(result.getStatus());
+
+    CursorResponse response = std::move(result.getValue());
+    ASSERT_EQ(response.getCursorId(), CursorId(123));
+    ASSERT_EQ(response.getNSS().ns(), "db.coll");
+    ASSERT_EQ(response.getBatch().size(), 2U);
+    ASSERT_BSONOBJ_EQ(response.getBatch()[0], BSON("_id" << 1));
+    ASSERT_BSONOBJ_EQ(response.getBatch()[1], BSON("_id" << 2));
+    ASSERT_TRUE(response.getVarsField());
+    ASSERT_BSONOBJ_EQ(response.getVarsField().get(), varsContents);
+}
+
+TEST(CursorResponseTest, parseFromBSONVarsFieldWrongType) {
+    StatusWith<CursorResponse> result = CursorResponse::parseFromBSON(BSON(
+        "cursor" << BSON("id" << CursorId(123) << "ns"
+                              << "db.coll"
+                              << "firstBatch" << BSON_ARRAY(BSON("_id" << 1) << BSON("_id" << 2)))
+                 << "vars" << 2 << "ok" << 1));
+    ASSERT_NOT_OK(result.getStatus());
+}
+
+TEST(CursorResponseTest, parseFromBSONMultipleVars) {
+    BSONObj varsContents = BSON("randomVar" << 7 << "otherVar" << BSON("nested" << 2));
+    StatusWith<CursorResponse> result = CursorResponse::parseFromBSON(BSON(
+        "cursor" << BSON("id" << CursorId(123) << "ns"
+                              << "db.coll"
+                              << "firstBatch" << BSON_ARRAY(BSON("_id" << 1) << BSON("_id" << 2)))
+                 << "vars" << varsContents << "ok" << 1));
+    ASSERT_OK(result.getStatus());
+
+    CursorResponse response = std::move(result.getValue());
+    ASSERT_EQ(response.getCursorId(), CursorId(123));
+    ASSERT_EQ(response.getNSS().ns(), "db.coll");
+    ASSERT_EQ(response.getBatch().size(), 2U);
+    ASSERT_BSONOBJ_EQ(response.getBatch()[0], BSON("_id" << 1));
+    ASSERT_BSONOBJ_EQ(response.getBatch()[1], BSON("_id" << 2));
+    ASSERT_TRUE(response.getVarsField());
+    ASSERT_BSONOBJ_EQ(response.getVarsField().get(), varsContents);
+}
+
 TEST(CursorResponseTest, roundTripThroughCursorResponseBuilderWithPartialResultsReturned) {
     CursorResponseBuilder::Options options;
     options.isInitialResponse = true;
@@ -300,6 +347,7 @@ TEST(CursorResponseTest, toBSONPartialResultsReturned) {
                             boost::none,
                             boost::none,
                             boost::none,
+                            boost::none,
                             true);
     BSONObj responseObj = response.toBSON(CursorResponse::ResponseType::InitialResponse);
     BSONObj expectedResponse = BSON(
diff --git a/src/mongo/db/query/query_knobs.idl b/src/mongo/db/query/query_knobs.idl
index 71148f35dae..ddccc7ff8a3 100644
--- a/src/mongo/db/query/query_knobs.idl
+++ b/src/mongo/db/query/query_knobs.idl
@@ -415,3 +415,10 @@ server_parameters:
     validator:
         gt: 0
         lte: { expr: BSONObjMaxInternalSize }
+
+  enableSearchMeta:
+    description: "If false queries with $$SEARCH_META or $searchMeta will fail."
+    set_at: [ startup, runtime ]
+    cpp_varname: "enableSearchMeta"
+    cpp_vartype: AtomicWord<bool>
+    default: false
diff --git a/src/mongo/executor/task_executor_cursor.cpp b/src/mongo/executor/task_executor_cursor.cpp
index ac2c8dc2423..924d4648b1f 100644
--- a/src/mongo/executor/task_executor_cursor.cpp
+++ b/src/mongo/executor/task_executor_cursor.cpp
@@ -158,6 +158,8 @@ void TaskExecutorCursor::_getNextBatch(OperationContext* opCtx) {
     if (_cursorId == kUnitializedCursorId) {
         _ns = cr.getNSS();
         _rcr.dbname = _ns.db().toString();
+        // 'vars' are only included in the first batch.
+        _cursorVars = cr.getVarsField();
     }
 
     _cursorId = cr.getCursorId();
diff --git a/src/mongo/executor/task_executor_cursor.h b/src/mongo/executor/task_executor_cursor.h
index 020d4986039..8d6bd5dc0a0 100644
--- a/src/mongo/executor/task_executor_cursor.h
+++ b/src/mongo/executor/task_executor_cursor.h
@@ -107,6 +107,10 @@ public:
         return toRet;
     }
 
+    boost::optional<BSONObj> getCursorVars() {
+        return _cursorVars;
+    }
+
 private:
     /**
      * Runs a remote command and pipes the output back to this object
@@ -138,6 +142,9 @@ private:
 
     CursorId _cursorId = kUnitializedCursorId;
 
+    // Variables sent alongside the results in the cursor.
+    boost::optional<BSONObj> _cursorVars = boost::none;
+
     // This is a sum of the time spent waiting on remote calls.
     Milliseconds _millisecondsWaiting = Milliseconds(0);
 
diff --git a/src/mongo/s/query/cluster_find.cpp b/src/mongo/s/query/cluster_find.cpp
index bdde52a7833..dcde42d33df 100644
--- a/src/mongo/s/query/cluster_find.cpp
+++ b/src/mongo/s/query/cluster_find.cpp
@@ -850,6 +850,7 @@ StatusWith<CursorResponse> ClusterFind::runGetMore(OperationContext* opCtx,
                           startingFrom,
                           postBatchResumeToken,
                           boost::none,
+                          boost::none,
                           partialResultsReturned);
 }
 
diff --git a/src/mongo/s/query/establish_cursors.cpp b/src/mongo/s/query/establish_cursors.cpp
index 769f45d0ca2..42cf13eae1f 100644
--- a/src/mongo/s/query/establish_cursors.cpp
+++ b/src/mongo/s/query/establish_cursors.cpp
@@ -248,7 +248,7 @@ void CursorEstablisher::_handleFailure(const AsyncRequestsSender::Response& resp
         // This exception is eligible to be swallowed. Add an entry with a cursorID of 0, an
         // empty HostAndPort, and which has the 'partialResultsReturned' flag set to true.
         _remoteCursors.push_back(
-            {response.shardId.toString(), {}, {_nss, CursorId{0}, {}, {}, {}, {}, true}});
+            {response.shardId.toString(), {}, {_nss, CursorId{0}, {}, {}, {}, {}, {}, true}});
         return;
     }