diff --git a/pymongo/_cmessagemodule.c b/pymongo/_cmessagemodule.c index a0bb3d81b..a8e29a914 100644 --- a/pymongo/_cmessagemodule.c +++ b/pymongo/_cmessagemodule.c @@ -59,7 +59,8 @@ static PyObject* _error(char* name) { /* add a lastError message on the end of the buffer. * returns 0 on failure */ -static int add_last_error(PyObject* self, buffer_t buffer, int request_id, PyObject* args) { +static int add_last_error(PyObject* self, buffer_t buffer, + int request_id, char* ns, int nslen, PyObject* args) { struct module_state *state = GETSTATE(self); int message_start; @@ -70,6 +71,9 @@ static int add_last_error(PyObject* self, buffer_t buffer, int request_id, PyObj PyObject* value; Py_ssize_t pos = 0; PyObject* one; + char *p = strchr(ns, '.'); + /* Length of the database portion of ns. */ + nslen = p ? (p - ns) : nslen; message_start = buffer_save_space(buffer, 4); if (message_start == -1) { @@ -80,11 +84,15 @@ static int add_last_error(PyObject* self, buffer_t buffer, int request_id, PyObj !buffer_write_bytes(buffer, "\x00\x00\x00\x00" /* responseTo */ "\xd4\x07\x00\x00" /* opcode */ - "\x00\x00\x00\x00" /* options */ - "admin.$cmd\x00" /* collection name */ + "\x00\x00\x00\x00", /* options */ + 12) || + !buffer_write_bytes(buffer, + ns, nslen) || /* database */ + !buffer_write_bytes(buffer, + ".$cmd\x00" /* collection name */ "\x00\x00\x00\x00" /* skip */ "\xFF\xFF\xFF\xFF", /* limit (-1) */ - 31)) { + 14)) { return 0; } @@ -184,14 +192,13 @@ static PyObject* _cbson_insert_message(PyObject* self, PyObject* args) { return NULL; } - PyMem_Free(collection_name); - iterator = PyObject_GetIter(docs); if (iterator == NULL) { PyObject* InvalidOperation = _error("InvalidOperation"); PyErr_SetString(InvalidOperation, "input is not iterable"); Py_DECREF(InvalidOperation); buffer_free(buffer); + PyMem_Free(collection_name); return NULL; } while ((doc = PyIter_Next(iterator)) != NULL) { @@ -200,6 +207,7 @@ static PyObject* _cbson_insert_message(PyObject* self, PyObject* args) { Py_DECREF(doc); Py_DECREF(iterator); buffer_free(buffer); + PyMem_Free(collection_name); return NULL; } Py_DECREF(doc); @@ -213,6 +221,7 @@ static PyObject* _cbson_insert_message(PyObject* self, PyObject* args) { PyErr_SetString(InvalidOperation, "cannot do an empty bulk insert"); Py_DECREF(InvalidOperation); buffer_free(buffer); + PyMem_Free(collection_name); return NULL; } @@ -220,12 +229,16 @@ static PyObject* _cbson_insert_message(PyObject* self, PyObject* args) { memcpy(buffer_get_buffer(buffer) + length_location, &message_length, 4); if (safe) { - if (!add_last_error(self, buffer, request_id, last_error_args)) { + if (!add_last_error(self, buffer, request_id, collection_name, + collection_name_length, last_error_args)) { buffer_free(buffer); + PyMem_Free(collection_name); return NULL; } } + PyMem_Free(collection_name); + /* objectify buffer */ result = Py_BuildValue("i" BYTES_FORMAT_STRING "i", request_id, buffer_get_buffer(buffer), @@ -318,18 +331,20 @@ static PyObject* _cbson_update_message(PyObject* self, PyObject* args) { cur_size = buffer_get_position(buffer) - before; max_size = (cur_size > max_size) ? cur_size : max_size; - PyMem_Free(collection_name); - message_length = buffer_get_position(buffer) - length_location; memcpy(buffer_get_buffer(buffer) + length_location, &message_length, 4); if (safe) { - if (!add_last_error(self, buffer, request_id, last_error_args)) { + if (!add_last_error(self, buffer, request_id, collection_name, + collection_name_length, last_error_args)) { buffer_free(buffer); + PyMem_Free(collection_name); return NULL; } } + PyMem_Free(collection_name); + /* objectify buffer */ result = Py_BuildValue("i" BYTES_FORMAT_STRING "i", request_id, buffer_get_buffer(buffer), diff --git a/pymongo/message.py b/pymongo/message.py index 69765bb36..2b0c9d553 100644 --- a/pymongo/message.py +++ b/pymongo/message.py @@ -45,12 +45,13 @@ MAX_INT32 = 2147483647 MIN_INT32 = -2147483648 -def __last_error(args): +def __last_error(namespace, args): """Data to send to do a lastError. """ cmd = SON([("getlasterror", 1)]) cmd.update(args) - return query(0, "admin.$cmd", 0, -1, cmd) + splitns = namespace.split('.', 1) + return query(0, splitns[0] + '.$cmd', 0, -1, cmd) def __pack_message(operation, data): @@ -83,7 +84,8 @@ def insert(collection_name, docs, check_keys, data += EMPTY.join(encoded) if safe: (_, insert_message) = __pack_message(2002, data) - (request_id, error_message, _) = __last_error(last_error_args) + (request_id, error_message, _) = __last_error(collection_name, + last_error_args) return (request_id, insert_message + error_message, max_bson_size) else: (request_id, insert_message) = __pack_message(2002, data) @@ -110,7 +112,8 @@ def update(collection_name, upsert, multi, data += encoded if safe: (_, update_message) = __pack_message(2001, data) - (request_id, error_message, _) = __last_error(last_error_args) + (request_id, error_message, _) = __last_error(collection_name, + last_error_args) return (request_id, update_message + error_message, len(encoded)) else: (request_id, update_message) = __pack_message(2001, data) @@ -163,7 +166,8 @@ def delete(collection_name, spec, safe, last_error_args, uuid_subtype): data += encoded if safe: (_, remove_message) = __pack_message(2006, data) - (request_id, error_message, _) = __last_error(last_error_args) + (request_id, error_message, _) = __last_error(collection_name, + last_error_args) return (request_id, remove_message + error_message, len(encoded)) else: (request_id, remove_message) = __pack_message(2006, data) diff --git a/test/test_connection.py b/test/test_connection.py index 6fb3ed042..955236321 100644 --- a/test/test_connection.py +++ b/test/test_connection.py @@ -238,7 +238,12 @@ class TestConnection(unittest.TestCase): c.admin.system.users.remove({}) c.pymongo_test.system.users.remove({}) - c.admin.add_user("admin", "pass") + try: + # First admin user add fails gle in MongoDB >= 2.1.2 + # See SERVER-4225 for more information. + c.admin.add_user("admin", "pass") + except OperationFailure: + pass c.admin.authenticate("admin", "pass") c.pymongo_test.add_user("user", "pass") diff --git a/test/test_database.py b/test/test_database.py index 023c4a056..4a152cb19 100644 --- a/test/test_database.py +++ b/test/test_database.py @@ -312,6 +312,28 @@ class TestDatabase(unittest.TestCase): db.logout() db.logout() + def test_authenticate_and_safe(self): + db = self.connection.auth_test + db.system.users.remove({}) + db.add_user("bernie", "password") + db.authenticate("bernie", "password") + + db.test.remove({}) + self.assertTrue(db.test.insert({"bim": "baz"}, safe=True)) + self.assertEqual(1, db.test.count()) + + self.assertEqual(1, + db.test.update({"bim": "baz"}, + {"$set": {"bim": "bar"}}, + safe=True).get('n')) + + self.assertEqual(1, + db.test.remove({}, safe=True).get('n')) + + self.assertEqual(0, db.test.count()) + self.connection.drop_database("auth_test") + + def test_authenticate_and_request(self): # Database.authenticate() needs to be in a request - check that it # always runs in a request, and that it restores the request state diff --git a/test/test_threads.py b/test/test_threads.py index 96de72141..934ef171a 100644 --- a/test/test_threads.py +++ b/test/test_threads.py @@ -387,7 +387,12 @@ class BaseTestThreadsAuth(object): raise SkipTest("Authentication is not enabled on server") self.conn = conn self.conn.admin.system.users.remove({}) - self.conn.admin.add_user('admin-user', 'password') + try: + # First admin user add fails gle in MongoDB >= 2.1.2 + # See SERVER-4225 for more information. + self.conn.admin.add_user('admin-user', 'password') + except OperationFailure: + pass self.conn.admin.authenticate("admin-user", "password") self.conn.auth_test.system.users.remove({}) self.conn.auth_test.add_user("test-user", "password")