From ccb995145c29f4ff6335ef31324f8b1b615c0aea Mon Sep 17 00:00:00 2001 From: Noah Stapp Date: Mon, 19 Oct 2026 14:56:08 -0400 Subject: [PATCH] PYTHON-6030 Consolidate _OIDCAuthenticator classes into shared base class --- pymongo/asynchronous/auth_oidc.py | 57 +++-------------------------- pymongo/auth_oidc_shared.py | 60 ++++++++++++++++++++++++++++++- pymongo/synchronous/auth_oidc.py | 57 +++-------------------------- 3 files changed, 67 insertions(+), 107 deletions(-) diff --git a/pymongo/asynchronous/auth_oidc.py b/pymongo/asynchronous/auth_oidc.py index 584dcc18c9..b0403d372a 100644 --- a/pymongo/asynchronous/auth_oidc.py +++ b/pymongo/asynchronous/auth_oidc.py @@ -24,7 +24,6 @@ from typing import TYPE_CHECKING, Any, Optional, Union import bson -from bson.binary import Binary from pymongo._csot import remaining from pymongo.auth_oidc_shared import ( CALLBACK_VERSION, @@ -36,10 +35,9 @@ OIDCCallbackResult, OIDCIdPInfo, _get_authenticator, - _OIDCProperties, + _OIDCAuthenticatorBase, ) from pymongo.errors import OperationFailure -from pymongo.helpers_shared import _AUTHENTICATION_FAILURE_CODE from pymongo.lock import Lock, _async_create_lock if TYPE_CHECKING: @@ -50,24 +48,16 @@ @dataclass -class _OIDCAuthenticator: - username: str - properties: _OIDCProperties - refresh_token: Optional[str] = field(default=None) - access_token: Optional[str] = field(default=None) - idp_info: Optional[OIDCIdPInfo] = field(default=None) - token_gen_id: int = field(default=0) +class _OIDCAuthenticator(_OIDCAuthenticatorBase): if not _IS_SYNC: lock: Lock = field(default_factory=_async_create_lock) # type: ignore[assignment] else: lock: threading.Lock = field(default_factory=_async_create_lock) # type: ignore[assignment, no-redef] - last_call_time: float = field(default=0) - async def reauthenticate(self, conn: AsyncConnection) -> Optional[Mapping[str, Any]]: """Handle a reauthenticate from the server.""" # Invalidate the token for the connection. - self._invalidate(conn) + self._invalidate(conn.oidc_token_gen_id) # Call the appropriate auth logic for the callback type. if self.properties.callback: return await self._authenticate_machine(conn) @@ -91,12 +81,6 @@ async def authenticate(self, conn: AsyncConnection) -> Optional[Mapping[str, Any return await self._authenticate_machine(conn) return await self._authenticate_human(conn) - def get_spec_auth_cmd(self) -> Optional[MutableMapping[str, Any]]: - """Get the appropriate speculative auth command.""" - if not self.access_token: - return None - return self._get_start_command({"jwt": self.access_token}) - async def _authenticate_machine(self, conn: AsyncConnection) -> Mapping[str, Any]: # If there is a cached access token, try to authenticate with it. If # authentication fails with error code 18, invalidate the access token, @@ -212,22 +196,9 @@ async def _run_command( return await conn.command("$external", cmd, no_reauth=True) # type: ignore[call-arg] except OperationFailure as e: if self._is_auth_error(e): - self._invalidate(conn) + self._invalidate(conn.oidc_token_gen_id) raise - def _is_auth_error(self, err: Exception) -> bool: - if not isinstance(err, OperationFailure): - return False - return err.code == _AUTHENTICATION_FAILURE_CODE - - def _invalidate(self, conn: AsyncConnection) -> None: - # Ignore the invalidation if a token gen id is given and is less than our - # current token gen id. - token_gen_id = conn.oidc_token_gen_id or 0 - if token_gen_id is not None and token_gen_id < self.token_gen_id: - return - self.access_token = None - async def _sasl_continue_jwt( self, conn: AsyncConnection, start_resp: Mapping[str, Any] ) -> Mapping[str, Any]: @@ -247,26 +218,6 @@ async def _sasl_start_jwt(self, conn: AsyncConnection) -> Mapping[str, Any]: cmd = self._get_start_command({"jwt": access_token}) return await self._run_command(conn, cmd) - def _get_start_command(self, payload: Optional[Mapping[str, Any]]) -> MutableMapping[str, Any]: - if payload is None: - principal_name = self.username - if principal_name: - payload = {"n": principal_name} - else: - payload = {} - bin_payload = Binary(bson.encode(payload)) - return {"saslStart": 1, "mechanism": "MONGODB-OIDC", "payload": bin_payload} - - def _get_continue_command( - self, payload: Mapping[str, Any], start_resp: Mapping[str, Any] - ) -> MutableMapping[str, Any]: - bin_payload = Binary(bson.encode(payload)) - return { - "saslContinue": 1, - "payload": bin_payload, - "conversationId": start_resp["conversationId"], - } - async def _authenticate_oidc( credentials: MongoCredential, conn: AsyncConnection, reauthenticate: bool diff --git a/pymongo/auth_oidc_shared.py b/pymongo/auth_oidc_shared.py index 6a307ec935..a774c85c0e 100644 --- a/pymongo/auth_oidc_shared.py +++ b/pymongo/auth_oidc_shared.py @@ -19,13 +19,17 @@ import abc import os +from collections.abc import Mapping, MutableMapping from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any, Callable, Optional from urllib.parse import quote +import bson +from bson.binary import Binary from pymongo._azure_helpers import _get_azure_response from pymongo._gcp_helpers import _get_gcp_response -from pymongo.errors import ConfigurationError +from pymongo.errors import ConfigurationError, OperationFailure +from pymongo.helpers_shared import _AUTHENTICATION_FAILURE_CODE if TYPE_CHECKING: from pymongo.auth_shared import MongoCredential @@ -137,6 +141,60 @@ def _get_k8s_token() -> str: return fid.read() +@dataclass +class _OIDCAuthenticatorBase: + """Shared base class for async and sync _OIDCAuthenticator. + + Subclasses contain all methods that perform I/O. + """ + + username: str + properties: _OIDCProperties + refresh_token: Optional[str] = field(default=None) + access_token: Optional[str] = field(default=None) + idp_info: Optional[OIDCIdPInfo] = field(default=None) + token_gen_id: int = field(default=0) + last_call_time: float = field(default=0) + + def get_spec_auth_cmd(self) -> Optional[MutableMapping[str, Any]]: + """Get the appropriate speculative auth command.""" + if not self.access_token: + return None + return self._get_start_command({"jwt": self.access_token}) + + def _is_auth_error(self, err: Exception) -> bool: + if not isinstance(err, OperationFailure): + return False + return err.code == _AUTHENTICATION_FAILURE_CODE + + def _invalidate(self, oidc_token_gen_id: Optional[int]) -> None: + # Ignore the invalidation if a token gen id is given and is less than our + # current token gen id. + if (oidc_token_gen_id or 0) < self.token_gen_id: + return + self.access_token = None + + def _get_start_command(self, payload: Optional[Mapping[str, Any]]) -> MutableMapping[str, Any]: + if payload is None: + principal_name = self.username + if principal_name: + payload = {"n": principal_name} + else: + payload = {} + bin_payload = Binary(bson.encode(payload)) + return {"saslStart": 1, "mechanism": "MONGODB-OIDC", "payload": bin_payload} + + def _get_continue_command( + self, payload: Mapping[str, Any], start_resp: Mapping[str, Any] + ) -> MutableMapping[str, Any]: + bin_payload = Binary(bson.encode(payload)) + return { + "saslContinue": 1, + "payload": bin_payload, + "conversationId": start_resp["conversationId"], + } + + def _get_authenticator( credentials: MongoCredential, address: tuple[str, int], authenticator_cls: Callable[..., Any] ) -> Any: diff --git a/pymongo/synchronous/auth_oidc.py b/pymongo/synchronous/auth_oidc.py index de7801bd05..7d242bc326 100644 --- a/pymongo/synchronous/auth_oidc.py +++ b/pymongo/synchronous/auth_oidc.py @@ -24,7 +24,6 @@ from typing import TYPE_CHECKING, Any, Optional, Union import bson -from bson.binary import Binary from pymongo._csot import remaining from pymongo.auth_oidc_shared import ( CALLBACK_VERSION, @@ -36,10 +35,9 @@ OIDCCallbackResult, OIDCIdPInfo, _get_authenticator, - _OIDCProperties, + _OIDCAuthenticatorBase, ) from pymongo.errors import OperationFailure -from pymongo.helpers_shared import _AUTHENTICATION_FAILURE_CODE from pymongo.lock import Lock, _create_lock if TYPE_CHECKING: @@ -50,24 +48,16 @@ @dataclass -class _OIDCAuthenticator: - username: str - properties: _OIDCProperties - refresh_token: Optional[str] = field(default=None) - access_token: Optional[str] = field(default=None) - idp_info: Optional[OIDCIdPInfo] = field(default=None) - token_gen_id: int = field(default=0) +class _OIDCAuthenticator(_OIDCAuthenticatorBase): if not _IS_SYNC: lock: Lock = field(default_factory=_create_lock) # type: ignore[assignment] else: lock: threading.Lock = field(default_factory=_create_lock) # type: ignore[assignment, no-redef] - last_call_time: float = field(default=0) - def reauthenticate(self, conn: Connection) -> Optional[Mapping[str, Any]]: """Handle a reauthenticate from the server.""" # Invalidate the token for the connection. - self._invalidate(conn) + self._invalidate(conn.oidc_token_gen_id) # Call the appropriate auth logic for the callback type. if self.properties.callback: return self._authenticate_machine(conn) @@ -91,12 +81,6 @@ def authenticate(self, conn: Connection) -> Optional[Mapping[str, Any]]: return self._authenticate_machine(conn) return self._authenticate_human(conn) - def get_spec_auth_cmd(self) -> Optional[MutableMapping[str, Any]]: - """Get the appropriate speculative auth command.""" - if not self.access_token: - return None - return self._get_start_command({"jwt": self.access_token}) - def _authenticate_machine(self, conn: Connection) -> Mapping[str, Any]: # If there is a cached access token, try to authenticate with it. If # authentication fails with error code 18, invalidate the access token, @@ -210,22 +194,9 @@ def _run_command(self, conn: Connection, cmd: MutableMapping[str, Any]) -> Mappi return conn.command("$external", cmd, no_reauth=True) # type: ignore[call-arg] except OperationFailure as e: if self._is_auth_error(e): - self._invalidate(conn) + self._invalidate(conn.oidc_token_gen_id) raise - def _is_auth_error(self, err: Exception) -> bool: - if not isinstance(err, OperationFailure): - return False - return err.code == _AUTHENTICATION_FAILURE_CODE - - def _invalidate(self, conn: Connection) -> None: - # Ignore the invalidation if a token gen id is given and is less than our - # current token gen id. - token_gen_id = conn.oidc_token_gen_id or 0 - if token_gen_id is not None and token_gen_id < self.token_gen_id: - return - self.access_token = None - def _sasl_continue_jwt( self, conn: Connection, start_resp: Mapping[str, Any] ) -> Mapping[str, Any]: @@ -245,26 +216,6 @@ def _sasl_start_jwt(self, conn: Connection) -> Mapping[str, Any]: cmd = self._get_start_command({"jwt": access_token}) return self._run_command(conn, cmd) - def _get_start_command(self, payload: Optional[Mapping[str, Any]]) -> MutableMapping[str, Any]: - if payload is None: - principal_name = self.username - if principal_name: - payload = {"n": principal_name} - else: - payload = {} - bin_payload = Binary(bson.encode(payload)) - return {"saslStart": 1, "mechanism": "MONGODB-OIDC", "payload": bin_payload} - - def _get_continue_command( - self, payload: Mapping[str, Any], start_resp: Mapping[str, Any] - ) -> MutableMapping[str, Any]: - bin_payload = Binary(bson.encode(payload)) - return { - "saslContinue": 1, - "payload": bin_payload, - "conversationId": start_resp["conversationId"], - } - def _authenticate_oidc( credentials: MongoCredential, conn: Connection, reauthenticate: bool