SERVER-14268: Improve server message parsing

1 files changed, 43 insertions, 104 deletions

diff --git a/src/mongo/db/dbmessage.h b/src/mongo/db/dbmessage.h
index 621624dd9be..5131e2d30a2 100644
--- a/src/mongo/db/dbmessage.h
+++ b/src/mongo/db/dbmessage.h
@@ -122,135 +122,74 @@ namespace mongo {
        See http://dochub.mongodb.org/core/mongowireprotocol
     */
     class DbMessage {
+    // Assume sizeof(int) == 4 bytes
+    BOOST_STATIC_ASSERT(sizeof(int) == 4);
+
     public:
-        DbMessage(const Message& _m) : m(_m) , mark(0) {
-            // for received messages, Message has only one buffer
-            theEnd = _m.singleData()->_data + _m.header()->dataLen();
-            char *r = _m.singleData()->_data;
-            reserved = (int *) r;
-            data = r + 4;
-            nextjsobj = data;
+        // Note: DbMessage constructor reads the first 4 bytes and stores it in reserved
+        DbMessage(const Message& msg);
+
+        // Indicates whether this message is expected to have a ns
+        // or in the case of dbMsg, a string in the same place as ns
+        bool messageShouldHaveNs() const {
+            return (_msg.operation() >= dbMsg) & (_msg.operation() <= dbDelete);
         }
 
-        /** the 32 bit field before the ns 
+        /** the 32 bit field before the ns
          * track all bit usage here as its cross op
          * 0: InsertOption_ContinueOnError
          * 1: fromWriteback
          */
-        int& reservedField() { return *reserved; }
-
-        const char * getns() const {
-            return data;
-        }
+        int reservedField() const { return _reserved; }
 
-        const char * afterNS() const {
-            return data + strlen( data ) + 1;
-        }
-
-        int getInt( int num ) const {
-            const int * foo = (const int*)afterNS();
-            return foo[num];
-        }
+        const char * getns() const;
+        int getQueryNToReturn() const;
 
-        int getQueryNToReturn() const {
-            return getInt( 1 );
-        }
+        int pullInt();
+        long long pullInt64();
+        const long long* getArray(size_t count) const;
 
-        /**
-         * get an int64 at specified offsetBytes after ns
-         */
-        long long getInt64( int offsetBytes ) const {
-            const char * x = afterNS();
-            x += offsetBytes;
-            const long long * ll = (const long long*)x;
-            return ll[0];
+        /* for insert and update msgs */
+        bool moreJSObjs() const {
+            return _nextjsobj != 0;
         }
 
-        void resetPull() { nextjsobj = data; }
-        int pullInt() const { return pullInt(); }
-        int& pullInt() {
-            if ( nextjsobj == data )
-                nextjsobj += strlen(data) + 1; // skip namespace
-            int& i = *((int *)nextjsobj);
-            nextjsobj += 4;
-            return i;
-        }
-        long long pullInt64() const {
-            return pullInt64();
-        }
-        long long &pullInt64() {
-            if ( nextjsobj == data )
-                nextjsobj += strlen(data) + 1; // skip namespace
-            long long &i = *((long long *)nextjsobj);
-            nextjsobj += 8;
-            return i;
-        }
+        BSONObj nextJsObj();
 
-        OID* getOID() const {
-            return (OID *) (data + strlen(data) + 1); // skip namespace
-        }
+        const Message& msg() const { return _msg; }
 
-        void getQueryStuff(const char *&query, int& ntoreturn) {
-            int *i = (int *) (data + strlen(data) + 1);
-            ntoreturn = *i;
-            i++;
-            query = (const char *) i;
+        const char * markGet() const {
+            return _nextjsobj;
         }
 
-        /* for insert and update msgs */
-        bool moreJSObjs() const {
-            return nextjsobj != 0;
+        void markSet() {
+            _mark = _nextjsobj;
         }
-        BSONObj nextJsObj() {
-            if ( nextjsobj == data ) {
-                nextjsobj += strlen(data) + 1; // skip namespace
-                massert( 13066 ,  "Message contains no documents", theEnd > nextjsobj );
-            }
-            massert( 10304,
-                     "Client Error: Remaining data too small for BSON object",
-                     theEnd - nextjsobj >= 5 );
-
-            if (serverGlobalParams.objcheck) {
-                Status status = validateBSON( nextjsobj, theEnd - nextjsobj );
-                massert( 10307,
-                         str::stream() << "Client Error: bad object in message: " << status.reason(),
-                         status.isOK() );
-            }
 
-            BSONObj js(nextjsobj);
-            verify( js.objsize() >= 5 );
-            verify( js.objsize() < ( theEnd - data ) );
+        void markReset(const char * toMark);
 
-            nextjsobj += js.objsize();
-            if ( nextjsobj >= theEnd )
-                nextjsobj = 0;
-            return js;
-        }
+    private:
+        // Check if we have enough data to read
+        template<typename T>
+        void checkRead(const char* start, size_t count = 0) const;
 
-        const Message& msg() const { return m; }
+        // Read some type without advancing our position
+        template<typename T>
+        T read() const;
 
-        const char * markGet() {
-            return nextjsobj;
-        }
+        // Read some type, and advance our position
+        template<typename T> T readAndAdvance();
 
-        void markSet() {
-            mark = nextjsobj;
-        }
+        const Message& _msg;
+        int _reserved; // flags or zero depending on packet, starts the packet
 
-        void markReset( const char * toMark = 0) {
-            if( toMark == 0 ) toMark = mark;
-            verify( toMark );
-            nextjsobj = toMark;
-        }
+        const char* _nsStart; // start of namespace string, +4 from message start
+        const char* _nextjsobj; // current position reading packet
+        const char* _theEnd; // end of packet
 
-    private:
-        const Message& m;
-        int* reserved;
-        const char *data;
-        const char *nextjsobj;
-        const char *theEnd;
+        const char* _mark;
 
-        const char * mark;
+        unsigned int _nsLen;
     };