PYTHON-4256 Clean up handling of TOKEN_RESOURCE (#1620)

pymongo/_azure_helpers.py

@@ -32,7 +32,6 @@ def _get_azure_response(
         url += f"&client_id={client_id}"
     headers = {"Metadata": "true", "Accept": "application/json"}
     request = Request(url, headers=headers)  # noqa: S310
-    print("fetching url", url)  # noqa: T201
     try:
         with urlopen(request, timeout=timeout) as response:  # noqa: S310
             status = response.status

pymongo/auth.py

@@ -33,7 +33,7 @@
     Optional,
     cast,
 )
-from urllib.parse import quote, unquote
+from urllib.parse import quote
 
 from bson.binary import Binary
 from pymongo.auth_aws import _authenticate_aws
@@ -138,7 +138,7 @@ def _build_credentials_tuple(
             raise ValueError("authentication source must be $external or None for GSSAPI")
         properties = extra.get("authmechanismproperties", {})
         service_name = properties.get("SERVICE_NAME", "mongodb")
-        canonicalize = properties.get("CANONICALIZE_HOST_NAME", False)
+        canonicalize = bool(properties.get("CANONICALIZE_HOST_NAME", False))
         service_realm = properties.get("SERVICE_REALM")
         props = GSSAPIProperties(
             service_name=service_name,
@@ -173,8 +173,6 @@ def _build_credentials_tuple(
         human_callback = properties.get("OIDC_HUMAN_CALLBACK")
         environ = properties.get("ENVIRONMENT")
         token_resource = properties.get("TOKEN_RESOURCE", "")
-        if unquote(token_resource) == token_resource:
-            token_resource = quote(token_resource)
         default_allowed = [
             "*.mongodb.net",
             "*.mongodb-dev.net",
@@ -227,6 +225,7 @@ def _build_credentials_tuple(
             human_callback=human_callback,
             environment=environ,
             allowed_hosts=allowed_hosts,
+            token_resource=token_resource,
             username=user,
         )
         return MongoCredential(mech, "$external", user, passwd, oidc_props, _Cache())

pymongo/auth_oidc.py

@@ -21,6 +21,7 @@
 import time
 from dataclasses import dataclass, field
 from typing import TYPE_CHECKING, Any, Mapping, MutableMapping, Optional, Union
+from urllib.parse import quote
 
 import bson
 from bson.binary import Binary
@@ -72,6 +73,7 @@ class _OIDCProperties:
     human_callback: Optional[OIDCCallback] = field(default=None)
     environment: Optional[str] = field(default=None)
     allowed_hosts: list[str] = field(default_factory=list)
+    token_resource: Optional[str] = field(default=None)
     username: str = ""
 
 
@@ -126,7 +128,7 @@ def fetch(self, context: OIDCCallbackContext) -> OIDCCallbackResult:
 
 class _OIDCAzureCallback(OIDCCallback):
     def __init__(self, token_resource: str) -> None:
-        self.token_resource = token_resource
+        self.token_resource = quote(token_resource)
 
     def fetch(self, context: OIDCCallbackContext) -> OIDCCallbackResult:
         resp = _get_azure_response(self.token_resource, context.username, context.timeout_seconds)
@@ -137,7 +139,7 @@ def fetch(self, context: OIDCCallbackContext) -> OIDCCallbackResult:
 
 class _OIDCGCPCallback(OIDCCallback):
     def __init__(self, token_resource: str) -> None:
-        self.token_resource = token_resource
+        self.token_resource = quote(token_resource)
 
     def fetch(self, context: OIDCCallbackContext) -> OIDCCallbackResult:
         resp = _get_gcp_response(self.token_resource, context.timeout_seconds)

pymongo/common.py

@@ -453,26 +453,21 @@ def validate_auth_mechanism_properties(option: str, value: Any) -> dict[str, Uni
 
     value = validate_string(option, value)
     for opt in value.split(","):
-        try:
-            key, val = opt.split(":")
-        except ValueError:
+        key, _, val = opt.partition(":")
+        if key not in _MECHANISM_PROPS:
             # Try not to leak the token.
-            if "AWS_SESSION_TOKEN" in opt:
-                opt = (  # noqa: PLW2901
-                    "AWS_SESSION_TOKEN:<redacted token>, did you forget "
-                    "to percent-escape the token with quote_plus?"
+            if "AWS_SESSION_TOKEN" in key:
+                raise ValueError(
+                    "auth mechanism properties must be "
+                    "key:value pairs like AWS_SESSION_TOKEN:<token>"
                 )
-            raise ValueError(
-                "auth mechanism properties must be "
-                "key:value pairs like SERVICE_NAME:"
-                f"mongodb, not {opt}."
-            ) from None
-        if key not in _MECHANISM_PROPS:
+
             raise ValueError(
                 f"{key} is not a supported auth "
                 "mechanism property. Must be one of "
                 f"{tuple(_MECHANISM_PROPS)}."
             )
+
         if key == "CANONICALIZE_HOST_NAME":
             props[key] = validate_boolean_or_string(key, val)
         else:

test/auth/legacy/connection-string.json

@@ -474,7 +474,7 @@
       }
     },
     {
-      "description": "should  throw an exception if username and password is specified for test environment (MONGODB-OIDC)",
+      "description": "should throw an exception if username and password is specified for test environment (MONGODB-OIDC)",
       "uri": "mongodb://user:pass@localhost/?authMechanism=MONGODB-OIDC&authMechanismProperties=ENVIRONMENT:test",
       "valid": false,
       "credential": null
@@ -486,23 +486,11 @@
       "credential": null
     },
     {
-      "description": "should throw an exception if specified provider is not supported (MONGODB-OIDC)",
+      "description": "should throw an exception if specified environment is not supported (MONGODB-OIDC)",
       "uri": "mongodb://localhost/?authMechanism=MONGODB-OIDC&authMechanismProperties=ENVIRONMENT:invalid",
       "valid": false,
       "credential": null
     },
-    {
-      "description": "should throw an exception custom callback is chosen but no callback is provided (MONGODB-OIDC)",
-      "uri": "mongodb://localhost/?authMechanism=MONGODB-OIDC&authMechanismProperties=ENVIRONMENT:custom",
-      "valid": false,
-      "credential": null
-    },
-    {
-      "description": "should throw an exception custom callback is chosen but no callback is provided (MONGODB-OIDC)",
-      "uri": "mongodb://localhost/?authMechanism=MONGODB-OIDC&authMechanismProperties=PROVIDER_NAME:custom",
-      "valid": false,
-      "credential": null
-    },
     {
       "description": "should throw an exception if neither provider nor callbacks specified (MONGODB-OIDC)",
       "uri": "mongodb://localhost/?authMechanism=MONGODB-OIDC",
@@ -541,7 +529,37 @@
     },
     {
       "description": "should accept a url-encoded TOKEN_RESOURCE (MONGODB-OIDC)",
-      "uri": "mongodb://user@localhost/?authMechanism=MONGODB-OIDC&authMechanismProperties=ENVIRONMENT:azure,TOKEN_RESOURCE:mongodb%253A//test-cluster",
+      "uri": "mongodb://user@localhost/?authMechanism=MONGODB-OIDC&authMechanismProperties=ENVIRONMENT:azure,TOKEN_RESOURCE:mongodb%3A%2F%2Ftest-cluster",
+      "valid": true,
+      "credential": {
+        "username": "user",
+        "password": null,
+        "source": "$external",
+        "mechanism": "MONGODB-OIDC",
+        "mechanism_properties": {
+          "ENVIRONMENT": "azure",
+          "TOKEN_RESOURCE": "mongodb://test-cluster"
+        }
+      }
+    },
+    {
+      "description": "should accept an un-encoded TOKEN_RESOURCE (MONGODB-OIDC)",
+      "uri": "mongodb://user@localhost/?authMechanism=MONGODB-OIDC&authMechanismProperties=ENVIRONMENT:azure,TOKEN_RESOURCE:mongodb://test-cluster",
+      "valid": true,
+      "credential": {
+        "username": "user",
+        "password": null,
+        "source": "$external",
+        "mechanism": "MONGODB-OIDC",
+        "mechanism_properties": {
+          "ENVIRONMENT": "azure",
+          "TOKEN_RESOURCE": "mongodb://test-cluster"
+        }
+      }
+    },
+    {
+      "description": "should handle a complicated url-encoded TOKEN_RESOURCE (MONGODB-OIDC)",
+      "uri": "mongodb://user@localhost/?authMechanism=MONGODB-OIDC&authMechanismProperties=ENVIRONMENT:azure,TOKEN_RESOURCE:abc%2Cd%25ef%3Ag%26hi",
       "valid": true,
       "credential": {
         "username": "user",
@@ -550,7 +568,7 @@
         "mechanism": "MONGODB-OIDC",
         "mechanism_properties": {
           "ENVIRONMENT": "azure",
-          "TOKEN_RESOURCE": "mongodb%253A//test-cluster"
+          "TOKEN_RESOURCE": "abc,d%ef:g&hi"
         }
       }
     },
@@ -565,7 +583,7 @@
         "mechanism": "MONGODB-OIDC",
         "mechanism_properties": {
           "ENVIRONMENT": "azure",
-          "TOKEN_RESOURCE": "a%24b"
+          "TOKEN_RESOURCE": "a$b"
         }
       }
     },

test/test_auth_spec.py

@@ -52,12 +52,7 @@ def run_test(self):
                 warnings.simplefilter("default")
                 self.assertRaises(Exception, MongoClient, uri, connect=False)
         else:
-            props = {}
-            if credential:
-                props = credential["mechanism_properties"] or {}
-                if props.get("CALLBACK"):
-                    props["callback"] = SampleHumanCallback()
-            client = MongoClient(uri, connect=False, authmechanismproperties=props)
+            client = MongoClient(uri, connect=False)
             credentials = client.options.pool_options._credentials
             if credential is None:
                 self.assertIsNone(credentials)
@@ -73,25 +68,8 @@ def run_test(self):
                 expected = credential["mechanism_properties"]
                 if expected is not None:
                     actual = credentials.mechanism_properties
-                    for key, _val in expected.items():
-                        if "SERVICE_NAME" in expected:
-                            self.assertEqual(actual.service_name, expected["SERVICE_NAME"])
-                        elif "CANONICALIZE_HOST_NAME" in expected:
-                            self.assertEqual(
-                                actual.canonicalize_host_name, expected["CANONICALIZE_HOST_NAME"]
-                            )
-                        elif "SERVICE_REALM" in expected:
-                            self.assertEqual(actual.service_realm, expected["SERVICE_REALM"])
-                        elif "AWS_SESSION_TOKEN" in expected:
-                            self.assertEqual(
-                                actual.aws_session_token, expected["AWS_SESSION_TOKEN"]
-                            )
-                        elif "ENVIRONMENT" in expected:
-                            self.assertEqual(actual.environment, expected["ENVIRONMENT"])
-                        elif "callback" in expected:
-                            self.assertEqual(actual.callback, expected["callback"])
-                        else:
-                            self.fail(f"Unhandled property: {key}")
+                    for key, value in expected.items():
+                        self.assertEqual(getattr(actual, key.lower()), value)
                 else:
                     if credential["mechanism"] == "MONGODB-AWS":
                         self.assertIsNone(credentials.mechanism_properties.aws_session_token)

test/test_uri_parser.py

@@ -504,20 +504,30 @@ def test_unquote_after_parsing(self):
         self.assertEqual(options, res["options"])
 
     def test_redact_AWS_SESSION_TOKEN(self):
-        unquoted_colon = "token:"
+        token = "token"
         uri = (
             "mongodb://user:password@localhost/?authMechanism=MONGODB-AWS"
-            "&authMechanismProperties=AWS_SESSION_TOKEN:" + unquoted_colon
+            "&authMechanismProperties=AWS_SESSION_TOKEN-" + token
         )
         with self.assertRaisesRegex(
             ValueError,
-            "auth mechanism properties must be key:value pairs like "
-            "SERVICE_NAME:mongodb, not AWS_SESSION_TOKEN:<redacted token>"
-            ", did you forget to percent-escape the token with "
-            "quote_plus?",
+            "auth mechanism properties must be key:value pairs like AWS_SESSION_TOKEN:<token>",
         ):
             parse_uri(uri)
 
+    def test_handle_colon(self):
+        token = "token:foo"
+        uri = (
+            "mongodb://user:password@localhost/?authMechanism=MONGODB-AWS"
+            "&authMechanismProperties=AWS_SESSION_TOKEN:" + token
+        )
+        res = parse_uri(uri)
+        options = {
+            "authmechanism": "MONGODB-AWS",
+            "authMechanismProperties": {"AWS_SESSION_TOKEN": token},
+        }
+        self.assertEqual(options, res["options"])
+
     def test_special_chars(self):
         user = "user@ /9+:?~!$&'()*+,;="
         pwd = "pwd@ /9+:?~!$&'()*+,;="