PYTHON-3845 OIDC: Implement Machine Callback Mechanism

.evergreen/run-mongodb-oidc-test.sh

@@ -45,4 +45,5 @@ fi
 export TEST_AUTH_OIDC=1
 export COVERAGE=1
 export AUTH="auth"
+export AWS_WEB_IDENTITY_TOKEN_FILE="$OIDC_TOKEN_DIR/test_user1"
 bash ./.evergreen/tox.sh -m test-eg

.evergreen/run-tests.sh

@@ -30,7 +30,7 @@ set -o xtrace
 
 AUTH=${AUTH:-noauth}
 SSL=${SSL:-nossl}
-TEST_ARGS="$1"
+TEST_ARGS="${*:1}"
 PYTHON=$(which python)
 export PIP_QUIET=1  # Quiet by default
 
@@ -52,6 +52,7 @@ if [ "$AUTH" != "noauth" ]; then
     elif [ ! -z "$TEST_AUTH_OIDC" ]; then
         export DB_USER=$OIDC_ALTAS_USER
         export DB_PASSWORD=$OIDC_ATLAS_PASSWORD
+        export DB_IP="$MONGODB_URI"
     else
         export DB_USER="bob"
         export DB_PASSWORD="pwd123"

pymongo/auth.py

@@ -38,7 +38,12 @@
 from bson.binary import Binary
 from bson.son import SON
 from pymongo.auth_aws import _authenticate_aws
-from pymongo.auth_oidc import _authenticate_oidc, _get_authenticator, _OIDCProperties
+from pymongo.auth_oidc import (
+    _authenticate_oidc,
+    _get_authenticator,
+    _OIDCAWSCallback,
+    _OIDCProperties,
+)
 from pymongo.errors import ConfigurationError, OperationFailure
 from pymongo.saslprep import saslprep
 
@@ -164,7 +169,8 @@ def _build_credentials_tuple(
     elif mech == "MONGODB-OIDC":
         properties = extra.get("authmechanismproperties", {})
         request_token_callback = properties.get("request_token_callback")
-        provider_name = properties.get("PROVIDER_NAME", "")
+        custom_token_callback = properties.get("custom_token_callback")
+        provider_name = properties.get("PROVIDER_NAME")
         default_allowed = [
             "*.mongodb.net",
             "*.mongodb-dev.net",
@@ -175,12 +181,25 @@ def _build_credentials_tuple(
             "::1",
         ]
         allowed_hosts = properties.get("allowed_hosts", default_allowed)
-        if not request_token_callback and provider_name != "aws":
-            raise ConfigurationError(
-                "authentication with MONGODB-OIDC requires providing an request_token_callback or a provider_name of 'aws'"
-            )
+        msg = "authentication with MONGODB-OIDC requires providing a request_token_callback, a provider_name, or a custom_token_callback"
+        if request_token_callback is not None:
+            if provider_name is not None or custom_token_callback is not None:
+                raise ConfigurationError(msg)
+        elif provider_name is not None:
+            if custom_token_callback is not None:
+                raise ConfigurationError(msg)
+            if provider_name == "aws":
+                custom_token_callback = _OIDCAWSCallback()
+            else:
+                raise ConfigurationError(
+                    f"unrecognized provider_name for MONGODB-OIDC: {provider_name}"
+                )
+        elif custom_token_callback is None:
+            raise ConfigurationError(msg)
+
         oidc_props = _OIDCProperties(
             request_token_callback=request_token_callback,
+            custom_token_callback=custom_token_callback,
             provider_name=provider_name,
             allowed_hosts=allowed_hosts,
         )

pymongo/auth_oidc.py

@@ -15,32 +15,93 @@
 """MONGODB-OIDC Authentication helpers."""
 from __future__ import annotations
 
+import abc
+import os
 import threading
 from dataclasses import dataclass, field
-from typing import TYPE_CHECKING, Any, Callable, Mapping, MutableMapping, Optional
+from typing import TYPE_CHECKING, Any, Mapping, MutableMapping, Optional, Union
 
 import bson
 from bson.binary import Binary
 from bson.son import SON
+from pymongo._csot import remaining
 from pymongo.errors import ConfigurationError, OperationFailure
 
 if TYPE_CHECKING:
     from pymongo.auth import MongoCredential
     from pymongo.pool import Connection
 
 
+@dataclass
+class OIDCIdPInfo:
+    issuer: str
+    clientId: str
+    requestScopes: Optional[list[str]] = field(default=None)
+
+
+@dataclass
+class OIDCHumanCallbackContext:
+    timeout_seconds: float
+    version: int
+    refresh_token: Optional[str] = field(default=None)
+
+
+@dataclass
+class OIDCMachineCallbackContext:
+    timeout_seconds: float
+    version: int
+
+
+@dataclass
+class OIDCHumanCallbackResult:
+    access_token: str
+    expires_in_seconds: Optional[float] = field(default=None)
+    refresh_token: Optional[str] = field(default=None)
+
+
+@dataclass
+class OIDCMachineCallbackResult:
+    access_token: str
+    expires_in_seconds: Optional[float] = field(default=None)
+
+
+class OIDCMachineCallback(abc.ABC):
+    """A base class for defining OIDC machine (workload federation)
+    callbacks.
+    """
+
+    @abc.abstractmethod
+    def fetch(self, context: OIDCMachineCallbackContext) -> OIDCMachineCallbackResult:
+        """Convert the given BSON value into our own type."""
+
+
+class OIDCHumanCallback(abc.ABC):
+    """A base class for defining OIDC human (workforce federation)
+    callbacks.
+    """
+
+    @abc.abstractmethod
+    def fetch(
+        self, idp_info: OIDCIdPInfo, context: OIDCHumanCallbackContext
+    ) -> OIDCHumanCallbackResult:
+        """Convert the given BSON value into our own type."""
+
+
 @dataclass
 class _OIDCProperties:
-    request_token_callback: Optional[Callable[..., dict]]
-    provider_name: Optional[str]
-    allowed_hosts: list[str]
+    request_token_callback: Optional[OIDCHumanCallback] = field(default=None)
+    custom_token_callback: Optional[OIDCMachineCallback] = field(default=None)
+    provider_name: Optional[str] = field(default=None)
+    allowed_hosts: list[str] = field(default_factory=list)
 
 
 """Mechanism properties for MONGODB-OIDC authentication."""
 
 TOKEN_BUFFER_MINUTES = 5
-CALLBACK_TIMEOUT_SECONDS = 5 * 60
-CALLBACK_VERSION = 1
+HUMAN_CALLBACK_TIMEOUT_SECONDS = 5 * 60
+HUMAN_CALLBACK_VERSION = 1
+MACHINE_CALLBACK_TIMEOUT_SECONDS = 60
+MACHINE_CALLBACK_VERSION = 1
 
 
 def _get_authenticator(
@@ -72,22 +133,37 @@ def _get_authenticator(
     return credentials.cache.data
 
 
+class _OIDCAWSCallback(OIDCMachineCallback):
+    def fetch(self, context: OIDCMachineCallbackContext) -> OIDCMachineCallbackResult:
+        token_file = os.environ.get("AWS_WEB_IDENTITY_TOKEN_FILE")
+        if not token_file:
+            raise RuntimeError(
+                'MONGODB-OIDC with an "aws" provider requires "AWS_WEB_IDENTITY_TOKEN_FILE" to be set'
+            )
+        with open(token_file) as fid:
+            return OIDCMachineCallbackResult(access_token=fid.read().strip())
+
+
 @dataclass
 class _OIDCAuthenticator:
     username: str
     properties: _OIDCProperties
     refresh_token: Optional[str] = field(default=None)
     access_token: Optional[str] = field(default=None)
-    idp_info: Optional[dict] = field(default=None)
+    idp_info: Optional[OIDCIdPInfo] = field(default=None)
     token_gen_id: int = field(default=0)
     lock: threading.Lock = field(default_factory=threading.Lock)
 
     def get_current_token(self, use_callback: bool = True) -> Optional[str]:
         properties = self.properties
-
-        # TODO: DRIVERS-2672, handle machine callback here as well.
-        cb = properties.request_token_callback if use_callback else None
-        cb_type = "human"
+        cb: Union[None, OIDCHumanCallback, OIDCMachineCallback]
+        resp: Union[None, OIDCHumanCallbackResult, OIDCMachineCallbackResult]
+        if not use_callback:
+            cb = None
+        elif properties.request_token_callback:
+            cb = properties.request_token_callback
+        elif properties.custom_token_callback:
+            cb = properties.custom_token_callback
 
         prev_token = self.access_token
         if prev_token:
@@ -104,37 +180,32 @@ def get_current_token(self, use_callback: bool = True) -> Optional[str]:
                 if new_token != prev_token:
                     return new_token
 
-                # TODO: DRIVERS-2672 handle machine callback here.
-                if cb_type == "human":
-                    context = {
-                        "timeout_seconds": CALLBACK_TIMEOUT_SECONDS,
-                        "version": CALLBACK_VERSION,
-                        "refresh_token": self.refresh_token,
-                    }
-                    resp = cb(self.idp_info, context)
-
-                    self.validate_request_token_response(resp)
-
+                if isinstance(cb, OIDCHumanCallback):
+                    human_context = OIDCHumanCallbackContext(
+                        timeout_seconds=HUMAN_CALLBACK_TIMEOUT_SECONDS,
+                        version=HUMAN_CALLBACK_VERSION,
+                        refresh_token=self.refresh_token,
+                    )
+                    assert self.idp_info is not None
+                    resp = cb.fetch(self.idp_info, human_context)
+                    if not isinstance(resp, OIDCHumanCallbackResult):
+                        raise ValueError("Callback result must be of type OIDCHumanCallbackResult")
+                    self.refresh_token = resp.refresh_token
+                else:
+                    machine_context = OIDCMachineCallbackContext(
+                        timeout_seconds=remaining() or MACHINE_CALLBACK_TIMEOUT_SECONDS,
+                        version=MACHINE_CALLBACK_VERSION,
+                    )
+                    resp = cb.fetch(machine_context)
+                    if not isinstance(resp, OIDCMachineCallbackResult):
+                        raise ValueError(
+                            "Callback result must be of type OIDCMachineCallbackResult"
+                        )
+                self.access_token = resp.access_token
                 self.token_gen_id += 1
 
         return self.access_token
 
-    def validate_request_token_response(self, resp: Mapping[str, Any]) -> None:
-        # Validate callback return value.
-        if not isinstance(resp, dict):
-            raise ValueError("OIDC callback returned invalid result")
-
-        if "access_token" not in resp:
-            raise ValueError("OIDC callback did not return an access_token")
-
-        expected = ["access_token", "refresh_token", "expires_in_seconds"]
-        for key in resp:
-            if key not in expected:
-                raise ValueError(f'Unexpected field in callback result "{key}"')
-
-        self.access_token = resp["access_token"]
-        self.refresh_token = resp.get("refresh_token")
-
     def principal_step_cmd(self) -> SON[str, Any]:
         """Get a SASL start command with an optional principal name"""
         # Send the SASL start with the optional principal name.
@@ -154,8 +225,7 @@ def principal_step_cmd(self) -> SON[str, Any]:
         )
 
     def auth_start_cmd(self, use_callback: bool = True) -> Optional[SON[str, Any]]:
-        # TODO: DRIVERS-2672, check for provider_name in self.properties here.
-        if self.idp_info is None:
+        if self.properties.request_token_callback is not None and self.idp_info is None:
             return self.principal_step_cmd()
 
         token = self.get_current_token(use_callback)
@@ -192,8 +262,10 @@ def reauthenticate(self, conn: Connection) -> Optional[Mapping[str, Any]]:
 
         self.access_token = None
 
-        # TODO: DRIVERS-2672, check for provider_name in self.properties here.
-        # If so, we clear the access token and return finish_auth.
+        # If we are using machine callbacks, clear the access token and
+        # re-authenticate.
+        if self.properties.provider_name:
+            return self.authenticate(conn)
 
         # Next see if the idp info has changed.
         prev_idp_info = self.idp_info
@@ -203,7 +275,7 @@ def reauthenticate(self, conn: Connection) -> Optional[Mapping[str, Any]]:
         assert resp is not None
         server_resp: dict = bson.decode(resp["payload"])
         if "issuer" in server_resp:
-            self.idp_info = server_resp
+            self.idp_info = OIDCIdPInfo(**server_resp)
 
         # Handle the case of changed idp info.
         if self.idp_info != prev_idp_info:
@@ -240,7 +312,7 @@ def authenticate(self, conn: Connection) -> Optional[Mapping[str, Any]]:
 
         server_resp: dict = bson.decode(resp["payload"])
         if "issuer" in server_resp:
-            self.idp_info = server_resp
+            self.idp_info = OIDCIdPInfo(**server_resp)
 
         return self.finish_auth(resp, conn)
 
@@ -273,4 +345,10 @@ def _authenticate_oidc(
     if reauthenticate:
         return authenticator.reauthenticate(conn)
     else:
-        return authenticator.authenticate(conn)
+        try:
+            return authenticator.authenticate(conn)
+        except Exception as e:
+            # Try one more time an an authentication failure.
+            if isinstance(e, OperationFailure) and e.code == 18:
+                return authenticator.authenticate(conn)
+            raise

pymongo/common.py

@@ -17,7 +17,6 @@
 from __future__ import annotations
 
 import datetime
-import inspect
 import warnings
 from collections import OrderedDict, abc
 from typing import (
@@ -41,6 +40,7 @@
 from bson.codec_options import CodecOptions, DatetimeConversion, TypeRegistry
 from bson.raw_bson import RawBSONDocument
 from pymongo.auth import MECHANISMS
+from pymongo.auth_oidc import OIDCHumanCallback, OIDCMachineCallback
 from pymongo.compression_support import (
     validate_compressors,
     validate_zlib_compression_level,
@@ -438,20 +438,16 @@ def validate_auth_mechanism_properties(option: str, value: Any) -> dict[str, Uni
                 props[key] = str(value).lower()
             elif key in ["allowed_hosts"] and isinstance(value, list):
                 props[key] = value
-            elif inspect.isfunction(value):
-                signature = inspect.signature(value)
-                if key == "request_token_callback":
-                    expected_params = 2
-                else:
-                    raise ValueError(f"Unrecognized Auth mechanism function {key}")
-                if len(signature.parameters) != expected_params:
-                    msg = f"{key} must accept {expected_params} parameters"
-                    raise ValueError(msg)
+            elif key == "request_token_callback":
+                if not isinstance(value, OIDCHumanCallback):
+                    raise ValueError("request_token_callback must be a OIDCHumanCallback object")
+                props[key] = value
+            elif key == "custom_token_callback":
+                if not isinstance(value, OIDCMachineCallback):
+                    raise ValueError("custom_token_callback must be a OIDCMachineCallback object")
                 props[key] = value
             else:
-                raise ValueError(
-                    "Auth mechanism property values must be strings or callback functions"
-                )
+                raise ValueError(f"Invalid type for auth mechanism property {key}, {type(value)}")
         return props
 
     value = validate_string(option, value)

pyproject.toml

@@ -219,6 +219,7 @@ exclude_lines = [
     "return NotImplemented",
     "_use_c = true",
     "if __name__ == '__main__':",
+    "if TYPE_CHECKING:"
     ]
 partial_branches = ["if (.*and +)*not _use_c( and.*)*:"]