TOOLS-754 prevent applyOps from overflowing maximum message size

2 files changed, 66 insertions, 1 deletions

diff --git a/mongorestore/oplog.go b/mongorestore/oplog.go
index de2014b4d65..95ed1bf8da1 100644
--- a/mongorestore/oplog.go
+++ b/mongorestore/oplog.go
@@ -14,7 +14,11 @@ import (
 	"time"
 )
 
-const oplogMaxCommandSize = 1024 * 1024 * 16.5
+// oplogMaxCommandSize sets the maximum size for multiple buffered ops in the
+// applyOps command. This is to prevent pathological cases where the array overhead
+// of many small operations can overflow the maximum command size.
+// Note that ops > 8MB will still be buffered, just as single elements.
+const oplogMaxCommandSize = 1024 * 1024 * 8
 
 // RestoreOplog attempts to restore a MongoDB oplog.
 func (restore *MongoRestore) RestoreOplog() error {
diff --git a/test/qa-tests/jstests/restore/oplog_replay_size_safety.js b/test/qa-tests/jstests/restore/oplog_replay_size_safety.js
new file mode 100644
index 00000000000..1b37b9f0d1a
--- /dev/null
+++ b/test/qa-tests/jstests/restore/oplog_replay_size_safety.js
@@ -0,0 +1,61 @@
+(function() {
+
+    if (typeof getToolTest === 'undefined') {
+        load('jstests/configs/plain_28.config.js');
+    }
+
+    var commonToolArgs = getCommonToolArguments();
+    var dumpTarget = 'oplog_replay_sizes';
+
+    // Helper for using mongorestore with --oplogReplay and a large oplog.bson
+    function tryOplogReplay(oplogSize, documentSize) {
+      var toolTest = getToolTest('oplog_replay_sizes');
+      // the test db and collections we'll be using
+      var testDB = toolTest.db.getSiblingDB('test_oplog');
+      var testColl = testDB.oplog;
+      var testRestoreDB = toolTest.db.getSiblingDB('test');
+      var testRestoreColl = testRestoreDB.op;
+      resetDbpath(dumpTarget);
+
+      var debugString = 'with ' + oplogSize + ' ops of size ' + documentSize;
+      jsTest.log('Testing --oplogReplay ' + debugString);
+
+
+      // create a fake oplog consisting of a large number of inserts
+      var xStr = new Array(documentSize).join("x"); // ~documentSize bytes string
+      for (var i = 0; i < oplogSize; i++) {
+        testColl.insert({ts:  new Timestamp(0,i), op: "i",
+          o: {_id:i, x: xStr}, "ns":"test.op"});
+      }
+
+      // dump the fake oplog
+      var ret = toolTest.runTool.apply(
+          toolTest,
+          ['dump', '--db', 'test_oplog', '-c', 'oplog', '--out', dumpTarget].concat(commonToolArgs)
+          );
+      assert.eq(0, ret, "dump operation failed " + debugString);
+
+      // create the test.op collection
+      testRestoreColl.drop();
+      testRestoreDB.createCollection("op");
+      assert.eq(0, testRestoreColl.count());
+
+      // trick restore into replaying the "oplog" we forged above
+      ret = toolTest.runTool.apply(
+          toolTest,
+          ['restore', '--oplogReplay', dumpTarget+'/test_oplog'].concat(commonToolArgs)
+          );
+      assert.eq(0, ret, "restore operation failed " + debugString);
+      assert.eq(oplogSize, testRestoreColl.count(),
+          "all oplog entries should be inserted " + debugString);
+      toolTest.stop();
+    }
+
+    // run the test on various oplog and op sizes
+    tryOplogReplay(1024, 1024);      // sanity check
+    tryOplogReplay(1024*1024, 1);    // millions of micro ops
+    tryOplogReplay(8, 16*1024*1023); // 8 ~16MB ops
+    tryOplogReplay(32, 1024*1024);   // 32 ~1MB ops
+    tryOplogReplay(32*1024, 1024);   // many ~1KB ops
+
+}());