SERVER-8183 Iterative implementation of validateBSON().

3 files changed, 193 insertions, 111 deletions

diff --git a/src/mongo/bson/bson_validate.cpp b/src/mongo/bson/bson_validate.cpp
index 7e4dfd10839..9ac0bcbbbfb 100644
--- a/src/mongo/bson/bson_validate.cpp
+++ b/src/mongo/bson/bson_validate.cpp
@@ -15,6 +15,8 @@
  *    limitations under the License.
  */
 
+#include <deque>
+
 #include "mongo/bson/bson_validate.h"
 #include "mongo/bson/oid.h"
 
@@ -92,143 +94,219 @@ namespace mongo {
             uint64_t _maxLength;
         };
 
-        Status validateBSONInternal( Buffer* buffer, int* bsonLength ) {
-            const int start = buffer->position();
+        struct ValidationState {
+            enum State {
+                BeginObj = 1,
+                WithinObj,
+                EndObj,
+                BeginCodeWScope,
+                EndCodeWScope,
+                Done
+            };
+        };
 
-            int supposedSize;
-            if ( !buffer->readNumber<int>(&supposedSize) )
-                return Status( ErrorCodes::InvalidBSON, "bson size is larger than buffer size" );
+        class ValidationObjectFrame {
+        public:
+            int startPosition() const { return _startPosition & ~(1 << 31); }
+            bool isCodeWithScope() const { return _startPosition & (1 << 31); }
+
+            void setStartPosition(int pos) {
+                _startPosition = (_startPosition & (1 << 31)) | (pos & ~(1 << 31));
+            }
+            void setIsCodeWithScope(bool isCodeWithScope) {
+                if (isCodeWithScope) {
+                    _startPosition |= 1 << 31;
+                }
+                else {
+                    _startPosition &= ~(1 << 31);
+                }
+            }
 
+            int expectedSize;
+        private:
+            int _startPosition;
+        };
+
+        Status validateElementInfo(Buffer* buffer, ValidationState::State* nextState) {
             Status status = Status::OK();
 
-            while ( true ) {
-                char type;
-                if ( !buffer->readNumber<char>(&type) )
-                    return Status( ErrorCodes::InvalidBSON, "invalid bson" );
+            char type;
+            if ( !buffer->readNumber<char>(&type) )
+                return Status( ErrorCodes::InvalidBSON, "invalid bson" );
 
-                if ( type == EOO )
-                    break;
+            if ( type == EOO ) {
+                *nextState = ValidationState::EndObj;
+                return Status::OK();
+            }
 
-                StringData name;
-                status = buffer->readCString( &name );
-                if ( !status.isOK() )
-                    return status;
+            StringData name;
+            status = buffer->readCString( &name );
+            if ( !status.isOK() )
+                return status;
 
-                switch ( type ) {
-                case MinKey:
-                case MaxKey:
-                case jstNULL:
-                case Undefined:
-                    break;
+            switch ( type ) {
+            case MinKey:
+            case MaxKey:
+            case jstNULL:
+            case Undefined:
+                return Status::OK();
 
-                case jstOID:
-                    if ( !buffer->skip( sizeof(OID) ) )
-                        return Status( ErrorCodes::InvalidBSON, "invalid bson" );
-                    break;
+            case jstOID:
+                if ( !buffer->skip( sizeof(OID) ) )
+                    return Status( ErrorCodes::InvalidBSON, "invalid bson" );
+                return Status::OK();
 
-                case NumberInt:
-                    if ( !buffer->skip( sizeof(int32_t) ) )
-                        return Status( ErrorCodes::InvalidBSON, "invalid bson" );
-                    break;
+            case NumberInt:
+                if ( !buffer->skip( sizeof(int32_t) ) )
+                    return Status( ErrorCodes::InvalidBSON, "invalid bson" );
+                return Status::OK();
 
-                case Bool:
-                    if ( !buffer->skip( sizeof(int8_t) ) )
-                        return Status( ErrorCodes::InvalidBSON, "invalid bson" );
-                    break;
+            case Bool:
+                if ( !buffer->skip( sizeof(int8_t) ) )
+                    return Status( ErrorCodes::InvalidBSON, "invalid bson" );
+                return Status::OK();
 
 
-                case NumberDouble:
-                case NumberLong:
-                case Timestamp:
-                case Date:
-                    if ( !buffer->skip( sizeof(int64_t) ) )
-                        return Status( ErrorCodes::InvalidBSON, "invalid bson" );
-                    break;
+            case NumberDouble:
+            case NumberLong:
+            case Timestamp:
+            case Date:
+                if ( !buffer->skip( sizeof(int64_t) ) )
+                    return Status( ErrorCodes::InvalidBSON, "invalid bson" );
+                return Status::OK();
 
-                case DBRef:
-                    status = buffer->readUTF8String( NULL );
-                    if ( !status.isOK() )
-                        return status;
-                    buffer->skip( sizeof(OID) );
-                    break;
+            case DBRef:
+                status = buffer->readUTF8String( NULL );
+                if ( !status.isOK() )
+                    return status;
+                buffer->skip( sizeof(OID) );
+                return Status::OK();
 
-                case CodeWScope: {
-                    int myStart = buffer->position();
-                    int sz;
+            case RegEx:
+                status = buffer->readCString( NULL );
+                if ( !status.isOK() )
+                    return status;
+                status = buffer->readCString( NULL );
+                if ( !status.isOK() )
+                    return status;
 
-                    if ( !buffer->readNumber<int>( &sz ) )
-                        return Status( ErrorCodes::InvalidBSON, "invalid bson" );
+                return Status::OK();
 
-                    status = buffer->readUTF8String( NULL );
-                    if ( !status.isOK() )
-                        return status;
+            case Code:
+            case Symbol:
+            case String:
+                status = buffer->readUTF8String( NULL );
+                if ( !status.isOK() )
+                    return status;
+                return Status::OK();
 
-                    status = validateBSONInternal( buffer, NULL );
-                    if ( !status.isOK() )
-                        return status;
+            case BinData: {
+                int sz;
+                if ( !buffer->readNumber<int>( &sz ) )
+                    return Status( ErrorCodes::InvalidBSON, "invalid bson" );
+                if ( !buffer->skip( 1 + sz ) )
+                    return Status( ErrorCodes::InvalidBSON, "invalid bson" );
+                return Status::OK();
+            }
+            case CodeWScope:
+                *nextState = ValidationState::BeginCodeWScope;
+                return Status::OK();
+            case Object:
+            case Array:
+                *nextState = ValidationState::BeginObj;
+                return Status::OK();
 
-                    if ( sz != static_cast<int>(buffer->position() - myStart) )
-                        return Status( ErrorCodes::InvalidBSON, "CodeWScope len is wrong" );
+            default:
+                return Status( ErrorCodes::InvalidBSON, "invalid bson type" );
+            }
+        }
 
+        Status validateBSONIterative(Buffer* buffer) {
+            std::deque<ValidationObjectFrame> frames;
+            ValidationObjectFrame* curr = NULL;
+            ValidationState::State state = ValidationState::BeginObj;
+
+            while (state != ValidationState::Done) {
+                switch (state) {
+                case ValidationState::BeginObj:
+                    frames.push_back(ValidationObjectFrame());
+                    curr = &frames.back();
+                    curr->setStartPosition(buffer->position());
+                    curr->setIsCodeWithScope(false);
+                    if (!buffer->readNumber<int>(&curr->expectedSize)) {
+                        return Status(ErrorCodes::InvalidBSON,
+                                      "bson size is larger than buffer size");
+                    }
+                    state = ValidationState::WithinObj;
+                    // fall through
+                case ValidationState::WithinObj: {
+                    Status status = validateElementInfo(buffer, &state);
+                    if (!status.isOK())
+                        return status;
                     break;
                 }
-                case RegEx:
-                    status = buffer->readCString( NULL );
-                    if ( !status.isOK() )
-                        return status;
-                    status = buffer->readCString( NULL );
-                    if ( !status.isOK() )
-                        return status;
-
+                case ValidationState::EndObj: {
+                    int actualLength = buffer->position() - curr->startPosition();
+                    if ( actualLength != curr->expectedSize ) {
+                        return Status( ErrorCodes::InvalidBSON,
+                                       "bson length doesn't match what we found" );
+                    }
+                    frames.pop_back();
+                    if (frames.empty()) {
+                        state = ValidationState::Done;
+                    }
+                    else {
+                        curr = &frames.back();
+                        if (curr->isCodeWithScope())
+                            state = ValidationState::EndCodeWScope;
+                        else
+                            state = ValidationState::WithinObj;
+                    }
                     break;
-
-                case Code:
-                case Symbol:
-                case String:
-                    status = buffer->readUTF8String( NULL );
+                }
+                case ValidationState::BeginCodeWScope: {
+                    frames.push_back(ValidationObjectFrame());
+                    curr = &frames.back();
+                    curr->setStartPosition(buffer->position());
+                    curr->setIsCodeWithScope(true);
+                    if ( !buffer->readNumber<int>( &curr->expectedSize ) )
+                        return Status( ErrorCodes::InvalidBSON, "invalid bson CodeWScope size" );
+                    Status status = buffer->readUTF8String( NULL );
                     if ( !status.isOK() )
                         return status;
+                    state = ValidationState::BeginObj;
                     break;
-
-                case BinData: {
-                    int sz;
-                    if ( !buffer->readNumber<int>( &sz ) )
-                        return Status( ErrorCodes::InvalidBSON, "invalid bson" );
-                    if ( !buffer->skip( 1 + sz ) )
-                        return Status( ErrorCodes::InvalidBSON, "invalid bson" );
+                }
+                case ValidationState::EndCodeWScope: {
+                    int actualLength = buffer->position() - curr->startPosition();
+                    if ( actualLength != curr->expectedSize ) {
+                        return Status( ErrorCodes::InvalidBSON,
+                                       "bson length for CodeWScope doesn't match what we found" );
+                    }
+                    frames.pop_back();
+                    if (frames.empty())
+                        return Status(ErrorCodes::InvalidBSON, "unnested CodeWScope");
+                    curr = &frames.back();
+                    state = ValidationState::WithinObj;
                     break;
                 }
-                case Object:
-                case Array:
-                    status = validateBSONInternal( buffer, NULL );
-                    if ( !status.isOK() )
-                        return status;
+                case ValidationState::Done:
                     break;
-
-                default:
-                    return Status( ErrorCodes::InvalidBSON, "invalid bson type" );
                 }
             }
 
-            const int end = buffer->position();
-
-            if ( end - start != supposedSize )
-                return Status( ErrorCodes::InvalidBSON, "bson length doesn't match what we found" );
-
-            if ( bsonLength )
-                *bsonLength = supposedSize;
-
             return Status::OK();
         }
-    }
 
-    Status validateBSON( const char* originalBuffer, uint64_t maxLength, int* bsonLength ) {
+    }  // namespace
+
+    Status validateBSON( const char* originalBuffer, uint64_t maxLength ) {
         if ( maxLength < 5 ) {
             return Status( ErrorCodes::InvalidBSON, "bson data has to be at least 5 bytes" );
         }
 
         Buffer buf( originalBuffer, maxLength );
-        return validateBSONInternal( &buf, bsonLength );
+        return validateBSONIterative( &buf );
     }
 
-}
+}  // namespace mongo
diff --git a/src/mongo/bson/bson_validate.h b/src/mongo/bson/bson_validate.h
index db9ddcfa7b4..fb1995d2bd2 100644
--- a/src/mongo/bson/bson_validate.h
+++ b/src/mongo/bson/bson_validate.h
@@ -28,9 +28,8 @@ namespace mongo {
      * @param buf - bson data
      * @param maxLength - maxLength of buffer
      *                    this is NOT the bson size, but how far we know the buffer is valid
-     * @param bsonLength - OUT size of bson, set if bsonLength != NULL _and_ data is valid
      */
-    Status validateBSON( const char* buf, uint64_t maxLength, int* bsonLength = NULL );
+    Status validateBSON( const char* buf, uint64_t maxLength );
 
 }
 
diff --git a/src/mongo/bson/bson_validate_test.cpp b/src/mongo/bson/bson_validate_test.cpp
index 3de58f5dd60..61e67f33c46 100644
--- a/src/mongo/bson/bson_validate_test.cpp
+++ b/src/mongo/bson/bson_validate_test.cpp
@@ -57,10 +57,10 @@ namespace {
             if ( o.valid() ) {
                 numValid++;
                 jsonSize += o.jsonString().size();
-                ASSERT( validateBSON( o.objdata(), o.objsize() ).isOK() );
+                ASSERT_OK( validateBSON( o.objdata(), o.objsize() ) );
             }
             else {
-                ASSERT( !validateBSON( o.objdata(), o.objsize() ).isOK() );
+                ASSERT_NOT_OK( validateBSON( o.objdata(), o.objsize() ) );
             }
 
             delete[] x;
@@ -104,10 +104,10 @@ namespace {
             if ( mine.valid() ) {
                 numValid++;
                 jsonSize += mine.jsonString().size();
-                ASSERT( validateBSON( mine.objdata(), mine.objsize() ).isOK() );
+                ASSERT_OK( validateBSON( mine.objdata(), mine.objsize() ) );
             }
             else {
-                ASSERT( !validateBSON( mine.objdata(), mine.objsize() ).isOK() );
+                ASSERT_NOT_OK( validateBSON( mine.objdata(), mine.objsize() ) );
             }
 
         }
@@ -161,22 +161,22 @@ namespace {
 
     TEST( BSONValidateFast, Empty ) {
         BSONObj x;
-        ASSERT( validateBSON( x.objdata(), x.objsize() ).isOK() );
+        ASSERT_OK( validateBSON( x.objdata(), x.objsize() ) );
     }
 
     TEST( BSONValidateFast, RegEx ) {
         BSONObjBuilder b;
         b.appendRegex( "foo", "i" );
         BSONObj x = b.obj();
-        ASSERT( validateBSON( x.objdata(), x.objsize() ).isOK() );
+        ASSERT_OK( validateBSON( x.objdata(), x.objsize() ) );
     }
 
     TEST(BSONValidateFast, Simple0 ) {
         BSONObj x;
-        ASSERT( validateBSON( x.objdata(), x.objsize() ).isOK() );
+        ASSERT_OK( validateBSON( x.objdata(), x.objsize() ) );
 
         x = BSON( "foo" << 17 << "bar" << "eliot" );
-        ASSERT( validateBSON( x.objdata(), x.objsize() ).isOK() );
+        ASSERT_OK( validateBSON( x.objdata(), x.objsize() ) );
 
     }
 
@@ -189,7 +189,7 @@ namespace {
             sprintf( buf, "bar%d", i );
             b.appendMaxForType( buf, i );
             BSONObj x = b.obj();
-            ASSERT( validateBSON( x.objdata(), x.objsize() ).isOK() );
+            ASSERT_OK( validateBSON( x.objdata(), x.objsize() ) );
         }
     }
 
@@ -204,8 +204,13 @@ namespace {
             b.appendMaxForType( buf, i );
         }
         BSONObj x = b.obj();
-        ASSERT( validateBSON( x.objdata(), x.objsize() ).isOK() );
+        ASSERT_OK( validateBSON( x.objdata(), x.objsize() ) );
     }
 
+    TEST(BSONValidateFast, NestedObject) {
+        BSONObj x = BSON( "a" << 1 << "b" << BSON("c" << 2 << "d" << BSONArrayBuilder().obj() << "e" << BSON_ARRAY("1" << 2 << 3)));
+        ASSERT_OK(validateBSON(x.objdata(), x.objsize()));
+        ASSERT_NOT_OK(validateBSON(x.objdata(), x.objsize() / 2));
+    }
 
 }