SERVER-7379 Prohibit shard key mutation in updates

9 files changed, 462 insertions, 13 deletions

diff --git a/src/mongo/base/error_codes.err b/src/mongo/base/error_codes.err
index d6a6a371b7c..97f06ffaf7a 100644
--- a/src/mongo/base/error_codes.err
+++ b/src/mongo/base/error_codes.err
@@ -65,5 +65,6 @@ error_code("OplogOperationUnsupported", 62)
 error_code("StaleShardVersion", 63)
 error_code("WriteConcernFailed", 64)
 error_code("MultipleErrorsOccurred", 65)
+error_code("ImmutableShardKeyField", 66)
 
 error_class("NetworkError", ["HostUnreachable", "HostNotFound"])
diff --git a/src/mongo/db/field_ref.h b/src/mongo/db/field_ref.h
index 06e60c2c165..6f68f019126 100644
--- a/src/mongo/db/field_ref.h
+++ b/src/mongo/db/field_ref.h
@@ -117,6 +117,8 @@ namespace mongo {
          */
         size_t numParts() const { return _size; }
 
+        bool empty() const { return numParts() == 0; }
+
     private:
         // Dotted fields are most often not longer than four parts. We use a mixed structure
         // here that will not require any extra memory allocation when that is the case. And
diff --git a/src/mongo/db/field_ref_set.cpp b/src/mongo/db/field_ref_set.cpp
index 65e81344300..c1a3c79cff5 100644
--- a/src/mongo/db/field_ref_set.cpp
+++ b/src/mongo/db/field_ref_set.cpp
@@ -56,6 +56,28 @@ namespace mongo {
     FieldRefSet::FieldRefSet() {
     }
 
+    void FieldRefSet::getConflicts(const FieldRef* toCheck, FieldRefSet* conflicts) const {
+
+        // If the set is empty, there is no work to do.
+        if (_fieldSet.empty())
+            return;
+
+        StringData prefixStr = safeFirstPart(toCheck);
+        FieldRef prefixField;
+        prefixField.parse(prefixStr);
+
+        FieldSet::iterator it = _fieldSet.lower_bound(&prefixField);
+        // Now, iterate over all the present fields in the set that have the same prefix.
+
+        while (it != _fieldSet.end() && safeFirstPart(*it) == prefixStr) {
+            size_t common = (*it)->commonPrefixSize(*toCheck);
+            if ((*it)->numParts() == common || toCheck->numParts() == common) {
+                conflicts->_fieldSet.insert(*it);
+            }
+            ++it;
+        }
+    }
+
     bool FieldRefSet::insert(const FieldRef* toInsert, const FieldRef** conflict) {
 
         // We can determine if two fields conflict by checking their common prefix.
diff --git a/src/mongo/db/field_ref_set.h b/src/mongo/db/field_ref_set.h
index 5abd0ad9892..c80e3714e69 100644
--- a/src/mongo/db/field_ref_set.h
+++ b/src/mongo/db/field_ref_set.h
@@ -42,9 +42,32 @@ namespace mongo {
      */
     class FieldRefSet {
         MONGO_DISALLOW_COPYING(FieldRefSet);
+
+        struct FieldRefPtrLessThan {
+            bool operator()(const FieldRef* lhs, const FieldRef* rhs) const;
+        };
+
+        typedef std::set<const FieldRef*, FieldRefPtrLessThan> FieldSet;
+
     public:
+        typedef FieldSet::iterator iterator;
+        typedef FieldSet::const_iterator const_iterator;
+
         FieldRefSet();
 
+        /** Returns 'true' if the set is empty */
+        bool empty() const {
+            return _fieldSet.empty();
+        }
+
+        inline const_iterator begin() const {
+            return _fieldSet.begin();
+        }
+
+        inline const_iterator end() const {
+            return _fieldSet.end();
+        }
+
         /**
          * Returns true if the field 'toInsert' can be added in the set without
          * conflicts. Otwerwise returns false and fill in '*conflict' with the field 'toInsert'
@@ -56,17 +79,17 @@ namespace mongo {
          */
         bool insert(const FieldRef* toInsert, const FieldRef** conflict);
 
-        /** Returns 'true' if the set is empty */
-        bool empty() const {
-            return _fieldSet.empty();
+        /**
+         * Find all inserted fields which conflict with the FieldRef 'toCheck' by the semantics
+         * of 'insert', and add those fields to the 'conflicts' set.
+         */
+        void getConflicts(const FieldRef* toCheck, FieldRefSet* conflicts) const;
+
+        void clear() {
+            _fieldSet.clear();
         }
 
     private:
-        struct FieldRefPtrLessThan {
-            bool operator()(const FieldRef* lhs, const FieldRef* rhs) const;
-        };
-        typedef std::set<const FieldRef*, FieldRefPtrLessThan> FieldSet;
-
         // A set of field_ref pointers, none of which is owned here.
         FieldSet _fieldSet;
     };
diff --git a/src/mongo/db/instance.cpp b/src/mongo/db/instance.cpp
index 82542ec3a55..38abe8bec15 100644
--- a/src/mongo/db/instance.cpp
+++ b/src/mongo/db/instance.cpp
@@ -653,6 +653,17 @@ namespace mongo {
                 if ( ! broadcast && handlePossibleShardedMessage( m , 0 ) )
                     return;
 
+                // See if we have any sharding keys, and if we find that we do, inject them
+                // into the driver. If we don't, the empty BSONObj will reset any shard key
+                // state in the driver.
+                BSONObj shardKeyPattern;
+                if (shardingState.needCollectionMetadata( ns ) ) {
+                    const CollectionMetadataPtr metadata = shardingState.getCollectionMetadata( ns );
+                    if ( metadata )
+                        shardKeyPattern = metadata->getKeyPattern();
+                }
+                driver.refreshShardKeyPattern( shardKeyPattern );
+
                 Client::Context ctx( ns );
 
                 const NamespaceString requestNs(ns);
diff --git a/src/mongo/db/ops/update.cpp b/src/mongo/db/ops/update.cpp
index acbbe492e7b..5d1677e93b5 100644
--- a/src/mongo/db/ops/update.cpp
+++ b/src/mongo/db/ops/update.cpp
@@ -346,6 +346,11 @@ namespace mongo {
             BSONObj newObj;
             const char* source = NULL;
             bool inPlace = doc.getInPlaceUpdates(&damages, &source);
+
+            // If something changed in the document, verify that no shard keys were altered.
+            if ((!inPlace || !damages.empty()) && driver->modsAffectShardKeys())
+                uassertStatusOK( driver->checkShardKeysUnaltered (oldObj, doc ) );
+
             if ( inPlace && !driver->modsAffectIndices() ) {
                 // If a set of modifiers were all no-ops, we are still 'in place', but there is
                 // no work to do, in which case we want to consider the object unchanged.
@@ -441,10 +446,8 @@ namespace mongo {
         driver->setLogOp( false );
         driver->setContext( ModifierInterface::ExecInfo::INSERT_CONTEXT );
 
-        BSONObj baseObj;
-
         // Reset the document we will be writing to
-        doc.reset( baseObj, mutablebson::Document::kInPlaceDisabled );
+        doc.reset();
         if ( request.getQuery().hasElement("_id") ) {
             uassertStatusOK(doc.root().appendElement(request.getQuery().getField("_id")));
         }
@@ -465,7 +468,12 @@ namespace mongo {
             uasserted( 16836, status.reason() );
         }
 
+        // Validate that the object replacement or modifiers resulted in a document
+        // that contains all the shard keys.
+        uassertStatusOK( driver->checkShardKeysUnaltered(BSONObj(), doc) );
+
         BSONObj newObj = doc.getObject();
+
         theDataFileMgr.insertWithObjMod( nsString.ns().c_str(), newObj, false, request.isGod() );
         if ( request.shouldUpdateOpLog() ) {
             logOp( "i", nsString.ns().c_str(), newObj,
diff --git a/src/mongo/db/ops/update_driver.cpp b/src/mongo/db/ops/update_driver.cpp
index 0133730847b..64d2616a2f5 100644
--- a/src/mongo/db/ops/update_driver.cpp
+++ b/src/mongo/db/ops/update_driver.cpp
@@ -32,7 +32,6 @@
 #include "mongo/base/string_data.h"
 #include "mongo/bson/mutable/document.h"
 #include "mongo/db/field_ref.h"
-#include "mongo/db/field_ref_set.h"
 #include "mongo/db/ops/log_builder.h"
 #include "mongo/db/ops/modifier_object_replace.h"
 #include "mongo/db/ops/modifier_table.h"
@@ -205,8 +204,12 @@ namespace mongo {
         // TODO: assert that update() is called at most once in a !_multi case.
 
         FieldRefSet targetFields;
+
         _affectIndices = false;
 
+        if (_shardKeyState)
+            _shardKeyState->affectedKeySet.clear();
+
         _logDoc.reset();
         LogBuilder logBuilder(_logDoc.root());
 
@@ -223,7 +226,7 @@ namespace mongo {
             // strict update), we should respect that. If a mod doesn't care, it would state
             // it is fine with ANY update context.
             const bool validContext = (execInfo.context == ModifierInterface::ExecInfo::ANY_CONTEXT ||
-                                 execInfo.context == _context);
+                                       execInfo.context == _context);
 
             // Nothing to do if not in a valid context.
             if (!validContext) {
@@ -239,6 +242,11 @@ namespace mongo {
                     break;
                 }
 
+                if (!execInfo.noOp && _shardKeyState)
+                    _shardKeyState->keySet.getConflicts(
+                        execInfo.fieldRef[i],
+                        &_shardKeyState->affectedKeySet);
+
                 if (!targetFields.empty() || _mods.size() > 1) {
                     const FieldRef* other;
                     if (!targetFields.insert(execInfo.fieldRef[i], &other)) {
@@ -304,6 +312,152 @@ namespace mongo {
         _indexedFields = indexedFields;
     }
 
+    void UpdateDriver::refreshShardKeyPattern(const BSONObj& shardKeyPattern) {
+
+        // An empty pattern object means no shard keys.
+        if (shardKeyPattern.isEmpty()) {
+            _shardKeyState.reset();
+            return;
+        }
+
+        // If we have already parsed an identical shard key pattern, don't do it again.
+        if (_shardKeyState && (_shardKeyState->pattern.woCompare(shardKeyPattern) == 0))
+            return;
+
+        // Reset the shard key state and capture the new pattern.
+        _shardKeyState.reset(new ShardKeyState);
+        _shardKeyState->pattern = shardKeyPattern;
+
+        // Parse the shard keys into the states 'keys' and 'keySet' members.
+        BSONObjIterator patternIter = _shardKeyState->pattern.begin();
+        while (patternIter.more())  {
+            BSONElement current = patternIter.next();
+
+            _shardKeyState->keys.mutableVector().push_back(new FieldRef);
+            FieldRef* const newFieldRef = _shardKeyState->keys.mutableVector().back();
+            newFieldRef->parse(current.fieldNameStringData());
+
+            // TODO: what about bad parse?
+
+            const FieldRef* conflict;
+            if ( !_shardKeyState->keySet.insert( newFieldRef, &conflict ) ) {
+                // This seems pretty unlikely in practice.
+                uasserted(
+                    17152, str::stream()
+                    << "Shard key '"
+                    << newFieldRef->dottedField()
+                    << "' conflicts with shard key '"
+                    << conflict->dottedField()
+                    << "'" );
+            }
+        }
+    }
+
+    bool UpdateDriver::modsAffectShardKeys() const {
+        // If we have no shard key state, the mods could not have affected them.
+        if (!_shardKeyState)
+            return false;
+
+        // In replacement mode, we always assume that all shard keys need to be checked. For
+        // upsert, we are inserting a new object, so it must have all shard keys. Otherwise, it
+        // depends on whether any shard keys were added to the affectedKeySet state.
+        return (_replacementMode || !_shardKeyState->affectedKeySet.empty());
+    }
+
+    Status UpdateDriver::checkShardKeysUnaltered(const BSONObj& original,
+                                                 const mutablebson::Document& updated) const {
+
+        if (!_shardKeyState)
+            return Status::OK();
+
+        // In replacement mode, we validate the values for all shard keys. Otherwise, only the
+        // ones tagged as being potentially invalidated. For an upsert, we check all keys.
+        const FieldRefSet& affected = (_replacementMode || original.isEmpty()) ?
+            _shardKeyState->keySet :
+            _shardKeyState->affectedKeySet;
+
+        const mutablebson::ConstElement id = updated.root()["_id"];
+
+        FieldRefSet::const_iterator where = affected.begin();
+        const FieldRefSet::const_iterator end = affected.end();
+        for( ; where != end; ++where ) {
+            const FieldRef& current = **where;
+
+            // Find the affected field in the updated document.
+            mutablebson::ConstElement elt = updated.root();
+            size_t currentPart = 0;
+            while (elt.ok() && currentPart < current.numParts())
+                elt = elt[current.getPart(currentPart++)];
+
+            if (!elt.ok()) {
+                if (original.isEmpty()) {
+                    return Status(ErrorCodes::NoSuchKey,
+                                  mongoutils::str::stream()
+                                  << "After applying modifiers of replacement in an upsert"
+                                  << "', the shard key field '" << current.dottedField() <<
+                                  "' was not found in the resulting document.");
+
+                }
+                else {
+                    return Status(ErrorCodes::ImmutableShardKeyField,
+                                  mongoutils::str::stream()
+                                  << "After applying updates to the object with _id '"
+                                  << id.toString()
+                                  << "', the shard key field '" << current.dottedField() <<
+                                  "' was found to have been removed from the document");
+                }
+            }
+
+            // For upserts, we don't have an original object to compare against. The existence
+            // check above must be sufficient for now. Potentially, we could do keyBelongsTo me
+            // here.
+            //
+            // TODO: Investigate calling keyBelongsToMe here. We'd need to do this if we want
+            // mongod to be a full backstop for incorrect updates in a forward compatible way,
+            // looking at the day where we open up the floodgates for updates through mongod.
+            if (!original.isEmpty()) {
+
+                // Find the potentially affected field in the original document.
+                const BSONElement foundInOld = original.getFieldDotted(current.dottedField());
+                if (foundInOld.eoo()) {
+
+                    // NOTE: These errors should really never occur. It would mean that the base
+                    // document on disk was missing the shard key, so we have no way to validate
+                    // that it wasn't changed.
+
+                    BSONElement id;
+                    if (original.getObjectID(id)) {
+                        return Status(ErrorCodes::NoSuchKey,
+                                      mongoutils::str::stream()
+                                      << "While updating object with _id '" << id << "', the field"
+                                      << " for shard key '"
+                                      << current.dottedField() << "' was not found "
+                                      << "in the original document");
+                    }
+                    else {
+                        return Status(ErrorCodes::NoSuchKey,
+                                      mongoutils::str::stream()
+                                      << "While updating an object with no '_id' field, the field"
+                                      << " for shard key '"
+                                      << current.dottedField() << "' was not found "
+                                      << "in the original document");
+                    }
+                }
+
+                if (elt.compareWithBSONElement(foundInOld, false) != 0) {
+                    return Status(ErrorCodes::ImmutableShardKeyField,
+                                  mongoutils::str::stream()
+                                  << "After applying updates to the object with _id '" <<
+                                  id.toString()
+                                  << "', the shard key field '" << current.dottedField() <<
+                                  "' was found to have been altered");
+                }
+            }
+        }
+
+        return Status::OK();
+    }
+
     bool UpdateDriver::multi() const {
         return _multi;
     }
@@ -369,6 +523,7 @@ namespace mongo {
         }
         _indexedFields.clear();
         _replacementMode = false;
+        _shardKeyState.reset();
     }
 
 } // namespace mongo
diff --git a/src/mongo/db/ops/update_driver.h b/src/mongo/db/ops/update_driver.h
index e49fcb71e52..972647b7580 100644
--- a/src/mongo/db/ops/update_driver.h
+++ b/src/mongo/db/ops/update_driver.h
@@ -32,7 +32,9 @@
 #include <vector>
 
 #include "mongo/base/status.h"
+#include "mongo/base/owned_pointer_vector.h"
 #include "mongo/bson/mutable/document.h"
+#include "mongo/db/field_ref_set.h"
 #include "mongo/db/index_set.h"
 #include "mongo/db/jsobj.h"
 #include "mongo/db/ops/modifier_interface.h"
@@ -97,6 +99,30 @@ namespace mongo {
         bool modsAffectIndices() const;
         void refreshIndexKeys(const IndexPathSet& indexedFields);
 
+        /** Inform the update driver of which fields are shard keys so that attempts to modify
+         *  those fields can be rejected by the driver. Pass an empty object to indicate that
+         *  no shard keys are in play.
+         */
+        void refreshShardKeyPattern(const BSONObj& shardKeyPattern);
+
+        /** After calling 'update' above, this will return true if it appears that the modifier
+         *  updates may have altered any shard keys. If this returns 'true',
+         *  'verifyShardKeysUnaltered' should be called with the original unmutated object so
+         *  field comparisons can be made and illegal mutations detected.
+         */
+        bool modsAffectShardKeys() const;
+
+        /** If the mods were detected to have potentially affected shard keys during a
+         *  non-upsert udpate, call this method, providing the original unaltered document so
+         *  that the apparently altered fields can be verified to have not actually changed. A
+         *  non-OK status indicates that at least one mutation to a shard key was detected, and
+         *  the update should be rejected rather than applied. You may pass an empty original
+         *  object on an upsert, since there is not an original object against which to
+         *  compare. In that case, only the existence of shard keys in 'updated' is verified.
+         */
+        Status checkShardKeysUnaltered(const BSONObj& original,
+                                       const mutablebson::Document& updated) const;
+
         bool multi() const;
         void setMulti(bool multi);
 
@@ -153,6 +179,25 @@ namespace mongo {
         // at each call to update.
         bool _affectIndices;
 
+        // Holds the fields relevant to any optional shard key state.
+        struct ShardKeyState {
+            // The current shard key pattern
+            BSONObj pattern;
+
+            // A vector owning the FieldRefs parsed from the pattern field names.
+            OwnedPointerVector<FieldRef> keys;
+
+            // A FieldRefSet containing pointers to the FieldRefs in 'keys'.
+            FieldRefSet keySet;
+
+            // The current set of keys known to be affected by the current update. This is
+            // reset on each call to 'update'.
+            FieldRefSet affectedKeySet;
+        };
+
+        // If shard keys have been set, holds the relevant state.
+        boost::scoped_ptr<ShardKeyState> _shardKeyState;
+
         // Is this update going to be an upsert?
         ModifierInterface::ExecInfo::UpdateContext _context;
 
diff --git a/src/mongo/db/ops/update_driver_test.cpp b/src/mongo/db/ops/update_driver_test.cpp
index d0f4e8603fd..cdc9baadd44 100644
--- a/src/mongo/db/ops/update_driver_test.cpp
+++ b/src/mongo/db/ops/update_driver_test.cpp
@@ -28,14 +28,20 @@
 
 #include "mongo/db/ops/update_driver.h"
 
+#include "mongo/db/field_ref_set.h"
 #include "mongo/db/index_set.h"
 #include "mongo/db/json.h"
 #include "mongo/unittest/unittest.h"
 
 namespace {
 
+    using mongo::BSONObj;
+    using mongo::FieldRef;
+    using mongo::FieldRefSet;
     using mongo::fromjson;
     using mongo::IndexPathSet;
+    using mongo::mutablebson::Document;
+    using mongo::StringData;
     using mongo::UpdateDriver;
 
     TEST(Parse, Normal) {
@@ -109,4 +115,180 @@ namespace {
         ASSERT_FALSE(driver.isDocReplacement());
     }
 
+    // A base class for shard key immutability tests. We construct a document (see 'setUp'
+    // below for the document structure), and declare the two subfields "s.a' and 's.b' to be
+    // the shard keys, then test that various mutations that affect (or don't) the shard keys
+    // are rejected (or permitted).
+    class ShardKeyTest : public mongo::unittest::Test {
+    public:
+        ShardKeyTest()
+            : _shardKeyPattern(fromjson("{ 's.a' : 1, 's.c' : 1 }")) {
+        }
+
+        void setUp() {
+            // All elements here are arrays so that we can perform a no-op that won't be
+            // detected as such by the update code, which would foil our testing. Instead, we
+            // use $push with $slice.
+            _obj.reset(new BSONObj(
+                           fromjson("{ x : [1], s : { a : [1], b : [2], c : [ 3, 3, 3 ] } }")));
+            _doc.reset(new Document(*_obj));
+            _driver.reset(new UpdateDriver(UpdateDriver::Options()));
+        }
+
+    protected:
+        BSONObj _shardKeyPattern;
+        boost::scoped_ptr<BSONObj> _obj;
+        boost::scoped_ptr<Document> _doc;
+        boost::scoped_ptr<UpdateDriver> _driver;
+    };
+
+    TEST_F(ShardKeyTest, NoOpsDoNotAffectShardKeys) {
+        BSONObj mod(fromjson("{ $set : { 's.a.0' : 1, 's.c.0' : 3 } }"));
+        ASSERT_OK(_driver->parse(mod));
+        _driver->refreshShardKeyPattern(_shardKeyPattern);
+        ASSERT_OK(_driver->update(StringData(), _doc.get(), NULL));
+        ASSERT_FALSE(_driver->modsAffectShardKeys());
+    }
+
+    TEST_F(ShardKeyTest, MutatingShardKeyFieldRejected) {
+        BSONObj mod(fromjson("{ $push : { 's.a' : { $each : [2], $slice : -1 } } }"));
+        ASSERT_OK(_driver->parse(mod));
+        _driver->refreshShardKeyPattern(_shardKeyPattern);
+        ASSERT_OK(_driver->update(StringData(), _doc.get(), NULL));
+
+        ASSERT_TRUE(_driver->modsAffectShardKeys());
+
+        // Should be rejected, we are changing the value of a shard key.
+        ASSERT_NOT_OK(_driver->checkShardKeysUnaltered(*_obj, *_doc));
+    }
+
+    TEST_F(ShardKeyTest, MutatingShardKeyFieldRejectedObjectReplace) {
+        BSONObj mod(fromjson("{ x : [1], s : { a : [2], b : [2], c : [ 3, 3, 3 ] } }"));
+        ASSERT_OK(_driver->parse(mod));
+        _driver->refreshShardKeyPattern(_shardKeyPattern);
+        ASSERT_OK(_driver->update(StringData(), _doc.get(), NULL));
+
+        ASSERT_TRUE(_driver->modsAffectShardKeys());
+
+        // Should be rejected, we are changing the value of a shard key.
+        ASSERT_NOT_OK(_driver->checkShardKeysUnaltered(*_obj, *_doc));
+    }
+
+    TEST_F(ShardKeyTest, SettingShardKeyFieldToSameValueIsNotRejected) {
+        BSONObj mod(fromjson("{ $push : { 's.a' : { $each : [1], $slice : -1 } } }"));
+        ASSERT_OK(_driver->parse(mod));
+        _driver->refreshShardKeyPattern(_shardKeyPattern);
+        ASSERT_OK(_driver->update(StringData(), _doc.get(), NULL));
+
+        // It is a no-op, so we don't see it as affecting.
+        ASSERT_TRUE(_driver->modsAffectShardKeys());
+
+        // Should not be rejected: 's.a' has the same value as it did originally.
+        ASSERT_OK(_driver->checkShardKeysUnaltered(*_obj, *_doc));
+    }
+
+    TEST_F(ShardKeyTest, UnsettingShardKeyFieldRejected) {
+        BSONObj mod(fromjson("{ $unset : { 's.a' : 1 } }"));
+        ASSERT_OK(_driver->parse(mod));
+        _driver->refreshShardKeyPattern(_shardKeyPattern);
+        ASSERT_OK(_driver->update(StringData(), _doc.get(), NULL));
+
+        ASSERT_TRUE(_driver->modsAffectShardKeys());
+
+        // Should be rejected, we are removing one of the shard key fields
+        ASSERT_NOT_OK(_driver->checkShardKeysUnaltered(*_obj, *_doc));
+    }
+
+    TEST_F(ShardKeyTest, SettingShardKeyChildrenRejected) {
+        BSONObj mod(fromjson("{ $set : { 's.c.0' : 0 } }"));
+        ASSERT_OK(_driver->parse(mod));
+        _driver->refreshShardKeyPattern(_shardKeyPattern);
+        ASSERT_OK(_driver->update(StringData(), _doc.get(), NULL));
+
+        ASSERT_TRUE(_driver->modsAffectShardKeys());
+
+        // Should be rejected, we are setting a value subsumed under one of the shard keys.
+        ASSERT_NOT_OK(_driver->checkShardKeysUnaltered(*_obj, *_doc));
+    }
+
+    TEST_F(ShardKeyTest, UnsettingShardKeyChildrenRejected) {
+        BSONObj mod(fromjson("{ $unset : { 's.c.0' : 1 } }"));
+        ASSERT_OK(_driver->parse(mod));
+        _driver->refreshShardKeyPattern(_shardKeyPattern);
+        ASSERT_OK(_driver->update(StringData(), _doc.get(), NULL));
+
+        ASSERT_TRUE(_driver->modsAffectShardKeys());
+
+        // Should be rejected, we are removing one of the shard key fields
+        ASSERT_NOT_OK(_driver->checkShardKeysUnaltered(*_obj, *_doc));
+    }
+
+    TEST_F(ShardKeyTest, SettingShardKeyChildrenToSameValueIsNotRejected) {
+        BSONObj mod(fromjson("{ $push : { 's.c' : { $each : [3], $slice : -3 } } }"));
+        ASSERT_OK(_driver->parse(mod));
+        _driver->refreshShardKeyPattern(_shardKeyPattern);
+        ASSERT_OK(_driver->update(StringData(), _doc.get(), NULL));
+
+        ASSERT_TRUE(_driver->modsAffectShardKeys());
+
+        // Should not be rejected, we are setting a value subsumed under one of the shard keys,
+        // but the set is a logical no-op.
+        ASSERT_OK(_driver->checkShardKeysUnaltered(*_obj, *_doc));
+    }
+
+    TEST_F(ShardKeyTest, AppendingToShardKeyChildrenRejected) {
+        BSONObj mod(fromjson("{ $push : { 's.c' : 4 } }"));
+        ASSERT_OK(_driver->parse(mod));
+        _driver->refreshShardKeyPattern(_shardKeyPattern);
+        ASSERT_OK(_driver->update(StringData(), _doc.get(), NULL));
+
+        ASSERT_TRUE(_driver->modsAffectShardKeys());
+
+        // Should be rejected, we are adding a new child under one of the shard keys.
+        ASSERT_NOT_OK(_driver->checkShardKeysUnaltered(*_obj, *_doc));
+    }
+
+    TEST_F(ShardKeyTest, ModificationsToUnrelatedFieldsAreOK) {
+        BSONObj mod(fromjson("{ $set : { x : 2, 's.b' : 'x' } }"));
+        ASSERT_OK(_driver->parse(mod));
+        _driver->refreshShardKeyPattern(_shardKeyPattern);
+        ASSERT_OK(_driver->update(StringData(), _doc.get(), NULL));
+
+        // Should not claim to have affected shard keys
+        ASSERT_FALSE(_driver->modsAffectShardKeys());
+    }
+
+    TEST_F(ShardKeyTest, RemovingUnrelatedFieldsIsOK) {
+        BSONObj mod(fromjson("{ $unset : { x : 1, 's.b' : 1 } }"));
+        ASSERT_OK(_driver->parse(mod));
+        _driver->refreshShardKeyPattern(_shardKeyPattern);
+        ASSERT_OK(_driver->update(StringData(), _doc.get(), NULL));
+
+        // Should not claim to have affected shard keys
+        ASSERT_FALSE(_driver->modsAffectShardKeys());
+    }
+
+    TEST_F(ShardKeyTest, AddingUnrelatedFieldsIsOK) {
+        BSONObj mod(fromjson("{ $set : { z : 1 } }"));
+        ASSERT_OK(_driver->parse(mod));
+        _driver->refreshShardKeyPattern(_shardKeyPattern);
+        ASSERT_OK(_driver->update(StringData(), _doc.get(), NULL));
+
+        // Should not claim to have affected shard keys
+        ASSERT_FALSE(_driver->modsAffectShardKeys());
+    }
+
+    TEST_F(ShardKeyTest, OverwriteShardKeyFieldWithSameValueIsNotAnErrorObjectReplace) {
+        BSONObj mod(fromjson("{ x : [1], s : { a : [1], b : [2], c : [ 3, 3, 3 ] } }"));
+        ASSERT_OK(_driver->parse(mod));
+        _driver->refreshShardKeyPattern(_shardKeyPattern);
+        ASSERT_OK(_driver->update(StringData(), _doc.get(), NULL));
+
+        ASSERT_TRUE(_driver->modsAffectShardKeys());
+
+        // Applying the above mod should be OK, since we didn't actually change any of the
+        // shard key values.
+        ASSERT_OK(_driver->checkShardKeysUnaltered(*_obj, *_doc));
+    }
+
 } // unnamed namespace