Automatic per-socket authentication now occurs in Connection after initial Database.authenticate call

Author: jmurty

pymongo/connection.py

@@ -135,14 +135,18 @@ class _Pool(threading.local):
  """
 
  # Non thread-locals
- __slots__ = ["sockets", "socket_factory", "pool_size"]
+ __slots__ = ["sockets", "socket_factory", "pool_size",
+ "connection", "auth_credentials"]
  sock = None
 
- def __init__(self, socket_factory):
+ def __init__(self, socket_factory, connection):
  self.pool_size = 10
  self.socket_factory = socket_factory
+ self.connection = connection
  if not hasattr(self, "sockets"):
  self.sockets = []
+ if not hasattr(self, "auth_credentials"):
+ self.auth_credentials = {}
 
  def socket(self):
  # we store the pid here to avoid issues with fork /
@@ -159,6 +163,15 @@ def socket(self):
  except IndexError:
  self.sock = (pid, self.socket_factory())
 
+ # Authenticate new socket for known DBs, 'admin' by preference
+ if 'admin' in self.auth_credentials:
+ username, password = self.auth_credentials['admin']
+ self.connection['admin'].authenticate(username, password)
+ else:
+ # Authenticate against all known databases
+ for db_name, (u, p) in self.auth_credentials.items():
+ self.connection[db_name].authenticate(u, p)
+
  return self.sock[1]
 
  def return_socket(self):
@@ -172,8 +185,12 @@ def return_socket(self):
  self.sock[1].close()
  self.sock = None
 
- def socket_ids(self):
- return [id(sock) for sock in self.sockets]
+ def add_db_auth(self, db_name, username, password):
+ self.auth_credentials[db_name] = (username, password)
+
+ def remove_db_auth(self, db_name):
+ if db_name in self.auth_credentials:
+ del(self.auth_credentials[db_name])
 
 
 class Connection(object):
@@ -294,7 +311,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)
  self.__last_checkout = time.time()
 
  self.__network_timeout = network_timeout
@@ -307,15 +324,10 @@ 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):
@@ -569,7 +581,7 @@ def disconnect(self):
  .. seealso:: :meth:`end_request`
  .. versionadded:: 1.3
  """
- self.__pool = _Pool(self.__connect)
+ self.__pool = _Pool(self.__connect, self)
  self.__host = None
  self.__port = None
 
@@ -622,30 +634,6 @@ def __check_response_to_last_error(self, response):
  else:
  raise OperationFailure(error["err"])
 
- 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.
@@ -663,14 +651,6 @@ 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
@@ -928,28 +908,8 @@ 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 _add_db_auth(self, db_name, username, password):
+ self.__pool.add_db_auth(db_name, username, password)
 
- 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()
+ def _remove_db_auth(self, db_name):
+ self.__pool.remove_db_auth(db_name)

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._add_db_auth(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._remove_db_auth(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")
@@ -450,7 +450,7 @@ 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):
+ def test_auth_from_database(self):
  conn = Connection(self.host, self.port)
 
  # Setup admin user
@@ -465,11 +465,6 @@ def test_auto_db_authentication(self):
 
  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:
@@ -482,53 +477,48 @@ def test_auto_db_authentication(self):
  # 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'))
+ # Logout admin
+ conn.admin.logout()
  self.assertRaises(OperationFailure, conn.pymongo_test.test.count)
 
- # Admin log in via add_db_auth
+ # Admin log in via Database.authenticate
  conn = Connection(self.host, self.port)
  conn.admin.system.users.find()
- conn.add_db_auth('admin', 'admin-user', 'password')
+ 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
- self.assertTrue(conn.has_db_auth('admin'))
- conn.remove_db_auth('admin')
- self.assertFalse(conn.has_db_auth('admin'))
+ conn.admin.logout()
  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.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.add_db_auth('pymongo_test', 'test-user', 'password')
- self.assertRaises(OperationFailure, conn.admin.system.users.find_one)
+ 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.add_db_auth('pymongo_test', 'wrong-user', 'password')
+ conn.pymongo_test.authenticate('wrong-user', 'password')
  self.assertRaises(OperationFailure, conn.pymongo_test.test.find_one)
  finally:
  # Remove auth users from databases

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)
 

test/test_threads.py

@@ -229,7 +229,7 @@ def test_auto_auth_login(self):
 
  # Admin auth
  conn = get_connection()
- conn.add_db_auth("admin", "admin-user", "password")
+ conn.admin.authenticate("admin-user", "password")
 
  threads = []
  for _ in range(10):
@@ -242,7 +242,7 @@ def test_auto_auth_login(self):
 
  # Database-specific auth
  conn = get_connection()
- conn.add_db_auth("auth_test", "test-user", "password")
+ conn.auth_test.authenticate("test-user", "password")
 
  threads = []
  for _ in range(10):