PYTHON-2858 Use OP_MSG to authenticate if server supports OP_MSG

pymongo/pool.py

@@ -552,7 +552,8 @@ def unpin(self):
     def hello_cmd(self):
         # Handshake spec requires us to use OP_MSG+hello command for the
         # initial handshake in load balanced or versioned api mode.
-        if self.opts.server_api or self.hello_ok or self.opts.load_balanced:
+        if (self.opts.server_api or self.hello_ok or self.opts.load_balanced
+            or self.max_wire_version >= 6):
             self.op_msg_enabled = True
             return SON([(HelloCompat.CMD, 1)])
         else:

test/mockupdb/test_handshake.py

@@ -21,6 +21,7 @@
 from bson.objectid import ObjectId
 
 import unittest
+from copy import deepcopy
 
 
 def test_hello_with_option(self, protocol, **kwargs):
@@ -44,13 +45,13 @@ def respond(r):
 
     # We need a special dict because MongoClient uses "server_api" and all
     # of the commands use "apiVersion".
-    k_map = {("apiVersion", "1"):("server_api", ServerApi(
-                                        ServerApiVersion.V1))}
+    k_map = {("apiVersion", "1"): ("server_api", ServerApi(
+                                         ServerApiVersion.V1))}
     client = MongoClient("mongodb://"+primary.address_string,
                          appname='my app', # For _check_handshake_data()
                          **dict([k_map.get((k, v), (k, v)) for k, v
                                  in kwargs.items()]))
-    
+
     self.addCleanup(client.close)
 
     # We have an autoresponder luckily, so no need for `go()`.
@@ -84,76 +85,78 @@ def test_client_handshake_data(self):
             self.addCleanup(server.stop)
 
         hosts = [server.address_string for server in (primary, secondary)]
-        primary_response = OpReply('ismaster', True,
-                                   setName='rs', hosts=hosts,
-                                   minWireVersion=2, maxWireVersion=6)
-        error_response = OpReply(
+        primary_response = {"setName": 'rs', "hosts": hosts,
+                            "minWireVersion": 2, "maxWireVersion": 13}
+        error_response = OpMsgReply(
             0, errmsg='Cache Reader No keys found for HMAC ...', code=211)
 
-        secondary_response = OpReply('ismaster', False,
-                                     setName='rs', hosts=hosts,
-                                     secondary=True,
-                                     minWireVersion=2, maxWireVersion=6)
+        secondary_response = {"setName": 'rs', "hosts": hosts,
+                              "minWireVersion": 2, "maxWireVersion": 13}
 
         client = MongoClient(primary.uri,
                              replicaSet='rs',
                              appname='my app',
                              heartbeatFrequencyMS=500)  # Speed up the test.
-
+    
         self.addCleanup(client.close)
 
         # New monitoring sockets send data during handshake.
-        heartbeat = primary.receives('ismaster')
+        heartbeat = primary.receives(Command('ismaster'))
         _check_handshake_data(heartbeat)
-        heartbeat.ok(primary_response)
+        heartbeat.ok(**primary_response)
 
-        heartbeat = secondary.receives('ismaster')
+        heartbeat = secondary.receives(Command('ismaster'))
         _check_handshake_data(heartbeat)
-        heartbeat.ok(secondary_response)
+        heartbeat.ok(**secondary_response)
 
         # Subsequent heartbeats have no client data.
-        primary.receives('ismaster', 1, client=absent).ok(error_response)
-        secondary.receives('ismaster', 1, client=absent).ok(error_response)
+        primary.receives(OpMsg('hello', 1, client=absent)).ok(error_response)
+        secondary.receives(OpMsg('hello', 1, client=absent)).ok(
+            error_response)
 
         # The heartbeat retry (on a new connection) does have client data.
-        heartbeat = primary.receives('ismaster')
+        heartbeat = primary.receives(Command('ismaster'))
         _check_handshake_data(heartbeat)
-        heartbeat.ok(primary_response)
+        heartbeat.reply(primary_response)
 
-        heartbeat = secondary.receives('ismaster')
+        heartbeat = secondary.receives(Command('ismaster'))
         _check_handshake_data(heartbeat)
-        heartbeat.ok(secondary_response)
+        heartbeat.reply(secondary_response)
 
         # Still no client data.
-        primary.receives('ismaster', 1, client=absent).ok(primary_response)
-        secondary.receives('ismaster', 1, client=absent).ok(secondary_response)
+        primary.receives(OpMsg('hello', 1, client=absent)).reply(**primary_response)
+        secondary.receives(OpMsg('hello', 1, client=absent)).reply(**secondary_response)
 
+        primary.receives(OpMsg('hello', 1, client=absent)).hangup()
+        secondary.receives(OpMsg('hello')).hangup()
         # After a disconnect, next ismaster has client data again.
-        primary.receives('ismaster', 1, client=absent).hangup()
         heartbeat = primary.receives('ismaster')
         _check_handshake_data(heartbeat)
-        heartbeat.ok(primary_response)
-
-        secondary.autoresponds('ismaster', secondary_response)
+        hb_port = deepcopy(heartbeat.client_port)
+        heartbeat.reply(primary_response)
 
         # Start a command, so the client opens an application socket.
-        future = go(client.db.command, 'whatever')
-
+        future = go(client.db.command, "whatever")
+        handshook = None
         for request in primary:
             if request.matches(Command('ismaster')):
-                if request.client_port == heartbeat.client_port:
-                    # This is the monitor again, keep going.
-                    request.ok(primary_response)
+                if handshook is None:
+                    handshook = True
                 else:
-                    # Handshaking a new application socket.
-                    _check_handshake_data(request)
-                    request.ok(primary_response)
-            else:
+                    handshook = False
+                # Handshaking a new application socket.
+                _check_handshake_data(request)
+                request.reply(**primary_response)
+            elif request.matches(OpMsg("whatever")):
                 # Command succeeds.
-                request.assert_matches(OpMsg('whatever'))
-                request.ok()
-                assert future()
+                request.reply(OpMsgReply(**primary_response))
+                # If handshook is false it means it tried to handshake
+                # multiple times. If it's None it means it never tried.
+                assert future() and handshook
                 return
+            else:
+                request.reply(isWritablePrimary=True,
+                                         **primary_response)
 
     def test_client_handshake_saslSupportedMechs(self):
         server = MockupDB()
@@ -217,5 +220,42 @@ def test_handshake_not_either(self):
         with self.assertRaisesRegex(AssertionError, "does not match"):
             test_hello_with_option(self, OpMsg)
 
+    def test_handshake_max_wire(self):
+        server = MockupDB()
+        primary_response = {"hello": 1, "ok": 1,
+                            "minWireVersion": 0, "maxWireVersion": 6}
+        self.found_auth_msg = False
+
+        def responder(request):
+            if request.matches(OpMsg, saslStart=1):
+                self.found_auth_msg = True
+                # Immediately closes the connection with
+                request.reply(OpMsgReply(**primary_response,
+                                         **{'payload':
+                                                b'r=wPleNM8S5p8gMaffMDF7Py4ru9bnmmoqb0'
+                                                b'1WNPsil6o=pAvr6B1garhlwc6MKNQ93ZfFky'
+                                                b'tXdF9r,'
+                                                b's=4dcxugMJq2P4hQaDbGXZR8uR3ei'
+                                                b'PHrSmh4uhkg==,i=15000',
+                                            "saslSupportedMechs": [
+                                                "SCRAM-SHA-1"]}))
+            else:
+                return request.reply(OpMsgReply(**primary_response))
+
+        server.autoresponds(responder)
+        self.addCleanup(server.stop)
+        server.run()
+        client = MongoClient(server.uri,
+                             username='username',
+                             password='password',
+                             appname='my app',
+                             )
+        self.addCleanup(client.close)
+        self.assertRaises(OperationFailure, client.db.collection.find_one,
+                          {"a": 1})
+        assert self.found_auth_msg, """Could not find authentication command 
+                                       with correct protocol"""
+
+
 if __name__ == '__main__':
     unittest.main()