SERVER-56516 Fix undefined behaviour in $slice arguments parsing

(partially cherry picked from commit 19764c1864c1c57d9274238146e6843591bc6ce6)

5 files changed, 139 insertions, 36 deletions

diff --git a/src/mongo/bson/bsonelement.h b/src/mongo/bson/bsonelement.h
index 946f99ee1d8..ed1ac10366a 100644
--- a/src/mongo/bson/bsonelement.h
+++ b/src/mongo/bson/bsonelement.h
@@ -338,6 +338,14 @@ public:
 
     /** Retrieve int value for the element safely.  Zero returned if not a number. */
     int numberInt() const;
+
+    /** Like numberInt() but with well-defined behavior for doubles that
+     *  are NaNs, or too large/small to be represented as int.
+     *  NaNs -> 0
+     *  very large doubles -> INT_MAX
+     *  very small doubles -> INT_MIN  */
+    int safeNumberInt() const;
+
     /** Retrieve long value for the element safely.  Zero returned if not a number.
      *  Behavior is not defined for double values that are NaNs, or too large/small
      *  to be represented by long longs */
@@ -822,6 +830,21 @@ inline int BSONElement::numberInt() const {
     }
 }
 
+/** Like numberInt() but with well-defined behavior for doubles that
+ *  are NaNs, or too large/small to be represented as int.
+ *  NaNs -> 0
+ *  very large doubles -> INT_MAX
+ *  very small doubles -> INT_MIN  */
+inline int BSONElement::safeNumberInt() const {
+    const long long minInt = std::numeric_limits<int>::min();
+    const long long maxInt = std::numeric_limits<int>::max();
+    const auto numberLong = safeNumberLong();
+    if (numberLong < minInt) {
+        return minInt;
+    }
+    return std::min(numberLong, maxInt);
+}
+
 /** Retrieve long value for the element safely.  Zero returned if not a number. */
 inline long long BSONElement::numberLong() const {
     switch (type()) {
diff --git a/src/mongo/db/exec/projection_exec.cpp b/src/mongo/db/exec/projection_exec.cpp
index 83e5ad721e9..e6e6f03a669 100644
--- a/src/mongo/db/exec/projection_exec.cpp
+++ b/src/mongo/db/exec/projection_exec.cpp
@@ -108,7 +108,7 @@ ProjectionExec::ProjectionExec(OperationContext* opCtx,
             BSONElement e2 = obj.firstElement();
             if (mongoutils::str::equals(e2.fieldName(), "$slice")) {
                 if (e2.isNumber()) {
-                    int i = e2.numberInt();
+                    long long i = e2.safeNumberInt();
                     if (i < 0) {
                         add(e.fieldName(), i, -i);  // limit is now positive
                     } else {
@@ -120,8 +120,8 @@ ProjectionExec::ProjectionExec(OperationContext* opCtx,
                     verify(2 == arr.nFields());
 
                     BSONObjIterator it(arr);
-                    int skip = it.next().numberInt();
-                    int limit = it.next().numberInt();
+                    int skip = it.next().safeNumberInt();
+                    int limit = it.next().safeNumberInt();
 
                     verify(limit > 0);
 
@@ -215,7 +215,7 @@ void ProjectionExec::add(const string& field, bool include) {
     }
 }
 
-void ProjectionExec::add(const string& field, int skip, int limit) {
+void ProjectionExec::add(const string& field, long long skip, long long limit) {
     _special = true;  // can't include or exclude whole object
 
     if (field.empty()) {  // this is the field the user referred to
@@ -447,11 +447,11 @@ Status ProjectionExec::transform(const BSONObj& in,
 }
 
 void ProjectionExec::appendArray(BSONObjBuilder* bob, const BSONObj& array, bool nested) const {
-    int skip = nested ? 0 : _skip;
-    int limit = nested ? -1 : _limit;
+    long long skip = nested ? 0 : _skip;
+    long long limit = nested ? -1 : _limit;
 
     if (skip < 0) {
-        skip = max(0, skip + array.nFields());
+        skip = max(0ll, skip + array.nFields());
     }
 
     int index = 0;
diff --git a/src/mongo/db/exec/projection_exec.h b/src/mongo/db/exec/projection_exec.h
index 93af3c083ed..7882fa8bc94 100644
--- a/src/mongo/db/exec/projection_exec.h
+++ b/src/mongo/db/exec/projection_exec.h
@@ -95,7 +95,7 @@ private:
     /**
      * Add 'field' as a field name that is sliced as part of the projection.
      */
-    void add(const std::string& field, int skip, int limit);
+    void add(const std::string& field, long long skip, long long limit);
 
     //
     // Execution
@@ -157,8 +157,8 @@ private:
     bool _includeID;
 
     // Arguments from the $slice operator.
-    int _skip;
-    int _limit;
+    long long _skip;
+    long long _limit;
 
     // Used for $elemMatch and positional operator ($)
     Matchers _matchers;
diff --git a/src/mongo/db/exec/projection_exec_test.cpp b/src/mongo/db/exec/projection_exec_test.cpp
index 9c9e4f5f090..04e6112fbc2 100644
--- a/src/mongo/db/exec/projection_exec_test.cpp
+++ b/src/mongo/db/exec/projection_exec_test.cpp
@@ -66,25 +66,22 @@ unique_ptr<MatchExpression> parseMatchExpression(const BSONObj& obj) {
  * test function to verify results of transform()
  * on a working set member.
  *
- * specStr - projection specification
- * queryStr - query
- * objStr - object to run projection on
+ * spec - projection specification
+ * query - query
+ * obj - object to run projection on
  * data - computed data. Owned by working set member created in this function if not null.
  * expectedStatusOK - expected status of transformation
- * expectedObjStr - expected object after successful projection.
- *                  Ignored if expectedStatusOK is false.
+ * expectedObj - expected object after successful projection.
+ *               Ignored if expectedStatusOK is false.
  */
-
-void testTransform(const char* specStr,
-                   const char* queryStr,
-                   const char* objStr,
+void testTransform(const BSONObj& spec,
+                   const BSONObj& query,
+                   const BSONObj& obj,
                    WorkingSetComputedData* data,
                    const CollatorInterface* collator,
                    bool expectedStatusOK,
-                   const char* expectedObjStr) {
+                   const BSONObj& expectedObj) {
     // Create projection exec object.
-    BSONObj spec = fromjson(specStr);
-    BSONObj query = fromjson(queryStr);
     unique_ptr<MatchExpression> queryExpression = parseMatchExpression(query);
     QueryTestServiceContext serviceCtx;
     auto opCtx = serviceCtx.makeOperationContext();
@@ -92,7 +89,7 @@ void testTransform(const char* specStr,
 
     // Create working set member.
     WorkingSetMember wsm;
-    wsm.obj = Snapshotted<BSONObj>(SnapshotId(), fromjson(objStr));
+    wsm.obj = Snapshotted<BSONObj>(SnapshotId(), obj);
     if (data) {
         wsm.addComputed(data);
     }
@@ -106,8 +103,8 @@ void testTransform(const char* specStr,
         if (status.isOK()) {
             mongoutils::str::stream ss;
             ss << "expected transform() to fail but got success instead."
-               << "\nprojection spec: " << specStr << "\nquery: " << queryStr
-               << "\nobject before projection: " << objStr;
+               << "\nprojection spec: " << spec.toString() << "\nquery: " << query.toString()
+               << "\nobject before projection: " << obj.toString();
             FAIL(ss);
         }
         return;
@@ -118,26 +115,41 @@ void testTransform(const char* specStr,
     if (!status.isOK()) {
         mongoutils::str::stream ss;
         ss << "transform() test failed: unexpected failed status: " << status.toString()
-           << "\nprojection spec: " << specStr << "\nquery: " << queryStr
-           << "\nobject before projection: " << objStr
-           << "\nexpected object after projection: " << expectedObjStr;
+           << "\nprojection spec: " << spec.toString() << "\nquery: " << query.toString()
+           << "\nobject before projection: " << obj.toString()
+           << "\nexpected object after projection: " << expectedObj.toString();
         FAIL(ss);
     }
 
     // Finally, we compare the projected object.
-    const BSONObj& obj = wsm.obj.value();
-    BSONObj expectedObj = fromjson(expectedObjStr);
-    if (SimpleBSONObjComparator::kInstance.evaluate(obj != expectedObj)) {
+    const BSONObj& resultObj = wsm.obj.value();
+    if (SimpleBSONObjComparator::kInstance.evaluate(resultObj != expectedObj)) {
         mongoutils::str::stream ss;
         ss << "transform() test failed: unexpected projected object."
-           << "\nprojection spec: " << specStr << "\nquery: " << queryStr
-           << "\nobject before projection: " << objStr
-           << "\nexpected object after projection: " << expectedObjStr
-           << "\nactual object after projection: " << obj.toString();
+           << "\nprojection spec: " << spec.toString() << "\nquery: " << query.toString()
+           << "\nobject before projection: " << obj.toString()
+           << "\nexpected object after projection: " << expectedObj.toString()
+           << "\nactual object after projection: " << resultObj.toString();
         FAIL(ss);
     }
 }
 
+void testTransform(const char* specStr,
+                   const char* queryStr,
+                   const char* objStr,
+                   WorkingSetComputedData* data,
+                   const CollatorInterface* collator,
+                   bool expectedStatusOK,
+                   const char* expectedObjStr) {
+    testTransform(fromjson(specStr),
+                  fromjson(queryStr),
+                  fromjson(objStr),
+                  data,
+                  collator,
+                  expectedStatusOK,
+                  fromjson(expectedObjStr));
+}
+
 /**
  * testTransform without computed data or collator arguments.
  */
@@ -266,6 +278,74 @@ TEST(ProjectionExecTest, TransformSliceSkipLimit) {
     testTransform("{a: {$slice: [10, 10]}}", "{}", "{a: [4, 6, 8]}", true, "{a: []}");
 }
 
+TEST(ProjectionExecTest, TransformSliceWithSpecialValuesAsArguments) {
+    auto assertSliceResult =
+        [&](const BSONObj& spec, const char* inputDocument, const char* expectedResult) {
+            testTransform(spec,
+                          BSONObj(),
+                          fromjson(inputDocument),
+                          nullptr,
+                          nullptr,
+                          true,
+                          fromjson(expectedResult));
+        };
+
+    const auto positiveClamping =
+        BSON_ARRAY((static_cast<long long>(std::numeric_limits<int>::max()) + 1)
+                   << std::numeric_limits<long long>::max()
+                   << std::numeric_limits<double>::max()
+                   << std::numeric_limits<double>::infinity()
+                   << Decimal128::kPositiveInfinity
+                   << Decimal128::kLargestPositive);
+    for (const auto& element : positiveClamping) {
+        assertSliceResult(
+            BSON("a" << BSON("$slice" << element)), "{ a: [ 1, 2, 3 ] }", "{ a: [ 1, 2, 3 ] }");
+
+        assertSliceResult(BSON("a" << BSON("$slice" << BSON_ARRAY(1 << element))),
+                          "{ a: [1, 2, 3] }",
+                          "{ a: [ 2, 3 ] }");
+
+        assertSliceResult(BSON("a" << BSON("$slice" << BSON_ARRAY(element << 1))),
+                          "{ a: [ 1, 2, 3 ] }",
+                          "{ a: [] }");
+
+        assertSliceResult(BSON("a" << BSON("$slice" << BSON_ARRAY(element << element))),
+                          "{ a: [ 1, 2, 3 ] }",
+                          "{ a: [] }");
+    }
+
+    const auto negativeClamping =
+        BSON_ARRAY((static_cast<long long>(std::numeric_limits<int>::min()) - 1)
+                   << std::numeric_limits<long long>::min()
+                   << std::numeric_limits<double>::lowest()
+                   << -std::numeric_limits<double>::infinity()
+                   << Decimal128::kNegativeInfinity
+                   << Decimal128::kLargestNegative);
+    for (const auto& element : negativeClamping) {
+        assertSliceResult(
+            BSON("a" << BSON("$slice" << element)), "{ a: [ 1, 2, 3 ] }", "{ a: [ 1, 2, 3 ] }");
+
+        assertSliceResult(BSON("a" << BSON("$slice" << BSON_ARRAY(element << 1))),
+                          "{ a: [ 1, 2, 3 ] }",
+                          "{ a: [ 1 ] }");
+    }
+
+    const auto convertionToZero =
+        BSON_ARRAY(0ll << 0.0 << 0.3 << -0.3 << std::numeric_limits<double>::quiet_NaN()
+                       << Decimal128::kNegativeNaN
+                       << Decimal128::kPositiveNaN
+                       << Decimal128::kSmallestPositive
+                       << Decimal128::kSmallestNegative);
+    for (const auto& element : convertionToZero) {
+        assertSliceResult(
+            BSON("a" << BSON("$slice" << element)), "{ a: [ 1, 2, 3 ] }", "{ a: [] }");
+
+        assertSliceResult(BSON("a" << BSON("$slice" << BSON_ARRAY(element << 1))),
+                          "{ a: [ 1, 2, 3 ] }",
+                          "{ a: [ 1 ] }");
+    }
+}
+
 //
 // Dotted projections.
 //
diff --git a/src/mongo/db/query/parsed_projection.cpp b/src/mongo/db/query/parsed_projection.cpp
index e3a18928d4b..ca9eee97186 100644
--- a/src/mongo/db/query/parsed_projection.cpp
+++ b/src/mongo/db/query/parsed_projection.cpp
@@ -91,7 +91,7 @@ Status ParsedProjection::make(OperationContext* opCtx,
                     BSONObjIterator it(arr);
                     // Skip over 'skip'.
                     it.next();
-                    int limit = it.next().numberInt();
+                    int limit = it.next().safeNumberInt();
                     if (limit <= 0) {
                         return Status(ErrorCodes::BadValue, "$slice limit must be positive");
                     }