Automatic per-socket authentication in Connection class

pymongo/collection.py

@@ -267,7 +267,8 @@ def insert(self, doc_or_docs,
             safe = True
         self.__database.connection._send_message(
             message.insert(self.__full_name, docs,
-                           check_keys, safe, kwargs), safe)
+                           check_keys, safe, kwargs), safe,
+                           collection_name=self.__full_name)
 
         ids = [doc.get("_id", None) for doc in docs]
         return return_one and ids[0] or ids
@@ -360,7 +361,8 @@ def update(self, spec, document, upsert=False, manipulate=False,
 
         return self.__database.connection._send_message(
             message.update(self.__full_name, upsert, multi,
-                           spec, document, safe, kwargs), safe)
+                           spec, document, safe, kwargs), safe,
+                           collection_name=self.__full_name)
 
     def drop(self):
         """Alias for :meth:`~pymongo.database.Database.drop_collection`.
@@ -433,7 +435,8 @@ def remove(self, spec_or_id=None, safe=False, **kwargs):
             safe = True
 
         return self.__database.connection._send_message(
-            message.delete(self.__full_name, spec_or_id, safe, kwargs), safe)
+            message.delete(self.__full_name, spec_or_id, safe, kwargs), safe,
+                           collection_name=self.__full_name)
 
     def find_one(self, spec_or_id=None, *args, **kwargs):
         """Get a single document from the database.

pymongo/connection.py

@@ -172,8 +172,11 @@ def return_socket(self):
                 self.sock[1].close()
         self.sock = None
 
+    def socket_ids(self):
+        return [id(sock) for sock in self.sockets]
 
-class Connection(object):  # TODO support auth for pooling
+
+class Connection(object):
     """Connection to MongoDB.
     """
 
@@ -304,10 +307,15 @@ def __init__(self, host=None, port=None, pool_size=None,
         if _connect:
             self.__find_master()
 
+        # cache of auth username/password credential keyed by DB name
+        self.__auth_credentials = {}
+        self.__sock_auths_by_id = {}
         if username:
             database = database or "admin"
             if not self[database].authenticate(username, password):
                 raise ConfigurationError("authentication failed")
+            # Add database auth credentials for auto-auth later
+            self.add_db_auth(database, username, password)
 
     @classmethod
     def from_uri(cls, uri="mongodb://localhost", **connection_args):
@@ -614,7 +622,32 @@ def __check_response_to_last_error(self, response):
         else:
             raise OperationFailure(error["err"])
 
-    def _send_message(self, message, with_last_error=False):
+    def _authenticate_socket_for_db(self, sock, db_name):
+        # Periodically remove cached auth flags of expired sockets
+        if len(self.__sock_auths_by_id) > self.pool_size:
+            cached_sock_ids = self.__sock_auths_by_id.keys()
+            current_sock_ids = self.__pool.socket_ids()
+            for sock_id in cached_sock_ids:
+                if not sock_id in current_sock_ids:
+                    del(self.__sock_auths_by_id[sock_id])
+        if not self.__auth_credentials:
+            return  # No credentials for any database
+        sock_id = id(sock)
+        if db_name in self.__sock_auths_by_id.get(sock_id, {}):
+            return  # Already authenticated for database
+        if not self.has_db_auth(db_name):
+            return  # No credentials for database
+        username, password = self.get_db_auth(db_name)
+        if not self[db_name].authenticate(username, password):
+            raise ConfigurationError("authentication to db %s failed for %s"
+                                     % (db_name, username))
+        if not sock_id in self.__sock_auths_by_id:
+            self.__sock_auths_by_id[sock_id] = {}
+        self.__sock_auths_by_id[sock_id][db_name] = 1
+        return True
+
+    def _send_message(self, message, with_last_error=False,
+                      collection_name=None):
         """Say something to Mongo.
 
         Raises ConnectionFailure if the message cannot be sent. Raises
@@ -630,6 +663,14 @@ def _send_message(self, message, with_last_error=False):
         """
         sock = self.__socket()
         try:
+            # Always authenticate for admin database, if possible
+            if self._authenticate_socket_for_db(sock, 'admin'):
+                pass  # No need for futher auth with admin login
+            elif collection_name and collection_name.split('.') >= 1:
+                # Authenticate for specific database
+                db_name = collection_name.split('.')[0]
+                self._authenticate_socket_for_db(sock, db_name)
+
             (request_id, data) = message
             sock.sendall(data)
             # Safe mode. We pack the message together with a lastError
@@ -886,3 +927,29 @@ def __iter__(self):
 
     def next(self):
         raise TypeError("'Connection' object is not iterable")
+
+    def add_db_auth(self, db_name, username, password):
+        if not username or not isinstance(username, basestring):
+            raise ConfigurationError('invalid username')
+        if not password or not isinstance(password, basestring):
+            raise ConfigurationError('invalid password')
+        self.__auth_credentials[db_name] = (username, password)
+
+    def has_db_auth(self, db_name):
+        return db_name in self.__auth_credentials
+
+    def get_db_auth(self, db_name):
+        if self.has_db_auth(db_name):
+            return self.__auth_credentials[db_name]
+        return None
+
+    def remove_db_auth(self, db_name):
+        if self.has_db_auth(db_name):
+            del(self.__auth_credentials[db_name])
+        # Force close any existing sockets to flush auths
+        self.disconnect()
+
+    def clear_db_auths(self):
+        self.__auth_credentials = {}  # Forget all credentials
+        # Force close any existing sockets to flush auths
+        self.disconnect()

test/test_connection.py

@@ -282,28 +282,46 @@ def test_from_uri(self):
         c.pymongo_test.system.users.remove({})
 
         c.admin.add_user("admin", "pass")
-        c.pymongo_test.add_user("user", "pass")
-
-        self.assertRaises(ConfigurationError, Connection,
-                          "mongodb://foo:bar@%s:%s" % (self.host, self.port))
-        self.assertRaises(ConfigurationError, Connection,
-                          "mongodb://admin:bar@%s:%s" % (self.host, self.port))
-        self.assertRaises(ConfigurationError, Connection,
-                          "mongodb://user:pass@%s:%s" % (self.host, self.port))
-        Connection("mongodb://admin:pass@%s:%s" % (self.host, self.port))
-
-        self.assertRaises(ConfigurationError, Connection,
-                          "mongodb://admin:pass@%s:%s/pymongo_test" %
-                          (self.host, self.port))
-        self.assertRaises(ConfigurationError, Connection,
-                          "mongodb://user:foo@%s:%s/pymongo_test" %
-                          (self.host, self.port))
-        Connection("mongodb://user:pass@%s:%s/pymongo_test" %
-                   (self.host, self.port))
-
-        self.assert_(Connection("mongodb://%s:%s" %
-                                (self.host, self.port),
-                                slave_okay=True).slave_okay)
+        try:
+            # Not yet logged in
+            try:
+                c.admin.system.users.find_one()
+                # If we get this far auth must not be enabled in server
+                raise SkipTest()
+            except OperationFailure:
+                pass
+
+            # Now we log in
+            c.admin.authenticate("admin", "pass")
+
+            c.pymongo_test.add_user("user", "pass")
+
+            self.assertRaises(ConfigurationError, Connection,
+                              "mongodb://foo:bar@%s:%s" % (self.host, self.port))
+            self.assertRaises(ConfigurationError, Connection,
+                              "mongodb://admin:bar@%s:%s" % (self.host, self.port))
+            self.assertRaises(ConfigurationError, Connection,
+                              "mongodb://user:pass@%s:%s" % (self.host, self.port))
+            Connection("mongodb://admin:pass@%s:%s" % (self.host, self.port))
+
+            self.assertRaises(ConfigurationError, Connection,
+                              "mongodb://admin:pass@%s:%s/pymongo_test" %
+                              (self.host, self.port))
+            self.assertRaises(ConfigurationError, Connection,
+                              "mongodb://user:foo@%s:%s/pymongo_test" %
+                              (self.host, self.port))
+            Connection("mongodb://user:pass@%s:%s/pymongo_test" %
+                       (self.host, self.port))
+
+            self.assert_(Connection("mongodb://%s:%s" %
+                                    (self.host, self.port),
+                                    slave_okay=True).slave_okay)
+        finally:
+            # Remove auth users from databases
+            c = Connection(self.host, self.port)
+            c.admin.authenticate("admin", "pass")
+            c.admin.system.users.remove({})
+            c.pymongo_test.system.users.remove({})
 
     def test_fork(self):
         """Test using a connection before and after a fork.
@@ -432,6 +450,93 @@ def test_tz_aware(self):
         self.assertEqual(aware.pymongo_test.test.find_one()["x"].replace(tzinfo=None),
                          naive.pymongo_test.test.find_one()["x"])
 
+    def test_auto_db_authentication(self):
+        conn = Connection(self.host, self.port)
+
+        # Setup admin user
+        conn.admin.system.users.remove({})
+        conn.admin.add_user("admin-user", "password")
+        conn.admin.authenticate("admin-user", "password")
+
+        try:  # try/finally to ensure we remove admin user
+            # Setup test database user
+            conn.pymongo_test.system.users.remove({})
+            conn.pymongo_test.add_user("test-user", "password")
+
+            conn.pymongo_test.drop_collection("test")
+
+            self.assertRaises(TypeError, conn.add_db_auth, "", "password")
+            self.assertRaises(TypeError, conn.add_db_auth, 5, "password")
+            self.assertRaises(TypeError, conn.add_db_auth, "test-user", "")
+            self.assertRaises(TypeError, conn.add_db_auth, "test-user", 5)
+
+            # Not yet logged in
+            conn = Connection(self.host, self.port)
+            try:
+                conn.admin.system.users.find_one()
+                # If we get this far auth must not be enabled in server
+                raise SkipTest()
+            except OperationFailure:
+                pass
+
+            # Not yet logged in
+            conn = Connection(self.host, self.port)
+            self.assertRaises(OperationFailure, conn.pymongo_test.test.count)
+            self.assertFalse(conn.has_db_auth('admin'))
+            self.assertEquals(None, conn.get_db_auth('admin'))
+
+            # Admin log in via URI
+            conn = Connection('admin-user:password@%s' % self.host, self.port)
+            self.assertTrue(conn.has_db_auth('admin'))
+            self.assertEquals('admin-user', conn.get_db_auth('admin')[0])
+            conn.admin.system.users.find()
+            conn.pymongo_test.test.insert({'_id':1, 'test':'data'}, safe=True)
+            self.assertEquals(1, conn.pymongo_test.test.find({'_id':1}).count())
+            conn.pymongo_test.test.remove({'_id':1})
+
+            # Clear and reset database authentication for all sockets
+            conn.clear_db_auths()
+            self.assertFalse(conn.has_db_auth('admin'))
+            self.assertRaises(OperationFailure, conn.pymongo_test.test.count)
+
+            # Admin log in via add_db_auth
+            conn = Connection(self.host, self.port)
+            conn.admin.system.users.find()
+            conn.add_db_auth('admin', 'admin-user', 'password')
+            conn.pymongo_test.test.insert({'_id':2, 'test':'data'}, safe=True)
+            self.assertEquals(1, conn.pymongo_test.test.find({'_id':2}).count())
+            conn.pymongo_test.test.remove({'_id':2})
+
+            # Remove database authentication for specific database
+            self.assertTrue(conn.has_db_auth('admin'))
+            conn.remove_db_auth('admin')
+            self.assertFalse(conn.has_db_auth('admin'))
+            self.assertRaises(OperationFailure, conn.pymongo_test.test.count)
+
+            # Incorrect admin credentials
+            conn = Connection(self.host, self.port)
+            conn.add_db_auth('admin', 'admin-user', 'wrong-password')
+            self.assertRaises(OperationFailure, conn.pymongo_test.test.count)
+
+            # Database-specific log in
+            conn = Connection(self.host, self.port)
+            conn.add_db_auth('pymongo_test', 'test-user', 'password')
+            self.assertRaises(OperationFailure, conn.admin.system.users.find_one)
+            conn.pymongo_test.test.insert({'_id':3, 'test':'data'}, safe=True)
+            self.assertEquals(1, conn.pymongo_test.test.find({'_id':3}).count())
+            conn.pymongo_test.test.remove({'_id':3})
+
+            # Incorrect database credentials
+            conn = Connection(self.host, self.port)
+            conn.add_db_auth('pymongo_test', 'wrong-user', 'password')
+            self.assertRaises(OperationFailure, conn.pymongo_test.test.find_one)
+        finally:
+            # Remove auth users from databases
+            conn = Connection(self.host, self.port)
+            conn.admin.authenticate("admin-user", "password")
+            conn.admin.system.users.remove({})
+            conn.pymongo_test.system.users.remove({})
+
 
 if __name__ == "__main__":
     unittest.main()