Skip to content
Browse files

Fix safe with auth for MongoDB >= 2.1.x PYTHON-371

The getLastError command requires authentication in
MongoDB 2.2. See the associated ticket for more
information.
  • Loading branch information...
1 parent a5432ea commit a33b0fce4dd5fb54020ee4eb87fef99bdad4eb42 @behackett behackett committed
Showing with 68 additions and 17 deletions.
  1. +25 −10 pymongo/_cmessagemodule.c
  2. +9 −5 pymongo/message.py
  3. +6 −1 test/test_connection.py
  4. +22 −0 test/test_database.py
  5. +6 −1 test/test_threads.py
View
35 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),
View
14 pymongo/message.py
@@ -45,12 +45,13 @@
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)
View
7 test/test_connection.py
@@ -238,7 +238,12 @@ def test_from_uri(self):
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")
View
22 test/test_database.py
@@ -312,6 +312,28 @@ def test_authenticate_add_remove_user(self):
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
View
7 test/test_threads.py
@@ -387,7 +387,12 @@ def setUp(self):
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")

0 comments on commit a33b0fc

Please sign in to comment.
Something went wrong with that request. Please try again.