SERVER-17224 Reserve room for EOO byte when starting BSONObj building

Since _done() is called from ~BSONObjBuilder we need to ensure that it cannot
fail. This prevents a double exception leading to a std::terminate call.

This also resolves SERVER-17226.

(cherry picked from commit a0db9321139e8da657638ddbe7e86d8bac9ea3cc)

3 files changed, 76 insertions, 7 deletions

diff --git a/jstests/aggregation/bugs/server17224.js b/jstests/aggregation/bugs/server17224.js
new file mode 100644
index 00000000000..33042abab53
--- /dev/null
+++ b/jstests/aggregation/bugs/server17224.js
@@ -0,0 +1,24 @@
+// SERVER-17224 An aggregation result with exactly the right size could crash the server rather than
+//              returning an error.
+(function() {
+    'use strict';
+
+    var t = db.server17224;
+    t.drop();
+
+    // first 63MB
+    for (var i = 0; i < 63; i++) {
+        t.insert({a: new Array(1024 * 1024 + 1).join('a')});
+    }
+
+    // the remaining ~1MB with room for field names and other overhead
+    t.insert({a: new Array(1024 * 1024 - 1105).join('a')});
+
+    // do not use cursor form, since it has a different workaroud for this issue.
+    assert.commandFailed(
+        db.runCommand({aggregate: t.getName(),
+                      pipeline: [{$match: {}}, {$group: {_id: null, arr: {$push: {a: '$a'}}}}]}));
+
+    // Make sure the server is still up.
+    assert.commandWorked(db.runCommand('ping'));
+}());
diff --git a/src/mongo/bson/bsonobjbuilder.h b/src/mongo/bson/bsonobjbuilder.h
index 65c2e8f0d66..e064c092de6 100644
--- a/src/mongo/bson/bsonobjbuilder.h
+++ b/src/mongo/bson/bsonobjbuilder.h
@@ -68,10 +68,13 @@ namespace mongo {
             , _s(this)
             , _tracker(0)
             , _doneCalled(false) {
-            // Reserve space for a holder object at the beginning of the buffer, followed by
+            // Skip over space for a holder object at the beginning of the buffer, followed by
             // space for the object length. The length is filled in by _done.
             _b.skip(sizeof(BSONObj::Holder));
             _b.skip(sizeof(int));
+
+            // Reserve space for the EOO byte. This means _done() can't fail.
+            _b.reserveBytes(1);
         }
 
         /** @param baseBuilder construct a BSONObjBuilder using an existing BufBuilder
@@ -84,9 +87,13 @@ namespace mongo {
             , _s(this)
             , _tracker(0)
             , _doneCalled(false) {
-            // Reserve space for the object length, which is filled in by _done. We don't need a holder
-            // since we are a sub-builder, and some parent builder has already made the reservation.
+            // Skip over space for the object length, which is filled in by _done. We don't need a
+            // holder since we are a sub-builder, and some parent builder has already made the
+            // reservation.
             _b.skip(sizeof(int));
+
+            // Reserve space for the EOO byte. This means _done() can't fail.
+            _b.reserveBytes(1);
         }
 
         BSONObjBuilder( const BSONSizeTracker & tracker )
@@ -99,6 +106,9 @@ namespace mongo {
             // See the comments in the first constructor for details.
             _b.skip(sizeof(BSONObj::Holder));
             _b.skip(sizeof(int));
+
+            // Reserve space for the EOO byte. This means _done() can't fail.
+            _b.reserveBytes(1);
         }
 
         ~BSONObjBuilder() {
@@ -622,6 +632,7 @@ namespace mongo {
         BSONObj asTempObj() {
             BSONObj temp(_done());
             _b.setlen(_b.len()-1); //next append should overwrite the EOO
+            _b.reserveBytes(1); // Rereserve room for the real EOO
             _doneCalled = false;
             return temp;
         }
@@ -703,8 +714,14 @@ namespace mongo {
                 return _b.buf() + _offset;
 
             _doneCalled = true;
+
+            // TODO remove this or find some way to prevent it from failing. Since this is intended
+            // for use with BSON() literal queries, it is less likely to result in oversized BSON.
             _s.endField();
+
+            _b.claimReservedBytes(1); // Prevents adding EOO from failing.
             _b.appendNum((char) EOO);
+
             char *data = _b.buf() + _offset;
             int size = _b.len() - _offset;
             DataView(data).writeLE(size);
diff --git a/src/mongo/bson/util/builder.h b/src/mongo/bson/util/builder.h
index 3833e6747db..be309ca652a 100644
--- a/src/mongo/bson/util/builder.h
+++ b/src/mongo/bson/util/builder.h
@@ -125,6 +125,7 @@ namespace mongo {
                 data = 0;
             }
             l = 0;
+            reservedBytes = 0;
         }
         ~_BufBuilder() { kill(); }
 
@@ -137,9 +138,11 @@ namespace mongo {
 
         void reset() {
             l = 0;
+            reservedBytes = 0;
         }
         void reset( int maxSize ) {
             l = 0;
+            reservedBytes = 0;
             if ( maxSize && size > maxSize ) {
                 al.Free(data);
                 data = (char*)al.Malloc(maxSize);
@@ -228,13 +231,36 @@ namespace mongo {
         inline char* grow(int by) {
             int oldlen = l;
             int newLen = l + by;
-            if ( newLen > size ) {
-                grow_reallocate(newLen);
+            int minSize = newLen + reservedBytes;
+            if ( minSize > size ) {
+                grow_reallocate(minSize);
             }
             l = newLen;
             return data + oldlen;
         }
 
+        /**
+         * Reserve room for some number of bytes to be claimed at a later time.
+         */
+        void reserveBytes(int bytes) {
+            int minSize = l + reservedBytes + bytes;
+            if (minSize > size)
+                grow_reallocate(minSize);
+
+            // This must happen *after* any attempt to grow.
+            reservedBytes += bytes;
+        }
+
+        /**
+         * Claim an earlier reservation of some number of bytes. These bytes must already have been
+         * reserved. Appends of up to this many bytes immediately following a claim are
+         * guaranteed to succeed without a need to reallocate.
+         */
+        void claimReservedBytes(int bytes) {
+            invariant(reservedBytes >= bytes);
+            reservedBytes -= bytes;
+        }
+
     private:
         template<typename T>
         void appendNumImpl(T t) {
@@ -247,10 +273,11 @@ namespace mongo {
 
 
         /* "slow" portion of 'grow()'  */
-        void NOINLINE_DECL grow_reallocate(int newLen) {
+        void NOINLINE_DECL grow_reallocate(int minSize) {
             int a = 64;
-            while( a < newLen ) 
+            while (a < minSize)
                 a = a * 2;
+
             if ( a > BufferMaxSize ) {
                 std::stringstream ss;
                 ss << "BufBuilder attempted to grow() to " << a << " bytes, past the 64MB limit.";
@@ -265,6 +292,7 @@ namespace mongo {
         char *data;
         int l;
         int size;
+        int reservedBytes; // eagerly grow_reallocate to keep this many bytes of spare room.
 
         friend class StringBuilderImpl<Allocator>;
     };