From 97eca5fe66c68f290e9c29535fce1e058988660d Mon Sep 17 00:00:00 2001
From: Eliot Horowitz <eliot@10gen.com>
Date: Wed, 9 Jan 2013 12:12:05 -0500
Subject: [PATCH] SERVER-7511: fix update with positional mod or index offset

---
 jstests/update_arraymatch8.js        | 113 +++++++++++++++++++++++++++
 src/mongo/db/ops/update_internal.cpp |  50 ++++++++++++
 src/mongo/db/ops/update_internal.h   |  59 +++++---------
 src/mongo/dbtests/updatetests.cpp    |  32 ++++++++
 4 files changed, 213 insertions(+), 41 deletions(-)
 create mode 100644 jstests/update_arraymatch8.js

diff --git a/jstests/update_arraymatch8.js b/jstests/update_arraymatch8.js
new file mode 100644
index 0000000000000..c9fe4bbcb1bcb
--- /dev/null
+++ b/jstests/update_arraymatch8.js
@@ -0,0 +1,113 @@
+// Checking for positional array updates with either .$ or .0 at the end
+// SERVER-7511
+
+// array.$.name
+t = db.jstests_update_arraymatch8;
+t.drop();
+t.ensureIndex( {'array.name': 1} );
+t.insert( {'array': [{'name': 'old'}]} );
+assert( t.findOne({'array.name': 'old'}) );
+t.update( {'array.name': 'old'}, {$set: {'array.$.name': 'new'}} );
+assert( t.findOne({'array.name': 'new'}) );
+assert( !t.findOne({'array.name': 'old'}) );
+
+// array.$   (failed in 2.2.2)
+t = db.jstests_update_arraymatch8;
+t.drop();
+t.ensureIndex( {'array.name': 1} );
+t.insert( {'array': [{'name': 'old'}]} );
+assert( t.findOne({'array.name': 'old'}) );
+t.update( {'array.name': 'old'}, {$set: {'array.$': {'name':'new'}}} );
+assert( t.findOne({'array.name': 'new'}) );
+assert( !t.findOne({'array.name': 'old'}) );
+
+// array.0.name
+t = db.jstests_update_arraymatch8;
+t.drop();
+t.ensureIndex( {'array.name': 1} );
+t.insert( {'array': [{'name': 'old'}]} );
+assert( t.findOne({'array.name': 'old'}) );
+t.update( {'array.name': 'old'}, {$set: {'array.0.name': 'new'}} );
+assert( t.findOne({'array.name': 'new'}) );
+assert( !t.findOne({'array.name': 'old'}) );
+
+// array.0   (failed in 2.2.2)
+t = db.jstests_update_arraymatch8;
+t.drop();
+t.ensureIndex( {'array.name': 1} );
+t.insert( {'array': [{'name': 'old'}]} );
+assert( t.findOne({'array.name': 'old'}) );
+t.update( {'array.name': 'old'}, {$set: {'array.0': {'name':'new'}}} );
+assert( t.findOne({'array.name': 'new'}) );
+assert( !t.findOne({'array.name': 'old'}) );
+
+// // array.12.name
+t = db.jstests_update_arraymatch8;
+t.drop();
+arr = new Array();
+for (var i=0; i<20; i++) {
+    arr.push({'name': 'old'});
+}
+t.ensureIndex( {'array.name': 1} );
+t.insert( {_id:0, 'array': arr} );
+assert( t.findOne({'array.name': 'old'}) );
+t.update( {_id:0}, {$set: {'array.12.name': 'new'}} );
+// note: both documents now have to be in the array
+assert( t.findOne({'array.name': 'new'}) );
+assert( t.findOne({'array.name': 'old'}) );
+
+// array.12   (failed in 2.2.2)
+t = db.jstests_update_arraymatch8;
+t.drop();
+arr = new Array();
+for (var i=0; i<20; i++) {
+    arr.push({'name': 'old'});
+}
+t.ensureIndex( {'array.name': 1} );
+t.insert( {_id:0, 'array': arr} );
+assert( t.findOne({'array.name': 'old'}) );
+t.update( {_id:0}, {$set: {'array.12': {'name':'new'}}} );
+// note: both documents now have to be in the array
+assert( t.findOne({'array.name': 'new'}) );
+assert( t.findOne({'array.name': 'old'}) );
+
+// array.$.123a.name
+t = db.jstests_update_arraymatch8;
+t.drop();
+t.ensureIndex( {'array.123a.name': 1} );
+t.insert( {'array': [{'123a':{'name': 'old'}}]} );
+assert( t.findOne({'array.123a.name': 'old'}) );
+t.update( {'array.123a.name': 'old'}, {$set: {'array.$.123a.name': 'new'}} );
+assert( t.findOne({'array.123a.name': 'new'}) );
+assert( !t.findOne({'array.123a.name': 'old'}) );
+
+// array.$.123a
+t = db.jstests_update_arraymatch8;
+t.drop();
+t.ensureIndex( {'array.name': 1} );
+t.insert( {'array': [{'123a':{'name': 'old'}}]} );
+assert( t.findOne({'array.123a.name': 'old'}) );
+t.update( {'array.123a.name': 'old'}, {$set: {'array.$.123a': {'name': 'new'}}} );
+assert( t.findOne({'array.123a.name': 'new'}) );
+assert( !t.findOne({'array.123a.name': 'old'}) );
+
+// array.0.123a.name
+t = db.jstests_update_arraymatch8;
+t.drop();
+t.ensureIndex( {'array.123a.name': 1} );
+t.insert( {'array': [{'123a':{'name': 'old'}}]} );
+assert( t.findOne({'array.123a.name': 'old'}) );
+t.update( {'array.123a.name': 'old'}, {$set: {'array.0.123a.name': 'new'}} );
+assert( t.findOne({'array.123a.name': 'new'}) );
+assert( !t.findOne({'array.123a.name': 'old'}) );
+
+// array.0.123a
+t = db.jstests_update_arraymatch8;
+t.drop();
+t.ensureIndex( {'array.name': 1} );
+t.insert( {'array': [{'123a':{'name': 'old'}}]} );
+assert( t.findOne({'array.123a.name': 'old'}) );
+t.update( {'array.123a.name': 'old'}, {$set: {'array.0.123a': {'name': 'new'}}} );
+assert( t.findOne({'array.123a.name': 'new'}) );
+assert( !t.findOne({'array.123a.name': 'old'}) );
+
diff --git a/src/mongo/db/ops/update_internal.cpp b/src/mongo/db/ops/update_internal.cpp
index 6f7b32e74a93c..19dab221ff677 100644
--- a/src/mongo/db/ops/update_internal.cpp
+++ b/src/mongo/db/ops/update_internal.cpp
@@ -1368,4 +1368,54 @@ namespace mongo {
             updateIsIndexed( i->second, idxKeys , backgroundKeys );
     }
 
+    bool getCanonicalIndexField( const StringData& fullName, string* out ) {
+        // check if fieldName contains ".$" or ".###" substrings (#=digit) and skip them
+        if ( fullName.find( '.' ) == string::npos )
+            return false;
+
+        bool modified = false;
+
+        StringBuilder buf;
+        for ( size_t i=0; i<fullName.size(); i++ ) {
+
+            char c = fullName[i];
+
+            if ( c != '.' ) {
+                buf << c;
+                continue;
+            }
+
+            // check for ".$", skip if present
+            if ( fullName[i+1] == '$' ) {
+                i++;
+                modified = true;
+                continue;
+            }
+
+            // check for ".###" for any number of digits (no letters)
+            if ( isdigit( fullName[i+1] ) ) {
+                size_t j = i;
+                // skip digits
+                while ( j+1 < fullName.size() && isdigit( fullName[j+1] ) )
+                    j++;
+
+                if ( j+1 == fullName.size() || fullName[j+1] == '.' ) {
+                    // only digits found, skip forward
+                    i = j;
+                    modified = true;
+                    continue;
+                }
+            }
+
+            buf << c;
+        }
+
+        if ( !modified )
+            return false;
+
+        *out = buf.str();
+        return true;
+    }
+
+
 } // namespace mongo
diff --git a/src/mongo/db/ops/update_internal.h b/src/mongo/db/ops/update_internal.h
index 9d1ab3f77e2ed..d49ddd04fe63b 100644
--- a/src/mongo/db/ops/update_internal.h
+++ b/src/mongo/db/ops/update_internal.h
@@ -30,6 +30,12 @@ namespace mongo {
     class ModState;
     class ModSetState;
 
+    /**
+     * a.$ -> a
+     * @return true if out is set and we made a change
+     */
+    bool getCanonicalIndexField( const StringData& fullName, string* out );
+
     /* Used for modifiers such as $inc, $set, $push, ...
      * stores the info about a single operation
      * once created should never be modified
@@ -143,6 +149,7 @@ namespace mongo {
             // check if there is an index key equal to mod
             if ( idxKeys.count(fullName) )
                 return true;
+
             // check if there is an index key that is a child of mod
             set< string >::const_iterator j = idxKeys.upper_bound( fullName );
             if ( j != idxKeys.end() && j->find( fullName ) == 0 && (*j)[fullName.size()] == '.' )
@@ -151,51 +158,21 @@ namespace mongo {
             return false;
         }
 
+        /**
+         * checks if mod is in the index by inspecting fieldName, and removing
+         * .$ or .### substrings (#=digit) with any number of digits.
+         *
+         * @return true iff the mod is indexed
+         */
         bool isIndexed( const set<string>& idxKeys ) const {
-            string fullName = fieldName;
 
-            if ( isIndexed( fullName , idxKeys ) )
+            // first, check if full name is in idxKeys
+            if ( isIndexed( fieldName , idxKeys ) )
                 return true;
 
-            if ( strstr( fieldName , "." ) ) {
-                // check for a.0.1
-                StringBuilder buf;
-                for ( size_t i=0; i<fullName.size(); i++ ) {
-                    char c = fullName[i];
-
-                    if ( c == '$' &&
-                            i > 0 && fullName[i-1] == '.' &&
-                            i+1<fullName.size() &&
-                            fullName[i+1] == '.' ) {
-                        i++;
-                        continue;
-                    }
-
-                    buf << c;
-
-                    if ( c != '.' )
-                        continue;
-
-                    if ( ! isdigit( fullName[i+1] ) )
-                        continue;
-
-                    bool possible = true;
-                    size_t j=i+2;
-                    for ( ; j<fullName.size(); j++ ) {
-                        char d = fullName[j];
-                        if ( d == '.' )
-                            break;
-                        if ( isdigit( d ) )
-                            continue;
-                        possible = false;
-                        break;
-                    }
-
-                    if ( possible )
-                        i = j;
-                }
-                string x = buf.str();
-                if ( isIndexed( x , idxKeys ) )
+            string x;
+            if ( getCanonicalIndexField( fieldName, &x ) ) {
+                if ( isIndexed( x, idxKeys ) )
                     return true;
             }
 
diff --git a/src/mongo/dbtests/updatetests.cpp b/src/mongo/dbtests/updatetests.cpp
index f3a951899be71..4cf5b5e0832a9 100644
--- a/src/mongo/dbtests/updatetests.cpp
+++ b/src/mongo/dbtests/updatetests.cpp
@@ -2509,6 +2509,36 @@ namespace UpdateTests {
 
     };
 
+    class IndexFieldNameTest {
+    public:
+        void run() {
+            string x;
+
+            ASSERT_FALSE( getCanonicalIndexField( "a", &x ) );
+            ASSERT_FALSE( getCanonicalIndexField( "aaa", &x ) );
+            ASSERT_FALSE( getCanonicalIndexField( "a.b", &x ) );
+
+            ASSERT_TRUE( getCanonicalIndexField( "a.$", &x ) );
+            ASSERT_EQUALS( x, "a" );
+            ASSERT_TRUE( getCanonicalIndexField( "a.0", &x ) );
+            ASSERT_EQUALS( x, "a" );
+            ASSERT_TRUE( getCanonicalIndexField( "a.123", &x ) );
+            ASSERT_EQUALS( x, "a" );
+
+            ASSERT_TRUE( getCanonicalIndexField( "a.$.b", &x ) );
+            ASSERT_EQUALS( x, "a.b" );
+            ASSERT_TRUE( getCanonicalIndexField( "a.0.b", &x ) );
+            ASSERT_EQUALS( x, "a.b" );
+            ASSERT_TRUE( getCanonicalIndexField( "a.123.b", &x ) );
+            ASSERT_EQUALS( x, "a.b" );
+
+            ASSERT_FALSE( getCanonicalIndexField( "a.123a", &x ) );
+            ASSERT_FALSE( getCanonicalIndexField( "a.a123", &x ) );
+            ASSERT_FALSE( getCanonicalIndexField( "a.123a.b", &x ) );
+            ASSERT_FALSE( getCanonicalIndexField( "a.a123.b", &x ) );
+        }
+    };
+
     class All : public Suite {
     public:
         All() : Suite( "update" ) {
@@ -2676,6 +2706,8 @@ namespace UpdateTests {
             add< basic::bit1 >();
             add< basic::unset >();
             add< basic::setswitchint >();
+
+            add< IndexFieldNameTest >();
         }
     } myall;