diff --git a/sdk/storage/azure-storage-blob/CHANGELOG.md b/sdk/storage/azure-storage-blob/CHANGELOG.md index 379ea2aa0d61..2a1d218f5dfd 100644 --- a/sdk/storage/azure-storage-blob/CHANGELOG.md +++ b/sdk/storage/azure-storage-blob/CHANGELOG.md @@ -4,6 +4,7 @@ ### Features Added - Added `list` support to `BlobSasPermissions` for use with directory-scoped SAS tokens. +- Added opt-in client session-based authentication via the new `use_session` keyword argument. When enabled, eligible blob download requests are authenticated with a short-lived, per-container session credential obtained from the service rather than the bearer token. Requires a `TokenCredential`. Sessions are managed by a session provider, which can be shared across clients via the `session_provider` keyword, and the account name used for signing can be set explicitly with `session_account_name`. ### Bugs Fixed - Fixed an issue where `destination_snapshot` on a blob's copy properties was always `None` when listing blobs with `response_format="arrow"`. diff --git a/sdk/storage/azure-storage-blob/azure/storage/blob/__init__.py b/sdk/storage/azure-storage-blob/azure/storage/blob/__init__.py index ef8b44a30790..d3bccf324401 100644 --- a/sdk/storage/azure-storage-blob/azure/storage/blob/__init__.py +++ b/sdk/storage/azure-storage-blob/azure/storage/blob/__init__.py @@ -26,6 +26,7 @@ UserDelegationKey, Services, ) +from ._shared.session import ContainerSessionProvider, Session, SessionProvider from ._generated.models import RehydratePriority from ._models import ( BlobType, @@ -267,4 +268,7 @@ def download_blob_from_url( "ObjectReplicationPolicy", "ObjectReplicationRule", "Services", + "Session", + "SessionProvider", + "ContainerSessionProvider", ] diff --git a/sdk/storage/azure-storage-blob/azure/storage/blob/_blob_service_client.py b/sdk/storage/azure-storage-blob/azure/storage/blob/_blob_service_client.py index d9a5974e238d..685dda8a95be 100644 --- a/sdk/storage/azure-storage-blob/azure/storage/blob/_blob_service_client.py +++ b/sdk/storage/azure-storage-blob/azure/storage/blob/_blob_service_client.py @@ -90,6 +90,19 @@ class BlobServiceClient(StorageAccountHostsMixin, StorageEncryptionMixin): :keyword str audience: The audience to use when requesting tokens for Azure Active Directory authentication. Only has an effect when credential is of type TokenCredential. The value could be https://storage.azure.com/ (default) or https://.blob.core.windows.net. + :keyword bool use_session: If True, enable session-based authentication for this container. + When enabled, eligible GET requests issued by this client will be authenticated using + a short-lived session credential obtained from the service instead + of the provided TokenCredential. Only supported with a TokenCredential; + ValueError is raised otherwise. Defaults to False. + :keyword session_provider: Creates, caches, and invalidates the session credentials used for + session-based authentication. Supply a shared instance to reuse the session cache across + multiple clients; when omitted, one is created and scoped to this client. Only has an + effect when `use_session` is True. + :paramtype session_provider: ~azure.storage.blob.SessionProvider + :keyword str session_account_name: The storage account name used to sign session-authenticated + requests. If omitted, it is derived from the account URL. Required when using a custom + endpoint, where the account name cannot be determined from the URL. .. admonition:: Example: diff --git a/sdk/storage/azure-storage-blob/azure/storage/blob/_blob_service_client.pyi b/sdk/storage/azure-storage-blob/azure/storage/blob/_blob_service_client.pyi index 526c2bfae18a..094e8f77e5e0 100644 --- a/sdk/storage/azure-storage-blob/azure/storage/blob/_blob_service_client.pyi +++ b/sdk/storage/azure-storage-blob/azure/storage/blob/_blob_service_client.pyi @@ -14,6 +14,7 @@ from typing import ( Optional, Union, ) +from azure.storage.blob._shared.session import SessionProvider from typing_extensions import Self from azure.core import MatchConditions @@ -56,6 +57,9 @@ class BlobServiceClient(StorageAccountHostsMixin, StorageEncryptionMixin): max_single_get_size: int = 32 * 1024 * 1024, max_chunk_get_size: int = 4 * 1024 * 1024, audience: Optional[str] = None, + use_session: bool = False, + session_provider: Optional[SessionProvider] = None, + session_account_name: Optional[str] = None, **kwargs: Any ) -> None: ... def __enter__(self) -> Self: ... diff --git a/sdk/storage/azure-storage-blob/azure/storage/blob/_container_client.py b/sdk/storage/azure-storage-blob/azure/storage/blob/_container_client.py index 05dcc2b49a12..8936b6d1f158 100644 --- a/sdk/storage/azure-storage-blob/azure/storage/blob/_container_client.py +++ b/sdk/storage/azure-storage-blob/azure/storage/blob/_container_client.py @@ -103,6 +103,19 @@ class ContainerClient(StorageAccountHostsMixin, StorageEncryptionMixin): # pyli :keyword str audience: The audience to use when requesting tokens for Azure Active Directory authentication. Only has an effect when credential is of type TokenCredential. The value could be https://storage.azure.com/ (default) or https://.blob.core.windows.net. + :keyword bool use_session: If True, enable session-based authentication for this container. + When enabled, eligible GET requests issued by this client will be authenticated using + a short-lived session credential obtained from the service instead + of the provided TokenCredential. Only supported with a TokenCredential; + ValueError is raised otherwise. Defaults to False. + :keyword session_provider: Creates, caches, and invalidates the session credentials used for + session-based authentication. Supply a shared instance to reuse the session cache across + multiple clients; when omitted, one is created and scoped to this client. Only has an + effect when `use_session` is True. + :paramtype session_provider: ~azure.storage.blob.SessionProvider + :keyword str session_account_name: The storage account name used to sign session-authenticated + requests. If omitted, it is derived from the account URL. Required when using a custom + endpoint, where the account name cannot be determined from the URL. .. admonition:: Example: diff --git a/sdk/storage/azure-storage-blob/azure/storage/blob/_container_client.pyi b/sdk/storage/azure-storage-blob/azure/storage/blob/_container_client.pyi index b956f375a2fe..ac6aaf49401d 100644 --- a/sdk/storage/azure-storage-blob/azure/storage/blob/_container_client.pyi +++ b/sdk/storage/azure-storage-blob/azure/storage/blob/_container_client.pyi @@ -28,6 +28,7 @@ from azure.core.credentials import AzureNamedKeyCredential, AzureSasCredential, from azure.core.paging import ItemPaged from azure.core.pipeline.transport import HttpResponse from azure.core.tracing.decorator import distributed_trace +from azure.storage.blob._shared.session import SessionProvider from ._blob_client import BlobClient from ._blob_service_client import BlobServiceClient from ._download import StorageStreamDownloader @@ -71,6 +72,9 @@ class ContainerClient(StorageAccountHostsMixin, StorageEncryptionMixin): max_single_get_size: int = 32 * 1024 * 1024, min_large_block_upload_threshold: int = 4 * 1024 * 1024 + 1, use_byte_buffer: Optional[bool] = None, + use_session: bool = False, + session_provider: Optional[SessionProvider] = None, + session_account_name: Optional[str] = None, **kwargs: Any, ) -> None: ... def __enter__(self) -> Self: ... diff --git a/sdk/storage/azure-storage-blob/azure/storage/blob/_shared/base_client.py b/sdk/storage/azure-storage-blob/azure/storage/blob/_shared/base_client.py index 8146c79d91c3..ea06fb6c15c2 100644 --- a/sdk/storage/azure-storage-blob/azure/storage/blob/_shared/base_client.py +++ b/sdk/storage/azure-storage-blob/azure/storage/blob/_shared/base_client.py @@ -57,12 +57,14 @@ StorageRequestHook, StorageResponseHook, StorageSensitiveHeaderCleanupPolicy, + StorageSessionPolicy, ) from .request_handlers import serialize_batch_body, _get_batch_request_delimiter from .response_handlers import PartialBatchErrorException, process_storage_error from .shared_access_signature import QueryStringConstants from .._version import VERSION from .._shared_access_signature import _is_credential_sastoken +from .session import ContainerSessionProvider if TYPE_CHECKING: from azure.core.credentials_async import AsyncTokenCredential @@ -295,12 +297,13 @@ def _create_pipeline( **kwargs: Any, ) -> Tuple[StorageConfiguration, Pipeline]: self._credential_policy: Any = None + audience = kwargs.pop("audience", None) if hasattr(credential, "get_token"): - if kwargs.get("audience"): - audience = str(kwargs.pop("audience")).rstrip("/") + DEFAULT_OAUTH_SCOPE + if audience: + scope = str(audience).rstrip("/") + DEFAULT_OAUTH_SCOPE else: - audience = STORAGE_OAUTH_SCOPE - self._credential_policy = StorageBearerTokenCredentialPolicy(cast(TokenCredential, credential), audience) + scope = STORAGE_OAUTH_SCOPE + self._credential_policy = StorageBearerTokenCredentialPolicy(cast(TokenCredential, credential), scope) elif isinstance(credential, SharedKeyCredentialPolicy): self._credential_policy = credential elif isinstance(credential, AzureSasCredential): @@ -330,12 +333,38 @@ def _create_pipeline( config.headers_policy, StorageRequestHook(**kwargs), self._credential_policy, - config.logging_policy, - StorageResponseHook(**kwargs), - DistributedTracingPolicy(**kwargs), - HttpLoggingPolicy(**kwargs), StorageSensitiveHeaderCleanupPolicy(**kwargs), ] + use_session = bool(kwargs.pop("use_session", False)) + session_provider = kwargs.pop("session_provider", None) + session_account_name = kwargs.pop("session_account_name", None) + if use_session: + if session_provider is None: + sub_kwargs = dict(kwargs) + sub_kwargs.pop("_configuration", None) + sub_kwargs.pop("pipeline", None) + sub_kwargs["transport"] = transport + session_provider = ContainerSessionProvider( + f"{self.scheme}://{self.primary_hostname}", + cast(TokenCredential, credential), + audience=audience, + **sub_kwargs, + ) + + policies.append( + StorageSessionPolicy( + account_name=session_account_name or self.account_name, + session_provider=session_provider, + ) + ) + policies.extend( + [ + config.logging_policy, + StorageResponseHook(**kwargs), + DistributedTracingPolicy(**kwargs), + HttpLoggingPolicy(**kwargs), + ] + ) if kwargs.get("_additional_pipeline_policies"): policies = policies + kwargs.get("_additional_pipeline_policies") # type: ignore config.transport = transport # type: ignore diff --git a/sdk/storage/azure-storage-blob/azure/storage/blob/_shared/models.py b/sdk/storage/azure-storage-blob/azure/storage/blob/_shared/models.py index 4ad6b1471cca..8e4be18b1235 100644 --- a/sdk/storage/azure-storage-blob/azure/storage/blob/_shared/models.py +++ b/sdk/storage/azure-storage-blob/azure/storage/blob/_shared/models.py @@ -34,6 +34,7 @@ class StorageErrorCode(str, Enum, metaclass=CaseInsensitiveEnumMeta): CONDITION_HEADERS_NOT_SUPPORTED = "ConditionHeadersNotSupported" CONDITION_NOT_MET = "ConditionNotMet" EMPTY_METADATA_KEY = "EmptyMetadataKey" + FEATURE_NOT_ENABLED = "FeatureNotEnabled" INSUFFICIENT_ACCOUNT_PERMISSIONS = "InsufficientAccountPermissions" INTERNAL_ERROR = "InternalError" INVALID_AUTHENTICATION_INFO = "InvalidAuthenticationInfo" @@ -64,6 +65,7 @@ class StorageErrorCode(str, Enum, metaclass=CaseInsensitiveEnumMeta): RESOURCE_ALREADY_EXISTS = "ResourceAlreadyExists" RESOURCE_NOT_FOUND = "ResourceNotFound" SERVER_BUSY = "ServerBusy" + SESSIONS_UNAVAILABLE = "SessionOperationsTemporarilyUnavailable" UNSUPPORTED_HEADER = "UnsupportedHeader" UNSUPPORTED_XML_NODE = "UnsupportedXmlNode" UNSUPPORTED_QUERY_PARAMETER = "UnsupportedQueryParameter" diff --git a/sdk/storage/azure-storage-blob/azure/storage/blob/_shared/policies.py b/sdk/storage/azure-storage-blob/azure/storage/blob/_shared/policies.py index a159b7dcae80..ac4203baf23d 100644 --- a/sdk/storage/azure-storage-blob/azure/storage/blob/_shared/policies.py +++ b/sdk/storage/azure-storage-blob/azure/storage/blob/_shared/policies.py @@ -9,11 +9,13 @@ import random import re import uuid +from datetime import timezone from io import BytesIO, SEEK_SET, UnsupportedOperation from time import time -from typing import Any, Dict, List, Optional, TYPE_CHECKING, Union +from typing import Any, Dict, List, Optional, Tuple, TYPE_CHECKING, Union from urllib.parse import ( parse_qsl, + unquote, urlencode, urlparse, urlunparse, @@ -30,9 +32,11 @@ SansIOHTTPPolicy, ) -from .authentication import AzureSigningError, StorageHttpChallenge +from . import sign_string +from .authentication import AzureSigningError, _storage_header_sort, StorageHttpChallenge from .constants import DEFAULT_OAUTH_SCOPE, DATA_BLOCK_SIZE from .models import LocationMode, StorageErrorCode +from .session import Session, SessionProvider from .streams import ( StructuredMessageDecoder, StructuredMessageEncodeStream, @@ -55,12 +59,26 @@ _LOGGER = logging.getLogger(__name__) +_SESSION_SIGNED_HEADERS = ( + "content-encoding", + "content-language", + "content-length", + "content-md5", + "content-type", + "date", + "if-modified-since", + "if-match", + "if-none-match", + "if-unmodified-since", + "byte_range", +) CONTENT_LENGTH_HEADER = "Content-Length" MD5_HEADER = "Content-MD5" CRC64_HEADER = "x-ms-content-crc64" SM_HEADER = "x-ms-structured-body" SM_HEADER_V1_CRC64 = "XSM/1.0; properties=crc64" SM_LENGTH_HEADER = "x-ms-structured-content-length" +SESSION_BEARER_AUTH_KEY = "_session_bearer_auth" def encode_base64(data: Union[bytes, str]) -> str: @@ -70,6 +88,50 @@ def encode_base64(data: Union[bytes, str]) -> str: return encoded.decode("utf-8") +def _apply_session_auth( + request: "PipelineRequest", session_token: str, session_key: str, account_name: str +) -> None: + """Sign an eligible request with the SharedKey protocol under the Session scheme. + + Shared by the sync and async session policies; ``account_name`` is passed in + rather than read from ``self`` so neither policy needs to instantiate the other. + + :param ~azure.core.pipeline.PipelineRequest request: The request to sign in place. + :param str session_token: The session token to embed in the Authorization header. + :param str session_key: The HMAC signing key for the session. + :param str account_name: Storage account name; the signer identity. + :raises ~azure.storage.blob._shared.authentication.AzureSigningError: if signing fails. + """ + http_request = request.http_request + http_request.headers["x-ms-date"] = format_date_time(time()) + + # 1) Standard headers. Storage omits content-length when it is "0". + headers = {name.lower(): value for name, value in http_request.headers.items() if value} + if headers.get("content-length") == "0": + del headers["content-length"] + signed_headers = "\n".join(headers.get(h, "") for h in _SESSION_SIGNED_HEADERS) + "\n" + + # 2) Canonicalized x-ms-* headers, sorted by the service-emulating comparator. + x_ms_headers = _storage_header_sort( + [(n.lower(), v) for n, v in http_request.headers.items() if n.lower().startswith("x-ms-")] + ) + canonicalized_headers = "".join(f"{n}:{v}\n" for n, v in x_ms_headers if v is not None) + + # 3) Canonicalized resource + query (query values must be url-decoded). + canonicalized_resource = "/" + account_name + urlparse(http_request.url).path + canonicalized_resource += "".join( + f"\n{n.lower()}:{unquote(v)}" for n, v in sorted(http_request.query.items()) if v is not None + ) + + string_to_sign = http_request.method + "\n" + signed_headers + canonicalized_headers + canonicalized_resource + + try: + signature = sign_string(session_key, string_to_sign) + except Exception as ex: # pylint: disable=broad-except + raise AzureSigningError(str(ex)) from ex + http_request.headers["Authorization"] = f"Session {session_token}:{signature}" + + # Are we out of retries? def is_exhausted(settings): retry_counts = ( @@ -927,3 +989,95 @@ def on_request(self, request: "PipelineRequest") -> None: # Clean up request headers for header in self._blocked_redirect_headers: request.http_request.headers.pop(header, None) + + +class StorageSessionPolicy(HTTPPolicy): + """ + A pipeline policy that selects between session token and bearer token authentication. + + Eligible requests are authenticated with a session token obtained from the + session provider. Everything else is left to the bearer token policy that + sits earlier in the pipeline. + """ + + def __init__( + self, + *, + account_name: Optional[str], + session_provider: SessionProvider, + ) -> None: + """Constructs a StorageSessionPolicy. + + :keyword str account_name: Storage account name; used as the signer + identity when signing session-authenticated requests. + :keyword session_provider: Creates, caches, and invalidates per-container sessions. + :paramtype session_provider: ~azure.storage.blob._shared.session.SessionProvider + :raises ValueError: if `account_name` is `None`. + """ + if account_name is None: + raise ValueError( + "Unable to determine the account name from the service URL. " + "Supply session_account_name when using a custom endpoint." + ) + super().__init__() + self._account_name = account_name + self._session_provider = session_provider + + def send(self, request: "PipelineRequest") -> "PipelineResponse": + """Orchestrate session auth. + + :param ~azure.core.pipeline.PipelineRequest request: The outgoing request. + :return: The pipeline response. + :rtype: ~azure.core.pipeline.PipelineResponse + """ + session = self.on_request(request) + response = self.next.send(request) + return self.on_response(request, response, session) + + def on_request(self, request: "PipelineRequest") -> Optional[Session]: + """Stamp session auth if eligible, otherwise leave the bearer header intact. + + :param ~azure.core.pipeline.PipelineRequest request: The request to (maybe) sign. + :return: The session that was applied, else None. + :rtype: ~azure.storage.blob._shared.session.Session or None + """ + session = self._session_provider.get_session(request) + if session is None or not session.session_token or not session.session_key: + return None + + # Kept so a 401 can serve this request with bearer instead of re-signing. + request.context.options[SESSION_BEARER_AUTH_KEY] = request.http_request.headers.get("Authorization") + _apply_session_auth(request, session.session_token, session.session_key, self._account_name) + return session + + def on_response( + self, + request: "PipelineRequest", + response: "PipelineResponse", + session: Optional[Session], + ) -> "PipelineResponse": + """On 401, invalidate the cached session and serve the request with bearer. + + :param ~azure.core.pipeline.PipelineRequest request: The original request. + :param ~azure.core.pipeline.PipelineResponse response: The response to inspect. + :param session: The session that signed the request, or `None` if bearer was used. + :type session: ~azure.storage.blob._shared.session.Session or None + :return: The final response. + :rtype: ~azure.core.pipeline.PipelineResponse + """ + if session is None: + return response # bearer was used; nothing session-related to react to + + status = response.http_response.status_code + + # 401 → drop the cached session and serve this request with bearer. + if status == 401: + _LOGGER.info("Session authentication: HTTP 401; invalidating session and retrying with bearer.") + self._session_provider.invalidate_session(request, session) + bearer = request.context.options.get(SESSION_BEARER_AUTH_KEY) + if bearer: + request.http_request.headers["Authorization"] = bearer + return self.next.send(request) + return response + + return response diff --git a/sdk/storage/azure-storage-blob/azure/storage/blob/_shared/session.py b/sdk/storage/azure-storage-blob/azure/storage/blob/_shared/session.py new file mode 100644 index 000000000000..05671c8cec12 --- /dev/null +++ b/sdk/storage/azure-storage-blob/azure/storage/blob/_shared/session.py @@ -0,0 +1,298 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for +# license information. +# -------------------------------------------------------------------------- + +import logging +from datetime import datetime, timedelta, timezone +from threading import Lock +from typing import Any, Dict, Optional, Tuple, TYPE_CHECKING +from typing_extensions import Protocol +from urllib.parse import urlparse + +from azure.core.exceptions import AzureError, HttpResponseError + +from .models import StorageErrorCode +from .._generated.models import CreateSessionConfiguration, CreateSessionResponse + +if TYPE_CHECKING: + from azure.core.credentials import TokenCredential + from azure.core.pipeline import PipelineRequest + +_LOGGER = logging.getLogger(__name__) +UTC = timezone.utc + + +def _extract_container(request: "PipelineRequest") -> Optional[str]: + http_request = request.http_request + if http_request.method != "GET": + return None + parsed = urlparse(http_request.url) + segments = [seg for seg in parsed.path.split("/") if seg] + if len(segments) < 2: + return None + query = http_request.query + if "comp" in query or query.get("restype") == "container": + return None + container_name = segments[0] + return container_name + + +def _extract_session(response: "CreateSessionResponse") -> Tuple[str, str, datetime]: + creds = getattr(response, "credentials", None) + if not creds or not getattr(creds, "session_token", None) or not getattr(creds, "session_key", None): + raise ValueError("CreateSession response missing SessionToken/SessionKey") + session_token: str = creds.session_token + session_key: str = creds.session_key + expires_at = getattr(response, "expiration", None) + if expires_at is None: + expires_at = datetime.now(UTC) + timedelta(minutes=5) + elif expires_at.tzinfo is None: + expires_at = expires_at.replace(tzinfo=UTC) + return session_token, session_key, expires_at + + +def _to_service_url(url: str) -> str: + parsed = urlparse(url) + return f"{parsed.scheme}://{parsed.netloc}" + + +def _is_cooldown_error(status: Optional[int], error_code: str) -> bool: + if status is None: + return False + if status >= 500 or status == 403: + return True + return status == 400 and error_code == StorageErrorCode.FEATURE_NOT_ENABLED + + +class Session: + """A session entry.""" + + __slots__ = ("session_token", "session_key", "expires_at", "is_fallback") + + REFRESH_BUFFER: timedelta = timedelta(seconds=30) + """Buffer before proactive refresh is initiated.""" + + def __init__( + self, + session_token: Optional[str], + session_key: Optional[str], + expires_at: datetime, + is_fallback: bool = False, + ) -> None: + self.session_token = session_token + self.session_key = session_key + self.expires_at = expires_at + self.is_fallback = is_fallback + + def expired(self) -> bool: + diff = timedelta(seconds=0) if self.is_fallback else Session.REFRESH_BUFFER + return datetime.now(UTC) >= self.expires_at - diff + + +class SessionProvider(Protocol): + """Creates, caches, and invalidates per-container sessions.""" + + def is_request_eligible(self, request: "PipelineRequest") -> bool: + ... + + def get_session(self, request: "PipelineRequest") -> Optional[Session]: + ... + + def invalidate_session(self, request: "PipelineRequest", current: Session) -> None: + ... + + +class SessionCache: + """Thread-safe, container-level storage for sessions on the sync stack. + + Concurrency model + ----------------- + * Reads (`get`) are lock-free. They perform a single dict.get and never + mutate the cache, so concurrent readers never need to coordinate. + * Writes (`put` / `put_fallback`) must be made under the lock returned by + :meth:`lock_container`, which callers also use to single-flight CreateSession. + * A single _locks_guard serializes only the *creation* of per-container + locks, so two threads racing on a brand-new container can't build two + different lock objects. + """ + + FALLBACK_COOLDOWN: timedelta = timedelta(minutes=5) + """Cooldown applied to the fallback-to-bearer sentinel after an eligible create session failure.""" + + def __init__(self) -> None: + self._locks: Dict[str, Lock] = {} + self._locks_guard: Lock = Lock() + self._entry: Dict[str, Session] = {} + + def lock_container(self, container_name: str) -> Lock: + """Return the per-container lock, creating it exactly once. + + :param str container_name: The container name to get the lock for. + :return: The single lock instance associated with the container. + :rtype: ~threading.Lock + """ + # Easy path: lock already exists, and on free threads it falls to slow path + existing_lock = self._locks.get(container_name) + if existing_lock is not None: + return existing_lock + # Slow path: create exactly one lock per container + with self._locks_guard: + return self._locks.setdefault(container_name, Lock()) + + def get(self, container_name: str) -> Optional[Session]: + """Return a live session for the container, or None. + + Lock-free and non-mutating. Expired entries are NOT deleted. + Instead, they are simply treated as a cache miss and overwritten on the next refresh. + + :param str container_name: The container name to look up. + :return: A live (non-expired) session, or None on miss/expiry. + :rtype: ~azure.storage.blob._shared.session.Session or None + """ + cached = self._entry.get(container_name, None) + if cached is None or cached.expired(): + return None + return cached + + def put(self, container_name: str, session: Session) -> None: + """Install a real session entry. + + Caller must hold the lock at the container-level. + + :param str container_name: The container name the session belongs to. + :param session: The session to cache. + :type session: ~azure.storage.blob._shared.session.Session + """ + self._entry[container_name] = session + + def put_fallback(self, container_name: str) -> None: + """Install a fallback-to-bearer sentinel for the cooldown window. + + Caller must hold the lock at the container-level. + + :param str container_name: The container name to mark for bearer fallback. + """ + self._entry[container_name] = Session( + None, None, datetime.now(UTC) + self.FALLBACK_COOLDOWN, is_fallback=True + ) + + def invalidate(self, container_name: str, session_token: Optional[str] = None) -> None: + """Drop the cached session if it still matches the rejected token. + + :param str container_name: The container name. + :param str session_token: The rejected token, or None if unknown. + """ + with self.lock_container(container_name): + cached = self._entry.get(container_name, None) + if cached is not None and cached.session_token == session_token: + self._entry.pop(container_name, None) + + +class ContainerSessionProvider: + """Creates, caches, and invalidates per-container sessions backed by a TokenCredential. + + A single provider may be shared across multiple clients to persist the session + cache beyond the lifetime of any one of them. When no provider is supplied, each + client creates one scoped to itself. + + :param str service_url: The blob service endpoint. Container and blob path segments + and all query parameters are stripped. + :param credential: The credential used to authorize CreateSession calls. + :type credential: ~azure.core.credentials.TokenCredential + :keyword str api_version: The Storage API version to use for CreateSession. + """ + + def __init__(self, service_url: str, credential: "TokenCredential", **kwargs: Any) -> None: + from .._blob_service_client import BlobServiceClient # module-level import would cycle + + if not hasattr(credential, "get_token"): + raise TypeError( + f"ContainerSessionProvider requires a TokenCredential; received {type(credential).__name__}." + ) + self._client = BlobServiceClient(_to_service_url(service_url), credential=credential, **kwargs) + self._cache = SessionCache() + + def is_request_eligible(self, request: "PipelineRequest") -> bool: + """Checks whether the request can be signed with a session token. + + :param ~azure.core.pipeline.PipelineRequest request: The outgoing request. + :return: True if the request is valid. + :rtype: bool + """ + return _extract_container(request) is not None + + def get_session(self, request: "PipelineRequest") -> Optional[Session]: + """Return a session, creating one on a miss. + + :param ~azure.core.pipeline.PipelineRequest request: The outgoing request. + :return: A session, or None if the caller should use bearer auth. + :rtype: ~azure.storage.blob._shared.session.Session or None + """ + container_name = _extract_container(request) + if container_name is None: + return None + + session = self._cache.get(container_name) + if session is None: + session = self._acquire(container_name) + if session is None or session.is_fallback: + return None + return session + + def invalidate_session(self, request: "PipelineRequest", current: Session) -> None: + """Drop the cached session if it still matches the rejected one. + + :param ~azure.core.pipeline.PipelineRequest request: The rejected request. + :param current: The session that was rejected. + :type current: ~azure.storage.blob._shared.session.Session + """ + container_name = _extract_container(request) + if container_name is not None: + self._cache.invalidate(container_name, current.session_token) + + def _acquire(self, container_name: str) -> Optional[Session]: + with self._cache.lock_container(container_name): + existing = self._cache.get(container_name) + if existing is not None: + return existing + try: + token, key, expires_at = self._create_session(container_name) + except HttpResponseError as error: + headers = getattr(error.response, "headers", {}) + error_code = headers.get("x-ms-error-code", "") + if _is_cooldown_error(error.status_code, error_code): + _LOGGER.warning( + "CreateSession failed for container '%s' (HTTP %s, %s); " + "falling back to bearer for %d seconds.", + container_name, + error.status_code, + error_code, + int(self._cache.FALLBACK_COOLDOWN.total_seconds()), + ) + self._cache.put_fallback(container_name) + else: + _LOGGER.warning( + "CreateSession failed for container '%s'; using bearer for this request.", + container_name, + exc_info=True, + ) + return None + except (AzureError, ValueError): + _LOGGER.warning( + "CreateSession failed for container '%s'; using bearer for this request.", + container_name, + exc_info=True, + ) + return None + session = Session(token, key, expires_at) + self._cache.put(container_name, session) + return session + + def _create_session(self, container_name: str) -> Tuple[str, str, datetime]: + container_client = self._client.get_container_client(container_name) + response = container_client._client.container.create_session( # pylint: disable=protected-access + create_session_configuration=CreateSessionConfiguration(authentication_type="HMAC") + ) + return _extract_session(response) diff --git a/sdk/storage/azure-storage-blob/tests/conftest.py b/sdk/storage/azure-storage-blob/tests/conftest.py index c5590814a4ed..d1df41d511f3 100644 --- a/sdk/storage/azure-storage-blob/tests/conftest.py +++ b/sdk/storage/azure-storage-blob/tests/conftest.py @@ -29,6 +29,11 @@ def add_sanitizers(test_proxy): add_general_regex_sanitizer(regex=tenant_id, value="00000000-0000-0000-0000-000000000000") add_header_regex_sanitizer(key="Set-Cookie", value="[set-cookie;]") add_header_regex_sanitizer(key="Cookie", value="cookie;") + add_header_regex_sanitizer(key="x-ms-session-token", value="Sanitized") + add_general_regex_sanitizer( + regex=r"[^<]*", value="Sanitized" + ) + add_general_regex_sanitizer(regex=r"[^<]*", value="U2FuaXRpemVk") add_oauth_response_sanitizer() add_header_regex_sanitizer(key="x-ms-copy-source-authorization", value="Sanitized") diff --git a/sdk/storage/azure-storage-blob/tests/test_container.py b/sdk/storage/azure-storage-blob/tests/test_container.py index 519bd4de862b..12501cf6da1f 100644 --- a/sdk/storage/azure-storage-blob/tests/test_container.py +++ b/sdk/storage/azure-storage-blob/tests/test_container.py @@ -14,6 +14,7 @@ from devtools_testutils import recorded_by_proxy, set_custom_default_matcher from devtools_testutils.storage import StorageRecordedTestCase from settings.testcase import BlobPreparer +from test_helpers import CaptureAuthHeader, _find_session_policy, _parse_session_token from azure.core import MatchConditions from azure.core.exceptions import HttpResponseError, ResourceExistsError, ResourceModifiedError, ResourceNotFoundError @@ -2744,3 +2745,116 @@ def recursive_walk(prefix): # Assert assert blobs is not None assert blobs == ["a/b/blob2", "a/b/blob3", "a/b/blob4", "a/blob1"] + + + @BlobPreparer() + @recorded_by_proxy + def test_create_session(self, **kwargs): + storage_account_name = kwargs.pop("storage_account_name") + + credential = self.get_credential(BlobServiceClient) + capture_auth_header = CaptureAuthHeader() + + service = BlobServiceClient( + self.account_url(storage_account_name, "blob"), + credential=credential, + use_session=True, + ) + container1_name = self.get_resource_name("utcontainer1") + container1 = service.get_container_client(container1_name) + try: + container1.create_container() + except ResourceExistsError: + pass + + blob1_name, blob1_data = self.get_resource_name("blob1"), b"abc123" + container1.upload_blob( + blob1_name, blob1_data, overwrite=True, raw_response_hook=capture_auth_header.hook("c1_upload") + ) + assert capture_auth_header["c1_upload"].startswith("Bearer ") + + blob1_actual = container1.download_blob( + blob1_name, raw_response_hook=capture_auth_header.hook("c1_download") + ).readall() + assert blob1_data == blob1_actual + assert capture_auth_header["c1_download"].startswith("Session ") + session1 = _parse_session_token(capture_auth_header["c1_download"]) + + container2_name = self.get_resource_name("utcontainer2") + container2 = service.get_container_client(container2_name) + try: + container2.create_container() + except ResourceExistsError: + pass + + blob2_name, blob2_data = self.get_resource_name("blob2"), b"def456" + container2.upload_blob( + blob2_name, blob2_data, overwrite=True, raw_response_hook=capture_auth_header.hook("c2_upload") + ) + assert capture_auth_header["c2_upload"].startswith("Bearer ") + + blob2_actual = container2.download_blob( + blob2_name, raw_response_hook=capture_auth_header.hook("c2_download") + ).readall() + assert blob2_data == blob2_actual + assert capture_auth_header["c2_download"].startswith("Session ") + session2 = _parse_session_token(capture_auth_header["c2_download"]) + + assert session1 != session2 + + blob1_actual = container1.download_blob( + blob1_name, raw_response_hook=capture_auth_header.hook("c1_download2") + ).readall() + assert blob1_data == blob1_actual + assert capture_auth_header["c1_download2"].startswith("Session ") + assert session1 == _parse_session_token(capture_auth_header["c1_download2"]) + + blob2_actual = container2.download_blob( + blob2_name, raw_response_hook=capture_auth_header.hook("c2_download2") + ).readall() + assert blob2_data == blob2_actual + assert capture_auth_header["c2_download2"].startswith("Session ") + assert session2 == _parse_session_token(capture_auth_header["c2_download2"]) + + policy = _find_session_policy(service._pipeline) + cached = policy._cache._entry[container1_name] + cached.expires_at = datetime.fromtimestamp(0, tz=cached.expires_at.tzinfo) + + blob1_actual = container1.download_blob( + blob1_name, raw_response_hook=capture_auth_header.hook("c1_download3") + ).readall() + assert blob1_data == blob1_actual + assert capture_auth_header["c1_download3"].startswith("Session ") + assert session1 != _parse_session_token(capture_auth_header["c1_download3"]) + assert session2 != _parse_session_token(capture_auth_header["c1_download3"]) + + @BlobPreparer() + @recorded_by_proxy + def test_sessions_disabled(self, **kwargs): + storage_account_name = kwargs.pop("storage_account_name") + + credential = self.get_credential(BlobServiceClient) + capture_auth_header = CaptureAuthHeader() + + service = BlobServiceClient( + self.account_url(storage_account_name, "blob"), + credential=credential, + use_session=False, + ) + container = service.get_container_client(self.get_resource_name("utcontainer")) + try: + container.create_container() + except ResourceExistsError: + pass + + blob_name, blob_data = self.get_resource_name("blob"), b"abc123" + container.upload_blob( + blob_name, blob_data, overwrite=True, raw_response_hook=capture_auth_header.hook("upload") + ) + assert capture_auth_header["upload"].startswith("Bearer ") + + blob_actual = container.download_blob( + blob_name, raw_response_hook=capture_auth_header.hook("download") + ).readall() + assert blob_data == blob_actual + assert capture_auth_header["download"].startswith("Bearer ") diff --git a/sdk/storage/azure-storage-blob/tests/test_helpers.py b/sdk/storage/azure-storage-blob/tests/test_helpers.py index cb7fc2d2bf6e..7c4aad539b40 100644 --- a/sdk/storage/azure-storage-blob/tests/test_helpers.py +++ b/sdk/storage/azure-storage-blob/tests/test_helpers.py @@ -64,6 +64,63 @@ def _create_file_share_oauth( return file_name, base_url +def _parse_session_token(auth: str) -> str: + """Extract the token from a "Session {token}:{signature}" Authorization header. + + :param str auth: The raw Authorization header value. + :return: The session token portion (before the ':'). + :rtype: str + """ + assert auth.startswith("Session ") + return auth[len("Session ") :].split(":", 1)[0] + + +def _find_session_policy(pipeline: Any, policy_name: str = "StorageSessionPolicy") -> Any: + """Return the session policy instance on a client pipeline, matched by class name. + + Matching by name avoids importing SDK internals into the test modules. + + :param pipeline: The client pipeline to search (e.g. ``client._pipeline``). + :type pipeline: Any + :param str policy_name: The policy class name to find. Use "StorageSessionPolicy" + for the sync stack and "AsyncStorageSessionPolicy" for the async stack. + :return: The matching policy instance. + :rtype: Any + """ + for policy in getattr(pipeline, "_impl_policies", []): + if type(policy).__name__ == policy_name: + return policy + raise AssertionError(f"{policy_name} not found on the pipeline") + + +class CaptureAuthHeader: + """Captures per-label Authorization headers via ``raw_response_hook`` callbacks. + + Encapsulates the captured-headers dict so the hook factory doesn't need a + closure over a test-local variable. Works for both sync and async clients, + since the response hook is invoked as a plain callable in both stacks. + """ + + def __init__(self) -> None: + self.captured: Dict[str, str] = {} + + def hook(self, label: str): + """Return a ``raw_response_hook`` that records the request's Authorization header. + + :param str label: The key under which to store the captured header. + :return: A callable suitable for ``raw_response_hook``. + :rtype: callable + """ + + def _hook(response): + self.captured[label] = response.http_request.headers.get("Authorization", "") + + return _hook + + def __getitem__(self, label: str) -> str: + return self.captured[label] + + class ProgressTracker: def __init__(self, total: int, step: int): self.total = total