PYTHON-732 Handle network errors when adding existing credentials to sockets.

Author: ajdavis

pymongo/mongo_client.py

@@ -673,13 +673,16 @@ def max_write_batch_size(self):
  'max_write_batch_size', common.MAX_WRITE_BATCH_SIZE)
 
  def __simple_command(self, sock_info, dbname, spec):
- """Send a command to the server.
+ """Send a command to the server. May raise AutoReconnect.
  """
  rqst_id, msg, _ = message.query(0, dbname + '.$cmd', 0, -1, spec)
  start = time.time()
  try:
  sock_info.sock.sendall(msg)
  response = self.__receive_message_on_socket(1, rqst_id, sock_info)
+ except socket.error, e:
+ sock_info.close()
+ raise AutoReconnect(e)
  except:
  sock_info.close()
  raise
@@ -916,7 +919,7 @@ def __socket(self, member):
  "%s %s" % (host_details, str(why)))
  try:
  self.__check_auth(sock_info)
- except OperationFailure:
+ except:
  connection_pool.maybe_return_socket(sock_info)
  raise
  return sock_info

test/test_client.py

@@ -34,7 +34,7 @@
 from pymongo.mongo_client import MongoClient
 from pymongo.database import Database
 from pymongo.pool import SocketInfo
-from pymongo import thread_util, common
+from pymongo import auth, thread_util, common
 from pymongo.errors import (AutoReconnect,
  ConfigurationError,
  ConnectionFailure,
@@ -54,7 +54,8 @@
  _TestLazyConnectMixin,
  lazy_client_trial,
  NTHREADS,
- get_pool)
+ get_pool,
+ one)
 
 
 def get_client(*args, **kwargs):
@@ -999,6 +1000,42 @@ def test_lazy_connect_w0(self):
  client = get_client(_connect=False)
  client.pymongo_test.test.remove(w=0)
 
+ def test_auth_network_error(self):
+ # Make sure there's no semaphore leak if we get a network error
+ # when authenticating a new socket with cached credentials.
+ auth_client = get_client()
+ if not server_started_with_auth(auth_client):
+ raise SkipTest('Authentication is not enabled on server')
+
+ auth_client.admin.add_user('admin', 'password')
+ auth_client.admin.authenticate('admin', 'password')
+ try:
+ # Get a client with one socket so we detect if it's leaked.
+ c = get_client(max_pool_size=1, waitQueueTimeoutMS=1)
+
+ # Simulate an authenticate() call on a different socket.
+ credentials = auth._build_credentials_tuple(
+ 'MONGODB-CR', 'admin',
+ unicode('admin'), unicode('password'),
+ {})
+
+ c._cache_credentials('test', credentials, connect=False)
+
+ # Cause a network error on the actual socket.
+ pool = get_pool(c)
+ socket_info = one(pool.sockets)
+ socket_info.sock.close()
+
+ # In __check_auth, the client authenticates its socket with the
+ # new credential, but gets a socket.error. Should be reraised as
+ # AutoReconnect.
+ self.assertRaises(AutoReconnect, c.test.collection.find_one)
+
+ # No semaphore leak, the pool is allowed to make a new socket.
+ c.test.collection.find_one()
+ finally:
+ remove_all_users(auth_client.admin)
+
 
 class TestClientLazyConnect(unittest.TestCase, _TestLazyConnectMixin):
  def _get_client(self, **kwargs):

test/test_replica_set_client.py

@@ -26,6 +26,7 @@
 import threading
 import traceback
 import unittest
+from pymongo import auth
 
 sys.path[0:0] = [""]
 
@@ -1126,6 +1127,42 @@ def test_alive(self):
 
  self.assertFalse(client.alive())
 
+ def test_auth_network_error(self):
+ # Make sure there's no semaphore leak if we get a network error
+ # when authenticating a new socket with cached credentials.
+ auth_client = self._get_client()
+ if not server_started_with_auth(auth_client):
+ raise SkipTest('Authentication is not enabled on server')
+
+ auth_client.admin.add_user('admin', 'password')
+ auth_client.admin.authenticate('admin', 'password')
+ try:
+ # Get a client with one socket so we detect if it's leaked.
+ c = self._get_client(max_pool_size=1, waitQueueTimeoutMS=1)
+
+ # Simulate an authenticate() call on a different socket.
+ credentials = auth._build_credentials_tuple(
+ 'MONGODB-CR', 'admin',
+ unicode('admin'), unicode('password'),
+ {})
+
+ c._cache_credentials('test', credentials, connect=False)
+
+ # Cause a network error on the actual socket.
+ pool = get_pool(c)
+ socket_info = one(pool.sockets)
+ socket_info.sock.close()
+
+ # In __check_auth, the client authenticates its socket with the
+ # new credential, but gets a socket.error. Should be reraised as
+ # AutoReconnect.
+ self.assertRaises(AutoReconnect, c.test.collection.find_one)
+
+ # No semaphore leak, the pool is allowed to make a new socket.
+ c.test.collection.find_one()
+ finally:
+ remove_all_users(auth_client.admin)
+
 
 class TestReplicaSetWireVersion(unittest.TestCase):
  def test_wire_version(self):