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

1 files changed, 20 insertions, 4 deletions

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);
         }
     }