PYTHON-764 SCRAM-SHA-1 automatic upgrade / downgrade.

ajdavis · ajdavis · commit ee11436675ea · 2014-10-23T16:30:55.000-04:00
diff --git a/pymongo/auth.py b/pymongo/auth.py
@@ -36,7 +36,7 @@
 
 
 MECHANISMS = frozenset(
-    ['GSSAPI', 'MONGODB-CR', 'MONGODB-X509', 'PLAIN', 'SCRAM-SHA-1'])
+    ['GSSAPI', 'MONGODB-CR', 'MONGODB-X509', 'PLAIN', 'SCRAM-SHA-1', 'DEFAULT'])
 """The authentication mechanisms supported by PyMongo."""
 
 
@@ -337,13 +337,21 @@ def _authenticate_mongo_cr(credentials, sock_info):
     sock_info.command(source, query)
 
 
+def _authenticate_default(credentials, sock_info):
+    if sock_info.max_wire_version >= 3:
+        return _authenticate_scram_sha1(credentials, sock_info)
+    else:
+        return _authenticate_mongo_cr(credentials, sock_info)
+
+
 _AUTH_MAP = {
     'CRAM-MD5': _authenticate_cram_md5,
     'GSSAPI': _authenticate_gssapi,
     'MONGODB-CR': _authenticate_mongo_cr,
     'MONGODB-X509': _authenticate_x509,
     'PLAIN': _authenticate_plain,
     'SCRAM-SHA-1': _authenticate_scram_sha1,
+    'DEFAULT': _authenticate_default,
 }
 
 
diff --git a/pymongo/client_options.py b/pymongo/client_options.py
@@ -28,7 +28,7 @@ def _parse_credentials(username, password, database, options):
     """Parse authentication credentials."""
     if username is None:
         return None
-    mechanism = options.get('authmechanism', 'MONGODB-CR')
+    mechanism = options.get('authmechanism', 'DEFAULT')
     source = options.get('authsource', database or 'admin')
     return _build_credentials_tuple(
         mechanism, source, _unicode(username), _unicode(password), options)
diff --git a/pymongo/database.py b/pymongo/database.py
@@ -851,7 +851,7 @@ def remove_user(self, name):
             raise
 
     def authenticate(self, name, password=None,
-                     source=None, mechanism='MONGODB-CR', **kwargs):
+                     source=None, mechanism='DEFAULT', **kwargs):
         """Authenticate to use this database.
 
         Authentication lasts for the life of the underlying client
@@ -883,11 +883,15 @@ def authenticate(self, name, password=None,
             specified the current database is used.
           - `mechanism` (optional): See
             :data:`~pymongo.auth.MECHANISMS` for options.
-            Defaults to MONGODB-CR (MongoDB Challenge Response protocol)
+            By default, use SCRAM-SHA-1 with MongoDB 2.8 and later,
+            MONGODB-CR (MongoDB Challenge Response protocol) for older servers.
           - `gssapiServiceName` (optional): Used with the GSSAPI mechanism
             to specify the service name portion of the service principal name.
             Defaults to 'mongodb'.
 
+        .. versionadded:: 2.8
+           Use SCRAM-SHA-1 with MongoDB 2.8 and later.
+
         .. versionchanged:: 2.5
            Added the `source` and `mechanism` parameters. :meth:`authenticate`
            now raises a subclass of :class:`~pymongo.errors.PyMongoError` if
diff --git a/pymongo/mongo_client.py b/pymongo/mongo_client.py
@@ -336,7 +336,7 @@ def _cache_credentials(self, source, credentials, connect=False):
 
             # get_socket() logs out of the database if logged in with old
             # credentials, and logs in with new ones.
-            with server.pool.get_socket(all_credentials) as sock_info:
+            with server.get_socket(all_credentials) as sock_info:
                 sock_info.authenticate(credentials)
 
         # If several threads run _cache_credentials at once, last one wins.
@@ -652,7 +652,7 @@ def alive(self):
 
             # Avoid race when other threads log in or out.
             all_credentials = self.__all_credentials.copy()
-            with server.pool.get_socket(all_credentials) as sock_info:
+            with server.get_socket(all_credentials) as sock_info:
                 return not pool._closed(sock_info.sock)
 
         except (socket.error, ConnectionFailure):
@@ -1095,7 +1095,7 @@ def copy_database(self, from_name, to_name,
 
         # Avoid race when other threads log in or out.
         all_credentials = self.__all_credentials.copy()
-        with server.pool.get_socket(all_credentials) as sock:
+        with server.get_socket(all_credentials) as sock:
             if username is not None:
                 get_nonce_cmd = SON([("copydbgetnonce", 1),
                                      ("fromhost", from_host)])
diff --git a/pymongo/monitor.py b/pymongo/monitor.py
@@ -138,7 +138,7 @@ def _check_once(self):
         Returns a ServerDescription, or None on error.
         """
         try:
-            with self._pool.get_socket(all_credentials={}) as sock_info:
+            with self._pool.get_socket({}, 0, 0) as sock_info:
                 response, round_trip_time = self._check_with_socket(sock_info)
                 old_rtts = self._server_description.round_trip_times
                 if old_rtts:
diff --git a/pymongo/pool.py b/pymongo/pool.py
@@ -150,6 +150,9 @@ def __init__(self, sock, pool, host):
         self.last_checkout = time.time()
         self.pool_ref = weakref.ref(pool)
 
+        self._min_wire_version = None
+        self._max_wire_version = None
+
         # The pool's pool_id changes with each reset() so we can close sockets
         # created before the last reset.
         self.pool_id = pool.pool_id
@@ -257,6 +260,20 @@ def close(self):
             self.sock.close()
         except:
             pass
+        
+    def set_wire_version_range(self, min_wire_version, max_wire_version):
+        self._min_wire_version = min_wire_version
+        self._max_wire_version = max_wire_version
+        
+    @property
+    def min_wire_version(self):
+        assert self._min_wire_version is not None
+        return self._min_wire_version
+        
+    @property
+    def max_wire_version(self):
+        assert self._max_wire_version is not None
+        return self._max_wire_version
 
     def exhaust(self, exhaust):
         self._exhaust = exhaust
@@ -425,19 +442,25 @@ def connect(self):
         sock.settimeout(self.opts.socket_timeout)
         return SocketInfo(sock, self, hostname)
 
-    def get_socket(self, all_credentials):
+    def get_socket(self, all_credentials, min_wire_version, max_wire_version):
         """Get a socket from the pool.
 
         Returns a :class:`SocketInfo` object wrapping a connected
-        :class:`socket.socket`, and a bool saying whether the socket was from
-        the pool or freshly created.
+        :class:`socket.socket`.
+
+        The socket is logged in or out as necessary to match ``all_credentials``
+        using the correct authentication mechanism for the server's wire
+        protocol version.
 
         :Parameters:
           - `all_credentials`: dict, maps auth source to MongoCredential.
+          - `min_wire_version`: int, minimum protocol the server supports.
+          - `max_wire_version`: int, maximum protocol the server supports.
         """
         # First get a socket, then attempt authentication. Simplifies
         # semaphore management in the face of network errors during auth.
         sock_info = self._get_socket_no_auth()
+        sock_info.set_wire_version_range(min_wire_version, max_wire_version)
         try:
             sock_info.check_auth(all_credentials)
             return sock_info
diff --git a/pymongo/server.py b/pymongo/server.py
@@ -60,7 +60,7 @@ def send_message(self, message, all_credentials):
         """
         request_id, data = self._check_bson_size(message)
         try:
-            with self._pool.get_socket(all_credentials) as sock_info:
+            with self.get_socket(all_credentials) as sock_info:
                 sock_info.send_message(data)
         except socket.error as exc:
             self._raise_connection_failure(exc)
@@ -82,7 +82,7 @@ def send_message_with_response(
         """
         request_id, data = self._check_bson_size(message)
         try:
-            with self._pool.get_socket(all_credentials) as sock_info:
+            with self.get_socket(all_credentials) as sock_info:
                 sock_info.exhaust(exhaust)
                 sock_info.send_message(data)
                 response_data = sock_info.receive_message(1, request_id)
@@ -99,6 +99,20 @@ def send_message_with_response(
         except socket.error as exc:
             self._raise_connection_failure(exc)
 
+    def get_socket(self, all_credentials):
+        sd = self.description
+        sock_info = self.pool.get_socket(all_credentials=all_credentials,
+                                         min_wire_version=sd.min_wire_version,
+                                         max_wire_version=sd.max_wire_version)
+
+        return sock_info
+
+    def maybe_return_socket(self, sock_info):
+        self.pool.maybe_return_socket(sock_info)
+
+    def discard_socket(self, sock_info):
+        self.pool.discard_socket(sock_info)
+
     def start_request(self):
         # TODO: Remove implicit threadlocal requests, use explicit requests.
         self.pool.start_request()
diff --git a/test/pymongo_mocks.py b/test/pymongo_mocks.py
@@ -38,7 +38,7 @@ def __init__(self, client, pair, *args, **kwargs):
         Pool.__init__(self,
             (default_host, default_port), PoolOptions(connect_timeout=20))
 
-    def get_socket(self, all_credentials):
+    def get_socket(self, all_credentials, min_wire_version, max_wire_version):
         client = self.client
         host_and_port = '%s:%s' % (self.mock_host, self.mock_port)
         if host_and_port in client.mock_down_hosts:
@@ -49,7 +49,9 @@ def get_socket(self, all_credentials):
             + client.mock_members
             + client.mock_mongoses), "bad host: %s" % host_and_port
 
-        sock_info = Pool.get_socket(self, all_credentials)
+        sock_info = Pool.get_socket(
+            self, all_credentials, min_wire_version, max_wire_version)
+
         sock_info.mock_host = self.mock_host
         sock_info.mock_port = self.mock_port
         return sock_info
diff --git a/test/test_auth.py b/test/test_auth.py
@@ -30,7 +30,7 @@
 from pymongo.auth import HAVE_KERBEROS, _build_credentials_tuple
 from pymongo.errors import OperationFailure, ConfigurationError
 from pymongo.read_preferences import ReadPreference
-from test import client_context, host, port, SkipTest, unittest
+from test import client_context, host, port, SkipTest, unittest, Version
 
 # YOU MUST RUN KINIT BEFORE RUNNING GSSAPI TESTS.
 GSSAPI_HOST = os.environ.get('GSSAPI_HOST')
@@ -256,11 +256,13 @@ class TestSCRAMSHA1(unittest.TestCase):
     def setUp(self):
         self.set_name = client_context.setname
 
-        cmd_line = client_context.cmd_line
-        if 'SCRAM-SHA-1' not in cmd_line.get(
-            'parsed', {}).get('setParameter',
-                            {}).get('authenticationMechanisms', ''):
-            raise SkipTest('SCRAM-SHA-1 mechanism not enabled')
+        # Before 2.7.7, SCRAM-SHA-1 had to be enabled from the command line.
+        if client_context.version < Version(2, 7, 7):
+            cmd_line = client_context.cmd_line
+            if 'SCRAM-SHA-1' not in cmd_line.get(
+                    'parsed', {}).get('setParameter',
+                    {}).get('authenticationMechanisms', ''):
+                raise SkipTest('SCRAM-SHA-1 mechanism not enabled')
 
         client = client_context.rs_or_standalone_client
         client.pymongo_test.add_user(
@@ -294,9 +296,7 @@ def test_scram_sha1(self):
             client.pymongo_test.command('dbstats')
 
     def tearDown(self):
-        client_context.rs_or_standalone_client.pymongo_test.remove_user(
-            'user',
-            w=client_context.w)
+        client_context.rs_or_standalone_client.pymongo_test.remove_user('user')
 
 
 class TestAuthURIOptions(unittest.TestCase):
diff --git a/test/test_client.py b/test/test_client.py
@@ -916,7 +916,7 @@ def test_auth_network_error(self):
 
         # Simulate an authenticate() call on a different socket.
         credentials = auth._build_credentials_tuple(
-            'MONGODB-CR', 'admin', db_user, db_pwd, {})
+            'DEFAULT', 'admin', db_user, db_pwd, {})
 
         c._cache_credentials('test', credentials, connect=False)
 
diff --git a/test/test_discovery_and_monitoring.py b/test/test_discovery_and_monitoring.py
@@ -53,12 +53,6 @@ def __init__(self, *args, **kwargs):
         self.pool_id = 0
         self._lock = threading.Lock()
 
-    def get_socket(self):
-        return MockSocketInfo()
-
-    def maybe_return_socket(self, _):
-        pass
-
     def reset(self):
         with self._lock:
             self.pool_id += 1
diff --git a/test/test_pooling.py b/test/test_pooling.py
diff --git a/test/test_topology.py b/test/test_topology.py
diff --git a/test/test_uri_parser.py b/test/test_uri_parser.py
diff --git a/test/utils.py b/test/utils.py
