SERVER-10338: mutablebson validation checking in update

3 files changed, 149 insertions, 50 deletions

diff --git a/jstests/update_dbref.js b/jstests/update_dbref.js
index d435b8c3416..bf31566fc28 100644
--- a/jstests/update_dbref.js
+++ b/jstests/update_dbref.js
@@ -4,6 +4,7 @@ t = db.jstests_update_dbref;
 t.drop();
 
 t.save({_id:1, a: new DBRef("a", "b")});
+assert.gleSuccess(db, "failed to save dbref");
 assert.docEq({_id:1, a: new DBRef("a", "b")}, t.findOne());
 
 t.update({}, {$set: {"a.$id": 2}});
diff --git a/jstests/updatea.js b/jstests/updatea.js
index 5b45d60f7bf..40b900d0c9d 100644
--- a/jstests/updatea.js
+++ b/jstests/updatea.js
@@ -5,18 +5,22 @@ t.drop();
 orig = { _id : 1 , a : [ { x : 1 , y : 2 } , { x : 10 , y : 11 } ] }
 
 t.save( orig )
+assert.gleSuccess(db, "orig");
 
 // SERVER-181
 t.update( {} , { $set : { "a.0.x" : 3 } } )
+assert.gleSuccess(db, "a.0.x");
 orig.a[0].x = 3;
 assert.eq( orig , t.findOne() , "A1" );
 
 t.update( {} , { $set : { "a.1.z" : 17 } } )
+assert.gleSuccess(db, "a.1.z");
 orig.a[1].z = 17;
 assert.eq( orig , t.findOne() , "A2" );
 
 // SERVER-273
 t.update( {} , { $unset : { "a.1.y" : 1 } } )
+assert.gleSuccess(db, "a.1.y");
 delete orig.a[1].y
 assert.eq( orig , t.findOne() , "A3" );
 
@@ -24,8 +28,11 @@ assert.eq( orig , t.findOne() , "A3" );
 t.drop();
 orig = { _id : 1 , comments : [ { name : "blah" , rate_up : 0 , rate_ups : [] } ] }
 t.save( orig );
+assert.gleSuccess(db, "save");
+
 
 t.update( {} , { $inc: { "comments.0.rate_up" : 1 } , $push: { "comments.0.rate_ups" : 99 } } )
+assert.gleSuccess(db, "comments.0.rate_up");
 orig.comments[0].rate_up++;
 orig.comments[0].rate_ups.push( 99 )
 assert.eq( orig , t.findOne() , "B1" )
@@ -37,13 +44,16 @@ for ( i=0; i<12; i++ )
 
 
 t.save( orig );
+assert.gleSuccess(db, "C1");
 assert.eq( orig , t.findOne() , "C1" );
 
 t.update( {} , { $inc: { "a.0" : 1 } } );
+assert.gleSuccess(db, "C2");
 orig.a[0]++;
 assert.eq( orig , t.findOne() , "C2" );
 
 t.update( {} , { $inc: { "a.10" : 1 } } );
+assert.gleSuccess(db, "a.10");
 orig.a[10]++;
 
 
@@ -51,6 +61,7 @@ orig.a[10]++;
 t.drop()
 t.insert({"a":{"c00":1}, 'c':2}) 
 t.update({"c":2}, {'$inc':{'a.c000':1}}) 
+assert.gleSuccess(db, "D1");
 
 assert.eq( { "c00" : 1 , "c000" : 1 } , t.findOne().a , "D1" )
 
diff --git a/src/mongo/db/ops/update.cpp b/src/mongo/db/ops/update.cpp
index 30a65eb855a..ed653c54679 100644
--- a/src/mongo/db/ops/update.cpp
+++ b/src/mongo/db/ops/update.cpp
@@ -78,6 +78,128 @@ namespace mongo {
             }
         }
 
+        Status storageValid(const mb::Document&, const bool);
+        Status storageValid(const mb::ConstElement&, const bool);
+        Status storageValidChildren(const mb::ConstElement&, const bool);
+
+        /**
+         * mutable::document storageValid check -- like BSONObj::_okForStorage
+         */
+        Status storageValid(const mb::Document& doc, const bool deep = true) {
+            mb::ConstElement currElem = doc.root().leftChild();
+            while (currElem.ok()) {
+                if (currElem.getFieldName() == idFieldName) {
+                    switch (currElem.getType()) {
+                        case RegEx:
+                        case Array:
+                        case Undefined:
+                            return Status(ErrorCodes::InvalidIdField,
+                                          str::stream() << "The '_id' value cannot be of type "
+                                                        << typeName(currElem.getType()));
+                        default:
+                            break;
+                    }
+                }
+                Status s = storageValid(currElem, deep);
+                if (!s.isOK())
+                    return s;
+                currElem = currElem.rightSibling();
+            }
+
+            return Status::OK();
+        }
+
+        Status storageValid(const mb::ConstElement& elem, const bool deep = true) {
+            if (!elem.ok())
+                return Status(ErrorCodes::BadValue, "Invalid elements cannot be stored.");
+
+            StringData fieldName = elem.getFieldName();
+            // Cannot start with "$", unless dbref which must start with ($ref, $id)
+            if (fieldName[0] == '$') {
+                // Check if it is a DBRef has this field {$ref, $id, [$db]}
+                mb::ConstElement curr = elem;
+                StringData currName = fieldName;
+
+                // Found a $db field
+                if (currName == "$db") {
+                    if (curr.getType() != String) {
+                        return Status(ErrorCodes::InvalidDBRef,
+                                str::stream() << "The DBRef $db field must be a String, not a "
+                                              << typeName(curr.getType()));
+                    }
+                    curr = curr.leftSibling();
+
+                    if (!curr.ok() || (curr.getFieldName() != "$id"))
+                        return Status(ErrorCodes::InvalidDBRef,
+                                      "Found $db field without a $id before it, which is invalid.");
+
+                    currName = curr.getFieldName();
+                }
+
+                // Found a $id field
+                if (currName == "$id") {
+                    Status s = storageValidChildren(curr, deep);
+                    if (!s.isOK())
+                        return s;
+
+                    curr = curr.leftSibling();
+                    if (!curr.ok() || (curr.getFieldName() != "$ref")) {
+                        return Status(ErrorCodes::InvalidDBRef,
+                                     "Found $id field without a $ref before it, which is invalid.");
+                    }
+
+                    currName = curr.getFieldName();
+                }
+
+                if (currName == "$ref") {
+                    if (curr.getType() != String) {
+                        return Status(ErrorCodes::InvalidDBRef,
+                                str::stream() << "The DBRef $ref field must be a String, not a "
+                                              << typeName(curr.getType()));
+                    }
+
+                    if (!curr.rightSibling().ok() || curr.rightSibling().getFieldName() != "$id")
+                        return Status(ErrorCodes::InvalidDBRef,
+                                str::stream() << "The DBRef $ref field must be "
+                                                 "following by a $id field");
+                }
+                else {
+                    // not an okay, $ prefixed field name.
+                    return Status(ErrorCodes::DollarPrefixedFieldName,
+                                  str::stream() << elem.getFieldName()
+                                                << " is not valid for storage.");
+                }
+            }
+
+            // Field name cannot have a "." in it.
+            if (fieldName.find(".") != string::npos) {
+                return Status(ErrorCodes::DottedFieldName,
+                              str::stream() << elem.getFieldName() << " is not valid for storage.");
+            }
+
+            // Check children if there are any.
+            Status s = storageValidChildren(elem, deep);
+            if (!s.isOK())
+                return s;
+
+            return Status::OK();
+        }
+
+        Status storageValidChildren(const mb::ConstElement& elem, const bool deep) {
+            if (!elem.hasChildren())
+                return Status::OK();
+
+            mb::ConstElement curr = elem.leftChild();
+            while (curr.ok()) {
+                Status s = storageValid(curr, deep);
+                if (!s.isOK())
+                    return s;
+                curr = curr.rightSibling();
+            }
+
+            return Status::OK();
+        }
+
         /**
          * This will verify that all updated fields are
          *   1.) Valid for storage (checking parent to make sure things like DBRefs are valid)
@@ -101,13 +223,7 @@ namespace mongo {
                    << " validate:" << opts.enforceOkForStorage;
 
             // 1.) Loop through each updated field and validate for storage
-            // and detect immutable updates
-
-            // Once we check the root once, there is no need to do it again
-            bool checkedRoot = false;
-
-            // TODO: Replace with mutablebson implementation -- this is wasteful to make a copy.
-            BSONObj doc = updated.getObject();
+            // and detect immutable field updates
 
             // The set of possibly changed immutable fields -- we will need to check their vals
             FieldRefSet changedImmutableFields;
@@ -117,7 +233,7 @@ namespace mongo {
             if (updatedFields.empty() || !opts.enforceOkForStorage) {
                 if (opts.enforceOkForStorage) {
                     // No specific fields were updated so the whole doc must be checked
-                    Status s = doc.storageValid(true);
+                    Status s = storageValid(updated, true);
                     if (!s.isOK())
                         return s;
                 }
@@ -134,54 +250,25 @@ namespace mongo {
                 if (immutableAndSingleValueFields)
                     immutableFieldRef.fillFrom(*immutableAndSingleValueFields);
 
-
                 FieldRefSet::const_iterator where = updatedFields.begin();
                 const FieldRefSet::const_iterator end = updatedFields.end();
                 for( ; where != end; ++where) {
                     const FieldRef& current = **where;
 
-                    //TODO: don't check paths more than once
-
-                    const bool isTopLevelField = (current.numParts() == 1);
-                    if (isTopLevelField) {
-                        // We check the top level since the current implementation checks BSONObj,
-                        // not BSONElements. NOTE: in the mutablebson version we can just check elem
-                        if (!checkedRoot) {
-                            // Check the root level (top level fields) only once, and don't go deep
-                            Status s = doc.storageValid(false);
-                            if (!s.isOK()) {
-                               return s;
-                            }
-                            checkedRoot = true;
-                        }
-
-                        // Check if the updated field conflicts with immutable fields
-                        immutableFieldRef.getConflicts(&current, &changedImmutableFields);
-
-                        // Traverse (deep)
-                        BSONElement elem = doc.getFieldDotted(current.dottedField());
-                        if (elem.isABSONObj()) {
-                            Status s = elem.embeddedObject().storageValid(true);
-                            if (!s.isOK()) {
-                               return s;
-                            }
-                        }
-                    }
-                    else {
-                        // Remove the right-most part, to get the parent.
-                        // "a.b.c" -> "a.b"
-                        StringData path = current.dottedField().substr(
-                                            0,
-                                            current.dottedField().size() -
-                                                current.getPart(current.numParts() - 1).size() - 1);
-
-                        BSONElement elem = doc.getFieldDotted(path);
-                        Status s = elem.embeddedObject().storageValid(true);
-                        if (!s.isOK()) {
-                           return s;
-                        }
-
+                    // Find the updated field in the updated document.
+                    mutablebson::ConstElement newElem = updated.root();
+                    size_t currentPart = 0;
+                    while (newElem.ok() && currentPart < current.numParts())
+                        newElem = newElem[current.getPart(currentPart++)];
+
+                    // newElem might be missing if $unset/$renamed-away
+                    if (newElem.ok()) {
+                        Status s = storageValid(newElem, true);
+                        if (!s.isOK())
+                            return s;
                     }
+                    // Check if the updated field conflicts with immutable fields
+                    immutableFieldRef.getConflicts(&current, &changedImmutableFields);
                 }
             }