From ee65d191b79609575192d0ec39edb07ac0213f3b Mon Sep 17 00:00:00 2001
From: David Percy <david.percy@mongodb.com>
Date: Mon, 24 May 2021 21:00:29 +0000
Subject: SERVER-57164 isVariableReference includes system variables

This fixes an invariant failure in DocumentSourceGroup.

(cherry picked from commit 2f4c2c6bdca301934d32a9f3a1f6fc76946f7b11)
---
 .../sources/group/group_by_system_variable.js      | 27 ++++++++++++++++++++++
 src/mongo/db/pipeline/expression.h                 | 18 +++++++++++++--
 src/mongo/db/pipeline/semantic_analysis.cpp        |  2 +-
 3 files changed, 44 insertions(+), 3 deletions(-)
 create mode 100644 jstests/aggregation/sources/group/group_by_system_variable.js

diff --git a/jstests/aggregation/sources/group/group_by_system_variable.js b/jstests/aggregation/sources/group/group_by_system_variable.js
new file mode 100644
index 00000000000..6c223343166
--- /dev/null
+++ b/jstests/aggregation/sources/group/group_by_system_variable.js
@@ -0,0 +1,27 @@
+/*
+ * Tests that the server doesn't crash when you group by a system variable.
+ * Reproduces SERVER-57164.
+ */
+(function() {
+"use strict";
+
+const coll = db.group_by_system_variable;
+coll.drop();
+assert.commandWorked(coll.insert({}));
+
+// These explains all should not throw.
+coll.explain().aggregate({$group: {_id: "$$IS_MR"}});
+coll.explain().aggregate({$group: {_id: "$$JS_SCOPE"}});
+coll.explain().aggregate({$group: {_id: "$$CLUSTER_TIME"}});
+
+// These queries throw, but the server should stay up.
+assert.throws(() => coll.aggregate({$group: {_id: "$$IS_MR"}}));
+assert.throws(() => coll.aggregate({$group: {_id: "$$JS_SCOPE"}}));
+try {
+    // This one may or may not throw: CLUSTER_TIME may or may not be defined,
+    // depending on what kind of cluster we're running against.
+    coll.aggregate({$group: {_id: "$$CLUSTER_TIME"}});
+} catch (e) {
+}
+db.hello();
+})();
diff --git a/src/mongo/db/pipeline/expression.h b/src/mongo/db/pipeline/expression.h
index d5f29a021b5..61479668339 100644
--- a/src/mongo/db/pipeline/expression.h
+++ b/src/mongo/db/pipeline/expression.h
@@ -1553,12 +1553,26 @@ public:
 
 class ExpressionFieldPath : public Expression {
 public:
-    bool isRootFieldPath() const {
+    /**
+     * Checks whether this field path is exactly "$$ROOT".
+     */
+    bool isROOT() const {
         return _variable == Variables::kRootId && _fieldPath.getPathLength() == 1;
     }
 
+    /**
+     * Checks whether this field path starts with a variable besides ROOT.
+     *
+     * For example, these are variable references:
+     *   "$$NOW"
+     *   "$$NOW.x"
+     * and these are not:
+     *   "$x"
+     *   "$$ROOT"
+     *   "$$ROOT.x"
+     */
     bool isVariableReference() const {
-        return Variables::isUserDefinedVariable(_variable);
+        return _variable != Variables::kRootId;
     }
 
     boost::intrusive_ptr<Expression> optimize() final;
diff --git a/src/mongo/db/pipeline/semantic_analysis.cpp b/src/mongo/db/pipeline/semantic_analysis.cpp
index 36f1985a913..c1613a6a85f 100644
--- a/src/mongo/db/pipeline/semantic_analysis.cpp
+++ b/src/mongo/db/pipeline/semantic_analysis.cpp
@@ -130,7 +130,7 @@ boost::optional<std::string> replaceRootNestsRoot(
     }
     auto&& [nestedName, expression] = children[0];
     if (!dynamic_cast<ExpressionFieldPath*>(expression.get()) ||
-        !dynamic_cast<ExpressionFieldPath*>(expression.get())->isRootFieldPath()) {
+        !dynamic_cast<ExpressionFieldPath*>(expression.get())->isROOT()) {
         return boost::none;
     }
     return nestedName;
-- 
cgit v1.2.1