Path for pymongo.Connection to perform automatic per-socket authentication for databases after first authentication

pymongo/connection.py

@@ -135,12 +135,14 @@ class _Pool(threading.local):
     """
 
     # Non thread-locals
-    __slots__ = ["sockets", "socket_factory", "pool_size"]
+    __slots__ = ["sockets", "socket_factory", "pool_size",
+                 "socket_authenticator"]
     sock = None
 
-    def __init__(self, socket_factory):
+    def __init__(self, socket_factory, socket_authenticator):
         self.pool_size = 10
         self.socket_factory = socket_factory
+        self.socket_authenticator = socket_authenticator
         if not hasattr(self, "sockets"):
             self.sockets = []
 
@@ -158,6 +160,7 @@ def socket(self):
             self.sock = (pid, self.sockets.pop())
         except IndexError:
             self.sock = (pid, self.socket_factory())
+            self.socket_authenticator()
 
         return self.sock[1]
 
@@ -173,7 +176,7 @@ def return_socket(self):
         self.sock = None
 
 
-class Connection(object):  # TODO support auth for pooling
+class Connection(object):
     """Connection to MongoDB.
     """
 
@@ -291,7 +294,7 @@ def __init__(self, host=None, port=None, pool_size=None,
 
         self.__cursor_manager = CursorManager(self)
 
-        self.__pool = _Pool(self.__connect)
+        self.__pool = _Pool(self.__connect, self.__authenticate_socket)
         self.__last_checkout = time.time()
 
         self.__network_timeout = network_timeout
@@ -300,6 +303,7 @@ def __init__(self, host=None, port=None, pool_size=None,
 
         # cache of existing indexes used by ensure_index ops
         self.__index_cache = {}
+        self.__auth_credentials = {}
 
         if _connect:
             self.__find_master()
@@ -391,6 +395,25 @@ def _purge_index(self, database_name,
         if index_name in self.__index_cache[database_name][collection_name]:
             del self.__index_cache[database_name][collection_name][index_name]
 
+    def _cache_database_credentials(self, db_name, username, password):
+        """Add credentials to the database authentication cache
+        for automatic login when a socket is created.
+
+        If credentials are already cached for the database they
+        will be replaced.
+        """
+        self.__auth_credentials[db_name] = (username, password)
+
+    def _purge_database_credentials(self, db_name):
+        """Purge credentials from the database authentication cache.
+
+        If `db_name` is None purge credentials for all databases.
+        """
+        if db_name is None:
+            self.__auth_credentials.clear()
+        elif db_name in self.__auth_credentials:
+            del(self.__auth_credentials[db_name])
+
     @property
     def host(self):
         """Current connected host.
@@ -528,6 +551,24 @@ def __connect(self):
             self.disconnect()
             raise AutoReconnect("could not connect to %r" % list(self.__nodes))
 
+    def __authenticate_socket(self):
+        """Authenticate using cached database credentials. If credentials for
+        the 'admin' database are available only this database is authenticated,
+        since this gives global access.
+
+        This method should be called by the socket pool when it
+        creates a new socket.
+        """
+        # Authenticate new socket with cached credentials
+        if 'admin' in self.__auth_credentials:
+            # Log in as 'admin' by preference, since it's basically root
+            username, password = self.__auth_credentials['admin']
+            self['admin'].authenticate(username, password)
+        else:
+            # Authenticate against all non-admin databases
+            for db_name, (u, p) in self.__auth_credentials.items():
+                self[db_name].authenticate(u, p)
+
     def __socket(self):
         """Get a socket from the pool.
 
@@ -561,7 +602,7 @@ def disconnect(self):
         .. seealso:: :meth:`end_request`
         .. versionadded:: 1.3
         """
-        self.__pool = _Pool(self.__connect)
+        self.__pool = _Pool(self.__connect, self.__authenticate_socket)
         self.__host = None
         self.__port = None
 

pymongo/database.py

@@ -464,28 +464,28 @@ def authenticate(self, name, password):
         Once authenticated, the user has full read and write access to
         this database. Raises :class:`TypeError` if either `name` or
         `password` is not an instance of ``(str,
-        unicode)``. Authentication lasts for the life of the database
-        connection, or until :meth:`logout` is called.
+        unicode)``. Authentication lasts for the life of the underlying
+        :class:`Connection`, or until :meth:`logout` is called.
 
         The "admin" database is special. Authenticating on "admin"
         gives access to *all* databases. Effectively, "admin" access
         means root access to the database.
 
-        .. note:: Currently, authentication is per
-           :class:`~socket.socket`. This means that there are a couple
-           of situations in which re-authentication is necessary:
-
-           - On failover (when an
-             :class:`~pymongo.errors.AutoReconnect` exception is
-             raised).
-
-           - After a call to
-             :meth:`~pymongo.connection.Connection.disconnect` or
-             :meth:`~pymongo.connection.Connection.end_request`.
+        .. note:: This method authenticates the current connection, and
+           will also cause all new :class:`~socket.socket` connections
+           in the underlying :class:`~pymongo.connection.Connection` to
+           be authenticated automatically.
 
            - When sharing a :class:`~pymongo.connection.Connection`
-             between multiple threads, each thread will need to
-             authenticate separately.
+             between multiple threads, all threads will share the
+             authentication. If you need different authentication profiles
+             for different purposes (e.g. admin users) you must use
+             distinct :class:`~pymongo.connection.Connection`s.
+
+           - To get authentication to apply immediately to all
+             connections including existing ones, you may need to
+             reset the connections sockets using
+             :meth:`~pymongo.connection.Connection.disconnect`.
 
         .. warning:: Currently, calls to
            :meth:`~pymongo.connection.Connection.end_request` will
@@ -511,16 +511,23 @@ def authenticate(self, name, password):
         try:
             self.command("authenticate", user=unicode(name),
                          nonce=nonce, key=key)
+            self.connection._cache_database_credentials(
+                self.name, unicode(name), unicode(password))
             return True
         except OperationFailure:
             return False
 
     def logout(self):
-        """Deauthorize use of this database for this connection.
+        """Deauthorize use of this database for this connection and
+        future connections.
 
-        Note that other databases may still be authorized.
+        Note that other databases may still be authenticated, and that other
+        existing :class:`~socket.socket` connections may remain
+        authenticated unless you reset all sockets with
+        :meth:`~pymongo.connection.Connection.disconnect`.
         """
         self.command("logout")
+        self.connection._purge_database_credentials(self.name)
 
     def dereference(self, dbref):
         """Dereference a DBRef, getting the SON object it points to.

test/test_connection.py

@@ -270,7 +270,7 @@ def test_parse_uri(self):
         self.assertEqual(([("localhost", 27017)], None, None, None),
                          _parse_uri("localhost/", 27017))
 
-    def test_from_uri(self):
+    def test_auth_from_uri(self):
         c = Connection(self.host, self.port)
 
         self.assertRaises(InvalidURI, Connection, "mongodb://localhost/baz")
@@ -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,83 @@ 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_auth_from_database(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")
+
+            # 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)
+
+            # Admin log in via URI
+            conn = Connection('admin-user:password@%s' % self.host, self.port)
+            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})
+
+            # Logout admin
+            conn.admin.logout()
+            self.assertRaises(OperationFailure, conn.pymongo_test.test.count)
+
+            # Admin log in via Database.authenticate
+            conn = Connection(self.host, self.port)
+            conn.admin.system.users.find()
+            conn.admin.authenticate('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
+            conn.admin.logout()
+            self.assertRaises(OperationFailure, conn.pymongo_test.test.count)
+
+            # Incorrect admin credentials
+            conn = Connection(self.host, self.port)
+            self.assertFalse(
+                conn.admin.authenticate('admin-user', 'wrong-password'))
+            self.assertRaises(OperationFailure, conn.pymongo_test.test.count)
+
+            # Database-specific log in
+            conn = Connection(self.host, self.port)
+            conn.pymongo_test.authenticate('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.pymongo_test.authenticate('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()

test/test_pooling.py

@@ -141,13 +141,13 @@ def test_disconnect(self):
         run_cases(self, [SaveAndFind, Disconnect, Unique])
 
     def test_independent_pools(self):
-        p = _Pool(None)
+        p = _Pool(None, None)
         self.assertEqual([], p.sockets)
         self.c.end_request()
         self.assertEqual([], p.sockets)
 
         # Sensical values aren't really important here
-        p1 = _Pool(5)
+        p1 = _Pool(5, 32)
         self.assertEqual(None, p.socket_factory)
         self.assertEqual(5, p1.socket_factory)