From 9660190555f684e8369afc6f24b3578bea6776bc Mon Sep 17 00:00:00 2001 From: Michal Hofman Date: Wed, 15 Jul 2026 14:23:42 +0200 Subject: [PATCH 1/6] feat: implement token cache key v2 with uniform SHA-256 hashing Extends TokenKey to five named fields (token_type, idp, snowflake, username, role) and introduces normalize_url/normalize_identifier helpers plus build_cache_key, which produces a stable SnowflakeTokenCache.v2. key used by both KeyringTokenCache and FileTokenCache. Threads the new fields through all auth call sites (_auth.py, _oauth_base.py, oauth_code.py, connection.py) and updates every affected test file, including a new test/unit/test_token_cache_key.py that validates the cross-driver golden hash vector. Co-authored-by: Cursor --- DESCRIPTION.md | 1 + src/snowflake/connector/auth/_auth.py | 83 +++-- src/snowflake/connector/auth/_oauth_base.py | 37 ++- src/snowflake/connector/auth/oauth_code.py | 3 + src/snowflake/connector/connection.py | 3 + src/snowflake/connector/token_cache.py | 159 +++++++--- .../sso_it/test_unit_mfa_cache_async.py | 8 +- test/integ/sso_it/test_unit_mfa_cache.py | 8 +- test/unit/aio/test_oauth_token_async.py | 68 ++++- test/unit/test_keyring_token_cache.py | 113 ++++--- test/unit/test_linux_local_file_cache.py | 92 +++--- test/unit/test_oauth_infinite_loop_fix.py | 4 +- test/unit/test_oauth_token.py | 68 ++++- test/unit/test_token_cache_key.py | 287 ++++++++++++++++++ 14 files changed, 736 insertions(+), 198 deletions(-) create mode 100644 test/unit/test_token_cache_key.py diff --git a/DESCRIPTION.md b/DESCRIPTION.md index b5eb2f6979..d44f750a14 100644 --- a/DESCRIPTION.md +++ b/DESCRIPTION.md @@ -8,6 +8,7 @@ Source code is also available at: https://github.com/snowflakedb/snowflake-conne # Release Notes - NEXT_RELEASE(TBD) + - Fixed token cache key collisions for multi-account (shared IdP) and multi-role scenarios by switching to a versioned, SHA256-hashed canonical-JSON key applied uniformly across macOS/Windows keyring and the Linux file backend (SNOW-3784431). - Added support for Python 3.14t (free-threaded). - **Note:** Python 3.14t CI testing excludes `win_arm64` (no `cryptography` wheels available) and `mitmproxy` proxy tests on all platforms (transitive dependencies `aioquic`/`pylsqpack` lack free-threaded-compatible wheels). diff --git a/src/snowflake/connector/auth/_auth.py b/src/snowflake/connector/auth/_auth.py index e60e607132..fd651e859f 100644 --- a/src/snowflake/connector/auth/_auth.py +++ b/src/snowflake/connector/auth/_auth.py @@ -410,7 +410,13 @@ def post_request_wrapper(self, url, headers, body) -> None: # raise an exception for reauth without id_token self._rest.id_token = None self._delete_temporary_credential( - self._rest._host, user, TokenType.ID_TOKEN + TokenKey( + token_type=TokenType.ID_TOKEN, + idp=self._rest._host, + snowflake=self._rest._host, + username=user, + role=role or "", + ) ) raise ReauthenticationRequest( ProgrammingError( @@ -446,7 +452,13 @@ def post_request_wrapper(self, url, headers, body) -> None: if isinstance(auth_instance, AuthByUsrPwdMfa): self._delete_temporary_credential( - self._rest._host, user, TokenType.MFA_TOKEN + TokenKey( + token_type=TokenType.MFA_TOKEN, + idp=self._rest._host, + snowflake=self._rest._host, + username=user, + role="", + ) ) Error.errorhandler_wrapper( self._rest._connection, @@ -514,7 +526,7 @@ def post_request_wrapper(self, url, headers, body) -> None: mfa_token=ret["data"].get("mfaToken"), ) self.write_temporary_credentials( - self._rest._host, user, session_parameters, ret + self._rest._host, user, session_parameters, ret, role=role or "" ) if ret["data"] and "sessionId" in ret["data"]: self._rest._connection._session_id = ret["data"].get("sessionId") @@ -531,19 +543,15 @@ def post_request_wrapper(self, url, headers, body) -> None: self._rest._connection._update_parameters(session_parameters) return session_parameters - def _read_temporary_credential( - self, - host: str, - user: str, - cred_type: TokenType, - ) -> str | None: - return self.get_token_cache().retrieve(TokenKey(host, user, cred_type)) + def _read_temporary_credential(self, key: TokenKey) -> str | None: + return self.get_token_cache().retrieve(key) def read_temporary_credentials( self, host: str, user: str, session_parameters: dict[str, Any], + role: str = "", ) -> None: """Attempt to load cached credentials to skip interactive authentication. @@ -557,31 +565,33 @@ def read_temporary_credentials( """ if session_parameters.get(PARAMETER_CLIENT_STORE_TEMPORARY_CREDENTIAL, False): self._rest.id_token = self._read_temporary_credential( - host, - user, - TokenType.ID_TOKEN, + TokenKey( + token_type=TokenType.ID_TOKEN, + idp=host, + snowflake=host, + username=user, + role=role, + ) ) if session_parameters.get(PARAMETER_CLIENT_REQUEST_MFA_TOKEN, False): self._rest.mfa_token = self._read_temporary_credential( - host, - user, - TokenType.MFA_TOKEN, + TokenKey( + token_type=TokenType.MFA_TOKEN, + idp=host, + snowflake=host, + username=user, + role="", + ) ) - def _write_temporary_credential( - self, - host: str, - user: str, - cred_type: TokenType, - cred: str | None, - ) -> None: + def _write_temporary_credential(self, key: TokenKey, cred: str | None) -> None: if not cred: logger.debug( "no credential is given when try to store temporary credential" ) return - self.get_token_cache().store(TokenKey(host, user, cred_type), cred) + self.get_token_cache().store(key, cred) def write_temporary_credentials( self, @@ -589,6 +599,7 @@ def write_temporary_credentials( user: str, session_parameters: dict[str, Any], response: dict[str, Any], + role: str = "", ) -> None: """Cache credentials received from successful authentication for future use. @@ -604,18 +615,30 @@ def write_temporary_credentials( ) ): self._write_temporary_credential( - host, user, TokenType.ID_TOKEN, response["data"].get("idToken") + TokenKey( + token_type=TokenType.ID_TOKEN, + idp=host, + snowflake=host, + username=user, + role=role, + ), + response["data"].get("idToken"), ) if session_parameters.get(PARAMETER_CLIENT_REQUEST_MFA_TOKEN, False): self._write_temporary_credential( - host, user, TokenType.MFA_TOKEN, response["data"].get("mfaToken") + TokenKey( + token_type=TokenType.MFA_TOKEN, + idp=host, + snowflake=host, + username=user, + role="", + ), + response["data"].get("mfaToken"), ) - def _delete_temporary_credential( - self, host: str, user: str, cred_type: TokenType - ) -> None: - self.get_token_cache().remove(TokenKey(host, user, cred_type)) + def _delete_temporary_credential(self, key: TokenKey) -> None: + self.get_token_cache().remove(key) def get_token_cache(self) -> TokenCache: if self._token_cache is None: diff --git a/src/snowflake/connector/auth/_oauth_base.py b/src/snowflake/connector/auth/_oauth_base.py index c75342c97c..6cd93e6968 100644 --- a/src/snowflake/connector/auth/_oauth_base.py +++ b/src/snowflake/connector/auth/_oauth_base.py @@ -40,7 +40,8 @@ class _OAuthTokensMixin: Access tokens: Short-lived (typically 10 minutes), cached to avoid immediate re-auth. Refresh tokens: Long-lived (hours/days), used to obtain new access tokens silently. - Tokens are cached per (user, IDP host) to support multiple OAuth providers/accounts. + Tokens are cached per (idp, snowflake, username, role) using the v2 cache key to + support multiple OAuth providers and multi-role scenarios without collision. """ def __init__( @@ -48,6 +49,8 @@ def __init__( token_cache: TokenCache | None, refresh_token_enabled: bool, idp_host: str, + snowflake_host: str = "", + role: str = "", ) -> None: self._access_token = None self._refresh_token_enabled = refresh_token_enabled @@ -55,6 +58,8 @@ def __init__( self._refresh_token = None self._token_cache = token_cache self._idp_host = idp_host + self._snowflake_host = snowflake_host + self._role = role self._tokens_loaded_from_cache = False # Prevents re-loading tokens from cache if self._token_cache: logger.debug("token cache is going to be used if needed") @@ -93,17 +98,25 @@ def _load_tokens_from_cache(self, user: str) -> bool: return self._access_token is not None def _get_access_token_cache_key(self) -> TokenKey | None: - return ( - TokenKey(self._user, self._idp_host, TokenType.OAUTH_ACCESS_TOKEN) - if self._token_cache and self._user - else None + if not (self._token_cache and self._user): + return None + return TokenKey( + token_type=TokenType.OAUTH_ACCESS_TOKEN, + idp=self._token_request_url, + snowflake=self._snowflake_host, + username=self._user, + role=self._role or "", ) def _get_refresh_token_cache_key(self) -> TokenKey | None: - return ( - TokenKey(self._user, self._idp_host, TokenType.OAUTH_REFRESH_TOKEN) - if self._refresh_token_enabled and self._token_cache and self._user - else None + if not (self._refresh_token_enabled and self._token_cache and self._user): + return None + return TokenKey( + token_type=TokenType.OAUTH_REFRESH_TOKEN, + idp=self._token_request_url, + snowflake=self._snowflake_host, + username=self._user, + role=self._role or "", ) def _invalidate_refresh_token(self) -> None: @@ -162,6 +175,8 @@ def __init__( token_cache: TokenCache | None, refresh_token_enabled: bool, is_snowflake_as_idp: bool = False, + snowflake_host: str = "", + role: str = "", **kwargs, ) -> None: super().__init__(**kwargs) @@ -169,7 +184,9 @@ def __init__( self, token_cache=token_cache, refresh_token_enabled=refresh_token_enabled, - idp_host=urllib.parse.urlparse(token_request_url).hostname, + idp_host=urllib.parse.urlparse(token_request_url).hostname or "", + snowflake_host=snowflake_host, + role=role, ) self._client_id = client_id self._client_secret = client_secret diff --git a/src/snowflake/connector/auth/oauth_code.py b/src/snowflake/connector/auth/oauth_code.py index d6d266169d..6ef6293ab1 100644 --- a/src/snowflake/connector/auth/oauth_code.py +++ b/src/snowflake/connector/auth/oauth_code.py @@ -66,6 +66,7 @@ def __init__( enable_single_use_refresh_tokens: bool = False, connection: SnowflakeConnection | None = None, uri: str | None = None, + role: str = "", **kwargs, ) -> None: authentication_url, redirect_uri = self._validate_oauth_code_uris( @@ -95,6 +96,8 @@ def __init__( is_snowflake_as_idp=self._is_snowflake_as_idp( authentication_url, token_request_url, host ), + snowflake_host=host, + role=role, **kwargs, ) self._application = application diff --git a/src/snowflake/connector/connection.py b/src/snowflake/connector/connection.py index 41fb8cc843..d41b1388cb 100644 --- a/src/snowflake/connector/connection.py +++ b/src/snowflake/connector/connection.py @@ -1492,6 +1492,7 @@ def __open_connection(self): self.host, self.user, self._session_parameters, + role=self._role or "", ) # Depending on whether self._rest.id_token is available we do different # auth_instance @@ -1564,6 +1565,7 @@ def __open_connection(self): refresh_token_enabled=self._oauth_enable_refresh_tokens, external_browser_timeout=self._external_browser_timeout, enable_single_use_refresh_tokens=self._oauth_enable_single_use_refresh_tokens, + role=self._role or "", ) elif self._authenticator == OAUTH_CLIENT_CREDENTIALS: if self._role and (self._oauth_scope == ""): @@ -1591,6 +1593,7 @@ def __open_connection(self): self.host, self.user, self._session_parameters, + role="", ) self.auth_class = AuthByUsrPwdMfa( password=self._password, diff --git a/src/snowflake/connector/token_cache.py b/src/snowflake/connector/token_cache.py index ed4fa406b3..54fd849459 100644 --- a/src/snowflake/connector/token_cache.py +++ b/src/snowflake/connector/token_cache.py @@ -5,6 +5,7 @@ import json import logging import os +import re import stat import sys from abc import ABC, abstractmethod @@ -42,21 +43,78 @@ class _InvalidTokenKeyError(Exception): @dataclass class TokenKey: - user: str - host: str - tokenType: TokenType + """Key identifying a cached token. + + All five fields are required. Raw (un-normalized) values are acceptable; + ``build_cache_key`` normalizes them before hashing. + + Fields: + token_type: The type of token being cached. + idp: IdP / token-endpoint URL. For MFA and external-browser flows this + equals the Snowflake server URL. + snowflake: Snowflake server URL. + username: Snowflake login name. + role: Snowflake role. Must be an empty string (not None) for MFA flows. + """ + + token_type: TokenType + idp: str + snowflake: str + username: str + role: str + + +def normalize_url(url: str) -> str: + """Strip scheme and userinfo, drop query/fragment, trim root slash, uppercase.""" + s = re.sub(r"^https?://", "", url) + at = s.find("@") + if at >= 0: + s = s[at + 1 :] + s = s.split("?")[0].split("#")[0] + s = s.rstrip("/") + return s.upper() + + +def normalize_identifier(identifier: str) -> str: + """Uppercase unquoted segments; preserve double-quoted segments verbatim.""" + result = [] + in_quotes = False + for ch in identifier: + if ch == '"': + in_quotes = not in_quotes + result.append(ch) + elif in_quotes: + result.append(ch) + else: + result.append(ch.upper()) + return "".join(result) + + +def build_cache_key(key: TokenKey) -> str: + """Build the versioned, uniformly-hashed v2 cache key. + + Format: ``SnowflakeTokenCache.v2.`` + + The canonical JSON is compact (no whitespace) with keys sorted + lexicographically, serialized to UTF-8. Hashing occurs exactly once here; + cache backends store and retrieve the returned string verbatim. + """ + if not key.snowflake: + raise _InvalidTokenKeyError("snowflake URL must not be empty") + if not key.username: + raise _InvalidTokenKeyError("username must not be empty") - def string_key(self) -> str: - if len(self.host) == 0: - raise _InvalidTokenKeyError("Invalid key, host is empty") - if len(self.user) == 0: - raise _InvalidTokenKeyError("Invalid key, user is empty") - return f"{self.host.upper()}:{self.user.upper()}:{self.tokenType.value}" + key_data = { + "idp": normalize_url(key.idp), + "role": normalize_identifier(key.role), + "snowflake": normalize_url(key.snowflake), + "token_type": key.token_type.value, + "username": normalize_identifier(key.username), + } - def hash_key(self) -> str: - m = hashlib.sha256() - m.update(self.string_key().encode(encoding="utf-8")) - return m.hexdigest() + canonical = json.dumps(key_data, sort_keys=True, separators=(",", ":")) + digest = hashlib.sha256(canonical.encode("utf-8")).hexdigest() + return f"SnowflakeTokenCache.v2.{digest}" def _warn(warning: str) -> None: @@ -72,7 +130,9 @@ class TokenCache(ABC): - Linux: Uses JSON file in ~/.cache/snowflake/ with 0o600 permissions - Fallback: NoopTokenCache (no caching) if secure storage unavailable - Tokens are keyed by (host, user, token_type) to support multiple accounts. + Tokens are keyed by a versioned, SHA-256-hashed canonical-JSON key (v2 format) + built from (token_type, idp, snowflake, username, role) to avoid collisions + across multi-account and multi-role scenarios. """ @staticmethod @@ -154,6 +214,12 @@ class FileTokenCache(TokenCache): Security: File must have 0o600 permissions and be owned by current user. Uses file locks to prevent concurrent access corruption. + + JSON map keys are the full ``SnowflakeTokenCache.v2.`` strings + produced by ``build_cache_key``; hashing is performed once before dispatch. + Note: the filename (``credential_cache_v1.json``) is unchanged for + backward compatibility; the ``v2`` in the key prefix refers to the + key-format version, not the file format. """ @staticmethod @@ -178,12 +244,13 @@ def __init__( def store(self, key: TokenKey, token: str) -> None: try: + final_key = build_cache_key(key) FileTokenCache.validate_cache_dir( self.cache_dir, self._skip_file_permissions_check ) with FileLock(self.lock_file()): cache = self._read_cache_file() - cache["tokens"][key.hash_key()] = token + cache["tokens"][final_key] = token self._write_cache_file(cache) except _FileTokenCacheError as e: self.logger.error(f"Failed to store token: {e=}") @@ -194,12 +261,13 @@ def store(self, key: TokenKey, token: str) -> None: def retrieve(self, key: TokenKey) -> str | None: try: + final_key = build_cache_key(key) FileTokenCache.validate_cache_dir( self.cache_dir, self._skip_file_permissions_check ) with FileLock(self.lock_file()): cache = self._read_cache_file() - token = cache["tokens"].get(key.hash_key(), None) + token = cache["tokens"].get(final_key, None) if isinstance(token, str): return token else: @@ -216,12 +284,13 @@ def retrieve(self, key: TokenKey) -> str | None: def remove(self, key: TokenKey) -> None: try: + final_key = build_cache_key(key) FileTokenCache.validate_cache_dir( self.cache_dir, self._skip_file_permissions_check ) with FileLock(self.lock_file()): cache = self._read_cache_file() - cache["tokens"].pop(key.hash_key(), None) + cache["tokens"].pop(final_key, None) self._write_cache_file(cache) except _FileTokenCacheError as e: self.logger.error(f"Failed to remove token: {e=}") @@ -398,14 +467,13 @@ class KeyringTokenCache(TokenCache): - macOS: Stores tokens in Keychain - Windows: Stores tokens in Windows Credential Manager - All tokens share a single keyring service name so that on macOS they fall - under one Keychain ACL entry, requiring only a single "Allow" prompt - instead of one per token. The keyring *account* field stores a SHA-256 hash - of ``{HOST}:{USER}:{TOKEN_TYPE}`` to avoid exposing plaintext identifiers - in the OS credential store. + The v2 cache key (``SnowflakeTokenCache.v2.``) is used as both + the keyring service name and the account field, ensuring a single Keychain + ACL entry per token and no plaintext identifiers in the OS credential store. For backward compatibility, :meth:`retrieve` also checks the legacy layout - where the service was the full string key and the account was the username. + where the service was ``{HOST}:{USER}:{TOKEN_TYPE}`` and the account was the + uppercase username; matching entries are silently migrated to v2 on first use. """ SERVICE_NAME = "com.snowflake.connector.python" @@ -415,61 +483,58 @@ def __init__(self) -> None: def store(self, key: TokenKey, token: str) -> None: try: - keyring.set_password( - self.SERVICE_NAME, - key.hash_key(), - token, - ) + final_key = build_cache_key(key) + keyring.set_password(final_key, key.username.upper(), token) except _InvalidTokenKeyError as e: - self.logger.error(f"Could not store {key.tokenType} in keyring, {e=}") + self.logger.error(f"Could not store {key.token_type} in keyring, {e=}") except keyring.errors.KeyringError as ke: self.logger.error("Could not store token in keyring, %s", str(ke)) def retrieve(self, key: TokenKey) -> str | None: try: - token = keyring.get_password( - self.SERVICE_NAME, - key.hash_key(), - ) + final_key = build_cache_key(key) + token = keyring.get_password(final_key, key.username.upper()) if token is not None: return token return self._retrieve_legacy(key) except keyring.errors.KeyringError as ke: self.logger.error( "Could not retrieve {} from secure storage : {}".format( - key.tokenType.value, str(ke) + key.token_type.value, str(ke) ) ) except _InvalidTokenKeyError as e: - self.logger.error(f"Could not retrieve {key.tokenType} from keyring, {e=}") + self.logger.error( + f"Could not retrieve {key.token_type} from keyring, {e=}" + ) def _retrieve_legacy(self, key: TokenKey) -> str | None: - """Try to read from the old per-token-type service layout and migrate.""" + """Try to read from the old per-token-type service layout and migrate to v2.""" + legacy_service = ( + f"{key.snowflake.upper()}:{key.username.upper()}:{key.token_type.value}" + ) try: - token = keyring.get_password( - key.string_key(), - key.user.upper(), - ) + token = keyring.get_password(legacy_service, key.username.upper()) except (keyring.errors.KeyringError, _InvalidTokenKeyError): return None if token is None: return None self.store(key, token) try: - keyring.delete_password(key.string_key(), key.user.upper()) + keyring.delete_password(legacy_service, key.username.upper()) except Exception: pass - self.logger.debug("migrated legacy keyring entry for %s", key.tokenType.value) + self.logger.debug( + "migrated legacy keyring entry for %s", key.token_type.value + ) return token def remove(self, key: TokenKey) -> None: try: - keyring.delete_password( - self.SERVICE_NAME, - key.hash_key(), - ) + final_key = build_cache_key(key) + keyring.delete_password(final_key, key.username.upper()) except _InvalidTokenKeyError as e: - self.logger.error(f"Could not remove {key.tokenType} from keyring, {e=}") + self.logger.error(f"Could not remove {key.token_type} from keyring, {e=}") except Exception as ex: self.logger.error( "Failed to delete credential in the keyring: err=[%s]", ex diff --git a/test/integ/aio_it/sso_it/test_unit_mfa_cache_async.py b/test/integ/aio_it/sso_it/test_unit_mfa_cache_async.py index eef35b96de..d760580b31 100644 --- a/test/integ/aio_it/sso_it/test_unit_mfa_cache_async.py +++ b/test/integ/aio_it/sso_it/test_unit_mfa_cache_async.py @@ -125,7 +125,13 @@ async def test_body(conn_cfg): from snowflake.connector.token_cache import TokenCache, TokenKey, TokenType TokenCache.make().remove( - TokenKey(conn_cfg["host"], conn_cfg["user"], TokenType.MFA_TOKEN) + TokenKey( + token_type=TokenType.MFA_TOKEN, + idp=conn_cfg["host"], + snowflake=conn_cfg["host"], + username=conn_cfg["user"], + role="", + ) ) # first connection, no mfa token cache diff --git a/test/integ/sso_it/test_unit_mfa_cache.py b/test/integ/sso_it/test_unit_mfa_cache.py index 8fe44db077..7ea355c120 100644 --- a/test/integ/sso_it/test_unit_mfa_cache.py +++ b/test/integ/sso_it/test_unit_mfa_cache.py @@ -121,7 +121,13 @@ def test_body(conn_cfg): from snowflake.connector.token_cache import TokenCache, TokenKey, TokenType TokenCache.make().remove( - TokenKey(conn_cfg["host"], conn_cfg["user"], TokenType.MFA_TOKEN) + TokenKey( + token_type=TokenType.MFA_TOKEN, + idp=conn_cfg["host"], + snowflake=conn_cfg["host"], + username=conn_cfg["user"], + role="", + ) ) # first connection, no mfa token cache diff --git a/test/unit/aio/test_oauth_token_async.py b/test/unit/aio/test_oauth_token_async.py index e54fd2dca5..0a3d312ea5 100644 --- a/test/unit/aio/test_oauth_token_async.py +++ b/test/unit/aio/test_oauth_token_async.py @@ -103,18 +103,20 @@ def webbrowser_mock_sync() -> Mock: def temp_cache_async(): """Async-compatible temporary cache.""" + from snowflake.connector.token_cache import build_cache_key + class TemporaryCache(TokenCache): def __init__(self): self._cache = {} def store(self, key: TokenKey, token: str) -> None: - self._cache[(key.user, key.host, key.tokenType)] = token + self._cache[build_cache_key(key)] = token def retrieve(self, key: TokenKey) -> str: - return self._cache.get((key.user, key.host, key.tokenType)) + return self._cache.get(build_cache_key(key)) def remove(self, key: TokenKey) -> None: - self._cache.pop((key.user, key.host, key.tokenType)) + self._cache.pop(build_cache_key(key), None) tmp_cache = TemporaryCache() # Patch both sync and async versions to be safe since async Auth inherits from sync Auth @@ -437,11 +439,20 @@ async def test_oauth_code_successful_refresh_token_flow_async( wiremock_generic_mappings_dir / "snowflake_disconnect_successful.json" ) user = "testUser" + token_request_url = f"http://{wiremock_client.wiremock_host}:{wiremock_client.wiremock_http_port}/oauth/token-request" access_token_key = TokenKey( - user, wiremock_client.wiremock_host, TokenType.OAUTH_ACCESS_TOKEN + token_type=TokenType.OAUTH_ACCESS_TOKEN, + idp=token_request_url, + snowflake=wiremock_client.wiremock_host, + username=user, + role="ANALYST", ) refresh_token_key = TokenKey( - user, wiremock_client.wiremock_host, TokenType.OAUTH_REFRESH_TOKEN + token_type=TokenType.OAUTH_REFRESH_TOKEN, + idp=token_request_url, + snowflake=wiremock_client.wiremock_host, + username=user, + role="ANALYST", ) temp_cache_async.store(access_token_key, "expired-access-token-123") temp_cache_async.store(refresh_token_key, "refresh-token-123") @@ -453,7 +464,7 @@ async def test_oauth_code_successful_refresh_token_flow_async( protocol="http", role="ANALYST", oauth_client_secret="testClientSecret", - oauth_token_request_url=f"http://{wiremock_client.wiremock_host}:{wiremock_client.wiremock_http_port}/oauth/token-request", + oauth_token_request_url=token_request_url, oauth_authorization_url=f"http://{wiremock_client.wiremock_host}:{wiremock_client.wiremock_http_port}/oauth/authorize", oauth_redirect_uri="http://localhost:8009/snowflake/oauth-redirect", host=wiremock_client.wiremock_host, @@ -505,11 +516,20 @@ async def test_oauth_code_expired_refresh_token_flow_async( ) user = "testUser" + token_request_url = f"http://{wiremock_client.wiremock_host}:{wiremock_client.wiremock_http_port}/oauth/token-request" access_token_key = TokenKey( - user, wiremock_client.wiremock_host, TokenType.OAUTH_ACCESS_TOKEN + token_type=TokenType.OAUTH_ACCESS_TOKEN, + idp=token_request_url, + snowflake=wiremock_client.wiremock_host, + username=user, + role="ANALYST", ) refresh_token_key = TokenKey( - user, wiremock_client.wiremock_host, TokenType.OAUTH_REFRESH_TOKEN + token_type=TokenType.OAUTH_REFRESH_TOKEN, + idp=token_request_url, + snowflake=wiremock_client.wiremock_host, + username=user, + role="ANALYST", ) temp_cache_async.store(access_token_key, "expired-access-token-123") temp_cache_async.store(refresh_token_key, "expired-refresh-token-123") @@ -523,7 +543,7 @@ async def test_oauth_code_expired_refresh_token_flow_async( protocol="http", role="ANALYST", oauth_client_secret="testClientSecret", - oauth_token_request_url=f"http://{wiremock_client.wiremock_host}:{wiremock_client.wiremock_http_port}/oauth/token-request", + oauth_token_request_url=token_request_url, oauth_authorization_url=f"http://{wiremock_client.wiremock_host}:{wiremock_client.wiremock_http_port}/oauth/authorize", oauth_redirect_uri="http://localhost:8009/snowflake/oauth-redirect", host=wiremock_client.wiremock_host, @@ -574,11 +594,20 @@ async def test_client_creds_successful_flow_async( wiremock_generic_mappings_dir / "snowflake_disconnect_successful.json" ) user = "testUser" + token_request_url = f"http://{wiremock_client.wiremock_host}:{wiremock_client.wiremock_http_port}/oauth/token-request" access_token_key = TokenKey( - user, wiremock_client.wiremock_host, TokenType.OAUTH_ACCESS_TOKEN + token_type=TokenType.OAUTH_ACCESS_TOKEN, + idp=token_request_url, + snowflake=wiremock_client.wiremock_host, + username=user, + role="ANALYST", ) refresh_token_key = TokenKey( - user, wiremock_client.wiremock_host, TokenType.OAUTH_REFRESH_TOKEN + token_type=TokenType.OAUTH_REFRESH_TOKEN, + idp=token_request_url, + snowflake=wiremock_client.wiremock_host, + username=user, + role="ANALYST", ) temp_cache_async.store(access_token_key, "unused-access-token-123") temp_cache_async.store(refresh_token_key, "unused-refresh-token-123") @@ -591,7 +620,7 @@ async def test_client_creds_successful_flow_async( account="testAccount", protocol="http", role="ANALYST", - oauth_token_request_url=f"http://{wiremock_client.wiremock_host}:{wiremock_client.wiremock_http_port}/oauth/token-request", + oauth_token_request_url=token_request_url, host=wiremock_client.wiremock_host, port=wiremock_client.wiremock_http_port, oauth_enable_refresh_tokens=True, @@ -671,11 +700,20 @@ async def test_client_creds_expired_refresh_token_flow_async( ) user = "testUser" + token_request_url = f"http://{wiremock_client.wiremock_host}:{wiremock_client.wiremock_http_port}/oauth/token-request" access_token_key = TokenKey( - user, wiremock_client.wiremock_host, TokenType.OAUTH_ACCESS_TOKEN + token_type=TokenType.OAUTH_ACCESS_TOKEN, + idp=token_request_url, + snowflake=wiremock_client.wiremock_host, + username=user, + role="ANALYST", ) refresh_token_key = TokenKey( - user, wiremock_client.wiremock_host, TokenType.OAUTH_REFRESH_TOKEN + token_type=TokenType.OAUTH_REFRESH_TOKEN, + idp=token_request_url, + snowflake=wiremock_client.wiremock_host, + username=user, + role="ANALYST", ) temp_cache_async.store(access_token_key, "expired-access-token-123") temp_cache_async.store(refresh_token_key, "expired-refresh-token-123") @@ -687,7 +725,7 @@ async def test_client_creds_expired_refresh_token_flow_async( protocol="http", role="ANALYST", oauth_client_secret="testClientSecret", - oauth_token_request_url=f"http://{wiremock_client.wiremock_host}:{wiremock_client.wiremock_http_port}/oauth/token-request", + oauth_token_request_url=token_request_url, host=wiremock_client.wiremock_host, port=wiremock_client.wiremock_http_port, oauth_enable_refresh_tokens=True, diff --git a/test/unit/test_keyring_token_cache.py b/test/unit/test_keyring_token_cache.py index 38aaf7cbc9..38ebfca3a9 100644 --- a/test/unit/test_keyring_token_cache.py +++ b/test/unit/test_keyring_token_cache.py @@ -5,7 +5,12 @@ import pytest from snowflake.connector.options import installed_keyring -from snowflake.connector.token_cache import KeyringTokenCache, TokenKey, TokenType +from snowflake.connector.token_cache import ( + KeyringTokenCache, + TokenKey, + TokenType, + build_cache_key, +) pytestmark = pytest.mark.skipif( not installed_keyring, @@ -26,53 +31,53 @@ def cache(): return KeyringTokenCache() -SERVICE = KeyringTokenCache.SERVICE_NAME KEY = TokenKey( - user="ALICE", - host="myhost.snowflakecomputing.com", - tokenType=TokenType.OAUTH_ACCESS_TOKEN, + token_type=TokenType.OAUTH_ACCESS_TOKEN, + idp="https://idp.example.com/oauth2", + snowflake="myhost.snowflakecomputing.com", + username="ALICE", + role="", ) -ACCOUNT = KEY.hash_key() +FINAL_KEY = build_cache_key(KEY) +ACCOUNT = KEY.username.upper() class TestStore: - def test_stores_under_unified_service_with_hashed_account( - self, cache, mock_keyring - ): + def test_stores_using_v2_key_as_service(self, cache, mock_keyring): cache.store(KEY, "tok123") - mock_keyring.set_password.assert_called_once_with( - SERVICE, - ACCOUNT, - "tok123", - ) - assert ACCOUNT != KEY.string_key(), "account should be a hash, not plaintext" + mock_keyring.set_password.assert_called_once_with(FINAL_KEY, ACCOUNT, "tok123") + + def test_v2_key_starts_with_prefix(self, cache, mock_keyring): + cache.store(KEY, "tok123") + service = mock_keyring.set_password.call_args.args[0] + assert service.startswith("SnowflakeTokenCache.v2.") + + def test_v2_key_differs_from_legacy_string_key(self, cache, mock_keyring): + cache.store(KEY, "tok123") + service = mock_keyring.set_password.call_args.args[0] + legacy_service = f"{KEY.snowflake.upper()}:{KEY.username.upper()}:{KEY.token_type.value}" + assert service != legacy_service, "v2 key should differ from legacy key" class TestRetrieve: - def test_retrieves_from_unified_service(self, cache, mock_keyring): + def test_retrieves_from_v2_key(self, cache, mock_keyring): mock_keyring.get_password.return_value = "tok123" assert cache.retrieve(KEY) == "tok123" - mock_keyring.get_password.assert_called_once_with(SERVICE, ACCOUNT) + mock_keyring.get_password.assert_called_once_with(FINAL_KEY, ACCOUNT) def test_falls_back_to_legacy_and_migrates(self, cache, mock_keyring): + legacy_service = f"{KEY.snowflake.upper()}:{KEY.username.upper()}:{KEY.token_type.value}" mock_keyring.get_password.side_effect = [None, "legacy_tok"] result = cache.retrieve(KEY) assert result == "legacy_tok" mock_keyring.get_password.assert_has_calls( [ - call(SERVICE, ACCOUNT), - call(KEY.string_key(), KEY.user.upper()), + call(FINAL_KEY, ACCOUNT), + call(legacy_service, ACCOUNT), ] ) - mock_keyring.set_password.assert_called_once_with( - SERVICE, - ACCOUNT, - "legacy_tok", - ) - mock_keyring.delete_password.assert_called_once_with( - KEY.string_key(), - KEY.user.upper(), - ) + mock_keyring.set_password.assert_called_once_with(FINAL_KEY, ACCOUNT, "legacy_tok") + mock_keyring.delete_password.assert_called_once_with(legacy_service, ACCOUNT) def test_returns_none_when_not_found_anywhere(self, cache, mock_keyring): mock_keyring.get_password.return_value = None @@ -87,16 +92,52 @@ def test_legacy_delete_failure_is_nonfatal(self, cache, mock_keyring): class TestRemove: - def test_removes_from_unified_service(self, cache, mock_keyring): + def test_removes_using_v2_key(self, cache, mock_keyring): cache.remove(KEY) - mock_keyring.delete_password.assert_called_once_with(SERVICE, ACCOUNT) + mock_keyring.delete_password.assert_called_once_with(FINAL_KEY, ACCOUNT) -class TestServiceNameConstant: - def test_all_token_types_share_service(self, cache, mock_keyring): +class TestMultiAccount: + def test_different_accounts_produce_different_keys(self, cache, mock_keyring): + """Keys for different accounts never collide.""" mock_keyring.get_password.return_value = None - for tt in TokenType: - k = TokenKey(user="BOB", host="host.com", tokenType=tt) - cache.store(k, "val") + key1 = TokenKey( + token_type=TokenType.OAUTH_ACCESS_TOKEN, + idp="https://idp.example.com/oauth2", + snowflake="account1.snowflakecomputing.com", + username="USER", + role="", + ) + key2 = TokenKey( + token_type=TokenType.OAUTH_ACCESS_TOKEN, + idp="https://idp.example.com/oauth2", + snowflake="account2.snowflakecomputing.com", + username="USER", + role="", + ) + cache.store(key1, "val1") + cache.store(key2, "val2") + services = [c.args[0] for c in mock_keyring.set_password.call_args_list] + assert services[0] != services[1], "different accounts must have different keys" + + def test_different_roles_produce_different_keys(self, cache, mock_keyring): + """Keys for different roles never collide.""" + mock_keyring.get_password.return_value = None + key1 = TokenKey( + token_type=TokenType.OAUTH_ACCESS_TOKEN, + idp="https://idp.example.com/oauth2", + snowflake="account.snowflakecomputing.com", + username="USER", + role="ANALYST", + ) + key2 = TokenKey( + token_type=TokenType.OAUTH_ACCESS_TOKEN, + idp="https://idp.example.com/oauth2", + snowflake="account.snowflakecomputing.com", + username="USER", + role="SYSADMIN", + ) + cache.store(key1, "val1") + cache.store(key2, "val2") services = [c.args[0] for c in mock_keyring.set_password.call_args_list] - assert all(s == SERVICE for s in services) + assert services[0] != services[1], "different roles must have different keys" diff --git a/test/unit/test_linux_local_file_cache.py b/test/unit/test_linux_local_file_cache.py index 56834ebd78..d2c384f22c 100644 --- a/test/unit/test_linux_local_file_cache.py +++ b/test/unit/test_linux_local_file_cache.py @@ -27,6 +27,16 @@ CRED_1 = "cred_1" +def make_key(host: str, user: str, cred_type) -> "TokenKey": + return TokenKey( + token_type=cred_type, + idp=host, + snowflake=host, + username=user, + role="", + ) + + @pytest.mark.skipolddriver def test_basic_store(tmpdir, monkeypatch): monkeypatch.setenv("SF_TEMPORARY_CREDENTIAL_CACHE_DIR", str(tmpdir)) @@ -35,13 +45,13 @@ def test_basic_store(tmpdir, monkeypatch): assert cache.cache_dir == pathlib.Path(tmpdir) cache.cache_file().unlink(missing_ok=True) - cache.store(TokenKey(HOST_0, USER_0, CRED_TYPE_0), CRED_0) - cache.store(TokenKey(HOST_1, USER_1, CRED_TYPE_1), CRED_1) - cache.store(TokenKey(HOST_0, USER_1, CRED_TYPE_1), CRED_1) + cache.store(make_key(HOST_0, USER_0, CRED_TYPE_0), CRED_0) + cache.store(make_key(HOST_1, USER_1, CRED_TYPE_1), CRED_1) + cache.store(make_key(HOST_0, USER_1, CRED_TYPE_1), CRED_1) - assert cache.retrieve(TokenKey(HOST_0, USER_0, CRED_TYPE_0)) == CRED_0 - assert cache.retrieve(TokenKey(HOST_1, USER_1, CRED_TYPE_1)) == CRED_1 - assert cache.retrieve(TokenKey(HOST_0, USER_1, CRED_TYPE_1)) == CRED_1 + assert cache.retrieve(make_key(HOST_0, USER_0, CRED_TYPE_0)) == CRED_0 + assert cache.retrieve(make_key(HOST_1, USER_1, CRED_TYPE_1)) == CRED_1 + assert cache.retrieve(make_key(HOST_0, USER_1, CRED_TYPE_1)) == CRED_1 cache.cache_file().unlink(missing_ok=True) @@ -51,15 +61,15 @@ def test_delete_specific_item(tmpdir, monkeypatch): cache = FileTokenCache.make() assert cache cache.cache_file().unlink(missing_ok=True) - cache.store(TokenKey(HOST_0, USER_0, CRED_TYPE_0), CRED_0) - cache.store(TokenKey(HOST_0, USER_0, CRED_TYPE_1), CRED_1) + cache.store(make_key(HOST_0, USER_0, CRED_TYPE_0), CRED_0) + cache.store(make_key(HOST_0, USER_0, CRED_TYPE_1), CRED_1) - assert cache.retrieve(TokenKey(HOST_0, USER_0, CRED_TYPE_0)) == CRED_0 - assert cache.retrieve(TokenKey(HOST_0, USER_0, CRED_TYPE_1)) == CRED_1 + assert cache.retrieve(make_key(HOST_0, USER_0, CRED_TYPE_0)) == CRED_0 + assert cache.retrieve(make_key(HOST_0, USER_0, CRED_TYPE_1)) == CRED_1 - cache.remove(TokenKey(HOST_0, USER_0, CRED_TYPE_0)) - assert not cache.retrieve(TokenKey(HOST_0, USER_0, CRED_TYPE_0)) - assert cache.retrieve(TokenKey(HOST_0, USER_0, CRED_TYPE_1)) == CRED_1 + cache.remove(make_key(HOST_0, USER_0, CRED_TYPE_0)) + assert not cache.retrieve(make_key(HOST_0, USER_0, CRED_TYPE_0)) + assert cache.retrieve(make_key(HOST_0, USER_0, CRED_TYPE_1)) == CRED_1 cache.cache_file().unlink(missing_ok=True) @@ -71,9 +81,9 @@ def test_malformed_json_cache(tmpdir, monkeypatch): cache.cache_file().touch(0o600) invalid_json = "[}" cache.cache_file().write_text(invalid_json) - assert cache.retrieve(TokenKey(HOST_0, USER_0, CRED_TYPE_0)) is None - cache.store(TokenKey(HOST_0, USER_0, CRED_TYPE_0), CRED_0) - assert cache.retrieve(TokenKey(HOST_0, USER_0, CRED_TYPE_0)) == CRED_0 + assert cache.retrieve(make_key(HOST_0, USER_0, CRED_TYPE_0)) is None + cache.store(make_key(HOST_0, USER_0, CRED_TYPE_0), CRED_0) + assert cache.retrieve(make_key(HOST_0, USER_0, CRED_TYPE_0)) == CRED_0 def test_malformed_utf_cache(tmpdir, monkeypatch): @@ -84,9 +94,9 @@ def test_malformed_utf_cache(tmpdir, monkeypatch): cache.cache_file().touch(0o600) invalid_utf_sequence = bytes.fromhex("c0af") cache.cache_file().write_bytes(invalid_utf_sequence) - assert cache.retrieve(TokenKey(HOST_0, USER_0, CRED_TYPE_0)) is None - cache.store(TokenKey(HOST_0, USER_0, CRED_TYPE_0), CRED_0) - assert cache.retrieve(TokenKey(HOST_0, USER_0, CRED_TYPE_0)) == CRED_0 + assert cache.retrieve(make_key(HOST_0, USER_0, CRED_TYPE_0)) is None + cache.store(make_key(HOST_0, USER_0, CRED_TYPE_0), CRED_0) + assert cache.retrieve(make_key(HOST_0, USER_0, CRED_TYPE_0)) == CRED_0 def test_cache_dir_is_not_a_directory(tmpdir, monkeypatch): @@ -159,9 +169,9 @@ def test_cache_file_incorrect_permissions(tmpdir, monkeypatch): assert cache cache.cache_file().unlink(missing_ok=True) cache.cache_file().touch(0o777) - assert cache.retrieve(TokenKey(HOST_0, USER_0, CRED_TYPE_0)) is None - cache.store(TokenKey(HOST_0, USER_0, CRED_TYPE_0), CRED_0) - assert cache.retrieve(TokenKey(HOST_0, USER_0, CRED_TYPE_0)) is None + assert cache.retrieve(make_key(HOST_0, USER_0, CRED_TYPE_0)) is None + cache.store(make_key(HOST_0, USER_0, CRED_TYPE_0), CRED_0) + assert cache.retrieve(make_key(HOST_0, USER_0, CRED_TYPE_0)) is None assert len(cache.cache_file().read_text("utf-8")) == 0 cache.cache_file().unlink() @@ -174,9 +184,9 @@ def test_cache_file_incorrect_permission_with_skip_file_permissions_check( assert cache cache.cache_file().unlink(missing_ok=True) cache.cache_file().touch(0o777) - assert cache.retrieve(TokenKey(HOST_0, USER_0, CRED_TYPE_0)) is None - cache.store(TokenKey(HOST_0, USER_0, CRED_TYPE_0), CRED_0) - assert cache.retrieve(TokenKey(HOST_0, USER_0, CRED_TYPE_0)) == CRED_0 + assert cache.retrieve(make_key(HOST_0, USER_0, CRED_TYPE_0)) is None + cache.store(make_key(HOST_0, USER_0, CRED_TYPE_0), CRED_0) + assert cache.retrieve(make_key(HOST_0, USER_0, CRED_TYPE_0)) == CRED_0 assert len(cache.cache_file().read_text("utf-8")) > 0 cache.cache_file().unlink() @@ -196,8 +206,8 @@ def test_cache_dir_xdg_cache_home(tmpdir, monkeypatch): cache.lock_file() == pathlib.Path(str(tmpdir)) / "snowflake" / "credential_cache_v1.json.lck" ) - cache.store(TokenKey(HOST_0, USER_0, CRED_TYPE_0), CRED_0) - assert cache.retrieve(TokenKey(HOST_0, USER_0, CRED_TYPE_0)) == CRED_0 + cache.store(make_key(HOST_0, USER_0, CRED_TYPE_0), CRED_0) + assert cache.retrieve(make_key(HOST_0, USER_0, CRED_TYPE_0)) == CRED_0 cache.cache_file().unlink() @@ -223,18 +233,18 @@ def test_cache_dir_home(tmpdir, monkeypatch): / "snowflake" / "credential_cache_v1.json.lck" ) - cache.store(TokenKey(HOST_0, USER_0, CRED_TYPE_0), CRED_0) - assert cache.retrieve(TokenKey(HOST_0, USER_0, CRED_TYPE_0)) == CRED_0 + cache.store(make_key(HOST_0, USER_0, CRED_TYPE_0), CRED_0) + assert cache.retrieve(make_key(HOST_0, USER_0, CRED_TYPE_0)) == CRED_0 def test_file_lock(tmpdir, monkeypatch): monkeypatch.setenv("SF_TEMPORARY_CREDENTIAL_CACHE_DIR", str(tmpdir)) cache = FileTokenCache.make() assert cache - cache.store(TokenKey(HOST_0, USER_0, CRED_TYPE_0), CRED_0) - assert cache.retrieve(TokenKey(HOST_0, USER_0, CRED_TYPE_0)) == CRED_0 + cache.store(make_key(HOST_0, USER_0, CRED_TYPE_0), CRED_0) + assert cache.retrieve(make_key(HOST_0, USER_0, CRED_TYPE_0)) == CRED_0 cache.lock_file().mkdir(0o700) - assert cache.retrieve(TokenKey(HOST_0, USER_0, CRED_TYPE_0)) is None + assert cache.retrieve(make_key(HOST_0, USER_0, CRED_TYPE_0)) is None assert cache.lock_file().exists() cache.lock_file().rmdir() @@ -243,11 +253,11 @@ def test_file_lock_stale(tmpdir, monkeypatch): monkeypatch.setenv("SF_TEMPORARY_CREDENTIAL_CACHE_DIR", str(tmpdir)) cache = FileTokenCache.make() assert cache - cache.store(TokenKey(HOST_0, USER_0, CRED_TYPE_0), CRED_0) - assert cache.retrieve(TokenKey(HOST_0, USER_0, CRED_TYPE_0)) == CRED_0 + cache.store(make_key(HOST_0, USER_0, CRED_TYPE_0), CRED_0) + assert cache.retrieve(make_key(HOST_0, USER_0, CRED_TYPE_0)) == CRED_0 cache.lock_file().mkdir(0o700) time.sleep(1) - assert cache.retrieve(TokenKey(HOST_0, USER_0, CRED_TYPE_0)) == CRED_0 + assert cache.retrieve(make_key(HOST_0, USER_0, CRED_TYPE_0)) == CRED_0 assert not cache.lock_file().exists() @@ -257,9 +267,9 @@ def test_file_missing_tokens_field(tmpdir, monkeypatch): assert cache cache.cache_file().touch(0o600) cache.cache_file().write_text("{}") - assert cache.retrieve(TokenKey(HOST_0, USER_0, CRED_TYPE_0)) is None - cache.store(TokenKey(HOST_0, USER_0, CRED_TYPE_0), CRED_0) - assert cache.retrieve(TokenKey(HOST_0, USER_0, CRED_TYPE_0)) == CRED_0 + assert cache.retrieve(make_key(HOST_0, USER_0, CRED_TYPE_0)) is None + cache.store(make_key(HOST_0, USER_0, CRED_TYPE_0), CRED_0) + assert cache.retrieve(make_key(HOST_0, USER_0, CRED_TYPE_0)) == CRED_0 cache.cache_file().unlink() @@ -269,7 +279,7 @@ def test_file_tokens_is_not_dict(tmpdir, monkeypatch): assert cache cache.cache_file().touch(0o600) cache.cache_file().write_text('{ "tokens": [] }') - assert cache.retrieve(TokenKey(HOST_0, USER_0, CRED_TYPE_0)) is None - cache.store(TokenKey(HOST_0, USER_0, CRED_TYPE_0), CRED_0) - assert cache.retrieve(TokenKey(HOST_0, USER_0, CRED_TYPE_0)) == CRED_0 + assert cache.retrieve(make_key(HOST_0, USER_0, CRED_TYPE_0)) is None + cache.store(make_key(HOST_0, USER_0, CRED_TYPE_0), CRED_0) + assert cache.retrieve(make_key(HOST_0, USER_0, CRED_TYPE_0)) == CRED_0 cache.cache_file().unlink() diff --git a/test/unit/test_oauth_infinite_loop_fix.py b/test/unit/test_oauth_infinite_loop_fix.py index a44bef8fce..2507dc056a 100644 --- a/test/unit/test_oauth_infinite_loop_fix.py +++ b/test/unit/test_oauth_infinite_loop_fix.py @@ -91,9 +91,9 @@ def test_load_tokens_from_cache_loads_both_tokens( """Verify both access and refresh tokens are loaded.""" def mock_retrieve(key: TokenKey): - if key.tokenType == TokenType.OAUTH_ACCESS_TOKEN: + if key.token_type == TokenType.OAUTH_ACCESS_TOKEN: return "test_access_token" - elif key.tokenType == TokenType.OAUTH_REFRESH_TOKEN: + elif key.token_type == TokenType.OAUTH_REFRESH_TOKEN: return "test_refresh_token" return None diff --git a/test/unit/test_oauth_token.py b/test/unit/test_oauth_token.py index b19d9415d6..e3cca4aa2c 100644 --- a/test/unit/test_oauth_token.py +++ b/test/unit/test_oauth_token.py @@ -81,18 +81,20 @@ def webbrowser_mock() -> Mock: @pytest.fixture() def temp_cache(): + from snowflake.connector.token_cache import build_cache_key + class TemporaryCache(TokenCache): def __init__(self): self._cache = {} def store(self, key: TokenKey, token: str) -> None: - self._cache[(key.user, key.host, key.tokenType)] = token + self._cache[build_cache_key(key)] = token def retrieve(self, key: TokenKey) -> str: - return self._cache.get((key.user, key.host, key.tokenType)) + return self._cache.get(build_cache_key(key)) def remove(self, key: TokenKey) -> None: - self._cache.pop((key.user, key.host, key.tokenType)) + self._cache.pop(build_cache_key(key), None) tmp_cache = TemporaryCache() with mock.patch( @@ -417,11 +419,20 @@ def test_oauth_code_successful_refresh_token_flow( wiremock_generic_mappings_dir / "snowflake_disconnect_successful.json" ) user = "testUser" + token_request_url = f"http://{wiremock_client.wiremock_host}:{wiremock_client.wiremock_http_port}/oauth/token-request" access_token_key = TokenKey( - user, wiremock_client.wiremock_host, TokenType.OAUTH_ACCESS_TOKEN + token_type=TokenType.OAUTH_ACCESS_TOKEN, + idp=token_request_url, + snowflake=wiremock_client.wiremock_host, + username=user, + role="ANALYST", ) refresh_token_key = TokenKey( - user, wiremock_client.wiremock_host, TokenType.OAUTH_REFRESH_TOKEN + token_type=TokenType.OAUTH_REFRESH_TOKEN, + idp=token_request_url, + snowflake=wiremock_client.wiremock_host, + username=user, + role="ANALYST", ) temp_cache.store(access_token_key, "expired-access-token-123") temp_cache.store(refresh_token_key, "refresh-token-123") @@ -433,7 +444,7 @@ def test_oauth_code_successful_refresh_token_flow( protocol="http", role="ANALYST", oauth_client_secret="testClientSecret", - oauth_token_request_url=f"http://{wiremock_client.wiremock_host}:{wiremock_client.wiremock_http_port}/oauth/token-request", + oauth_token_request_url=token_request_url, oauth_authorization_url=f"http://{wiremock_client.wiremock_host}:{wiremock_client.wiremock_http_port}/oauth/authorize", oauth_redirect_uri="http://localhost:8009/snowflake/oauth-redirect", host=wiremock_client.wiremock_host, @@ -485,11 +496,20 @@ def test_oauth_code_expired_refresh_token_flow( ) user = "testUser" + token_request_url = f"http://{wiremock_client.wiremock_host}:{wiremock_client.wiremock_http_port}/oauth/token-request" access_token_key = TokenKey( - user, wiremock_client.wiremock_host, TokenType.OAUTH_ACCESS_TOKEN + token_type=TokenType.OAUTH_ACCESS_TOKEN, + idp=token_request_url, + snowflake=wiremock_client.wiremock_host, + username=user, + role="ANALYST", ) refresh_token_key = TokenKey( - user, wiremock_client.wiremock_host, TokenType.OAUTH_REFRESH_TOKEN + token_type=TokenType.OAUTH_REFRESH_TOKEN, + idp=token_request_url, + snowflake=wiremock_client.wiremock_host, + username=user, + role="ANALYST", ) temp_cache.store(access_token_key, "expired-access-token-123") temp_cache.store(refresh_token_key, "expired-refresh-token-123") @@ -503,7 +523,7 @@ def test_oauth_code_expired_refresh_token_flow( protocol="http", role="ANALYST", oauth_client_secret="testClientSecret", - oauth_token_request_url=f"http://{wiremock_client.wiremock_host}:{wiremock_client.wiremock_http_port}/oauth/token-request", + oauth_token_request_url=token_request_url, oauth_authorization_url=f"http://{wiremock_client.wiremock_host}:{wiremock_client.wiremock_http_port}/oauth/authorize", oauth_redirect_uri="http://localhost:8009/snowflake/oauth-redirect", host=wiremock_client.wiremock_host, @@ -556,11 +576,20 @@ def test_client_creds_successful_flow( wiremock_generic_mappings_dir / "snowflake_disconnect_successful.json" ) user = "testUser" + token_request_url = f"http://{wiremock_client.wiremock_host}:{wiremock_client.wiremock_http_port}/oauth/token-request" access_token_key = TokenKey( - user, wiremock_client.wiremock_host, TokenType.OAUTH_ACCESS_TOKEN + token_type=TokenType.OAUTH_ACCESS_TOKEN, + idp=token_request_url, + snowflake=wiremock_client.wiremock_host, + username=user, + role="ANALYST", ) refresh_token_key = TokenKey( - user, wiremock_client.wiremock_host, TokenType.OAUTH_REFRESH_TOKEN + token_type=TokenType.OAUTH_REFRESH_TOKEN, + idp=token_request_url, + snowflake=wiremock_client.wiremock_host, + username=user, + role="ANALYST", ) temp_cache.store(access_token_key, "unused-access-token-123") temp_cache.store(refresh_token_key, "unused-refresh-token-123") @@ -573,7 +602,7 @@ def test_client_creds_successful_flow( account="testAccount", protocol="http", role="ANALYST", - oauth_token_request_url=f"http://{wiremock_client.wiremock_host}:{wiremock_client.wiremock_http_port}/oauth/token-request", + oauth_token_request_url=token_request_url, host=wiremock_client.wiremock_host, port=wiremock_client.wiremock_http_port, oauth_enable_refresh_tokens=True, @@ -654,11 +683,20 @@ def test_client_creds_expired_refresh_token_flow( ) user = "testUser" + token_request_url = f"http://{wiremock_client.wiremock_host}:{wiremock_client.wiremock_http_port}/oauth/token-request" access_token_key = TokenKey( - user, wiremock_client.wiremock_host, TokenType.OAUTH_ACCESS_TOKEN + token_type=TokenType.OAUTH_ACCESS_TOKEN, + idp=token_request_url, + snowflake=wiremock_client.wiremock_host, + username=user, + role="ANALYST", ) refresh_token_key = TokenKey( - user, wiremock_client.wiremock_host, TokenType.OAUTH_REFRESH_TOKEN + token_type=TokenType.OAUTH_REFRESH_TOKEN, + idp=token_request_url, + snowflake=wiremock_client.wiremock_host, + username=user, + role="ANALYST", ) temp_cache.store(access_token_key, "expired-access-token-123") temp_cache.store(refresh_token_key, "expired-refresh-token-123") @@ -670,7 +708,7 @@ def test_client_creds_expired_refresh_token_flow( protocol="http", role="ANALYST", oauth_client_secret="testClientSecret", - oauth_token_request_url=f"http://{wiremock_client.wiremock_host}:{wiremock_client.wiremock_http_port}/oauth/token-request", + oauth_token_request_url=token_request_url, host=wiremock_client.wiremock_host, port=wiremock_client.wiremock_http_port, oauth_enable_refresh_tokens=True, diff --git a/test/unit/test_token_cache_key.py b/test/unit/test_token_cache_key.py new file mode 100644 index 0000000000..00bce719d7 --- /dev/null +++ b/test/unit/test_token_cache_key.py @@ -0,0 +1,287 @@ +"""Tests for the v2 token cache key: normalization, building, and golden hash.""" +from __future__ import annotations + +import hashlib +import json + +import pytest + +from snowflake.connector.token_cache import ( + TokenKey, + TokenType, + _InvalidTokenKeyError, + build_cache_key, + normalize_identifier, + normalize_url, +) + +# --------------------------------------------------------------------------- +# normalize_url +# --------------------------------------------------------------------------- + + +def test_normalize_url_strips_https_scheme(): + assert normalize_url("https://example.com") == "EXAMPLE.COM" + + +def test_normalize_url_strips_http_scheme(): + assert normalize_url("http://example.com") == "EXAMPLE.COM" + + +def test_normalize_url_no_scheme(): + assert normalize_url("example.com") == "EXAMPLE.COM" + + +def test_normalize_url_preserves_port_and_path(): + assert ( + normalize_url("https://login.microsoftonline.com:443/tenant-id/oauth2/v2.0") + == "LOGIN.MICROSOFTONLINE.COM:443/TENANT-ID/OAUTH2/V2.0" + ) + + +def test_normalize_url_strips_userinfo(): + assert normalize_url("https://user:pass@example.com/path") == "EXAMPLE.COM/PATH" + + +def test_normalize_url_drops_query_and_fragment(): + assert normalize_url("https://example.com/path?q=1#frag") == "EXAMPLE.COM/PATH" + + +def test_normalize_url_trims_root_trailing_slash(): + assert normalize_url("https://example.com/") == "EXAMPLE.COM" + + +def test_normalize_url_keeps_non_root_trailing_slash_stripped(): + assert normalize_url("https://example.com/path/") == "EXAMPLE.COM/PATH" + + +def test_normalize_url_uppercases(): + assert ( + normalize_url("https://myorg-myaccount.privatelink.snowflakecomputing.com") + == "MYORG-MYACCOUNT.PRIVATELINK.SNOWFLAKECOMPUTING.COM" + ) + + +# --------------------------------------------------------------------------- +# normalize_identifier +# --------------------------------------------------------------------------- + + +def test_normalize_identifier_unquoted_uppercased(): + assert normalize_identifier("north_america") == "NORTH_AMERICA" + + +def test_normalize_identifier_quoted_segment_verbatim(): + assert normalize_identifier('"First Last"') == '"First Last"' + + +def test_normalize_identifier_mixed(): + assert ( + normalize_identifier('"First Last"@long-corporate-domain.example.com') + == '"First Last"@LONG-CORPORATE-DOMAIN.EXAMPLE.COM' + ) + + +def test_normalize_identifier_role_with_quoted_spaces(): + assert ( + normalize_identifier('"Analyst Role With Spaces":north_america:prod:readonly') + == '"Analyst Role With Spaces":NORTH_AMERICA:PROD:READONLY' + ) + + +def test_normalize_identifier_empty(): + assert normalize_identifier("") == "" + + +# --------------------------------------------------------------------------- +# build_cache_key — validation +# --------------------------------------------------------------------------- + + +def test_build_cache_key_rejects_empty_snowflake(): + key = TokenKey( + token_type=TokenType.MFA_TOKEN, + idp="https://example.com", + snowflake="", + username="user", + role="", + ) + with pytest.raises(_InvalidTokenKeyError): + build_cache_key(key) + + +def test_build_cache_key_rejects_empty_username(): + key = TokenKey( + token_type=TokenType.MFA_TOKEN, + idp="https://example.com", + snowflake="https://example.snowflakecomputing.com", + username="", + role="", + ) + with pytest.raises(_InvalidTokenKeyError): + build_cache_key(key) + + +# --------------------------------------------------------------------------- +# Golden hash (LOCK — must not change) +# --------------------------------------------------------------------------- + + +def test_golden_hash(): + """Assert byte-exact parity with the cross-driver golden vector in 00-INDEX.md §3. + + The golden vector uses uppercase content inside the double-quoted identifier + segments (e.g. ``"FIRST LAST"``) because those represent Snowflake quoted + identifiers whose content was already uppercased before caching. + ``normalize_identifier`` preserves quoted segments verbatim, so the content + inside quotes must already be in the correct case before normalization. + """ + idp_raw = "https://login.microsoftonline.com:443/tenant-id/oauth2/v2.0" + snowflake_raw = "https://myorg-myaccount.privatelink.snowflakecomputing.com" + # Quoted segments have uppercase content — this matches the Rust golden key. + username_raw = '"FIRST LAST"@long-corporate-domain.example.com' + role_raw = '"ANALYST ROLE WITH SPACES":north_america:prod:readonly' + + canonical = json.dumps( + { + "idp": normalize_url(idp_raw), + "role": normalize_identifier(role_raw), + "snowflake": normalize_url(snowflake_raw), + "token_type": "DPOP_BUNDLED_ACCESS_TOKEN", + "username": normalize_identifier(username_raw), + }, + sort_keys=True, + separators=(",", ":"), + ) + digest = hashlib.sha256(canonical.encode("utf-8")).hexdigest() + assert f"SnowflakeTokenCache.v2.{digest}" == ( + "SnowflakeTokenCache.v2." + "75ff2ad65a68afb402f125f62894697673c5ef3d863aba466d16b7a81053d1f4" + ) + + +# --------------------------------------------------------------------------- +# build_cache_key — prefix and format +# --------------------------------------------------------------------------- + + +def test_build_cache_key_prefix(): + key = TokenKey( + token_type=TokenType.OAUTH_ACCESS_TOKEN, + idp="https://login.example.com/oauth2", + snowflake="https://org.snowflakecomputing.com", + username="alice", + role="analyst", + ) + result = build_cache_key(key) + assert result.startswith("SnowflakeTokenCache.v2.") + + +def test_build_cache_key_hash_is_lowercase_hex(): + key = TokenKey( + token_type=TokenType.ID_TOKEN, + idp="https://host.example.com", + snowflake="https://host.example.com", + username="bob", + role="", + ) + suffix = build_cache_key(key).split(".")[-1] + assert suffix == suffix.lower() + assert len(suffix) == 64 + + +# --------------------------------------------------------------------------- +# Dimension isolation — different field → different key +# --------------------------------------------------------------------------- + + +def _base_key(**overrides) -> TokenKey: + defaults = dict( + token_type=TokenType.OAUTH_ACCESS_TOKEN, + idp="https://idp.example.com/oauth2", + snowflake="https://org.snowflakecomputing.com", + username="alice", + role="analyst", + ) + defaults.update(overrides) + return TokenKey(**defaults) + + +def test_different_snowflake_host_yields_different_key(): + k1 = build_cache_key(_base_key(snowflake="https://org1.snowflakecomputing.com")) + k2 = build_cache_key(_base_key(snowflake="https://org2.snowflakecomputing.com")) + assert k1 != k2 + + +def test_different_idp_yields_different_key(): + k1 = build_cache_key(_base_key(idp="https://idp1.example.com/oauth2")) + k2 = build_cache_key(_base_key(idp="https://idp2.example.com/oauth2")) + assert k1 != k2 + + +def test_different_role_yields_different_key(): + k1 = build_cache_key(_base_key(role="analyst")) + k2 = build_cache_key(_base_key(role="sysadmin")) + assert k1 != k2 + + +def test_different_token_type_yields_different_key(): + k1 = build_cache_key(_base_key(token_type=TokenType.OAUTH_ACCESS_TOKEN)) + k2 = build_cache_key(_base_key(token_type=TokenType.OAUTH_REFRESH_TOKEN)) + assert k1 != k2 + + +def test_mfa_empty_role_yields_stable_distinct_key(): + mfa_key = TokenKey( + token_type=TokenType.MFA_TOKEN, + idp="https://org.snowflakecomputing.com", + snowflake="https://org.snowflakecomputing.com", + username="alice", + role="", + ) + id_token_key = TokenKey( + token_type=TokenType.ID_TOKEN, + idp="https://org.snowflakecomputing.com", + snowflake="https://org.snowflakecomputing.com", + username="alice", + role="", + ) + mfa_result = build_cache_key(mfa_key) + assert mfa_result == build_cache_key(mfa_key) + assert mfa_result != build_cache_key(id_token_key) + + +def test_different_username_yields_different_key(): + k1 = build_cache_key(_base_key(username="alice")) + k2 = build_cache_key(_base_key(username="bob")) + assert k1 != k2 + + +# --------------------------------------------------------------------------- +# Normalization is applied consistently +# --------------------------------------------------------------------------- + + +def test_case_insensitive_for_url_fields(): + """Uppercase and lowercase URLs produce the same key.""" + k_lower = build_cache_key( + _base_key( + idp="https://idp.example.com/oauth2", + snowflake="https://org.snowflakecomputing.com", + ) + ) + k_upper = build_cache_key( + _base_key( + idp="https://IDP.EXAMPLE.COM/OAUTH2", + snowflake="https://ORG.SNOWFLAKECOMPUTING.COM", + ) + ) + assert k_lower == k_upper + + +def test_scheme_stripped_from_url(): + k_with_scheme = build_cache_key( + _base_key(snowflake="https://org.snowflakecomputing.com") + ) + k_no_scheme = build_cache_key(_base_key(snowflake="org.snowflakecomputing.com")) + assert k_with_scheme == k_no_scheme From e3f46bb4630407f1e5fe2687c3f3f3985dedf7a9 Mon Sep 17 00:00:00 2001 From: Michal Hofman Date: Thu, 16 Jul 2026 09:00:05 +0200 Subject: [PATCH 2/6] SNOW-3784431: remove prompt cross-references and simplify verbose comments Co-authored-by: Cursor --- test/unit/test_token_cache_key.py | 12 +++++------- 1 file changed, 5 insertions(+), 7 deletions(-) diff --git a/test/unit/test_token_cache_key.py b/test/unit/test_token_cache_key.py index 00bce719d7..8da7b2fd20 100644 --- a/test/unit/test_token_cache_key.py +++ b/test/unit/test_token_cache_key.py @@ -128,17 +128,15 @@ def test_build_cache_key_rejects_empty_username(): def test_golden_hash(): - """Assert byte-exact parity with the cross-driver golden vector in 00-INDEX.md §3. + """Assert that the cache key hash is stable and must not change between releases. - The golden vector uses uppercase content inside the double-quoted identifier - segments (e.g. ``"FIRST LAST"``) because those represent Snowflake quoted - identifiers whose content was already uppercased before caching. - ``normalize_identifier`` preserves quoted segments verbatim, so the content - inside quotes must already be in the correct case before normalization. + Quoted identifier segments (e.g. ``"FIRST LAST"``) contain uppercase content + because ``normalize_identifier`` preserves them verbatim — the content inside + quotes must already be in the correct case before normalization is called. """ idp_raw = "https://login.microsoftonline.com:443/tenant-id/oauth2/v2.0" snowflake_raw = "https://myorg-myaccount.privatelink.snowflakecomputing.com" - # Quoted segments have uppercase content — this matches the Rust golden key. + # Quoted segments have uppercase content because normalize_identifier preserves them verbatim. username_raw = '"FIRST LAST"@long-corporate-domain.example.com' role_raw = '"ANALYST ROLE WITH SPACES":north_america:prod:readonly' From f8717993e5f04c89e4b029439dc75f24b031347b Mon Sep 17 00:00:00 2001 From: Michal Hofman Date: Fri, 17 Jul 2026 13:00:01 +0200 Subject: [PATCH 3/6] fix: address token cache key v2 review findings Threads the role field through the async auth path (aio/_connection.py and aio/auth/_auth.py) so async SSO/OAuth/MFA flows build the same v2 cache keys as their sync counterparts, and replaces the stale positional TokenKey call sites in the async authenticate override. Extends legacy migration so both KeyringTokenCache and FileTokenCache read and migrate entries from the two prior key layouts (hashed-account and string-key), deriving the legacy host from the IdP hostname for OAuth tokens and the Snowflake host otherwise via shared _legacy_string_key/_legacy_hash_key helpers. Corrects the KeyringTokenCache docstring and updates the manual async SSO test to the five-field TokenKey. Adds migration coverage for both backends. Co-authored-by: Cursor --- DESCRIPTION.md | 2 +- src/snowflake/connector/aio/_connection.py | 3 + src/snowflake/connector/aio/auth/_auth.py | 20 ++- src/snowflake/connector/token_cache.py | 121 ++++++++++++++---- .../sso_it/test_connection_manual_async.py | 8 +- test/unit/test_keyring_token_cache.py | 43 ++++++- test/unit/test_linux_local_file_cache.py | 33 ++++- 7 files changed, 190 insertions(+), 40 deletions(-) diff --git a/DESCRIPTION.md b/DESCRIPTION.md index d44f750a14..a79454b01c 100644 --- a/DESCRIPTION.md +++ b/DESCRIPTION.md @@ -8,7 +8,7 @@ Source code is also available at: https://github.com/snowflakedb/snowflake-conne # Release Notes - NEXT_RELEASE(TBD) - - Fixed token cache key collisions for multi-account (shared IdP) and multi-role scenarios by switching to a versioned, SHA256-hashed canonical-JSON key applied uniformly across macOS/Windows keyring and the Linux file backend (SNOW-3784431). + - Fixed token cache key collisions for multi-account (shared IdP) and multi-role scenarios by switching to a versioned, SHA256-hashed canonical-JSON key applied uniformly across macOS/Windows keyring and the Linux file backend (SNOW-3784431). The new key is threaded through both the sync and async authentication paths, and existing cache entries from the two prior key layouts are transparently migrated to the new format on first use. - Added support for Python 3.14t (free-threaded). - **Note:** Python 3.14t CI testing excludes `win_arm64` (no `cryptography` wheels available) and `mitmproxy` proxy tests on all platforms (transitive dependencies `aioquic`/`pylsqpack` lack free-threaded-compatible wheels). diff --git a/src/snowflake/connector/aio/_connection.py b/src/snowflake/connector/aio/_connection.py index f62da43069..7becdac3ff 100644 --- a/src/snowflake/connector/aio/_connection.py +++ b/src/snowflake/connector/aio/_connection.py @@ -323,6 +323,7 @@ async def __open_connection(self): self.host, self.user, self._session_parameters, + role=self._role or "", ) # Depending on whether self._rest.id_token is available we do different # auth_instance @@ -395,6 +396,7 @@ async def __open_connection(self): refresh_token_enabled=self._oauth_enable_refresh_tokens, external_browser_timeout=self._external_browser_timeout, enable_single_use_refresh_tokens=self._oauth_enable_single_use_refresh_tokens, + role=self._role or "", ) elif self._authenticator == OAUTH_CLIENT_CREDENTIALS: if self._role and (self._oauth_scope == ""): @@ -428,6 +430,7 @@ async def __open_connection(self): self.host, self.user, self._session_parameters, + role="", ) self.auth_class = AuthByUsrPwdMfa( password=self._password, diff --git a/src/snowflake/connector/aio/auth/_auth.py b/src/snowflake/connector/aio/auth/_auth.py index a618a0ed41..6a0f4398a0 100644 --- a/src/snowflake/connector/aio/auth/_auth.py +++ b/src/snowflake/connector/aio/auth/_auth.py @@ -35,7 +35,7 @@ ReauthenticationRequest, ) from ...sqlstate import SQLSTATE_CONNECTION_WAS_NOT_ESTABLISHED -from ...token_cache import TokenType +from ...token_cache import TokenKey, TokenType from ._no_auth import AuthNoAuth if TYPE_CHECKING: @@ -288,7 +288,13 @@ async def post_request_wrapper(self, url, headers, body) -> None: # raise an exception for reauth without id_token self._rest.id_token = None self._delete_temporary_credential( - self._rest._host, user, TokenType.ID_TOKEN + TokenKey( + token_type=TokenType.ID_TOKEN, + idp=self._rest._host, + snowflake=self._rest._host, + username=user, + role=role or "", + ) ) raise ReauthenticationRequest( ProgrammingError( @@ -320,7 +326,13 @@ async def post_request_wrapper(self, url, headers, body) -> None: if isinstance(auth_instance, AuthByUsrPwdMfa): self._delete_temporary_credential( - self._rest._host, user, TokenType.MFA_TOKEN + TokenKey( + token_type=TokenType.MFA_TOKEN, + idp=self._rest._host, + snowflake=self._rest._host, + username=user, + role="", + ) ) Error.errorhandler_wrapper( self._rest._connection, @@ -388,7 +400,7 @@ async def post_request_wrapper(self, url, headers, body) -> None: mfa_token=ret["data"].get("mfaToken"), ) self.write_temporary_credentials( - self._rest._host, user, session_parameters, ret + self._rest._host, user, session_parameters, ret, role=role or "" ) if ret["data"] and "sessionId" in ret["data"]: self._rest._connection._session_id = ret["data"].get("sessionId") diff --git a/src/snowflake/connector/token_cache.py b/src/snowflake/connector/token_cache.py index 54fd849459..4407479803 100644 --- a/src/snowflake/connector/token_cache.py +++ b/src/snowflake/connector/token_cache.py @@ -8,6 +8,7 @@ import re import stat import sys +import urllib.parse from abc import ABC, abstractmethod from dataclasses import dataclass from enum import Enum @@ -117,6 +118,32 @@ def build_cache_key(key: TokenKey) -> str: return f"SnowflakeTokenCache.v2.{digest}" +def _legacy_string_key(key: TokenKey) -> str: + """Reconstruct the pre-v2 ``{HOST}:{USER}:{TOKEN_TYPE}`` string key. + + OAuth tokens historically keyed on the IdP hostname + (``urlparse(token_request_url).hostname``); all other flows keyed on the + Snowflake host. Used only to locate and migrate legacy cache entries. + """ + if key.token_type in ( + TokenType.OAUTH_ACCESS_TOKEN, + TokenType.OAUTH_REFRESH_TOKEN, + ): + host = urllib.parse.urlparse(key.idp).hostname or key.idp + else: + host = key.snowflake + if not host: + raise _InvalidTokenKeyError("Invalid key, host is empty") + if not key.username: + raise _InvalidTokenKeyError("Invalid key, user is empty") + return f"{host.upper()}:{key.username.upper()}:{key.token_type.value}" + + +def _legacy_hash_key(key: TokenKey) -> str: + """SHA-256 hex of the legacy string key (the pre-v2 ``hash_key`` layout).""" + return hashlib.sha256(_legacy_string_key(key).encode("utf-8")).hexdigest() + + def _warn(warning: str) -> None: logger.warning(warning) print("Warning: " + warning, file=sys.stderr) @@ -220,6 +247,10 @@ class FileTokenCache(TokenCache): Note: the filename (``credential_cache_v1.json``) is unchanged for backward compatibility; the ``v2`` in the key prefix refers to the key-format version, not the file format. + + For backward compatibility, :meth:`retrieve` also checks the legacy layout + where the map key was ``sha256("{HOST}:{USER}:{TOKEN_TYPE}")``; matching + entries are silently migrated to the v2 key on first use. """ @staticmethod @@ -267,11 +298,23 @@ def retrieve(self, key: TokenKey) -> str | None: ) with FileLock(self.lock_file()): cache = self._read_cache_file() - token = cache["tokens"].get(final_key, None) + tokens = cache["tokens"] + token = tokens.get(final_key, None) if isinstance(token, str): return token - else: - return None + # Legacy v1 fallback: entries keyed by sha256("{HOST}:{USER}:{TYPE}"). + legacy_key = _legacy_hash_key(key) + legacy_token = tokens.get(legacy_key, None) + if isinstance(legacy_token, str): + tokens[final_key] = legacy_token + tokens.pop(legacy_key, None) + self._write_cache_file(cache) + self.logger.debug( + "migrated legacy file cache entry for %s", + key.token_type.value, + ) + return legacy_token + return None except _FileTokenCacheError as e: self.logger.error(f"Failed to retrieve token: {e=}") return None @@ -467,13 +510,21 @@ class KeyringTokenCache(TokenCache): - macOS: Stores tokens in Keychain - Windows: Stores tokens in Windows Credential Manager - The v2 cache key (``SnowflakeTokenCache.v2.``) is used as both - the keyring service name and the account field, ensuring a single Keychain - ACL entry per token and no plaintext identifiers in the OS credential store. + The v2 cache key (``SnowflakeTokenCache.v2.``) is used as the + keyring service name, and the uppercase username is used as the account + field. This ensures a distinct entry per token dimension while still + letting related tokens share Keychain visibility per account. - For backward compatibility, :meth:`retrieve` also checks the legacy layout - where the service was ``{HOST}:{USER}:{TOKEN_TYPE}`` and the account was the - uppercase username; matching entries are silently migrated to v2 on first use. + For backward compatibility, :meth:`retrieve` also checks two legacy layouts + and silently migrates matching entries to v2 on first use: + + - hash layout (immediately prior): service ``com.snowflake.connector.python`` + with account ``sha256("{HOST}:{USER}:{TOKEN_TYPE}")`` + - string layout (oldest): service ``{HOST}:{USER}:{TOKEN_TYPE}`` with account + equal to the uppercase username + + For OAuth tokens the legacy ``HOST`` is the IdP hostname; for other flows it + is the Snowflake host. """ SERVICE_NAME = "com.snowflake.connector.python" @@ -509,25 +560,43 @@ def retrieve(self, key: TokenKey) -> str | None: ) def _retrieve_legacy(self, key: TokenKey) -> str | None: - """Try to read from the old per-token-type service layout and migrate to v2.""" - legacy_service = ( - f"{key.snowflake.upper()}:{key.username.upper()}:{key.token_type.value}" - ) + """Read from pre-v2 keyring layouts and migrate matching entries to v2. + + Two historical layouts are checked, newest first: + - hash layout: service ``SERVICE_NAME``, account ``sha256(string_key)`` + - string layout: service ``string_key``, account uppercase username + + where ``string_key`` is ``{HOST}:{USER}:{TOKEN_TYPE}`` (HOST being the + IdP hostname for OAuth tokens and the Snowflake host otherwise). + """ try: - token = keyring.get_password(legacy_service, key.username.upper()) - except (keyring.errors.KeyringError, _InvalidTokenKeyError): + legacy_string_key = _legacy_string_key(key) + legacy_hash_key = _legacy_hash_key(key) + except _InvalidTokenKeyError: return None - if token is None: - return None - self.store(key, token) - try: - keyring.delete_password(legacy_service, key.username.upper()) - except Exception: - pass - self.logger.debug( - "migrated legacy keyring entry for %s", key.token_type.value - ) - return token + + account = key.username.upper() + lookups = [ + (self.SERVICE_NAME, legacy_hash_key), + (legacy_string_key, account), + ] + for service, acct in lookups: + try: + token = keyring.get_password(service, acct) + except (keyring.errors.KeyringError, _InvalidTokenKeyError): + continue + if token is None: + continue + self.store(key, token) + try: + keyring.delete_password(service, acct) + except Exception: + pass + self.logger.debug( + "migrated legacy keyring entry for %s", key.token_type.value + ) + return token + return None def remove(self, key: TokenKey) -> None: try: diff --git a/test/integ/aio_it/sso_it/test_connection_manual_async.py b/test/integ/aio_it/sso_it/test_connection_manual_async.py index bfe5482604..7dcb821198 100644 --- a/test/integ/aio_it/sso_it/test_connection_manual_async.py +++ b/test/integ/aio_it/sso_it/test_connection_manual_async.py @@ -89,9 +89,11 @@ async def test_connect_externalbrowser(token_validity_test_values): TokenCache.make().remove( TokenKey( - CONNECTION_PARAMETERS_SSO["host"], - CONNECTION_PARAMETERS_SSO["user"], - TokenType.ID_TOKEN, + token_type=TokenType.ID_TOKEN, + idp=CONNECTION_PARAMETERS_SSO["host"], + snowflake=CONNECTION_PARAMETERS_SSO["host"], + username=CONNECTION_PARAMETERS_SSO["user"], + role=CONNECTION_PARAMETERS_SSO.get("role", "") or "", ) ) # delete existing temporary credential diff --git a/test/unit/test_keyring_token_cache.py b/test/unit/test_keyring_token_cache.py index 38ebfca3a9..92e561b729 100644 --- a/test/unit/test_keyring_token_cache.py +++ b/test/unit/test_keyring_token_cache.py @@ -9,6 +9,8 @@ KeyringTokenCache, TokenKey, TokenType, + _legacy_hash_key, + _legacy_string_key, build_cache_key, ) @@ -65,19 +67,50 @@ def test_retrieves_from_v2_key(self, cache, mock_keyring): assert cache.retrieve(KEY) == "tok123" mock_keyring.get_password.assert_called_once_with(FINAL_KEY, ACCOUNT) - def test_falls_back_to_legacy_and_migrates(self, cache, mock_keyring): - legacy_service = f"{KEY.snowflake.upper()}:{KEY.username.upper()}:{KEY.token_type.value}" + def test_falls_back_to_hash_layout_and_migrates(self, cache, mock_keyring): + """Immediately-prior layout: SERVICE_NAME + sha256(string_key) account.""" + legacy_hash = _legacy_hash_key(KEY) + # v2 miss, then hash-layout hit on the first legacy lookup. mock_keyring.get_password.side_effect = [None, "legacy_tok"] result = cache.retrieve(KEY) assert result == "legacy_tok" mock_keyring.get_password.assert_has_calls( [ call(FINAL_KEY, ACCOUNT), - call(legacy_service, ACCOUNT), + call(cache.SERVICE_NAME, legacy_hash), ] ) - mock_keyring.set_password.assert_called_once_with(FINAL_KEY, ACCOUNT, "legacy_tok") - mock_keyring.delete_password.assert_called_once_with(legacy_service, ACCOUNT) + mock_keyring.set_password.assert_called_once_with( + FINAL_KEY, ACCOUNT, "legacy_tok" + ) + mock_keyring.delete_password.assert_called_once_with( + cache.SERVICE_NAME, legacy_hash + ) + + def test_falls_back_to_string_layout_and_migrates(self, cache, mock_keyring): + """Oldest layout: string_key service + uppercase username account.""" + legacy_hash = _legacy_hash_key(KEY) + legacy_string = _legacy_string_key(KEY) + # v2 miss, hash-layout miss, then string-layout hit. + mock_keyring.get_password.side_effect = [None, None, "legacy_tok"] + result = cache.retrieve(KEY) + assert result == "legacy_tok" + mock_keyring.get_password.assert_has_calls( + [ + call(FINAL_KEY, ACCOUNT), + call(cache.SERVICE_NAME, legacy_hash), + call(legacy_string, ACCOUNT), + ] + ) + mock_keyring.set_password.assert_called_once_with( + FINAL_KEY, ACCOUNT, "legacy_tok" + ) + mock_keyring.delete_password.assert_called_once_with(legacy_string, ACCOUNT) + + def test_oauth_legacy_string_key_uses_idp_host(self, cache, mock_keyring): + """OAuth legacy keys must be built from the IdP hostname, not the SF host.""" + legacy_string = _legacy_string_key(KEY) + assert legacy_string == "IDP.EXAMPLE.COM:ALICE:OAUTH_ACCESS_TOKEN" def test_returns_none_when_not_found_anywhere(self, cache, mock_keyring): mock_keyring.get_password.return_value = None diff --git a/test/unit/test_linux_local_file_cache.py b/test/unit/test_linux_local_file_cache.py index d2c384f22c..9cad72199f 100644 --- a/test/unit/test_linux_local_file_cache.py +++ b/test/unit/test_linux_local_file_cache.py @@ -12,7 +12,12 @@ pytestmark = pytest.mark.skipif(not IS_LINUX, reason="Testing on linux only") try: - from snowflake.connector.token_cache import FileTokenCache, TokenKey, TokenType + from snowflake.connector.token_cache import ( + FileTokenCache, + TokenKey, + TokenType, + _legacy_hash_key, + ) CRED_TYPE_0 = TokenType.ID_TOKEN CRED_TYPE_1 = TokenType.MFA_TOKEN @@ -283,3 +288,29 @@ def test_file_tokens_is_not_dict(tmpdir, monkeypatch): cache.store(make_key(HOST_0, USER_0, CRED_TYPE_0), CRED_0) assert cache.retrieve(make_key(HOST_0, USER_0, CRED_TYPE_0)) == CRED_0 cache.cache_file().unlink() + + +def test_retrieve_migrates_legacy_hash_key(tmpdir, monkeypatch): + """Tokens stored under the pre-v2 sha256(string_key) key are migrated to v2.""" + import json + + from snowflake.connector.token_cache import build_cache_key + + monkeypatch.setenv("SF_TEMPORARY_CREDENTIAL_CACHE_DIR", str(tmpdir)) + cache = FileTokenCache.make() + assert cache + cache.cache_file().unlink(missing_ok=True) + + key = make_key(HOST_0, USER_0, CRED_TYPE_0) + legacy_key = _legacy_hash_key(key) + cache.cache_file().touch(0o600) + cache.cache_file().write_text(json.dumps({"tokens": {legacy_key: CRED_0}})) + + # Retrieval finds the legacy entry and returns it. + assert cache.retrieve(key) == CRED_0 + + # The entry is migrated: now stored under the v2 key and legacy key removed. + tokens = json.loads(cache.cache_file().read_text("utf-8"))["tokens"] + assert tokens.get(build_cache_key(key)) == CRED_0 + assert legacy_key not in tokens + cache.cache_file().unlink() From c97321cd6826a9ee44195d8f803b2daf29bbf4f5 Mon Sep 17 00:00:00 2001 From: Michal Hofman Date: Tue, 8 Sep 2026 13:21:39 +0200 Subject: [PATCH 4/6] fix: correct legacy string key order for ID/MFA tokens and resolve merge conflict The old TokenKey call sites for ID/MFA passed args positionally as TokenKey(host, user, type) into a (user, host, tokenType) dataclass, swapping the fields. The resulting string_key() produced USER:HOST:TYPE, not HOST:USER:TYPE. _legacy_string_key now replicates this swap for ID_TOKEN and MFA_TOKEN so existing pre-v2 cache entries are found and migrated correctly. OAuth tokens were called in the correct order and are unaffected. Also resolves the DESCRIPTION.md merge conflict from the main merge. Co-authored-by: Cursor --- DESCRIPTION.md | 4 +--- src/snowflake/connector/token_cache.py | 32 ++++++++++++++++++-------- test/unit/test_keyring_token_cache.py | 14 +++++++++++ 3 files changed, 37 insertions(+), 13 deletions(-) diff --git a/DESCRIPTION.md b/DESCRIPTION.md index 290b02091f..fec49400bd 100644 --- a/DESCRIPTION.md +++ b/DESCRIPTION.md @@ -7,12 +7,10 @@ https://docs.snowflake.com/ Source code is also available at: https://github.com/snowflakedb/snowflake-connector-python # Release Notes -<<<<<<< michal.hofman/SNOW-3784431-token-cache-key-v2 - NEXT_RELEASE(TBD) - Fixed token cache key collisions for multi-account (shared IdP) and multi-role scenarios by switching to a versioned, SHA256-hashed canonical-JSON key applied uniformly across macOS/Windows keyring and the Linux file backend (SNOW-3784431). The new key is threaded through both the sync and async authentication paths, and existing cache entries from the two prior key layouts are transparently migrated to the new format on first use. -======= + - v4.7.1(Jul 15,2026) ->>>>>>> main - Added support for Python 3.14t (free-threaded). - **Note:** Python 3.14t CI testing excludes `win_arm64` (no `cryptography` wheels available) and `mitmproxy` proxy tests on all platforms (transitive dependencies `aioquic`/`pylsqpack` lack free-threaded-compatible wheels). - Improved verification of TLS connections (SNOW-3675579). diff --git a/src/snowflake/connector/token_cache.py b/src/snowflake/connector/token_cache.py index 4407479803..d2bbe22d33 100644 --- a/src/snowflake/connector/token_cache.py +++ b/src/snowflake/connector/token_cache.py @@ -119,24 +119,36 @@ def build_cache_key(key: TokenKey) -> str: def _legacy_string_key(key: TokenKey) -> str: - """Reconstruct the pre-v2 ``{HOST}:{USER}:{TOKEN_TYPE}`` string key. + """Reconstruct the pre-v2 string key used by the old ``hash_key()`` / ``string_key()`` methods. - OAuth tokens historically keyed on the IdP hostname - (``urlparse(token_request_url).hostname``); all other flows keyed on the - Snowflake host. Used only to locate and migrate legacy cache entries. + The historical ``TokenKey(host, user, tokenType)`` call sites had the + **positional** arguments in that order, but the dataclass was declared as + ``(user, host, tokenType)``. This swap means: + - OAuth (correctly called ``TokenKey(user, idp_host, ...)``): + ``string_key()`` produced ``{IDP_HOST}:{USER}:{TYPE}`` + - ID / MFA (incorrectly called ``TokenKey(host, user, ...)``): + ``string_key()`` produced ``{USER}:{HOST}:{TYPE}`` + + Used only to locate and migrate pre-v2 cache entries. """ if key.token_type in ( TokenType.OAUTH_ACCESS_TOKEN, TokenType.OAUTH_REFRESH_TOKEN, ): + # OAuth was called correctly: key order is HOST:USER:TYPE. host = urllib.parse.urlparse(key.idp).hostname or key.idp + if not host: + raise _InvalidTokenKeyError("Invalid key, host is empty") + if not key.username: + raise _InvalidTokenKeyError("Invalid key, user is empty") + return f"{host.upper()}:{key.username.upper()}:{key.token_type.value}" else: - host = key.snowflake - if not host: - raise _InvalidTokenKeyError("Invalid key, host is empty") - if not key.username: - raise _InvalidTokenKeyError("Invalid key, user is empty") - return f"{host.upper()}:{key.username.upper()}:{key.token_type.value}" + # ID/MFA was called with swapped positional args: key order is USER:HOST:TYPE. + if not key.snowflake: + raise _InvalidTokenKeyError("Invalid key, host is empty") + if not key.username: + raise _InvalidTokenKeyError("Invalid key, user is empty") + return f"{key.username.upper()}:{key.snowflake.upper()}:{key.token_type.value}" def _legacy_hash_key(key: TokenKey) -> str: diff --git a/test/unit/test_keyring_token_cache.py b/test/unit/test_keyring_token_cache.py index 92e561b729..d09cb7f4f2 100644 --- a/test/unit/test_keyring_token_cache.py +++ b/test/unit/test_keyring_token_cache.py @@ -112,6 +112,20 @@ def test_oauth_legacy_string_key_uses_idp_host(self, cache, mock_keyring): legacy_string = _legacy_string_key(KEY) assert legacy_string == "IDP.EXAMPLE.COM:ALICE:OAUTH_ACCESS_TOKEN" + def test_mfa_legacy_string_key_is_user_host_swapped(self, cache, mock_keyring): + """MFA/ID legacy keys are USER:HOST:TYPE because the old call site passed args + in the wrong positional order (TokenKey(host, user, type) into (user, host, type)).""" + mfa_key = TokenKey( + token_type=TokenType.MFA_TOKEN, + idp="myhost.snowflakecomputing.com", + snowflake="myhost.snowflakecomputing.com", + username="ALICE", + role="", + ) + legacy_string = _legacy_string_key(mfa_key) + # Old storage: USER:HOST:TYPE (swapped), not HOST:USER:TYPE + assert legacy_string == "ALICE:MYHOST.SNOWFLAKECOMPUTING.COM:MFA_TOKEN" + def test_returns_none_when_not_found_anywhere(self, cache, mock_keyring): mock_keyring.get_password.return_value = None assert cache.retrieve(KEY) is None From cd5e378ea36267e346c75a2fcaaa7695eca008be Mon Sep 17 00:00:00 2001 From: Michal Hofman Date: Wed, 22 Jul 2026 12:34:41 +0200 Subject: [PATCH 5/6] =?UTF-8?q?SNOW-3784431:=20token=20cache=20key=20v2=20?= =?UTF-8?q?fixup=20=E2=80=94=20type=20in=20prefix,=20flow-specific=20keyDa?= =?UTF-8?q?ta?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Move token_type out of keyData into the key prefix: SnowflakeTokenCache.v2.. - Use flow-specific keyData: * OAuth (OAUTH_ACCESS_TOKEN, OAUTH_REFRESH_TOKEN, DPOP_BUNDLED_ACCESS_TOKEN): 4 fields — idp, role, snowflake, username * MFA / ID token (MFA_TOKEN, ID_TOKEN): 2 fields — snowflake, username only - Make TokenKey.idp and TokenKey.role optional (default ""); required args are now token_type, snowflake, username — matching the MFA/ID flow. - Update all MFA and ID_TOKEN call sites in _auth.py to omit idp and role. - Replace old 5-field golden hash test with vectors A (OAuth/DPoP) and B (MFA). - Add test_mfa_key_has_no_idp_or_role and test_mfa_vs_oauth_key_differ_... - Update DESCRIPTION.md changelog entry. Co-authored-by: Cursor --- .../token-cache-updated/00-INDEX-upd.md | 432 ++++++++++++++++++ .../07-connector-python-upd.md | 300 ++++++++++++ DESCRIPTION.md | 2 +- src/snowflake/connector/auth/_auth.py | 12 - src/snowflake/connector/token_cache.py | 81 ++-- test/unit/test_token_cache_key.py | 103 +++-- 6 files changed, 861 insertions(+), 69 deletions(-) create mode 100644 ---prompts---/token-cache-updated/00-INDEX-upd.md create mode 100644 ---prompts---/token-cache-updated/07-connector-python-upd.md diff --git a/---prompts---/token-cache-updated/00-INDEX-upd.md b/---prompts---/token-cache-updated/00-INDEX-upd.md new file mode 100644 index 0000000000..c58dc2308a --- /dev/null +++ b/---prompts---/token-cache-updated/00-INDEX-upd.md @@ -0,0 +1,432 @@ +# Token Cache Key v2 — Follow-up Shared Contract & Rollout Index + +This is the authoritative cross-driver contract for the **second PR** in the token cache +key redesign. Read this file before making any changes in a per-repo prompt. + +--- + +## 0. Background: what the first PR implemented + +The first round of changes (already landed or in review under the branches below) fixed +the original `{host}:{username}:{type}` key by implementing a v2 format: + +``` +SnowflakeTokenCache.v2. +``` + +where `keyData` was a **5-field** JSON object — compact, keys sorted lexicographically: + +```json +{"idp":"…","role":"…","snowflake":"…","token_type":"…","username":"…"} +``` + +That first PR also: +- Added `normalize_url` (strips scheme/userinfo/query/fragment, uppercases remainder) and + `normalize_identifier` (uppercases outside `"…"` segments, preserves inside verbatim). +- Applied SHA-256 uniformly — one hash before dispatch, no per-backend hashing. +- Wired the same final key string to both OS keystore and JSON file backends. +- Removed legacy separator-injection guards (`;` / `:` in inputs). + +--- + +## 0.1 What this follow-up PR changes (three surgical fixes) + +### Fix 1 — Token type moves from `keyData` into the key prefix + +Old key: `SnowflakeTokenCache.v2.` +New key: `SnowflakeTokenCache.v2.MFA_TOKEN.` + +Putting the token type in the readable prefix lets keystore/keyring tooling identify +and remove specific token classes without decoding the opaque hash. + +### Fix 2 — MFA and ID token keys use only `snowflake` + `username` + +For MFA and ID token flows, `idp` and `role` are removed from `keyData`: +- `role` is absent because role is **not** embedded in MFA or external-browser + authentication calls. Including it would cause misses every time. +- `idp` is absent because MFA/ID token authentication always targets the Snowflake + host directly — there is no separate identity provider endpoint. + +OAuth flows are unaffected: they keep `idp`, `role`, `snowflake`, and `username`. + +### Fix 3 — Golden test updated to prove lowercase preservation inside quotes + +The previous golden test used all-uppercase strings inside double quotes +(e.g., `"FIRST LAST"`), which failed to exercise the quote-preservation branch of +`normalize_identifier`. The new test uses mixed case (e.g., `"First Last"`) and +asserts the quoted portion is preserved verbatim after normalization. + +--- + +## 1. Repositories, tickets, and follow-up branches + +The `` segment is your local git user handle +(`git config user.email | cut -d'@' -f1` or `git config user.name`). + +For each repo: **branch from the first PR's branch** (if it has not yet merged); or +**from the default branch** (if the first PR already merged). + +| # | Repo | Local path | Public remote | Default branch | Ticket | Follow-up branch | +|---|------|-----------|---------------|----------------|--------|-----------------| +| 01 | libsnowflakeclient (C++) | `/Users/mhofman/Projects/libsnowflakeclient` | `ssh://git@github.com/snowflakedb/libsnowflakeclient` | `master` | [SNOW-3784428](https://snowflakecomputing.atlassian.net/browse/SNOW-3784428) | `/SNOW-3784428-token-cache-key-v2-fixup` | +| 02 | snowflake-connector-net (C#) | `/Users/mhofman/Projects/snowflake-connector-net` | `git@github.com:snowflakedb/snowflake-connector-net.git` | `master` | [SNOW-3784414](https://snowflakecomputing.atlassian.net/browse/SNOW-3784414) | `/SNOW-3784414-token-cache-key-v2-fixup` | +| 03 | snowflake-jdbc (Java) | `/Users/mhofman/Projects/snowflake-jdbc` | `ssh://git@github.com/snowflakedb/snowflake-jdbc` | `master` | [SNOW-3784426](https://snowflakecomputing.atlassian.net/browse/SNOW-3784426) | `/SNOW-3784426-token-cache-key-v2-fixup` | +| 04 | snowflake-connector-nodejs (JS/TS) | `/Users/mhofman/Projects/snowflake-connector-nodejs` | `git@github.com:snowflakedb/snowflake-connector-nodejs.git` | `master` | [SNOW-3784415](https://snowflakecomputing.atlassian.net/browse/SNOW-3784415) | `/SNOW-3784415-token-cache-key-v2-fixup` | +| 05 | gosnowflake (Go) | `/Users/mhofman/Projects/gosnowflake` | `ssh://git@github.com/snowflakedb/gosnowflake.git` | `master` | [SNOW-3784429](https://snowflakecomputing.atlassian.net/browse/SNOW-3784429) | `/SNOW-3784429-token-cache-key-v2-fixup` | +| 06 | pdo_snowflake (PHP) | `/Users/mhofman/Projects/pdo_snowflake` | `ssh://git@github.com/snowflakedb/pdo_snowflake` | `master` | [SNOW-3784428](https://snowflakecomputing.atlassian.net/browse/SNOW-3784428) | `/SNOW-3784428-token-cache-key-v2-fixup` | +| 07 | snowflake-connector-python (Python) | `/Users/mhofman/Projects/snowflake-connector-python` | `git@github.com:snowflakedb/snowflake-connector-python.git` | **`main`** | [SNOW-3784431](https://snowflakecomputing.atlassian.net/browse/SNOW-3784431) | `/SNOW-3784431-token-cache-key-v2-fixup` | +| 08 | universal-driver (Rust) | `/Users/mhofman/Projects/universal-driver` | `git@github.com:snowflake-eng/universal-driver.git` | `main` | SNOW-TBD | `/SNOW-TBD-token-cache-key-v2-fixup` | + +**Not the targets** (wrong clones — never touch these): +- Any `/Users/mhofman/Projects/*-private` mirror +- `snowflake-connector-net-playground`, `snowflake-jdbc-playground` +- `pdo_snowflake/libsnowflakeclient/` (vendored tree inside pdo — changes go to repo 01 only) +- Any `/Users/mhofman/Projects/snowdrivers-analysis/drivers/*` copy + +**Rollout order**: (01, 08 in parallel) → (02–05, 07 in parallel) → 06. +Repo 06 (pdo) is blocked on 01 landing first (it vendors a prebuilt `libsnowflakeclient.a`). +Repo 08 (universal-driver) is the Rust reference implementation and can run in parallel with 01. + +--- + +## 2. The final v2 key contract + +### 2.1 Key format + +``` +SnowflakeTokenCache.v2.. +``` + +- Prefix: `SnowflakeTokenCache` (not `Snowflake`) +- Version: `v2` +- Token type: canonical uppercase string, e.g. `MFA_TOKEN`, `OAUTH_ACCESS_TOKEN` +- Hash: **lowercase** hex SHA-256 of the canonical JSON bytes of `keyData` + +The **identical** key string is used for both the OS keystore backend and the JSON file +fallback. Hashing happens **exactly once** when the key is built; backends store verbatim. + +### 2.2 `keyData` fields — differ by flow + +`keyData` does **not** contain `token_type`; the token type appears in the key prefix. + +**OAuth flows** (`OAUTH_ACCESS_TOKEN`, `OAUTH_REFRESH_TOKEN`, `DPOP_BUNDLED_ACCESS_TOKEN`, …): + +| Field | Value | +|-------|-------| +| `idp` | Normalized IdP/token-endpoint URL | +| `role` | Normalized role | +| `snowflake` | Normalized Snowflake server URL | +| `username` | Normalized Snowflake username | + +Sorted field order in canonical JSON: **`idp`, `role`, `snowflake`, `username`**. + +**MFA and ID token flows** (`MFA_TOKEN`, `ID_TOKEN`): + +| Field | Value | +|-------|-------| +| `snowflake` | Normalized Snowflake server URL | +| `username` | Normalized Snowflake username | + +Sorted field order in canonical JSON: **`snowflake`, `username`**. + +`role` is absent for MFA/ID because role is not embedded in those authentication calls. +`idp` is absent because MFA/ID authentication always targets the Snowflake host directly. + +### 2.3 Canonical JSON serialization + +The JSON hashed **must** be byte-for-byte identical across all drivers: + +- **Compact**: no spaces after `:` or `,`, no newlines, no indentation. +- **Keys sorted lexicographically** (Unicode code-point ascending). +- **Standard JSON string escaping**: `"` → `\"`, `\` → `\\`, control chars escaped. +- Serialize to UTF-8 bytes → SHA-256 → **lowercase** hex. + +Idiomatic per language: +- **C++**: emit manually (picojson does not sort keys — do not use it for key serialization). +- **C# / .NET**: `JsonConvert.SerializeObject(new SortedDictionary{…}, Formatting.None)`. +- **Java**: `new ObjectMapper().writeValueAsString(new TreeMap<>(map))`. +- **Node.js**: build the object, then sort keys explicitly — `JSON.stringify` does NOT sort. +- **Go**: `json.Marshal(map[string]string{…})` — Go maps are marshaled in sorted key order. +- **Python**: `json.dumps(keyData, sort_keys=True, separators=(",", ":"))`. + +### 2.4 Normalization rules + +#### `normalize_url` — applies to `idp` and `snowflake` +1. Strip scheme (`https://` or `http://`). +2. Strip optional userinfo (`user:pass@`). +3. Drop query string and fragment. +4. Trim a root-only trailing slash. +5. **Uppercase the entire remainder** (host + optional `:port` + optional `/path`). + +``` +https://login.microsoftonline.com:443/tenant-id/oauth2/v2.0 + → LOGIN.MICROSOFTONLINE.COM:443/TENANT-ID/OAUTH2/V2.0 + +https://myorg-myaccount.privatelink.snowflakecomputing.com + → MYORG-MYACCOUNT.PRIVATELINK.SNOWFLAKECOMPUTING.COM +``` + +#### `normalize_identifier` — applies to `username` and `role` +- Uppercase every character **outside** double-quoted segments. +- Preserve the **entire** `"…"` segment verbatim — surrounding `"`, spaces, and lowercase + letters inside are all kept exactly as-is. + +``` +"First Last"@long-corporate-domain.example.com + → "First Last"@LONG-CORPORATE-DOMAIN.EXAMPLE.COM + ↑↑↑↑ lowercase preserved inside quotes + +"Analyst Role With Spaces":north_america:prod:readonly + → "Analyst Role With Spaces":NORTH_AMERICA:PROD:READONLY +``` + +### 2.5 Field wiring per flow + +| Flow | `idp` | `snowflake` | `role` | +|------|-------|-------------|--------| +| OAuth (auth code / refresh / DPoP) | normalized token-endpoint URL | normalized Snowflake server URL | normalized role from login params | +| MFA | *(absent from keyData)* | normalized Snowflake server URL | *(absent from keyData)* | +| External-browser ID token | *(absent from keyData)* | normalized Snowflake server URL | *(absent from keyData)* | + +### 2.6 Canonical token type strings + +| Flow / enum member | Correct `token_type` in key prefix | +|--------------------|-------------------------------------| +| ID token / `ID_TOKEN` | `ID_TOKEN` | +| MFA token / `MFA_TOKEN` | `MFA_TOKEN` | +| OAuth access / `OAUTH_ACCESS_TOKEN` | `OAUTH_ACCESS_TOKEN` | +| OAuth refresh / `OAUTH_REFRESH_TOKEN` | `OAUTH_REFRESH_TOKEN` | +| DPoP bundled access | `DPOP_BUNDLED_ACCESS_TOKEN` | + +Driver-specific mapping hazards (these were wrong before the first PR too — verify they +are fixed in the first PR branch, or fix them here): + +- **JDBC**: `CachedCredentialType.MFA_TOKEN.getValue()` previously returned `"MFATOKEN"` + (no underscore). Must be `"MFA_TOKEN"`. +- **.NET**: `TokenType.MFAToken.ToString()` returns `"MFAToken"`. Use the `StringAttr` + wire value `"MFA_TOKEN"` instead. +- **Node.js**: production used `AuthenticationTypes` strings such as + `USERNAME_PASSWORD_MFA` (for MFA) and `OAUTH_AUTHORIZATION_CODE_ACCESS_TOKEN` (for + OAuth). Map to canonical: `"MFA_TOKEN"`, `"OAUTH_ACCESS_TOKEN"`, `"OAUTH_REFRESH_TOKEN"`. +- **Python**: `TokenType.MFA_TOKEN.value = "MFA_TOKEN"` is correct; double-check others. +- **Go**: existing constants `idToken = "ID_TOKEN"`, `mfaToken = "MFA_TOKEN"`, + `oauthAccessToken = "OAUTH_ACCESS_TOKEN"`, `oauthRefreshToken = "OAUTH_REFRESH_TOKEN"` — + all correct; verify they were not changed. + +### 2.7 Validation + +- Reject (error or return empty) if `username` or `snowflake` is empty. +- `role` may be an empty string for OAuth when no role is configured. +- Old separator-injection guards (rejecting `;` or `:` in inputs) must be removed. + +--- + +## 3. Golden test vectors (LOCK — do not change) + +Every driver must reproduce **both** of these exact outputs. Add unit tests asserting them. + +### Vector A — OAuth flow + +**Raw inputs** (pre-normalization): + +``` +token_type : DPOP_BUNDLED_ACCESS_TOKEN (key prefix only — not in keyData) +idp : https://login.microsoftonline.com:443/tenant-id/oauth2/v2.0 +snowflake : https://myorg-myaccount.privatelink.snowflakecomputing.com +username : "First Last"@long-corporate-domain.example.com +role : "Analyst Role With Spaces":north_america:prod:readonly +``` + +**After normalization** (note: `First Last` and `Analyst Role With Spaces` preserved +verbatim because they are inside double quotes): + +``` +idp : LOGIN.MICROSOFTONLINE.COM:443/TENANT-ID/OAUTH2/V2.0 +snowflake : MYORG-MYACCOUNT.PRIVATELINK.SNOWFLAKECOMPUTING.COM +username : "First Last"@LONG-CORPORATE-DOMAIN.EXAMPLE.COM +role : "Analyst Role With Spaces":NORTH_AMERICA:PROD:READONLY +``` + +**Canonical JSON** (compact, sorted keys, 4 OAuth fields — no `token_type`): + +``` +{"idp":"LOGIN.MICROSOFTONLINE.COM:443/TENANT-ID/OAUTH2/V2.0","role":"\"Analyst Role With Spaces\":NORTH_AMERICA:PROD:READONLY","snowflake":"MYORG-MYACCOUNT.PRIVATELINK.SNOWFLAKECOMPUTING.COM","username":"\"First Last\"@LONG-CORPORATE-DOMAIN.EXAMPLE.COM"} +``` + +**Expected final key**: + +``` +SnowflakeTokenCache.v2.DPOP_BUNDLED_ACCESS_TOKEN.be782aa7c9abf8698adc9e6de61b954ccec7d9202899b44c2eb4e1dfa4313d5f +``` + +### Vector B — MFA flow + +**Raw inputs** (pre-normalization): + +``` +token_type : MFA_TOKEN (key prefix only — not in keyData) +snowflake : https://myorg-myaccount.privatelink.snowflakecomputing.com +username : "First Last"@long-corporate-domain.example.com +``` + +**After normalization**: + +``` +snowflake : MYORG-MYACCOUNT.PRIVATELINK.SNOWFLAKECOMPUTING.COM +username : "First Last"@LONG-CORPORATE-DOMAIN.EXAMPLE.COM +``` + +**Canonical JSON** (compact, sorted keys, 2 MFA/ID fields — no `idp`, `role`, `token_type`): + +``` +{"snowflake":"MYORG-MYACCOUNT.PRIVATELINK.SNOWFLAKECOMPUTING.COM","username":"\"First Last\"@LONG-CORPORATE-DOMAIN.EXAMPLE.COM"} +``` + +**Expected final key**: + +``` +SnowflakeTokenCache.v2.MFA_TOKEN.a508fa2858a6e22e9fdbc90b4149a3ff666d1acbb286c85ff179499ac92d75c8 +``` + +`DPOP_BUNDLED_ACCESS_TOKEN` in Vector A is used purely as a test literal. Do **not** add +DPoP flows to repos that don't have them today (only JDBC has this type). Pass the string +`"DPOP_BUNDLED_ACCESS_TOKEN"` directly in the golden test without a named constant. + +--- + +## 4. Cross-cutting pitfalls + +### 4.1 `token_type` must not appear in `keyData` + +The first PR serialized `token_type` inside the JSON object. Remove it. The type goes in +the key prefix only. Double-check no serialization helper or sort utility accidentally +re-inserts it. + +### 4.2 Two `keyData` shapes — flow-dispatch must be correct + +OAuth uses 4 fields: `idp`, `role`, `snowflake`, `username`. +MFA/ID uses 2 fields: `snowflake`, `username`. + +Common mistakes: +- Leaving `idp`/`role` in the MFA key (gives wrong hash). +- Omitting `idp`/`role` from an OAuth key (gives wrong hash). +- Treating ID token the same as OAuth (ID token = MFA path: snowflake + username only). + +### 4.3 Lowercase hex + +- **JDBC** `HexUtil.byteToHexString()` returns **UPPERCASE** → add `byteToHexStringLower`. +- **.NET** `StringUtils.ToSha256Hash()` calls `BitConverter.ToString(…).Replace("-","")` → + **UPPERCASE** → add `ToSha256HashLower`. +- **Go** `hex.EncodeToString` → lowercase ✓ +- **Python** `hashlib.sha256(…).hexdigest()` → lowercase ✓ +- **C++** — verify `Sha256.cpp` output; add `tolower` transform if it produces uppercase. + +### 4.4 Sorted JSON keys + +- **Node.js**: `JSON.stringify` uses insertion order, **not** sorted. Sort keys explicitly + before building the JSON string. +- **Java**: `ObjectMapper` + `TreeMap` sorts ✓ +- **.NET**: `SortedDictionary` sorts ✓ +- **Go**: `json.Marshal(map[string]string{…})` sorts ✓ +- **Python**: `json.dumps(…, sort_keys=True)` sorts ✓ +- **C++**: emit manually in sorted order — do **not** use picojson (it preserves insertion + order, not lexicographic). + +### 4.5 Single hashing point + +The final key is built once. Remove any secondary hashing inside backends: +- **Node.js**: `JsonCredentialManager.hashKey(key)` applied SHA-256 inside the JSON + file backend on top of whatever the caller passed. Remove `hashKey()` entirely. +- **Python**: `FileTokenCache` called `key.hash_key()` internally; `KeyringTokenCache` + stored the raw `string_key()`. Both must call `build_cache_key(key)` to get the + pre-built final key. +- **C++**: confirm no platform backend (Apple, Windows, Linux) calls `sha256()` again + after receiving the key from `convertTarget()`. + +### 4.6 OAuth idp must be the full token-endpoint URL, not just the host + +Several first-PR implementations may have accidentally stored only the hostname in `idp`: +- **Python** `_oauth_base.py`: `_idp_host` was already just a hostname. Replace with + `_token_request_url` (full URL). Key change: `"LOGIN.MICROSOFTONLINE.COM"` (hostname + only) → `"LOGIN.MICROSOFTONLINE.COM:443/TENANT-ID/OAUTH2/V2.0"` (full path). +- **JDBC** `getHostForOAuthCacheKey()` returned `.getHost()` (bare hostname). Rename and + return the full token-request URL. +- **Node.js**: `authorizationUrl.host` / `tokenUrl.host` (hostname-only) → use the full + URL string. +- **Go**: `oa.tokenURL()` returns the full URL already — verify no caller truncates it. +- **.NET**: `new Uri(GetTokenEndpoint()).Host` (hostname-only) → pass `GetTokenEndpoint()` + directly and let `NormalizeUrl` handle it. + +### 4.7 Cache filename stays `credential_cache_v1.json` + +Do not rename the JSON file. The `v2` in the key prefix is the key-format version, +independent of the filename. + +--- + +## 5. Tests required (every repo) + +- **Golden hash A** — assert `SnowflakeTokenCache.v2.DPOP_BUNDLED_ACCESS_TOKEN.be782aa7…` exactly. +- **Golden hash B** — assert `SnowflakeTokenCache.v2.MFA_TOKEN.a508fa28…` exactly. +- **`normalize_url`** — scheme stripping, port/path uppercasing, no trailing slash on bare host. +- **`normalize_identifier`** — unquoted uppercased; quoted segment preserved **verbatim including lowercase**; mixed-case input. +- **Dimension isolation**: + - Same IdP + different Snowflake host → different OAuth keys. + - Same host/user + different role → different OAuth keys. + - MFA and OAuth for same user/host → different keys (different prefix + field set). + - Different `token_type` prefix → different keys by construction. +- **File backend** — stored key equals `SnowflakeTokenCache.v2..`; no double-hash; round-trip set/get/delete. +- **OS keystore** — multi-account no-collision, multi-role no-collision (OAuth; where applicable). +- **Integration/E2E** — update any test that seeds the cache; add multi-account (shared IdP) and multi-role OAuth scenarios. + +--- + +## 6. Docs and changelog + +- Update doc comments that still describe the old `{host}:{user}:{type}` format or the 5-field keyData format. +- Add a **Bug fixes:** changelog entry: + > Fixed token cache key collisions for multi-account (shared IdP) and multi-role scenarios + > by switching to a versioned, SHA256-hashed canonical-JSON key with the token type in the + > key prefix, applied uniformly across OS keystore and file backends. +- Entry must end with the repo-appropriate PR link (see each driver prompt). + +--- + +## 7. Migration and compatibility + +- No reader looks up first-PR v2 keys (or original v1 keys). Old entries become orphaned; + the next connect transparently re-authenticates and writes a fresh entry. +- Active cleanup of orphaned entries is out of scope. + +--- + +## 8. Definition of done (per repo) + +- [ ] Golden hash A (OAuth) passes byte-exact (§3 Vector A). +- [ ] Golden hash B (MFA) passes byte-exact (§3 Vector B). +- [ ] Both backends use the final key verbatim; hashing occurs exactly once. +- [ ] OAuth call sites thread `idp` (full token-endpoint URL) + `snowflake` + `username` + `role`. +- [ ] MFA/ID call sites thread only `snowflake` + `username`; no `idp` or `role`. +- [ ] `token_type` does NOT appear in `keyData`; it is the third segment of the key prefix. +- [ ] Normalization + dimension-isolation tests pass. +- [ ] Integration/E2E seed keys updated; multi-account + multi-role no-collision scenarios added. +- [ ] Docs updated; changelog bug-fix entry with PR link added. +- [ ] Repo's own lint/format/test gates pass. +- [ ] Self-review checklist completed. + +--- + +## 9. Self-review checklist (run before committing) + +- [ ] Key format is `SnowflakeTokenCache.v2..` — four dot-separated segments. +- [ ] `token_type` is **not** present in any serialized `keyData` JSON. +- [ ] Hash is **lowercase** hex. +- [ ] OAuth `keyData` JSON emits exactly 4 keys in sorted order: `idp, role, snowflake, username`. +- [ ] MFA/ID `keyData` JSON emits exactly 2 keys in sorted order: `snowflake, username`. +- [ ] JSON is compact (no extra whitespace). +- [ ] Token type value in the key prefix is the canonical string (e.g., `MFA_TOKEN` not `MFATOKEN`). +- [ ] Hashing occurs exactly once — no leftover `hashKey` / `ToSha256Hash` / `sha256` inside backends. +- [ ] OAuth `idp` is the full token-endpoint URL (not hostname-only). +- [ ] Empty-`username` and empty-`snowflake` validation present. +- [ ] Old separator-injection guards removed. +- [ ] Lint/format/test gate passes. diff --git a/---prompts---/token-cache-updated/07-connector-python-upd.md b/---prompts---/token-cache-updated/07-connector-python-upd.md new file mode 100644 index 0000000000..309d77a455 --- /dev/null +++ b/---prompts---/token-cache-updated/07-connector-python-upd.md @@ -0,0 +1,300 @@ +# Token Cache Key v2 Fixup — snowflake-connector-python (Python) + +You are implementing a **follow-up PR** on top of the token-cache-key v2 change in +**snowflake-connector-python**. This is repo **07 of 07** and can run in parallel with +repos 02–05. + +Read `00-INDEX-upd.md` (same directory as this file) before proceeding. This file +contains only the Python-specific implementation details. + +--- + +## 1. Context: what the first PR already implemented + +The first PR introduced: +- `TokenKey` dataclass (5 named fields): `token_type: TokenType`, `idp: str`, + `snowflake: str`, `username: str`, `role: str`. +- `normalize_url(url: str) -> str` and `normalize_identifier(identifier: str) -> str` + module-level functions in `token_cache.py`. +- `build_cache_key(key: TokenKey) -> str`: builds 5-field compact sorted JSON via + `json.dumps(…, sort_keys=True, separators=(',', ':'))`, SHA-256-hashes it via + `hashlib.sha256(…).hexdigest()` (lowercase), returns + `SnowflakeTokenCache.v2.`. +- `KeyringTokenCache.store/retrieve/remove` updated to call `build_cache_key(key)`. +- `FileTokenCache.store/retrieve/remove` updated to call `build_cache_key(key)` (no more + `key.hash_key()` inside the backend). +- `_auth.py` call sites updated: `TokenKey` constructed with named fields for ID/MFA. +- `_oauth_base.py`: `_get_access_token_cache_key()` uses `_token_request_url` (full URL) + instead of the old `_idp_host` (hostname only). `_token_request_url`, `_snowflake_host`, + and `_role` threaded through `_OAuthTokensMixin.__init__`. +- No positional `TokenKey(a, b, c)` construction anywhere — all keyword arguments. + +--- + +## 2. Repo setup — do this first + +```bash +cd /Users/mhofman/Projects/snowflake-connector-python + +REMOTE=$(git remote get-url origin) +echo "Remote: $REMOTE" +# Expected: git@github.com:snowflakedb/snowflake-connector-python.git +# The private mirror must NOT be used. + +git fetch origin +git branch -r | grep "SNOW-3784431-token-cache-key-v2$" && \ + BASE="origin/$(git config user.email | cut -d'@' -f1)/SNOW-3784431-token-cache-key-v2" || \ + BASE="origin/main" +echo "Branching from: $BASE" + +USER=$(git config user.email | cut -d'@' -f1) +git switch -c "${USER}/SNOW-3784431-token-cache-key-v2-fixup" --track $BASE +``` + +Ticket: [SNOW-3784431](https://snowflakecomputing.atlassian.net/browse/SNOW-3784431) + +--- + +## 3. Pre-flight drift check + +Verify the first PR's changes exist: + +- `src/snowflake/connector/token_cache.py` — `TokenKey` has 5 named fields including + `idp` and `role`; `build_cache_key`, `normalize_url`, `normalize_identifier` present. +- `KeyringTokenCache.store/retrieve/remove` call `build_cache_key`. +- `FileTokenCache.store/retrieve/remove` call `build_cache_key` (no `hash_key()` call). +- `auth/_auth.py` — `TokenKey` for ID/MFA uses named fields, `idp` and `role` threaded. +- `auth/_oauth_base.py` — `_token_request_url` field present; `_get_access_token_cache_key` + uses it. +- Existing golden hash test asserting the 5-field format (will be replaced). + +--- + +## 4. Implementation checklist + +### 4.1 Update `build_cache_key` for flow-specific `keyData` + +The new format is: +``` +SnowflakeTokenCache.v2.. +``` + +`keyData` is flow-dependent and **never contains `token_type`**: + +```python +_OAUTH_TYPES = frozenset({ + 'OAUTH_ACCESS_TOKEN', + 'OAUTH_REFRESH_TOKEN', + 'DPOP_BUNDLED_ACCESS_TOKEN', +}) + + +def build_cache_key(key: TokenKey) -> str: + """ + Build the versioned, uniformly-hashed v2 cache key. + + Format: SnowflakeTokenCache.v2.. + OAuth flows include idp/role; MFA and ID token flows include only + snowflake/username. + """ + if not key.snowflake: + raise ValueError("snowflake URL must not be empty") + if not key.username: + raise ValueError("username must not be empty") + + token_type_value = key.token_type.value + + if token_type_value in _OAUTH_TYPES: + key_data = { + 'idp': normalize_url(key.idp or ''), + 'role': normalize_identifier(key.role or ''), + 'snowflake': normalize_url(key.snowflake), + 'username': normalize_identifier(key.username), + } + else: + # MFA_TOKEN, ID_TOKEN — no idp or role + key_data = { + 'snowflake': normalize_url(key.snowflake), + 'username': normalize_identifier(key.username), + } + + # sort_keys=True + no whitespace = canonical JSON required by spec + canonical = json.dumps(key_data, sort_keys=True, separators=(',', ':')) + digest = hashlib.sha256(canonical.encode('utf-8')).hexdigest() + return f'SnowflakeTokenCache.v2.{token_type_value}.{digest}' +``` + +Key changes from the first PR: +1. `token_type` removed from `key_data`. +2. MFA/ID path uses only `snowflake` + `username`. +3. `token_type_value` inserted between `v2.` and the hash. + +### 4.2 Update MFA and ID token call sites — omit `idp` and `role` + +In `auth/_auth.py`, `read_temporary_credentials`, `write_temporary_credentials`: + +```python +# External browser / ID token +id_token_key = TokenKey( + token_type=TokenType.ID_TOKEN, + snowflake=self.host, + username=self.user, + # idp and role default to "" — omit them; build_cache_key will skip them +) + +# MFA +mfa_key = TokenKey( + token_type=TokenType.MFA_TOKEN, + snowflake=self.host, + username=self.user, +) +``` + +If `TokenKey` was defined with required positional fields for `idp` and `role` in the +first PR, change them to keyword-only with `""` defaults so MFA/ID call sites can omit them: + +```python +@dataclass +class TokenKey: + token_type: TokenType + snowflake: str + username: str + idp: str = "" + role: str = "" +``` + +### 4.3 Ensure OAuth `idp` is the full token-endpoint URL + +If the first PR stored only the hostname in `_token_request_url` or `_idp_host`, correct +it. The full URL must be passed to `TokenKey.idp`: + +```python +# _oauth_base.py +def _get_access_token_cache_key(self) -> TokenKey | None: + if not (self._token_cache and self._user): + return None + return TokenKey( + token_type=TokenType.OAUTH_ACCESS_TOKEN, + snowflake=self._snowflake_host, + username=self._user, + idp=self._token_request_url, # full URL, e.g. https://login.microsoftonline.com:443/… + role=self._role or '', + ) +``` + +Eviction paths must build the same `TokenKey`. Verify they use `_token_request_url` (not +a hostname-only variant). + +### 4.4 `TokenKey` field layout — no positional construction + +All `TokenKey` construction must use keyword arguments. Grep for any remaining +positional `TokenKey(a, b, c)` patterns and convert them. + +### 4.5 Update the golden test + +Replace the old 5-field golden hash test with the two new vectors (see §5). + +--- + +## 5. Test plan + +- [ ] **Golden hash A (OAuth)** in `test/unit/test_token_cache_key.py`: + ```python + def test_oauth_golden_hash(): + # Build keyData manually to inject the DPOP literal as token_type prefix + canonical = json.dumps({ + 'idp': normalize_url('https://login.microsoftonline.com:443/tenant-id/oauth2/v2.0'), + 'role': normalize_identifier('"Analyst Role With Spaces":north_america:prod:readonly'), + 'snowflake': normalize_url('https://myorg-myaccount.privatelink.snowflakecomputing.com'), + 'username': normalize_identifier('"First Last"@long-corporate-domain.example.com'), + }, sort_keys=True, separators=(',', ':')) + digest = hashlib.sha256(canonical.encode('utf-8')).hexdigest() + assert f'SnowflakeTokenCache.v2.DPOP_BUNDLED_ACCESS_TOKEN.{digest}' == \ + 'SnowflakeTokenCache.v2.DPOP_BUNDLED_ACCESS_TOKEN.be782aa7c9abf8698adc9e6de61b954ccec7d9202899b44c2eb4e1dfa4313d5f' + ``` +- [ ] **Golden hash B (MFA)**: + ```python + def test_mfa_golden_hash(): + key = TokenKey( + token_type=TokenType.MFA_TOKEN, + snowflake='https://myorg-myaccount.privatelink.snowflakecomputing.com', + username='"First Last"@long-corporate-domain.example.com', + ) + assert build_cache_key(key) == \ + 'SnowflakeTokenCache.v2.MFA_TOKEN.a508fa2858a6e22e9fdbc90b4149a3ff666d1acbb286c85ff179499ac92d75c8' + ``` +- [ ] **`test_normalize_identifier`** — `'"First Last"@example.com'` → + `'"First Last"@EXAMPLE.COM'` (lowercase inside quotes preserved, not uppercased). +- [ ] **`test_mfa_key_has_no_idp_or_role`** — assert MFA `key_data` JSON is + `'{"snowflake":"…","username":"…"}'` (no `idp`, `role`, or `token_type`). +- [ ] **Dimension isolation**: different Snowflake URL → different OAuth keys; different role → + different OAuth keys; MFA ≠ OAuth key for same user/host. +- [ ] **File backend** (`test_linux_local_file_cache.py`) — stored key equals + `SnowflakeTokenCache.v2..`; no double-hash; round-trip. +- [ ] **Keyring backend** — uses `build_cache_key`; multi-account no-collision. +- [ ] Update `test_oauth_token.py` and `test_auth_mfa.py`. + +--- + +## 6. Build and test commands + +```bash +cd /Users/mhofman/Projects/snowflake-connector-python +pip install -e ".[development]" +pytest test/unit/test_token_cache_key.py \ + test/unit/test_linux_local_file_cache.py \ + test/unit/test_oauth_token.py \ + test/unit/test_auth_mfa.py -v +pytest test/unit/ -v +``` + +--- + +## 7. Docs and changelog + +- Update the module-level docstring in `token_cache.py` to describe the updated format. +- Update any docstring referencing `{host}:{user}:{type}`, `string_key()`, or `hash_key()`. +- Add a bug-fix entry to `DESCRIPTION.md`: + +```markdown +- vX.Y(TBD) + - Fixed token cache key collisions for multi-account (shared IdP) and multi-role + scenarios by switching to a versioned, SHA256-hashed canonical-JSON key with the + token type in the key prefix, applied uniformly across macOS/Windows keyring and + the Linux file backend. +``` + +--- + +## 8. Self-review pass + +Run through `00-INDEX-upd.md §9`, plus Python-specific items: + +- [ ] `json.dumps(…, sort_keys=True, separators=(',', ':'))` — compact and sorted. +- [ ] `hashlib.sha256(…).hexdigest()` → lowercase hex (Python default). +- [ ] `KeyringTokenCache` calls `build_cache_key` — raw key not stored in keyring. +- [ ] `FileTokenCache` stores `build_cache_key` result (no `hash_key()` call). +- [ ] `token_type` NOT a key in `key_data` dict. +- [ ] MFA/ID `key_data` has exactly 2 keys (`snowflake`, `username`). +- [ ] OAuth `key_data` has exactly 4 keys (`idp`, `role`, `snowflake`, `username`). +- [ ] No positional `TokenKey(a, b, c)` construction — all keyword arguments. +- [ ] `_idp_host` (hostname only) no longer used to build cache keys. +- [ ] `mypy` / `ruff` / `flake8` checks pass. + +--- + +## 9. Commit and report back + +```bash +cd /Users/mhofman/Projects/snowflake-connector-python +git add -A +git commit -m "SNOW-3784431: token cache key v2 fixup — type in prefix, flow-specific keyData" +git log --oneline -3 +``` + +Reply with: +1. Branch name. +2. One-paragraph change summary. +3. Golden test A and B results (pass/fail + actual output if fail). +4. Self-review verdict. +5. Deviations from spec. diff --git a/DESCRIPTION.md b/DESCRIPTION.md index 8089d78d35..d21d50ddaa 100644 --- a/DESCRIPTION.md +++ b/DESCRIPTION.md @@ -8,7 +8,7 @@ Source code is also available at: https://github.com/snowflakedb/snowflake-conne # Release Notes - NEXT_RELEASE(TBD) - - Fixed token cache key collisions for multi-account (shared IdP) and multi-role scenarios by switching to a versioned, SHA256-hashed canonical-JSON key applied uniformly across macOS/Windows keyring and the Linux file backend (SNOW-3784431). The new key is threaded through both the sync and async authentication paths, and existing cache entries from the two prior key layouts are transparently migrated to the new format on first use. + - Fixed token cache key collisions for multi-account (shared IdP) and multi-role scenarios by switching to a versioned, SHA256-hashed canonical-JSON key with the token type in the key prefix, applied uniformly across macOS/Windows keyring and the Linux file backend (SNOW-3784431). MFA and ID token keys use only the Snowflake host and username; OAuth keys additionally include the full token-endpoint URL and role. Existing cache entries from prior key layouts are transparently migrated on first use. - Fixed a bug where a TLS handshake terminated by the peer (`SSLError` containing `SysCallError(-1, 'Unexpected EOF')`) was classified as non-retryable and surfaced as an `OperationalError`, unlike `ECONNRESET`. Such handshake `Unexpected EOF` errors are now retried, on both the sync and async request paths (SNOW-4058589). - v4.7.3(Sep 3,2026) diff --git a/src/snowflake/connector/auth/_auth.py b/src/snowflake/connector/auth/_auth.py index 9d1ac4ad8e..60cb7cb743 100644 --- a/src/snowflake/connector/auth/_auth.py +++ b/src/snowflake/connector/auth/_auth.py @@ -417,10 +417,8 @@ def post_request_wrapper(self, url, headers, body) -> None: self._delete_temporary_credential( TokenKey( token_type=TokenType.ID_TOKEN, - idp=self._rest._host, snowflake=self._rest._host, username=user, - role=role or "", ) ) raise ReauthenticationRequest( @@ -472,10 +470,8 @@ def post_request_wrapper(self, url, headers, body) -> None: self._delete_temporary_credential( TokenKey( token_type=TokenType.MFA_TOKEN, - idp=self._rest._host, snowflake=self._rest._host, username=user, - role="", ) ) Error.errorhandler_wrapper( @@ -589,10 +585,8 @@ def read_temporary_credentials( self._rest.id_token = self._read_temporary_credential( TokenKey( token_type=TokenType.ID_TOKEN, - idp=host, snowflake=host, username=user, - role=role, ) ) @@ -600,10 +594,8 @@ def read_temporary_credentials( self._rest.mfa_token = self._read_temporary_credential( TokenKey( token_type=TokenType.MFA_TOKEN, - idp=host, snowflake=host, username=user, - role="", ) ) @@ -639,10 +631,8 @@ def write_temporary_credentials( self._write_temporary_credential( TokenKey( token_type=TokenType.ID_TOKEN, - idp=host, snowflake=host, username=user, - role=role, ), response["data"].get("idToken"), ) @@ -651,10 +641,8 @@ def write_temporary_credentials( self._write_temporary_credential( TokenKey( token_type=TokenType.MFA_TOKEN, - idp=host, snowflake=host, username=user, - role="", ), response["data"].get("mfaToken"), ) diff --git a/src/snowflake/connector/token_cache.py b/src/snowflake/connector/token_cache.py index d2bbe22d33..f924680ad3 100644 --- a/src/snowflake/connector/token_cache.py +++ b/src/snowflake/connector/token_cache.py @@ -46,23 +46,25 @@ class _InvalidTokenKeyError(Exception): class TokenKey: """Key identifying a cached token. - All five fields are required. Raw (un-normalized) values are acceptable; - ``build_cache_key`` normalizes them before hashing. + ``snowflake`` and ``username`` are required for all flows. + ``idp`` and ``role`` are used only for OAuth flows; they default to ``""`` + and are ignored by ``build_cache_key`` for MFA and ID token flows. + Raw (un-normalized) values are acceptable; ``build_cache_key`` normalizes + them before hashing. Fields: token_type: The type of token being cached. - idp: IdP / token-endpoint URL. For MFA and external-browser flows this - equals the Snowflake server URL. snowflake: Snowflake server URL. username: Snowflake login name. - role: Snowflake role. Must be an empty string (not None) for MFA flows. + idp: IdP / token-endpoint URL (OAuth flows only). + role: Snowflake role (OAuth flows only). """ token_type: TokenType - idp: str snowflake: str username: str - role: str + idp: str = "" + role: str = "" def normalize_url(url: str) -> str: @@ -91,10 +93,27 @@ def normalize_identifier(identifier: str) -> str: return "".join(result) +_OAUTH_TYPES: frozenset[str] = frozenset( + { + "OAUTH_ACCESS_TOKEN", + "OAUTH_REFRESH_TOKEN", + "DPOP_BUNDLED_ACCESS_TOKEN", + } +) + + def build_cache_key(key: TokenKey) -> str: """Build the versioned, uniformly-hashed v2 cache key. - Format: ``SnowflakeTokenCache.v2.`` + Format: ``SnowflakeTokenCache.v2..`` + + ``keyData`` is flow-dependent and never contains ``token_type``: + + - OAuth (``OAUTH_ACCESS_TOKEN``, ``OAUTH_REFRESH_TOKEN``, + ``DPOP_BUNDLED_ACCESS_TOKEN``): 4 fields — ``idp``, ``role``, + ``snowflake``, ``username``. + - MFA / ID token (``MFA_TOKEN``, ``ID_TOKEN``): 2 fields — + ``snowflake``, ``username`` only. The canonical JSON is compact (no whitespace) with keys sorted lexicographically, serialized to UTF-8. Hashing occurs exactly once here; @@ -105,17 +124,25 @@ def build_cache_key(key: TokenKey) -> str: if not key.username: raise _InvalidTokenKeyError("username must not be empty") - key_data = { - "idp": normalize_url(key.idp), - "role": normalize_identifier(key.role), - "snowflake": normalize_url(key.snowflake), - "token_type": key.token_type.value, - "username": normalize_identifier(key.username), - } + token_type_value = key.token_type.value + + if token_type_value in _OAUTH_TYPES: + key_data: dict[str, str] = { + "idp": normalize_url(key.idp or ""), + "role": normalize_identifier(key.role or ""), + "snowflake": normalize_url(key.snowflake), + "username": normalize_identifier(key.username), + } + else: + # MFA_TOKEN, ID_TOKEN — idp and role are not part of the key + key_data = { + "snowflake": normalize_url(key.snowflake), + "username": normalize_identifier(key.username), + } canonical = json.dumps(key_data, sort_keys=True, separators=(",", ":")) digest = hashlib.sha256(canonical.encode("utf-8")).hexdigest() - return f"SnowflakeTokenCache.v2.{digest}" + return f"SnowflakeTokenCache.v2.{token_type_value}.{digest}" def _legacy_string_key(key: TokenKey) -> str: @@ -169,9 +196,10 @@ class TokenCache(ABC): - Linux: Uses JSON file in ~/.cache/snowflake/ with 0o600 permissions - Fallback: NoopTokenCache (no caching) if secure storage unavailable - Tokens are keyed by a versioned, SHA-256-hashed canonical-JSON key (v2 format) - built from (token_type, idp, snowflake, username, role) to avoid collisions - across multi-account and multi-role scenarios. + Tokens are keyed by a versioned, SHA-256-hashed canonical-JSON key (v2 format): + ``SnowflakeTokenCache.v2..``. OAuth flows include + ``idp`` and ``role`` in the hashed JSON; MFA and ID token flows use only + ``snowflake`` and ``username``. """ @staticmethod @@ -254,8 +282,9 @@ class FileTokenCache(TokenCache): Security: File must have 0o600 permissions and be owned by current user. Uses file locks to prevent concurrent access corruption. - JSON map keys are the full ``SnowflakeTokenCache.v2.`` strings - produced by ``build_cache_key``; hashing is performed once before dispatch. + JSON map keys are the full ``SnowflakeTokenCache.v2..`` + strings produced by ``build_cache_key``; hashing is performed once before + dispatch. Note: the filename (``credential_cache_v1.json``) is unchanged for backward compatibility; the ``v2`` in the key prefix refers to the key-format version, not the file format. @@ -522,9 +551,9 @@ class KeyringTokenCache(TokenCache): - macOS: Stores tokens in Keychain - Windows: Stores tokens in Windows Credential Manager - The v2 cache key (``SnowflakeTokenCache.v2.``) is used as the - keyring service name, and the uppercase username is used as the account - field. This ensures a distinct entry per token dimension while still + The v2 cache key (``SnowflakeTokenCache.v2..``) is + used as the keyring service name, and the uppercase username is used as the + account field. This ensures a distinct entry per token dimension while still letting related tokens share Keychain visibility per account. For backward compatibility, :meth:`retrieve` also checks two legacy layouts @@ -535,8 +564,8 @@ class KeyringTokenCache(TokenCache): - string layout (oldest): service ``{HOST}:{USER}:{TOKEN_TYPE}`` with account equal to the uppercase username - For OAuth tokens the legacy ``HOST`` is the IdP hostname; for other flows it - is the Snowflake host. + where ``string_key`` is ``{HOST}:{USER}:{TOKEN_TYPE}`` (HOST being the + IdP hostname for OAuth tokens and the Snowflake host otherwise). """ SERVICE_NAME = "com.snowflake.connector.python" diff --git a/test/unit/test_token_cache_key.py b/test/unit/test_token_cache_key.py index 8da7b2fd20..17bc28133b 100644 --- a/test/unit/test_token_cache_key.py +++ b/test/unit/test_token_cache_key.py @@ -101,10 +101,8 @@ def test_normalize_identifier_empty(): def test_build_cache_key_rejects_empty_snowflake(): key = TokenKey( token_type=TokenType.MFA_TOKEN, - idp="https://example.com", snowflake="", username="user", - role="", ) with pytest.raises(_InvalidTokenKeyError): build_cache_key(key) @@ -113,10 +111,8 @@ def test_build_cache_key_rejects_empty_snowflake(): def test_build_cache_key_rejects_empty_username(): key = TokenKey( token_type=TokenType.MFA_TOKEN, - idp="https://example.com", snowflake="https://example.snowflakecomputing.com", username="", - role="", ) with pytest.raises(_InvalidTokenKeyError): build_cache_key(key) @@ -127,35 +123,88 @@ def test_build_cache_key_rejects_empty_username(): # --------------------------------------------------------------------------- -def test_golden_hash(): - """Assert that the cache key hash is stable and must not change between releases. +def test_oauth_golden_hash(): + """Vector A — OAuth (DPoP) flow. Hash must never change between releases.""" + canonical = json.dumps( + { + "idp": normalize_url( + "https://login.microsoftonline.com:443/tenant-id/oauth2/v2.0" + ), + "role": normalize_identifier( + '"Analyst Role With Spaces":north_america:prod:readonly' + ), + "snowflake": normalize_url( + "https://myorg-myaccount.privatelink.snowflakecomputing.com" + ), + "username": normalize_identifier( + '"First Last"@long-corporate-domain.example.com' + ), + }, + sort_keys=True, + separators=(",", ":"), + ) + digest = hashlib.sha256(canonical.encode("utf-8")).hexdigest() + assert f"SnowflakeTokenCache.v2.DPOP_BUNDLED_ACCESS_TOKEN.{digest}" == ( + "SnowflakeTokenCache.v2.DPOP_BUNDLED_ACCESS_TOKEN." + "be782aa7c9abf8698adc9e6de61b954ccec7d9202899b44c2eb4e1dfa4313d5f" + ) + + +def test_mfa_golden_hash(): + """Vector B — MFA flow. Hash must never change between releases.""" + key = TokenKey( + token_type=TokenType.MFA_TOKEN, + snowflake="https://myorg-myaccount.privatelink.snowflakecomputing.com", + username='"First Last"@long-corporate-domain.example.com', + ) + assert build_cache_key(key) == ( + "SnowflakeTokenCache.v2.MFA_TOKEN." + "a508fa2858a6e22e9fdbc90b4149a3ff666d1acbb286c85ff179499ac92d75c8" + ) + - Quoted identifier segments (e.g. ``"FIRST LAST"``) contain uppercase content - because ``normalize_identifier`` preserves them verbatim — the content inside - quotes must already be in the correct case before normalization is called. - """ - idp_raw = "https://login.microsoftonline.com:443/tenant-id/oauth2/v2.0" - snowflake_raw = "https://myorg-myaccount.privatelink.snowflakecomputing.com" - # Quoted segments have uppercase content because normalize_identifier preserves them verbatim. - username_raw = '"FIRST LAST"@long-corporate-domain.example.com' - role_raw = '"ANALYST ROLE WITH SPACES":north_america:prod:readonly' +# --------------------------------------------------------------------------- +# build_cache_key — structural assertions +# --------------------------------------------------------------------------- + +def test_mfa_key_has_no_idp_or_role(): + """MFA keyData must contain exactly snowflake and username — no idp, role, or token_type.""" + key = TokenKey( + token_type=TokenType.MFA_TOKEN, + snowflake="https://myorg.snowflakecomputing.com", + username="alice", + ) + # Reconstruct the canonical JSON that build_cache_key hashes and verify its shape. canonical = json.dumps( { - "idp": normalize_url(idp_raw), - "role": normalize_identifier(role_raw), - "snowflake": normalize_url(snowflake_raw), - "token_type": "DPOP_BUNDLED_ACCESS_TOKEN", - "username": normalize_identifier(username_raw), + "snowflake": normalize_url("https://myorg.snowflakecomputing.com"), + "username": normalize_identifier("alice"), }, sort_keys=True, separators=(",", ":"), ) digest = hashlib.sha256(canonical.encode("utf-8")).hexdigest() - assert f"SnowflakeTokenCache.v2.{digest}" == ( - "SnowflakeTokenCache.v2." - "75ff2ad65a68afb402f125f62894697673c5ef3d863aba466d16b7a81053d1f4" + assert build_cache_key(key) == f"SnowflakeTokenCache.v2.MFA_TOKEN.{digest}" + + +def test_mfa_vs_oauth_key_differ_for_same_user_and_host(): + """MFA and OAuth produce different keys for the same user/host (different prefix + field set).""" + snowflake = "https://org.snowflakecomputing.com" + username = "alice" + mfa_key = TokenKey( + token_type=TokenType.MFA_TOKEN, + snowflake=snowflake, + username=username, + ) + oauth_key = TokenKey( + token_type=TokenType.OAUTH_ACCESS_TOKEN, + snowflake=snowflake, + username=username, + idp="https://idp.example.com/oauth2", + role="analyst", ) + assert build_cache_key(mfa_key) != build_cache_key(oauth_key) # --------------------------------------------------------------------------- @@ -172,16 +221,14 @@ def test_build_cache_key_prefix(): role="analyst", ) result = build_cache_key(key) - assert result.startswith("SnowflakeTokenCache.v2.") + assert result.startswith("SnowflakeTokenCache.v2.OAUTH_ACCESS_TOKEN.") def test_build_cache_key_hash_is_lowercase_hex(): key = TokenKey( token_type=TokenType.ID_TOKEN, - idp="https://host.example.com", snowflake="https://host.example.com", username="bob", - role="", ) suffix = build_cache_key(key).split(".")[-1] assert suffix == suffix.lower() @@ -232,17 +279,13 @@ def test_different_token_type_yields_different_key(): def test_mfa_empty_role_yields_stable_distinct_key(): mfa_key = TokenKey( token_type=TokenType.MFA_TOKEN, - idp="https://org.snowflakecomputing.com", snowflake="https://org.snowflakecomputing.com", username="alice", - role="", ) id_token_key = TokenKey( token_type=TokenType.ID_TOKEN, - idp="https://org.snowflakecomputing.com", snowflake="https://org.snowflakecomputing.com", username="alice", - role="", ) mfa_result = build_cache_key(mfa_key) assert mfa_result == build_cache_key(mfa_key) From c3164920366fa94dbee4ad5ddd09ba093d250c24 Mon Sep 17 00:00:00 2001 From: Michal Hofman Date: Thu, 23 Jul 2026 16:19:50 +0200 Subject: [PATCH 6/6] =?UTF-8?q?fix:=20token=20cache=20key=20v2=20=E2=80=94?= =?UTF-8?q?=20lowercase=20normalization=20and=20PascalCase=20token=20type?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - normalize_url now lowercases instead of uppercases; normalize_identifier returns verbatim for any value containing a double-quote, lowercases otherwise. - TokenType enum values changed to PascalCase (MfaToken, OauthAccessToken, …) matching the cross-driver v2 key contract update. - _OAUTH_TYPES frozenset updated to PascalCase values. - _legacy_string_key retains SCREAMING_SNAKE_CASE via _LEGACY_TOKEN_TYPE_VALUES to preserve backward-compatible migration of pre-v2 cache entries. - All affected unit tests updated; both golden hash vectors (A and B) pass. Co-authored-by: Cursor --- src/snowflake/connector/token_cache.py | 61 +++++++++++++------------- test/unit/test_token_cache_key.py | 57 ++++++++++++++---------- 2 files changed, 65 insertions(+), 53 deletions(-) diff --git a/src/snowflake/connector/token_cache.py b/src/snowflake/connector/token_cache.py index f924680ad3..4d855f9616 100644 --- a/src/snowflake/connector/token_cache.py +++ b/src/snowflake/connector/token_cache.py @@ -32,10 +32,18 @@ class TokenType(Enum): - OAUTH_REFRESH_TOKEN: Long-lived OAuth token to obtain new access tokens """ - ID_TOKEN = "ID_TOKEN" - MFA_TOKEN = "MFA_TOKEN" - OAUTH_ACCESS_TOKEN = "OAUTH_ACCESS_TOKEN" - OAUTH_REFRESH_TOKEN = "OAUTH_REFRESH_TOKEN" + ID_TOKEN = "IdToken" + MFA_TOKEN = "MfaToken" + OAUTH_ACCESS_TOKEN = "OauthAccessToken" + OAUTH_REFRESH_TOKEN = "OauthRefreshToken" + + +_LEGACY_TOKEN_TYPE_VALUES: dict[TokenType, str] = { + TokenType.ID_TOKEN: "ID_TOKEN", + TokenType.MFA_TOKEN: "MFA_TOKEN", + TokenType.OAUTH_ACCESS_TOKEN: "OAUTH_ACCESS_TOKEN", + TokenType.OAUTH_REFRESH_TOKEN: "OAUTH_REFRESH_TOKEN", +} class _InvalidTokenKeyError(Exception): @@ -68,36 +76,28 @@ class TokenKey: def normalize_url(url: str) -> str: - """Strip scheme and userinfo, drop query/fragment, trim root slash, uppercase.""" + """Strip scheme and userinfo, drop query/fragment, trim root slash, lowercase.""" s = re.sub(r"^https?://", "", url) at = s.find("@") if at >= 0: s = s[at + 1 :] s = s.split("?")[0].split("#")[0] s = s.rstrip("/") - return s.upper() + return s.lower() def normalize_identifier(identifier: str) -> str: - """Uppercase unquoted segments; preserve double-quoted segments verbatim.""" - result = [] - in_quotes = False - for ch in identifier: - if ch == '"': - in_quotes = not in_quotes - result.append(ch) - elif in_quotes: - result.append(ch) - else: - result.append(ch.upper()) - return "".join(result) + """Return verbatim if the value contains any double-quote character; otherwise lowercase.""" + if '"' in identifier: + return identifier + return identifier.lower() _OAUTH_TYPES: frozenset[str] = frozenset( { - "OAUTH_ACCESS_TOKEN", - "OAUTH_REFRESH_TOKEN", - "DPOP_BUNDLED_ACCESS_TOKEN", + "OauthAccessToken", + "OauthRefreshToken", + "DpopBundledAccessToken", } ) @@ -105,14 +105,14 @@ def normalize_identifier(identifier: str) -> str: def build_cache_key(key: TokenKey) -> str: """Build the versioned, uniformly-hashed v2 cache key. - Format: ``SnowflakeTokenCache.v2..`` + Format: ``SnowflakeTokenCache.v2..`` ``keyData`` is flow-dependent and never contains ``token_type``: - - OAuth (``OAUTH_ACCESS_TOKEN``, ``OAUTH_REFRESH_TOKEN``, - ``DPOP_BUNDLED_ACCESS_TOKEN``): 4 fields — ``idp``, ``role``, + - OAuth (``OauthAccessToken``, ``OauthRefreshToken``, + ``DpopBundledAccessToken``): 4 fields — ``idp``, ``role``, ``snowflake``, ``username``. - - MFA / ID token (``MFA_TOKEN``, ``ID_TOKEN``): 2 fields — + - MFA / ID token (``MfaToken``, ``IdToken``): 2 fields — ``snowflake``, ``username`` only. The canonical JSON is compact (no whitespace) with keys sorted @@ -158,6 +158,7 @@ def _legacy_string_key(key: TokenKey) -> str: Used only to locate and migrate pre-v2 cache entries. """ + legacy_type = _LEGACY_TOKEN_TYPE_VALUES.get(key.token_type, key.token_type.value) if key.token_type in ( TokenType.OAUTH_ACCESS_TOKEN, TokenType.OAUTH_REFRESH_TOKEN, @@ -168,14 +169,14 @@ def _legacy_string_key(key: TokenKey) -> str: raise _InvalidTokenKeyError("Invalid key, host is empty") if not key.username: raise _InvalidTokenKeyError("Invalid key, user is empty") - return f"{host.upper()}:{key.username.upper()}:{key.token_type.value}" + return f"{host.upper()}:{key.username.upper()}:{legacy_type}" else: # ID/MFA was called with swapped positional args: key order is USER:HOST:TYPE. if not key.snowflake: raise _InvalidTokenKeyError("Invalid key, host is empty") if not key.username: raise _InvalidTokenKeyError("Invalid key, user is empty") - return f"{key.username.upper()}:{key.snowflake.upper()}:{key.token_type.value}" + return f"{key.username.upper()}:{key.snowflake.upper()}:{legacy_type}" def _legacy_hash_key(key: TokenKey) -> str: @@ -197,7 +198,7 @@ class TokenCache(ABC): - Fallback: NoopTokenCache (no caching) if secure storage unavailable Tokens are keyed by a versioned, SHA-256-hashed canonical-JSON key (v2 format): - ``SnowflakeTokenCache.v2..``. OAuth flows include + ``SnowflakeTokenCache.v2..``. OAuth flows include ``idp`` and ``role`` in the hashed JSON; MFA and ID token flows use only ``snowflake`` and ``username``. """ @@ -282,7 +283,7 @@ class FileTokenCache(TokenCache): Security: File must have 0o600 permissions and be owned by current user. Uses file locks to prevent concurrent access corruption. - JSON map keys are the full ``SnowflakeTokenCache.v2..`` + JSON map keys are the full ``SnowflakeTokenCache.v2..`` strings produced by ``build_cache_key``; hashing is performed once before dispatch. Note: the filename (``credential_cache_v1.json``) is unchanged for @@ -551,7 +552,7 @@ class KeyringTokenCache(TokenCache): - macOS: Stores tokens in Keychain - Windows: Stores tokens in Windows Credential Manager - The v2 cache key (``SnowflakeTokenCache.v2..``) is + The v2 cache key (``SnowflakeTokenCache.v2..``) is used as the keyring service name, and the uppercase username is used as the account field. This ensures a distinct entry per token dimension while still letting related tokens share Keychain visibility per account. diff --git a/test/unit/test_token_cache_key.py b/test/unit/test_token_cache_key.py index 17bc28133b..b821aa6f54 100644 --- a/test/unit/test_token_cache_key.py +++ b/test/unit/test_token_cache_key.py @@ -21,44 +21,44 @@ def test_normalize_url_strips_https_scheme(): - assert normalize_url("https://example.com") == "EXAMPLE.COM" + assert normalize_url("https://example.com") == "example.com" def test_normalize_url_strips_http_scheme(): - assert normalize_url("http://example.com") == "EXAMPLE.COM" + assert normalize_url("http://example.com") == "example.com" def test_normalize_url_no_scheme(): - assert normalize_url("example.com") == "EXAMPLE.COM" + assert normalize_url("example.com") == "example.com" def test_normalize_url_preserves_port_and_path(): assert ( normalize_url("https://login.microsoftonline.com:443/tenant-id/oauth2/v2.0") - == "LOGIN.MICROSOFTONLINE.COM:443/TENANT-ID/OAUTH2/V2.0" + == "login.microsoftonline.com:443/tenant-id/oauth2/v2.0" ) def test_normalize_url_strips_userinfo(): - assert normalize_url("https://user:pass@example.com/path") == "EXAMPLE.COM/PATH" + assert normalize_url("https://user:pass@example.com/path") == "example.com/path" def test_normalize_url_drops_query_and_fragment(): - assert normalize_url("https://example.com/path?q=1#frag") == "EXAMPLE.COM/PATH" + assert normalize_url("https://example.com/path?q=1#frag") == "example.com/path" def test_normalize_url_trims_root_trailing_slash(): - assert normalize_url("https://example.com/") == "EXAMPLE.COM" + assert normalize_url("https://example.com/") == "example.com" def test_normalize_url_keeps_non_root_trailing_slash_stripped(): - assert normalize_url("https://example.com/path/") == "EXAMPLE.COM/PATH" + assert normalize_url("https://example.com/path/") == "example.com/path" -def test_normalize_url_uppercases(): +def test_normalize_url_lowercases(): assert ( normalize_url("https://myorg-myaccount.privatelink.snowflakecomputing.com") - == "MYORG-MYACCOUNT.PRIVATELINK.SNOWFLAKECOMPUTING.COM" + == "myorg-myaccount.privatelink.snowflakecomputing.com" ) @@ -67,28 +67,39 @@ def test_normalize_url_uppercases(): # --------------------------------------------------------------------------- -def test_normalize_identifier_unquoted_uppercased(): - assert normalize_identifier("north_america") == "NORTH_AMERICA" +def test_normalize_identifier_unquoted_lowercased(): + assert normalize_identifier("north_america") == "north_america" + + +def test_normalize_identifier_unquoted_uppercased_input_lowercased(): + assert normalize_identifier("ANALYST_ROLE") == "analyst_role" def test_normalize_identifier_quoted_segment_verbatim(): assert normalize_identifier('"First Last"') == '"First Last"' -def test_normalize_identifier_mixed(): +def test_normalize_identifier_mixed_verbatim(): + """Value containing a double-quote anywhere is returned entirely verbatim.""" assert ( normalize_identifier('"First Last"@long-corporate-domain.example.com') - == '"First Last"@LONG-CORPORATE-DOMAIN.EXAMPLE.COM' + == '"First Last"@long-corporate-domain.example.com' ) -def test_normalize_identifier_role_with_quoted_spaces(): +def test_normalize_identifier_role_with_quoted_spaces_verbatim(): + """Value containing a double-quote anywhere is returned entirely verbatim.""" assert ( normalize_identifier('"Analyst Role With Spaces":north_america:prod:readonly') - == '"Analyst Role With Spaces":NORTH_AMERICA:PROD:READONLY' + == '"Analyst Role With Spaces":north_america:prod:readonly' ) +def test_normalize_identifier_quote_not_at_position_zero_verbatim(): + """A quote that does NOT appear at position 0 still triggers the verbatim path.""" + assert normalize_identifier('prefix-"segment"') == 'prefix-"segment"' + + def test_normalize_identifier_empty(): assert normalize_identifier("") == "" @@ -144,9 +155,9 @@ def test_oauth_golden_hash(): separators=(",", ":"), ) digest = hashlib.sha256(canonical.encode("utf-8")).hexdigest() - assert f"SnowflakeTokenCache.v2.DPOP_BUNDLED_ACCESS_TOKEN.{digest}" == ( - "SnowflakeTokenCache.v2.DPOP_BUNDLED_ACCESS_TOKEN." - "be782aa7c9abf8698adc9e6de61b954ccec7d9202899b44c2eb4e1dfa4313d5f" + assert f"SnowflakeTokenCache.v2.DpopBundledAccessToken.{digest}" == ( + "SnowflakeTokenCache.v2.DpopBundledAccessToken." + "741b6d66d252666d6821bfd19e0151511cf4efdaaeba2b3c87673aa4de6d2c0b" ) @@ -158,8 +169,8 @@ def test_mfa_golden_hash(): username='"First Last"@long-corporate-domain.example.com', ) assert build_cache_key(key) == ( - "SnowflakeTokenCache.v2.MFA_TOKEN." - "a508fa2858a6e22e9fdbc90b4149a3ff666d1acbb286c85ff179499ac92d75c8" + "SnowflakeTokenCache.v2.MfaToken." + "10c5dde84bb8f584c0df06ea826d418c4f580e08f9db10187c0cb5e2a732a0d6" ) @@ -185,7 +196,7 @@ def test_mfa_key_has_no_idp_or_role(): separators=(",", ":"), ) digest = hashlib.sha256(canonical.encode("utf-8")).hexdigest() - assert build_cache_key(key) == f"SnowflakeTokenCache.v2.MFA_TOKEN.{digest}" + assert build_cache_key(key) == f"SnowflakeTokenCache.v2.MfaToken.{digest}" def test_mfa_vs_oauth_key_differ_for_same_user_and_host(): @@ -221,7 +232,7 @@ def test_build_cache_key_prefix(): role="analyst", ) result = build_cache_key(key) - assert result.startswith("SnowflakeTokenCache.v2.OAUTH_ACCESS_TOKEN.") + assert result.startswith("SnowflakeTokenCache.v2.OauthAccessToken.") def test_build_cache_key_hash_is_lowercase_hex():