-
Notifications
You must be signed in to change notification settings - Fork 211
Added withFmi method for cca app #876
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: dev
Are you sure you want to change the base?
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -15,7 +15,7 @@ | |
| from .mex import send_request as mex_send_request | ||
| from .wstrust_request import send_request as wst_send_request | ||
| from .wstrust_response import * | ||
| from .token_cache import TokenCache, _get_username, _GRANT_TYPE_BROKER | ||
| from .token_cache import TokenCache, _get_username, _GRANT_TYPE_BROKER, _compute_ext_cache_key | ||
| import msal.telemetry | ||
| from .region import _detect_region | ||
| from .throttled_http_client import ThrottledHttpClient | ||
|
|
@@ -1571,6 +1571,9 @@ def _acquire_token_silent_from_cache_and_possibly_refresh_it( | |
| key_id = kwargs.get("data", {}).get("key_id") | ||
| if key_id: # Some token types (SSH-certs, POP) are bound to a key | ||
| query["key_id"] = key_id | ||
| ext_cache_key = _compute_ext_cache_key(kwargs.get("data", {})) | ||
| if ext_cache_key: # FMI tokens need cache isolation by path | ||
| query["ext_cache_key"] = ext_cache_key | ||
| now = time.time() | ||
| refresh_reason = msal.telemetry.AT_ABSENT | ||
| for entry in self.token_cache.search( # A generator allows us to | ||
|
|
@@ -2491,6 +2494,32 @@ def remove_tokens_for_client(self): | |
| self.token_cache.remove_at(at) | ||
| # acquire_token_for_client() obtains no RTs, so we have no RT to remove | ||
|
|
||
| def acquire_token_for_client_with_fmi_path(self, scopes, fmi_path, claims_challenge=None, **kwargs): | ||
| """Acquires token for the current confidential client with a Federated Managed Identity (FMI) path. | ||
|
|
||
| This is a convenience wrapper around :func:`~acquire_token_for_client` | ||
| that attaches the ``fmi_path`` parameter to the token request body. | ||
|
|
||
| :param list[str] scopes: (Required) | ||
| Scopes requested to access a protected API (a resource). | ||
| :param str fmi_path: (Required) | ||
| The Federated Managed Identity path to attach to the request. | ||
| :param claims_challenge: | ||
| The claims_challenge parameter requests specific claims requested by the resource provider | ||
| in the form of a claims_challenge directive in the www-authenticate header to be | ||
| returned from the UserInfo Endpoint and/or in the ID Token and/or Access Token. | ||
| It is a string of a JSON object which contains lists of claims being requested from these locations. | ||
|
|
||
| :return: A dict representing the json response from Microsoft Entra: | ||
|
|
||
| - A successful response would contain "access_token" key, | ||
| - an error response would contain "error" and usually "error_description". | ||
| """ | ||
| data = kwargs.pop("data", {}) | ||
| data["fmi_path"] = fmi_path | ||
|
Member
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Are the tokens cached like in MSAL .NET and MSAL GO? I don't think they are. I see a PR here that look at proper caching #759 but I am not sure it is compliant with the rest.
Contributor
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Added a cache enhancement so now we create a hash key for some data keyvalues |
||
| return self.acquire_token_for_client( | ||
| scopes, claims_challenge=claims_challenge, data=data, **kwargs) | ||
|
|
||
| def acquire_token_on_behalf_of(self, user_assertion, scopes, claims_challenge=None, **kwargs): | ||
| """Acquires token using on-behalf-of (OBO) flow. | ||
|
|
||
|
|
||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,4 +1,6 @@ | ||
| import json | ||
| import base64 | ||
| import hashlib | ||
| import json | ||
| import threading | ||
| import time | ||
| import logging | ||
|
|
@@ -12,6 +14,63 @@ | |
| logger = logging.getLogger(__name__) | ||
| _GRANT_TYPE_BROKER = "broker" | ||
|
|
||
| # Fields in the request data dict that should NOT be included in the extended | ||
|
Member
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. The .NET spec mentioned using "atext" in the cache key. Maybe we are missing a test that attempts to load a token from a file. |
||
| # cache key hash. Everything else in data IS included, because those are extra | ||
| # body parameters going on the wire and must differentiate cached tokens. | ||
| # | ||
| # Excluded fields and reasons: | ||
| # - "key_id" : Already handled as a separate cache lookup field | ||
| # - "token_type" : Used for SSH-cert/POP detection; AT entry stores it separately | ||
| # - "req_cnf" : Ephemeral proof-of-possession nonce, changes per request | ||
| # - "claims" : Handled separately; its presence forces a token refresh | ||
| # - "scope" : Already represented as "target" in the AT cache key; | ||
| # also added to data only at wire-time, not at cache-lookup time | ||
| # - "username" : Standard ROPC grant parameter, not an extra body parameter | ||
| # - "password" : Standard ROPC grant parameter, not an extra body parameter | ||
| # | ||
| # Included fields (examples — anything NOT in this set is included): | ||
| # - "fmi_path" : Federated Managed Identity credential path | ||
| # - any future extra body parameter that should isolate cache entries | ||
| _EXT_CACHE_KEY_EXCLUDED_FIELDS = frozenset({ | ||
|
Member
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Pls add a comment that these fields are already taken care of separately. |
||
| "key_id", | ||
| "token_type", | ||
| "req_cnf", | ||
| "claims", | ||
|
Member
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. I think you are missing a few things:
We probably need to enhance existing tests to assert that no "unknown" key is used. |
||
| "scope", | ||
| "username", | ||
| "password", | ||
| }) | ||
|
|
||
|
|
||
| def _compute_ext_cache_key(data): | ||
| """Compute an extended cache key hash from extra body parameters in *data*. | ||
|
|
||
| All fields in *data* that go on the wire are included in the hash, | ||
| EXCEPT those listed in ``_EXT_CACHE_KEY_EXCLUDED_FIELDS``. | ||
| This ensures tokens acquired with different parameter values | ||
| (e.g., different FMI paths) are cached separately. | ||
|
|
||
| Returns an empty string when *data* has no hashable fields. | ||
|
|
||
| The algorithm matches the Go MSAL implementation (CacheExtKeyGenerator): | ||
| sorted key+value pairs are concatenated and SHA256 hashed, then base64url encoded. | ||
| """ | ||
| if not data: | ||
| return "" | ||
| cache_components = { | ||
| k: str(v) for k, v in data.items() | ||
| if k not in _EXT_CACHE_KEY_EXCLUDED_FIELDS and v | ||
| } | ||
| if not cache_components: | ||
| return "" | ||
| # Sort keys for consistent hashing (matches Go implementation) | ||
| key_str = "".join( | ||
| k + cache_components[k] for k in sorted(cache_components.keys()) | ||
| ) | ||
| hash_bytes = hashlib.sha256(key_str.encode("utf-8")).digest() | ||
| return base64.urlsafe_b64encode(hash_bytes).rstrip(b"=").decode("ascii").lower() | ||
|
|
||
|
|
||
| def is_subdict_of(small, big): | ||
| return dict(big, **small) == big | ||
|
|
||
|
|
@@ -59,6 +118,7 @@ def __init__(self): | |
| self.CredentialType.ACCESS_TOKEN: | ||
| lambda home_account_id=None, environment=None, client_id=None, | ||
| realm=None, target=None, | ||
| ext_cache_key=None, | ||
| # Note: New field(s) can be added here | ||
| #key_id=None, | ||
| **ignored_payload_from_a_real_token: | ||
|
|
@@ -70,7 +130,8 @@ def __init__(self): | |
| realm or "", | ||
| target or "", | ||
| #key_id or "", # So ATs of different key_id can coexist | ||
| ]).lower(), | ||
| ] + ([ext_cache_key] if ext_cache_key else []) | ||
| ).lower(), | ||
| self.CredentialType.ID_TOKEN: | ||
| lambda home_account_id=None, environment=None, client_id=None, | ||
| realm=None, **ignored_payload_from_a_real_token: | ||
|
|
@@ -98,6 +159,7 @@ def __init__(self): | |
| def _get_access_token( | ||
| self, | ||
| home_account_id, environment, client_id, realm, target, # Together they form a compound key | ||
| ext_cache_key=None, | ||
| default=None, | ||
| ): # O(1) | ||
| return self._get( | ||
|
|
@@ -108,6 +170,7 @@ def _get_access_token( | |
| client_id=client_id, | ||
| realm=realm, | ||
| target=" ".join(target), | ||
| ext_cache_key=ext_cache_key, | ||
| ), | ||
| default=default) | ||
|
|
||
|
|
@@ -153,7 +216,8 @@ def search(self, credential_type, target=None, query=None, *, now=None): # O(n) | |
| ): # Special case for O(1) AT lookup | ||
| preferred_result = self._get_access_token( | ||
| query["home_account_id"], query["environment"], | ||
| query["client_id"], query["realm"], target) | ||
| query["client_id"], query["realm"], target, | ||
| ext_cache_key=query.get("ext_cache_key")) | ||
| if preferred_result and self._is_matching( | ||
| preferred_result, query, | ||
| # Needs no target_set here because it is satisfied by dict key | ||
|
|
@@ -179,6 +243,13 @@ def search(self, credential_type, target=None, query=None, *, now=None): # O(n) | |
| if (entry != preferred_result # Avoid yielding the same entry twice | ||
| and self._is_matching(entry, query, target_set=target_set) | ||
| ): | ||
| # Cache isolation for extended cache keys (e.g., FMI path). | ||
| # Entries with ext_cache_key must not match queries without one. | ||
| if (credential_type == self.CredentialType.ACCESS_TOKEN | ||
| and "ext_cache_key" in entry | ||
| and "ext_cache_key" not in (query or {}) | ||
| ): | ||
| continue | ||
| yield entry | ||
| for at in expired_access_tokens: | ||
| self.remove_at(at) | ||
|
|
@@ -278,6 +349,12 @@ def __add(self, event, now=None): | |
| # So that we won't accidentally store a user's password etc. | ||
| "key_id", # It happens in SSH-cert or POP scenario | ||
| }}) | ||
| # Compute and store extended cache key for cache isolation | ||
| # (e.g., different FMI paths should have separate cache entries) | ||
| ext_cache_key = _compute_ext_cache_key(data) | ||
|
|
||
| if ext_cache_key: | ||
| at["ext_cache_key"] = ext_cache_key | ||
| if "refresh_in" in response: | ||
| refresh_in = response["refresh_in"] # It is an integer | ||
| at["refresh_on"] = str(now + refresh_in) # Schema wants a string | ||
|
|
||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Should we have FMI path as one of the kwargs instead? Creating an entire new API seems pretty heavy.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This comment still stands. Should we "an overload" of AcquireTokenForClient existing API?