Skip to content

Commit 6584dd2

Browse files
blink1073blink1073
authored
PYTHON-4256 Clean up handling of TOKEN_RESOURCE (#1620)
1 parent b83fd99 commit 6584dd2

File tree

7 files changed

+69
-68
lines changed

7 files changed

+69
-68
lines changed

‎pymongo/_azure_helpers.py

-1
Original file line numberDiff line numberDiff line change
@@ -32,7 +32,6 @@ def _get_azure_response(
3232
url += f"&client_id={client_id}"
3333
headers = {"Metadata": "true", "Accept": "application/json"}
3434
request = Request(url, headers=headers) # noqa: S310
35-
print("fetching url", url) # noqa: T201
3635
try:
3736
with urlopen(request, timeout=timeout) as response: # noqa: S310
3837
status = response.status

‎pymongo/auth.py

+3-4
Original file line numberDiff line numberDiff line change
@@ -33,7 +33,7 @@
3333
Optional,
3434
cast,
3535
)
36-
from urllib.parse import quote, unquote
36+
from urllib.parse import quote
3737

3838
from bson.binary import Binary
3939
from pymongo.auth_aws import _authenticate_aws
@@ -138,7 +138,7 @@ def _build_credentials_tuple(
138138
raise ValueError("authentication source must be $external or None for GSSAPI")
139139
properties = extra.get("authmechanismproperties", {})
140140
service_name = properties.get("SERVICE_NAME", "mongodb")
141-
canonicalize = properties.get("CANONICALIZE_HOST_NAME", False)
141+
canonicalize = bool(properties.get("CANONICALIZE_HOST_NAME", False))
142142
service_realm = properties.get("SERVICE_REALM")
143143
props = GSSAPIProperties(
144144
service_name=service_name,
@@ -173,8 +173,6 @@ def _build_credentials_tuple(
173173
human_callback = properties.get("OIDC_HUMAN_CALLBACK")
174174
environ = properties.get("ENVIRONMENT")
175175
token_resource = properties.get("TOKEN_RESOURCE", "")
176-
if unquote(token_resource) == token_resource:
177-
token_resource = quote(token_resource)
178176
default_allowed = [
179177
"*.mongodb.net",
180178
"*.mongodb-dev.net",
@@ -227,6 +225,7 @@ def _build_credentials_tuple(
227225
human_callback=human_callback,
228226
environment=environ,
229227
allowed_hosts=allowed_hosts,
228+
token_resource=token_resource,
230229
username=user,
231230
)
232231
return MongoCredential(mech, "$external", user, passwd, oidc_props, _Cache())

‎pymongo/auth_oidc.py

+4-2
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
import time
2222
from dataclasses import dataclass, field
2323
from typing import TYPE_CHECKING, Any, Mapping, MutableMapping, Optional, Union
24+
from urllib.parse import quote
2425

2526
import bson
2627
from bson.binary import Binary
@@ -72,6 +73,7 @@ class _OIDCProperties:
7273
human_callback: Optional[OIDCCallback] = field(default=None)
7374
environment: Optional[str] = field(default=None)
7475
allowed_hosts: list[str] = field(default_factory=list)
76+
token_resource: Optional[str] = field(default=None)
7577
username: str = ""
7678

7779

@@ -126,7 +128,7 @@ def fetch(self, context: OIDCCallbackContext) -> OIDCCallbackResult:
126128

127129
class _OIDCAzureCallback(OIDCCallback):
128130
def __init__(self, token_resource: str) -> None:
129-
self.token_resource = token_resource
131+
self.token_resource = quote(token_resource)
130132

131133
def fetch(self, context: OIDCCallbackContext) -> OIDCCallbackResult:
132134
resp = _get_azure_response(self.token_resource, context.username, context.timeout_seconds)
@@ -137,7 +139,7 @@ def fetch(self, context: OIDCCallbackContext) -> OIDCCallbackResult:
137139

138140
class _OIDCGCPCallback(OIDCCallback):
139141
def __init__(self, token_resource: str) -> None:
140-
self.token_resource = token_resource
142+
self.token_resource = quote(token_resource)
141143

142144
def fetch(self, context: OIDCCallbackContext) -> OIDCCallbackResult:
143145
resp = _get_gcp_response(self.token_resource, context.timeout_seconds)

‎pymongo/common.py

+8-13
Original file line numberDiff line numberDiff line change
@@ -453,26 +453,21 @@ def validate_auth_mechanism_properties(option: str, value: Any) -> dict[str, Uni
453453

454454
value = validate_string(option, value)
455455
for opt in value.split(","):
456-
try:
457-
key, val = opt.split(":")
458-
except ValueError:
456+
key, _, val = opt.partition(":")
457+
if key not in _MECHANISM_PROPS:
459458
# Try not to leak the token.
460-
if "AWS_SESSION_TOKEN" in opt:
461-
opt = ( # noqa: PLW2901
462-
"AWS_SESSION_TOKEN:<redacted token>, did you forget "
463-
"to percent-escape the token with quote_plus?"
459+
if "AWS_SESSION_TOKEN" in key:
460+
raise ValueError(
461+
"auth mechanism properties must be "
462+
"key:value pairs like AWS_SESSION_TOKEN:<token>"
464463
)
465-
raise ValueError(
466-
"auth mechanism properties must be "
467-
"key:value pairs like SERVICE_NAME:"
468-
f"mongodb, not {opt}."
469-
) from None
470-
if key not in _MECHANISM_PROPS:
464+
471465
raise ValueError(
472466
f"{key} is not a supported auth "
473467
"mechanism property. Must be one of "
474468
f"{tuple(_MECHANISM_PROPS)}."
475469
)
470+
476471
if key == "CANONICALIZE_HOST_NAME":
477472
props[key] = validate_boolean_or_string(key, val)
478473
else:

‎test/auth/legacy/connection-string.json

+35-17
Original file line numberDiff line numberDiff line change
@@ -474,7 +474,7 @@
474474
}
475475
},
476476
{
477-
"description": "should throw an exception if username and password is specified for test environment (MONGODB-OIDC)",
477+
"description": "should throw an exception if username and password is specified for test environment (MONGODB-OIDC)",
478478
"uri": "mongodb://user:pass@localhost/?authMechanism=MONGODB-OIDC&authMechanismProperties=ENVIRONMENT:test",
479479
"valid": false,
480480
"credential": null
@@ -486,23 +486,11 @@
486486
"credential": null
487487
},
488488
{
489-
"description": "should throw an exception if specified provider is not supported (MONGODB-OIDC)",
489+
"description": "should throw an exception if specified environment is not supported (MONGODB-OIDC)",
490490
"uri": "mongodb://localhost/?authMechanism=MONGODB-OIDC&authMechanismProperties=ENVIRONMENT:invalid",
491491
"valid": false,
492492
"credential": null
493493
},
494-
{
495-
"description": "should throw an exception custom callback is chosen but no callback is provided (MONGODB-OIDC)",
496-
"uri": "mongodb://localhost/?authMechanism=MONGODB-OIDC&authMechanismProperties=ENVIRONMENT:custom",
497-
"valid": false,
498-
"credential": null
499-
},
500-
{
501-
"description": "should throw an exception custom callback is chosen but no callback is provided (MONGODB-OIDC)",
502-
"uri": "mongodb://localhost/?authMechanism=MONGODB-OIDC&authMechanismProperties=PROVIDER_NAME:custom",
503-
"valid": false,
504-
"credential": null
505-
},
506494
{
507495
"description": "should throw an exception if neither provider nor callbacks specified (MONGODB-OIDC)",
508496
"uri": "mongodb://localhost/?authMechanism=MONGODB-OIDC",
@@ -541,7 +529,37 @@
541529
},
542530
{
543531
"description": "should accept a url-encoded TOKEN_RESOURCE (MONGODB-OIDC)",
544-
"uri": "mongodb://user@localhost/?authMechanism=MONGODB-OIDC&authMechanismProperties=ENVIRONMENT:azure,TOKEN_RESOURCE:mongodb%253A//test-cluster",
532+
"uri": "mongodb://user@localhost/?authMechanism=MONGODB-OIDC&authMechanismProperties=ENVIRONMENT:azure,TOKEN_RESOURCE:mongodb%3A%2F%2Ftest-cluster",
533+
"valid": true,
534+
"credential": {
535+
"username": "user",
536+
"password": null,
537+
"source": "$external",
538+
"mechanism": "MONGODB-OIDC",
539+
"mechanism_properties": {
540+
"ENVIRONMENT": "azure",
541+
"TOKEN_RESOURCE": "mongodb://test-cluster"
542+
}
543+
}
544+
},
545+
{
546+
"description": "should accept an un-encoded TOKEN_RESOURCE (MONGODB-OIDC)",
547+
"uri": "mongodb://user@localhost/?authMechanism=MONGODB-OIDC&authMechanismProperties=ENVIRONMENT:azure,TOKEN_RESOURCE:mongodb://test-cluster",
548+
"valid": true,
549+
"credential": {
550+
"username": "user",
551+
"password": null,
552+
"source": "$external",
553+
"mechanism": "MONGODB-OIDC",
554+
"mechanism_properties": {
555+
"ENVIRONMENT": "azure",
556+
"TOKEN_RESOURCE": "mongodb://test-cluster"
557+
}
558+
}
559+
},
560+
{
561+
"description": "should handle a complicated url-encoded TOKEN_RESOURCE (MONGODB-OIDC)",
562+
"uri": "mongodb://user@localhost/?authMechanism=MONGODB-OIDC&authMechanismProperties=ENVIRONMENT:azure,TOKEN_RESOURCE:abc%2Cd%25ef%3Ag%26hi",
545563
"valid": true,
546564
"credential": {
547565
"username": "user",
@@ -550,7 +568,7 @@
550568
"mechanism": "MONGODB-OIDC",
551569
"mechanism_properties": {
552570
"ENVIRONMENT": "azure",
553-
"TOKEN_RESOURCE": "mongodb%253A//test-cluster"
571+
"TOKEN_RESOURCE": "abc,d%ef:g&hi"
554572
}
555573
}
556574
},
@@ -565,7 +583,7 @@
565583
"mechanism": "MONGODB-OIDC",
566584
"mechanism_properties": {
567585
"ENVIRONMENT": "azure",
568-
"TOKEN_RESOURCE": "a%24b"
586+
"TOKEN_RESOURCE": "a$b"
569587
}
570588
}
571589
},

‎test/test_auth_spec.py

+3-25
Original file line numberDiff line numberDiff line change
@@ -52,12 +52,7 @@ def run_test(self):
5252
warnings.simplefilter("default")
5353
self.assertRaises(Exception, MongoClient, uri, connect=False)
5454
else:
55-
props = {}
56-
if credential:
57-
props = credential["mechanism_properties"] or {}
58-
if props.get("CALLBACK"):
59-
props["callback"] = SampleHumanCallback()
60-
client = MongoClient(uri, connect=False, authmechanismproperties=props)
55+
client = MongoClient(uri, connect=False)
6156
credentials = client.options.pool_options._credentials
6257
if credential is None:
6358
self.assertIsNone(credentials)
@@ -73,25 +68,8 @@ def run_test(self):
7368
expected = credential["mechanism_properties"]
7469
if expected is not None:
7570
actual = credentials.mechanism_properties
76-
for key, _val in expected.items():
77-
if "SERVICE_NAME" in expected:
78-
self.assertEqual(actual.service_name, expected["SERVICE_NAME"])
79-
elif "CANONICALIZE_HOST_NAME" in expected:
80-
self.assertEqual(
81-
actual.canonicalize_host_name, expected["CANONICALIZE_HOST_NAME"]
82-
)
83-
elif "SERVICE_REALM" in expected:
84-
self.assertEqual(actual.service_realm, expected["SERVICE_REALM"])
85-
elif "AWS_SESSION_TOKEN" in expected:
86-
self.assertEqual(
87-
actual.aws_session_token, expected["AWS_SESSION_TOKEN"]
88-
)
89-
elif "ENVIRONMENT" in expected:
90-
self.assertEqual(actual.environment, expected["ENVIRONMENT"])
91-
elif "callback" in expected:
92-
self.assertEqual(actual.callback, expected["callback"])
93-
else:
94-
self.fail(f"Unhandled property: {key}")
71+
for key, value in expected.items():
72+
self.assertEqual(getattr(actual, key.lower()), value)
9573
else:
9674
if credential["mechanism"] == "MONGODB-AWS":
9775
self.assertIsNone(credentials.mechanism_properties.aws_session_token)

‎test/test_uri_parser.py

+16-6
Original file line numberDiff line numberDiff line change
@@ -504,20 +504,30 @@ def test_unquote_after_parsing(self):
504504
self.assertEqual(options, res["options"])
505505

506506
def test_redact_AWS_SESSION_TOKEN(self):
507-
unquoted_colon = "token:"
507+
token = "token"
508508
uri = (
509509
"mongodb://user:password@localhost/?authMechanism=MONGODB-AWS"
510-
"&authMechanismProperties=AWS_SESSION_TOKEN:" + unquoted_colon
510+
"&authMechanismProperties=AWS_SESSION_TOKEN-" + token
511511
)
512512
with self.assertRaisesRegex(
513513
ValueError,
514-
"auth mechanism properties must be key:value pairs like "
515-
"SERVICE_NAME:mongodb, not AWS_SESSION_TOKEN:<redacted token>"
516-
", did you forget to percent-escape the token with "
517-
"quote_plus?",
514+
"auth mechanism properties must be key:value pairs like AWS_SESSION_TOKEN:<token>",
518515
):
519516
parse_uri(uri)
520517

518+
def test_handle_colon(self):
519+
token = "token:foo"
520+
uri = (
521+
"mongodb://user:password@localhost/?authMechanism=MONGODB-AWS"
522+
"&authMechanismProperties=AWS_SESSION_TOKEN:" + token
523+
)
524+
res = parse_uri(uri)
525+
options = {
526+
"authmechanism": "MONGODB-AWS",
527+
"authMechanismProperties": {"AWS_SESSION_TOKEN": token},
528+
}
529+
self.assertEqual(options, res["options"])
530+
521531
def test_special_chars(self):
522532
user = "user@ /9+:?~!$&'()*+,;="
523533
pwd = "pwd@ /9+:?~!$&'()*+,;="

0 commit comments

Comments
 (0)