clean up reauth handling

Author: blink1073

pymongo/auth.py

@@ -589,8 +589,6 @@ def authenticate(credentials, sock_info, reauthenticate=False):
     """Authenticate sock_info."""
     mechanism = credentials.mechanism
     auth_func = _AUTH_MAP[mechanism]
-    if reauthenticate:
-        sock_info.handle_reauthenticate()
     if mechanism == "MONGODB-OIDC":
         _authenticate_oidc(credentials, sock_info, reauthenticate)
     else:

pymongo/auth_oidc.py

@@ -100,7 +100,6 @@ class _OIDCAuthenticator:
 
     def get_current_token(self, use_callbacks=True):
         properties = self.properties
-        principal_name = self.username
 
         request_cb = properties.request_token_callback
         refresh_cb = properties.refresh_token_callback
@@ -238,29 +237,26 @@ def clear(self):
 
     def run_command(self, sock_info, cmd):
         try:
-            return sock_info.command("$external", cmd)
+            return sock_info.command("$external", cmd, no_reauth=True)
         except OperationFailure as exc:
             self.clear()
             if exc.code == _REAUTHENTICATION_REQUIRED_CODE:
                 if "jwt" in bson.decode(cmd["payload"]):
                     if self.idp_info_gen_id > self.reauth_gen_id:
                         raise
-                    self.handle_reauth(sock_info)
-                    return self.authenticate(sock_info)
+                    return self.authenticate(sock_info, reauthenticate=True)
             raise
 
-    def handle_reauth(self, sock_info):
-        prev_id = getattr(sock_info, "oidc_token_gen_id", None)
-        if prev_id != self.token_gen_id:
-            # No need to preemptively clear, we've already changed tokens.
-            return
-
-        self.reauth_gen_id = self.idp_info_gen_id
-        self.token_exp_utc = None
-        if not self.properties.refresh_token_callback:
-            self.clear()
+    def authenticate(self, sock_info, reauthenticate=False):
+        if reauthenticate:
+            prev_id = getattr(sock_info, "oidc_token_gen_id", None)
+            # Check if we've already changed tokens.
+            if prev_id == self.token_gen_id:
+                self.reauth_gen_id = self.idp_info_gen_id
+                self.token_exp_utc = None
+                if not self.properties.refresh_token_callback:
+                    self.clear()
 
-    def authenticate(self, sock_info):
         ctx = sock_info.auth_ctx
         cmd = None
 
@@ -300,6 +296,4 @@ def authenticate(self, sock_info):
 def _authenticate_oidc(credentials, sock_info, reauthenticate):
     """Authenticate using MONGODB-OIDC."""
     authenticator = _get_authenticator(credentials, sock_info.address)
-    if reauthenticate:
-        authenticator.handle_reauth(sock_info)
-    return authenticator.authenticate(sock_info)
+    return authenticator.authenticate(sock_info, reauthenticate=reauthenticate)

pymongo/helpers.py

@@ -270,3 +270,35 @@ def _handle_exception():
             pass
         finally:
             del einfo
+
+
+def _handle_reauth(func):
+    def inner(*args, **kwargs):
+        no_reauth = kwargs.pop("no_reauth", False)
+        from pymongo.pool import SocketInfo
+
+        try:
+            return func(*args, **kwargs)
+        except OperationFailure as exc:
+            if no_reauth:
+                raise
+            if exc.code == _REAUTHENTICATION_REQUIRED_CODE:
+                # Look for an argument that either is a SocketInfo
+                # or has a socket_info attribute, so we can trigger
+                # a reauth.
+                sock_info = None
+                for arg in args:
+                    if isinstance(arg, SocketInfo):
+                        sock_info = arg
+                        break
+                    if hasattr(arg, "sock_info"):
+                        sock_info = arg.sock_info
+                        break
+                if sock_info:
+                    sock_info.authenticate(reauthenticate=True)
+                else:
+                    raise
+                return func(*args, **kwargs)
+            raise
+
+    return inner

pymongo/message.py

@@ -54,6 +54,7 @@
     ProtocolError,
 )
 from pymongo.hello import HelloCompat
+from pymongo.helpers import _handle_reauth
 from pymongo.read_preferences import ReadPreference
 from pymongo.write_concern import WriteConcern
 
@@ -909,6 +910,7 @@ def unack_write(self, cmd, request_id, msg, max_doc_size, docs):
             self.start_time = datetime.datetime.now()
         return result
 
+    @_handle_reauth
     def write_command(self, cmd, request_id, msg, docs):
         """A proxy for SocketInfo.write_command that handles event publishing."""
         if self.publish:

pymongo/mongo_client.py

@@ -1397,14 +1397,7 @@ def is_retrying():
                             assert last_error is not None
                             raise last_error
                         retryable = False
-                    # Handle re-authentication.
-                    try:
-                        return func(session, sock_info, retryable)
-                    except OperationFailure as exc:
-                        if exc.code == helpers._REAUTHENTICATION_REQUIRED_CODE:
-                            sock_info.authenticate(reauthenticate=True)
-                            return func(session, sock_info, retryable)
-                        raise
+                    return func(session, sock_info, retryable)
             except ServerSelectionTimeoutError:
                 if is_retrying():
                     # The application may think the write was never attempted
@@ -1468,14 +1461,7 @@ def _retryable_read(self, func, read_pref, session, address=None, retryable=True
                         # not support retryable reads, raise the last error.
                         assert last_error is not None
                         raise last_error
-                    # Handle re-authentication.
-                    try:
-                        return func(session, server, sock_info, read_pref)
-                    except OperationFailure as exc:
-                        if exc.code == helpers._REAUTHENTICATION_REQUIRED_CODE:
-                            sock_info.authenticate(reauthenticate=True)
-                            return func(session, server, sock_info, read_pref)
-                        raise
+                    return func(session, server, sock_info, read_pref)
             except ServerSelectionTimeoutError:
                 if retrying:
                     # The application may think the write was never attempted

pymongo/pool.py

@@ -56,6 +56,7 @@
     _CertificateError,
 )
 from pymongo.hello import Hello, HelloCompat
+from pymongo.helpers import _handle_reauth
 from pymongo.lock import _create_lock
 from pymongo.monitoring import ConnectionCheckOutFailedReason, ConnectionClosedReason
 from pymongo.network import command, receive_message
@@ -704,6 +705,7 @@ def _next_reply(self):
         helpers._check_command_response(response_doc, self.max_wire_version)
         return response_doc
 
+    @_handle_reauth
     def command(
         self,
         dbname,
@@ -788,7 +790,7 @@ def command(
                 exhaust_allowed=exhaust_allowed,
                 write_concern=write_concern,
             )
-        except (OperationFailure, NotPrimaryError):
+        except (OperationFailure, NotPrimaryError) as exc:
             raise
         # Catch socket.error, KeyboardInterrupt, etc. and close ourselves.
         except BaseException as error:
@@ -864,7 +866,12 @@ def authenticate(self, reauthenticate=False):
         """
         # CMAP spec says to publish the ready event only after authenticating
         # the connection.
-        if not self.ready or reauthenticate:
+        if reauthenticate:
+            if self.performed_handshake:
+                # Existing auth_ctx is stale, remove it.
+                self.auth_ctx = None
+            self.ready = False
+        if not self.ready:
             creds = self.opts._credentials
             if creds:
                 auth.authenticate(creds, self, reauthenticate=reauthenticate)
@@ -927,12 +934,6 @@ def idle_time_seconds(self):
         """Seconds since this socket was last checked into its pool."""
         return time.monotonic() - self.last_checkin_time
 
-    def handle_reauthenticate(self):
-        """Handle a reauthentication."""
-        if self.performed_handshake:
-            # Existing auth_ctx is stale, remove it.
-            self.auth_ctx = None
-
     def _raise_connection_failure(self, error):
         # Catch *all* exceptions from socket methods and close the socket. In
         # regular Python, socket operations only raise socket.error, even if

pymongo/server.py

@@ -18,7 +18,7 @@
 
 from bson import _decode_all_selective
 from pymongo.errors import NotPrimaryError, OperationFailure
-from pymongo.helpers import _check_command_response
+from pymongo.helpers import _check_command_response, _handle_reauth
 from pymongo.message import _convert_exception, _OpMsg
 from pymongo.response import PinnedResponse, Response
 
@@ -73,6 +73,7 @@ def request_check(self):
         """Check the server's state soon."""
         self._monitor.request_check()
 
+    @_handle_reauth
     def run_operation(self, sock_info, operation, read_preference, listeners, unpack_res):
         """Run a _Query or _GetMore operation and return a Response object.
 

test/auth_aws/test_auth_oidc.py

@@ -28,9 +28,10 @@
 
 from bson import SON
 from pymongo import MongoClient
-from pymongo.auth import MongoCredential
 from pymongo.auth_oidc import _CACHE as _oidc_cache
+from pymongo.cursor import CursorType
 from pymongo.errors import ConfigurationError, OperationFailure
+from pymongo.operations import InsertOne
 
 
 class TestAuthOIDC(unittest.TestCase):
@@ -496,8 +497,8 @@ def test_reauthenticate_succeeds(self):
 
         with self.fail_point(
             {
-                "mode": {"times": 2},
-                "data": {"failCommands": ["find", "saslStart"], "errorCode": 391},
+                "mode": {"times": 1},
+                "data": {"failCommands": ["find"], "errorCode": 391},
             }
         ):
             # Perform a find operation.
@@ -529,7 +530,152 @@ def test_reauthenticate_succeeds(self):
         self.assertEqual(self.refresh_called, 1)
         client.close()
 
-    def test_reauthenticate_retries_and_succees_with_cache(self):
+    def test_reauthenticate_succeeds_bulk_write(self):
+        request_cb = self.create_request_cb()
+        refresh_cb = self.create_refresh_cb()
+
+        # Create a client with the callbacks.
+        props: Dict = dict(request_token_callback=request_cb, refresh_token_callback=refresh_cb)
+        client = MongoClient(self.uri_single, authmechanismproperties=props)
+
+        # Perform a find operation.
+        client.test.test.find_one()
+
+        # Assert that the refresh callback has not been called.
+        self.assertEqual(self.refresh_called, 0)
+
+        with self.fail_point(
+            {
+                "mode": {"times": 1},
+                "data": {"failCommands": ["insert"], "errorCode": 391},
+            }
+        ):
+            # Perform a bulk write operation.
+            client.test.test.bulk_write([InsertOne({})])
+
+        # Assert that the refresh callback has been called.
+        self.assertEqual(self.refresh_called, 1)
+        client.close()
+
+    def test_reauthenticate_succeeds_bulk_read(self):
+        request_cb = self.create_request_cb()
+        refresh_cb = self.create_refresh_cb()
+
+        # Create a client with the callbacks.
+        props: Dict = dict(request_token_callback=request_cb, refresh_token_callback=refresh_cb)
+        client = MongoClient(self.uri_single, authmechanismproperties=props)
+
+        # Perform a find operation.
+        client.test.test.find_one()
+
+        # Perform a bulk write operation.
+        client.test.test.bulk_write([InsertOne({})])
+
+        # Assert that the refresh callback has not been called.
+        self.assertEqual(self.refresh_called, 0)
+
+        with self.fail_point(
+            {
+                "mode": {"times": 1},
+                "data": {"failCommands": ["find"], "errorCode": 391},
+            }
+        ):
+            # Perform a bulk read operation.
+            cursor = client.test.test.find_raw_batches({})
+            list(cursor)
+
+        # Assert that the refresh callback has been called.
+        self.assertEqual(self.refresh_called, 1)
+        client.close()
+
+    def test_reauthenticate_succeeds_cursor(self):
+        request_cb = self.create_request_cb()
+        refresh_cb = self.create_refresh_cb()
+
+        # Create a client with the callbacks.
+        props: Dict = dict(request_token_callback=request_cb, refresh_token_callback=refresh_cb)
+        client = MongoClient(self.uri_single, authmechanismproperties=props)
+
+        # Perform an insert operation.
+        client.test.test.insert_one({"a": 1})
+
+        # Assert that the refresh callback has not been called.
+        self.assertEqual(self.refresh_called, 0)
+
+        with self.fail_point(
+            {
+                "mode": {"times": 1},
+                "data": {"failCommands": ["find"], "errorCode": 391},
+            }
+        ):
+            # Perform a find operation.
+            cursor = client.test.test.find({"a": 1})
+            self.assertGreaterEqual(len(list(cursor)), 1)
+
+        # Assert that the refresh callback has been called.
+        self.assertEqual(self.refresh_called, 1)
+        client.close()
+
+    def test_reauthenticate_succeeds_get_more(self):
+        request_cb = self.create_request_cb()
+        refresh_cb = self.create_refresh_cb()
+
+        # Create a client with the callbacks.
+        props: Dict = dict(request_token_callback=request_cb, refresh_token_callback=refresh_cb)
+        client = MongoClient(self.uri_single, authmechanismproperties=props)
+
+        # Perform an insert operation.
+        client.test.test.insert_many([{"a": 1}, {"a": 1}])
+
+        # Assert that the refresh callback has not been called.
+        self.assertEqual(self.refresh_called, 0)
+
+        with self.fail_point(
+            {
+                "mode": {"times": 1},
+                "data": {"failCommands": ["find"], "errorCode": 391},
+            }
+        ):
+            # Perform a find operation.
+            cursor = client.test.test.find({"a": 1}, batch_size=1, cursor_type=CursorType.EXHAUST)
+            self.assertGreaterEqual(len(list(cursor)), 1)
+
+        # Assert that the refresh callback has been called.
+        self.assertEqual(self.refresh_called, 1)
+        client.close()
+
+    def test_reauthenticate_succeeds_command(self):
+        request_cb = self.create_request_cb()
+        refresh_cb = self.create_refresh_cb()
+
+        # Create a client with the callbacks.
+        props: Dict = dict(request_token_callback=request_cb, refresh_token_callback=refresh_cb)
+
+        print("start of test")
+        client = MongoClient(self.uri_single, authmechanismproperties=props)
+
+        # Perform an insert operation.
+        client.test.test.insert_one({"a": 1})
+
+        # Assert that the refresh callback has not been called.
+        self.assertEqual(self.refresh_called, 0)
+
+        with self.fail_point(
+            {
+                "mode": {"times": 1},
+                "data": {"failCommands": ["count"], "errorCode": 391},
+            }
+        ):
+            # Perform a count operation.
+            cursor = client.test.command(dict(count="test"))
+
+        self.assertGreaterEqual(len(list(cursor)), 1)
+
+        # Assert that the refresh callback has been called.
+        self.assertEqual(self.refresh_called, 1)
+        client.close()
+
+    def test_reauthenticate_retries_and_succeeds_with_cache(self):
         listener = EventListener()
 
         # Create request and refresh callbacks that return valid credentials