SERVER-39788 Large values in $skip and $limit stages may cause an arithmetic overflow

6 files changed, 172 insertions, 21 deletions

diff --git a/jstests/aggregation/bugs/skip_limit_overflow.js b/jstests/aggregation/bugs/skip_limit_overflow.js
new file mode 100644
index 00000000000..f0d7e0b27c7
--- /dev/null
+++ b/jstests/aggregation/bugs/skip_limit_overflow.js
@@ -0,0 +1,123 @@
+/**
+ * SERVER-39788 Test that the values in the $limit and $skip stages do not overflow when the
+ * pipeline is optimized and the $sort stage doesn't crash the server due to unreasonable memory
+ * requirements.
+ *
+ * This test makes assumptions about how the explain output will be formatted, so cannot be
+ * transformed to be put inside a $facet stage or in a sharded explain output.
+ * @tags: [do_not_wrap_aggregations_in_facets, assumes_unsharded_collection]
+ */
+(function() {
+    "use strict";
+
+    load("jstests/libs/analyze_plan.js");  // For 'aggPlanHasStages' and other explain helpers.
+
+    const coll = db.server39788;
+    coll.drop();
+
+    function testPipeline(pipeline, expectedResult, optimizedAwayStages) {
+        const explainOutput = coll.explain().aggregate(pipeline);
+
+        assert(explainOutput.hasOwnProperty("stages"),
+               "Expected pipeline " + tojsononeline(pipeline) +
+                   " to use an aggregation framework in the explain output: " +
+                   tojson(explainOutput));
+
+        if (optimizedAwayStages) {
+            optimizedAwayStages.forEach(
+                (stage) =>
+                    assert(!aggPlanHasStage(explainOutput, stage),
+                           "Expected pipeline " + tojsononeline(pipeline) + " to *not* include a " +
+                               stage + " stage in the explain output: " + tojson(explainOutput)));
+        }
+
+        for (let path in expectedResult) {
+            const subPaths = path.split(".");
+            const stageName = subPaths[0];
+            const stages = getAggPlanStages(explainOutput, stageName);
+            assert(stages !== null,
+                   "Expected pipeline " + tojsononeline(pipeline) + " to include a " + stageName +
+                       " stage in the explain output: " + tojson(explainOutput));
+            assert(stages.length == expectedResult[path].length,
+                   "Expected pipeline " + tojsononeline(pipeline) + " to include " +
+                       expectedResult[path].length + stageName + " stages in the explain output: " +
+                       tojson(explainOutput));
+            assert.eq(
+                stages.reduce(
+                    (res, stage) => {
+                        res.push(subPaths.reduce((res, cur) => res[cur], stage));
+                        return res;
+                    },
+                    []),
+                expectedResult[path],
+                "Stage: " + stageName + ", path: " + path + ", explain: " + tojson(explainOutput));
+        }
+
+        // Ensure the aggregate command doesn't fail.
+        assert.eq(coll.aggregate(pipeline).toArray(), []);
+    }
+
+    // Case where overflow of limit + skip prevents limit stage from being absorbed. Values
+    // are specified as integrals > MAX_LONG. Note that we cannot specify this huge value as
+    // a NumberLong, as we get a number conversion error (even if it's passed as a string).
+    testPipeline([{$sort: {x: -1}}, {$skip: 18446744073709552000}, {$limit: 6}],
+                 {"$limit": [NumberLong(6)], "$skip": [NumberLong("9223372036854775807")]});
+    testPipeline([{$sort: {x: -1}}, {$skip: 6}, {$limit: 18446744073709552000}],
+                 {"$limit": [NumberLong("9223372036854775807")], "$skip": [NumberLong(6)]});
+
+    // Case where overflow of limit + skip prevents limit stage from being absorbed. One of the
+    // values == MAX_LONG, another one is 1.
+    testPipeline([{$sort: {x: -1}}, {$skip: NumberLong("9223372036854775807")}, {$limit: 1}],
+                 {"$limit": [NumberLong(1)], "$skip": [NumberLong("9223372036854775807")]});
+    testPipeline([{$sort: {x: -1}}, {$skip: 1}, {$limit: NumberLong("9223372036854775807")}],
+                 {"$limit": [NumberLong("9223372036854775807")], "$skip": [NumberLong(1)]});
+
+    // Case where limit + skip do not overflow. Limit == MAX_LONG and skip is 0. Should be able to
+    // absorb the limit and skip stages.
+    // Note that we cannot specify limit == 0, so we expect an error in this case.
+    testPipeline([{$sort: {x: -1}}, {$skip: 0}, {$limit: NumberLong("9223372036854775807")}],
+                 {"$cursor.limit": [NumberLong("9223372036854775807")]},
+                 ["$skip", "$limit"]);
+
+    // Case where limit + skip do not overflow. One value is MAX_LONG - 1 and another one is 1.
+    // Should be able to absorb the limit stage.
+    testPipeline([{$sort: {x: -1}}, {$skip: NumberLong("9223372036854775806")}, {$limit: 1}],
+                 {
+                   "$cursor.limit": [NumberLong("9223372036854775807")],
+                   "$skip": [NumberLong("9223372036854775806")]
+                 },
+                 ["$limit"]);
+    testPipeline([{$sort: {x: -1}}, {$skip: 1}, {$limit: NumberLong("9223372036854775806")}],
+                 {"$cursor.limit": [NumberLong("9223372036854775807")], "$skip": [NumberLong(1)]},
+                 ["$limit"]);
+
+    // Case where limit + skip do not overflow. Both values are < MAX_LONG.
+    testPipeline([{$sort: {x: -1}}, {$skip: 674761616283}, {$limit: 35361718}],
+                 {"$cursor.limit": [NumberLong(674796978001)], "$skip": [NumberLong(674761616283)]},
+                 ["$limit"]);
+    testPipeline([{$sort: {x: -1}}, {$skip: 35361718}, {$limit: 674761616283}],
+                 {"$cursor.limit": [NumberLong(674796978001)], "$skip": [NumberLong(35361718)]},
+                 ["$limit"]);
+
+    // Case where where overflow of limit + skip + skip prevents limit stage from being absorbed.
+    // One skip == MAX_LONG - 1, another one is 1. Should merge two skip stages into one.
+    testPipeline(
+        [{$sort: {x: -1}}, {$skip: 1}, {$skip: NumberLong("9223372036854775806")}, {$limit: 1}],
+        {"$limit": [NumberLong(1)], "$skip": [NumberLong("9223372036854775807")]});
+
+    // Case where where overflow of limit + skip + skip prevents limit stage from being absorbed.
+    // One skip == MAX_LONG, another one is 1. Should not absorb or merge any stages.
+    testPipeline(
+        [{$sort: {x: -1}}, {$skip: 1}, {$skip: NumberLong("9223372036854775807")}, {$limit: 1}],
+        {"$limit": [NumberLong(1)], "$skip": [NumberLong(1), NumberLong("9223372036854775807")]});
+
+    // Case where sample size is > MAX_LONG.
+    testPipeline([{$sample: {size: 18446744073709552000}}],
+                 {"$sample.size": [NumberLong("9223372036854775807")]});
+    // Case where sample size is == MAX_LONG.
+    testPipeline([{$sample: {size: NumberLong("9223372036854775807")}}],
+                 {"$sample.size": [NumberLong("9223372036854775807")]});
+    // Case where sample size is == MAX_LONG - 1.
+    testPipeline([{$sample: {size: NumberLong("9223372036854775806")}}],
+                 {"$sample.size": [NumberLong("9223372036854775806")]});
+})();
diff --git a/src/mongo/db/pipeline/document_source_limit.cpp b/src/mongo/db/pipeline/document_source_limit.cpp
index abe70f0d3db..bcc10f4f892 100644
--- a/src/mongo/db/pipeline/document_source_limit.cpp
+++ b/src/mongo/db/pipeline/document_source_limit.cpp
@@ -99,7 +99,7 @@ intrusive_ptr<DocumentSource> DocumentSourceLimit::createFromBson(
     BSONElement elem, const intrusive_ptr<ExpressionContext>& pExpCtx) {
     uassert(15957, "the limit must be specified as a number", elem.isNumber());
 
-    long long limit = elem.numberLong();
+    long long limit = elem.safeNumberLong();
     return DocumentSourceLimit::create(pExpCtx, limit);
 }
 }  // namespace mongo
diff --git a/src/mongo/db/pipeline/document_source_sample.cpp b/src/mongo/db/pipeline/document_source_sample.cpp
index 154fbe60174..b0f709366a1 100644
--- a/src/mongo/db/pipeline/document_source_sample.cpp
+++ b/src/mongo/db/pipeline/document_source_sample.cpp
@@ -102,8 +102,9 @@ intrusive_ptr<DocumentSource> DocumentSourceSample::createFromBson(
 
         if (fieldName == "size") {
             uassert(28746, "size argument to $sample must be a number", elem.isNumber());
-            uassert(28747, "size argument to $sample must not be negative", elem.numberLong() >= 0);
-            sample->_size = elem.numberLong();
+            auto size = elem.safeNumberLong();
+            uassert(28747, "size argument to $sample must not be negative", size >= 0);
+            sample->_size = size;
             sizeSpecified = true;
         } else {
             uasserted(28748, str::stream() << "unrecognized option to $sample: " << fieldName);
diff --git a/src/mongo/db/pipeline/document_source_skip.cpp b/src/mongo/db/pipeline/document_source_skip.cpp
index e58259232c7..364f148162e 100644
--- a/src/mongo/db/pipeline/document_source_skip.cpp
+++ b/src/mongo/db/pipeline/document_source_skip.cpp
@@ -107,7 +107,7 @@ intrusive_ptr<DocumentSource> DocumentSourceSkip::createFromBson(
     uassert(15972,
             str::stream() << "Argument to $skip must be a number not a " << typeName(elem.type()),
             elem.isNumber());
-    auto nToSkip = elem.numberLong();
+    auto nToSkip = elem.safeNumberLong();
     uassert(15956, "Argument to $skip cannot be negative", nToSkip >= 0);
 
     return DocumentSourceSkip::create(pExpCtx, nToSkip);
diff --git a/src/mongo/db/pipeline/document_source_sort.cpp b/src/mongo/db/pipeline/document_source_sort.cpp
index 0d5f0b331c8..9f9832fd961 100644
--- a/src/mongo/db/pipeline/document_source_sort.cpp
+++ b/src/mongo/db/pipeline/document_source_sort.cpp
@@ -40,6 +40,7 @@
 #include "mongo/db/pipeline/lite_parsed_document_source.h"
 #include "mongo/db/pipeline/value.h"
 #include "mongo/db/query/collation/collation_index_key.h"
+#include "mongo/platform/overflow_arithmetic.h"
 #include "mongo/s/query/document_source_merge_cursors.h"
 
 namespace mongo {
@@ -200,24 +201,34 @@ Pipeline::SourceContainer::iterator DocumentSourceSort::doOptimizeAt(
     Pipeline::SourceContainer::iterator itr, Pipeline::SourceContainer* container) {
     invariant(*itr == this);
 
-    auto sortItr = std::next(itr);
-    long long skipSum = 0;
-    while (sortItr != container->end()) {
-        auto nextStage = (*sortItr).get();
-
-        if (auto nextSkip = dynamic_cast<DocumentSourceSkip*>(nextStage)) {
-            skipSum += nextSkip->getSkip();
-            ++sortItr;
-        } else if (auto nextLimit = dynamic_cast<DocumentSourceLimit*>(nextStage)) {
-            nextLimit->setLimit(nextLimit->getLimit() + skipSum);
+    using SumType = long long;
+    // Just in case, since we're passing an address of a signed 64bit below.
+    static_assert(std::is_signed<SumType>::value && sizeof(SumType) == 8);
+
+    auto stageItr = std::next(itr);
+    SumType skipSum = 0;
+    while (stageItr != container->end()) {
+        auto nextStage = (*stageItr).get();
+        auto nextSkip = dynamic_cast<DocumentSourceSkip*>(nextStage);
+        auto nextLimit = dynamic_cast<DocumentSourceLimit*>(nextStage);
+        SumType safeSum = 0;
+
+        // The skip and limit values can be very large, so we need to make sure the sum doesn't
+        // overflow before applying an optimiztion to pull the limit into the sort stage.
+        if (nextSkip && !mongoSignedAddOverflow64(skipSum, nextSkip->getSkip(), &safeSum)) {
+            skipSum = safeSum;
+            ++stageItr;
+        } else if (nextLimit &&
+                   !mongoSignedAddOverflow64(nextLimit->getLimit(), skipSum, &safeSum)) {
+            nextLimit->setLimit(safeSum);
             setLimitSrc(nextLimit);
-            container->erase(sortItr);
-            sortItr = std::next(itr);
+            container->erase(stageItr);
+            stageItr = std::next(itr);
             skipSum = 0;
         } else if (!nextStage->constraints().canSwapWithLimitAndSample) {
             return std::next(itr);
         } else {
-            ++sortItr;
+            ++stageItr;
         }
     }
 
diff --git a/src/mongo/db/sorter/sorter.cpp b/src/mongo/db/sorter/sorter.cpp
index a6e038191fe..0e3e1ad7f38 100644
--- a/src/mongo/db/sorter/sorter.cpp
+++ b/src/mongo/db/sorter/sorter.cpp
@@ -59,6 +59,7 @@
 #include "mongo/db/storage/encryption_hooks.h"
 #include "mongo/db/storage/storage_options.h"
 #include "mongo/platform/atomic_word.h"
+#include "mongo/platform/overflow_arithmetic.h"
 #include "mongo/s/is_mongos.h"
 #include "mongo/util/assert_util.h"
 #include "mongo/util/bufreader.h"
@@ -655,10 +656,25 @@ public:
             _fileName = _opts.tempDir + "/" + nextFileName();
         }
 
-        // Preallocate a fixed sized vector of the required size if we
-        // don't expect it to have a major impact on our memory budget.
-        // This is the common case with small limits.
-        if ((sizeof(Data) * opts.limit) < opts.maxMemoryUsageBytes / 10) {
+        // Preallocate a fixed sized vector of the required size if we don't expect it to have a
+        // major impact on our memory budget. This is the common case with small limits. If the
+        // limit is really large, we need to take care when doing the check below. Both 'opts.limit'
+        // and 'sizeof(Data)' are unsigned long's, it's always clearly defined behaviour to multiply
+        // these two numbers, but the result may be wrapped around (truncated) due to modular
+        // arithmetics. So, the result of the multiplication may be < maxMemoryUsageBytes / 10, but
+        // 'opts.limit' may be still large enough to trigger an std::length_error exception and
+        // abnormally terminate the program. So, we'll check here if the multiply operation was safe
+        // and whether the limit fits into the maximum allowed vector size.
+        using MultiplicationType = size_t;
+        // Just in case, since we're passing an address of an unsigned 64bit below.
+        static_assert(std::is_unsigned<MultiplicationType>::value &&
+                      sizeof(MultiplicationType) == 8);
+
+        MultiplicationType requiredBytes;
+        auto isSafeMultiplication = mongoUnsignedMultiplyOverflow64(
+            sizeof(typename decltype(_data)::value_type), opts.limit, &requiredBytes);
+        if (isSafeMultiplication && requiredBytes < opts.maxMemoryUsageBytes / 10 &&
+            opts.limit <= _data.max_size()) {
             _data.reserve(opts.limit);
         }
     }