From ab9618c3962ed1d1df42c5f9f6ad57b269df0327 Mon Sep 17 00:00:00 2001 From: Paul Van Eck Date: Wed, 8 Oct 2025 01:54:50 +0000 Subject: [PATCH 01/12] [Identity] Implement binding mode support in WorkloadIdentityCredential - Implement a new http.client transport that will leverage Proxy, CA, SNI settings. Signed-off-by: Paul Van Eck --- .../azure/identity/_constants.py | 5 + .../_credentials/workload_identity.py | 28 + .../_internal/http_client_transport.py | 351 ++++++++++ .../aio/_credentials/workload_identity.py | 29 + .../azure-identity/tests/proxy_server.py | 302 +++++++++ .../tests/test_http_client_transport.py | 606 ++++++++++++++++++ .../test_workload_identity_credential.py | 207 ++++++ ...test_workload_identity_credential_async.py | 208 ++++++ 8 files changed, 1736 insertions(+) create mode 100644 sdk/identity/azure-identity/azure/identity/_internal/http_client_transport.py create mode 100644 sdk/identity/azure-identity/tests/proxy_server.py create mode 100644 sdk/identity/azure-identity/tests/test_http_client_transport.py diff --git a/sdk/identity/azure-identity/azure/identity/_constants.py b/sdk/identity/azure-identity/azure/identity/_constants.py index 09732f3de9ec..4f86f3d94c48 100644 --- a/sdk/identity/azure-identity/azure/identity/_constants.py +++ b/sdk/identity/azure-identity/azure/identity/_constants.py @@ -68,5 +68,10 @@ class EnvironmentVariables: AZURE_REGIONAL_AUTHORITY_NAME = "AZURE_REGIONAL_AUTHORITY_NAME" AZURE_FEDERATED_TOKEN_FILE = "AZURE_FEDERATED_TOKEN_FILE" + AZURE_KUBERNETES_SNI_NAME = "AZURE_KUBERNETES_SNI_NAME" + AZURE_KUBERNETES_TOKEN_PROXY = "AZURE_KUBERNETES_TOKEN_PROXY" + AZURE_KUBERNETES_CA_FILE = "AZURE_KUBERNETES_CA_FILE" + AZURE_KUBERNETES_CA_DATA = "AZURE_KUBERNETES_CA_DATA" + AZURE_TOKEN_CREDENTIALS = "AZURE_TOKEN_CREDENTIALS" WORKLOAD_IDENTITY_VARS = (AZURE_AUTHORITY_HOST, AZURE_TENANT_ID, AZURE_FEDERATED_TOKEN_FILE) diff --git a/sdk/identity/azure-identity/azure/identity/_credentials/workload_identity.py b/sdk/identity/azure-identity/azure/identity/_credentials/workload_identity.py index db15c4ed633c..0dc44294a8cc 100644 --- a/sdk/identity/azure-identity/azure/identity/_credentials/workload_identity.py +++ b/sdk/identity/azure-identity/azure/identity/_credentials/workload_identity.py @@ -9,6 +9,7 @@ from .client_assertion import ClientAssertionCredential from .._constants import EnvironmentVariables +from .._internal.http_client_transport import HttpClientTransport WORKLOAD_CONFIG_ERROR = ( @@ -57,6 +58,8 @@ class WorkloadIdentityCredential(ClientAssertionCredential, TokenFileMixin): :keyword str client_id: The client ID of a Microsoft Entra app registration. :keyword str token_file_path: The path to a file containing a Kubernetes service account token that authenticates the identity. + :keyword str use_token_proxy: Whether or not to to read token proxy configuration from environment variables and use + a token proxy to acquire tokens. Defaults to False. .. admonition:: Example: @@ -74,6 +77,7 @@ def __init__( tenant_id: Optional[str] = None, client_id: Optional[str] = None, token_file_path: Optional[str] = None, + use_token_proxy: bool = False, **kwargs: Any, ) -> None: tenant_id = tenant_id or os.environ.get(EnvironmentVariables.AZURE_TENANT_ID) @@ -99,6 +103,30 @@ def __init__( assert token_file_path is not None self._token_file_path = token_file_path + + if use_token_proxy: + token_proxy_endpoint = os.environ.get(EnvironmentVariables.AZURE_KUBERNETES_TOKEN_PROXY) + if not token_proxy_endpoint: + raise ValueError( + "use_token_proxy is True, but no token proxy endpoint was found. " + f"Ensure that the {EnvironmentVariables.AZURE_KUBERNETES_TOKEN_PROXY} environment variable is set." + ) + + sni = os.environ.get(EnvironmentVariables.AZURE_KUBERNETES_SNI_NAME) + ca_file = os.environ.get(EnvironmentVariables.AZURE_KUBERNETES_CA_FILE) + ca_data = os.environ.get(EnvironmentVariables.AZURE_KUBERNETES_CA_DATA) + + if ca_file and ca_data: + raise ValueError( + "Both AZURE_KUBERNETES_CA_FILE and AZURE_KUBERNETES_CA_DATA are set. Only one should be set" + ) + + kwargs["transport"] = HttpClientTransport( + sni=sni, + proxy_endpoint=token_proxy_endpoint, + ca_file=ca_file, + ca_data=ca_data, + ) super(WorkloadIdentityCredential, self).__init__( tenant_id=tenant_id, client_id=client_id, diff --git a/sdk/identity/azure-identity/azure/identity/_internal/http_client_transport.py b/sdk/identity/azure-identity/azure/identity/_internal/http_client_transport.py new file mode 100644 index 000000000000..4677deb28bae --- /dev/null +++ b/sdk/identity/azure-identity/azure/identity/_internal/http_client_transport.py @@ -0,0 +1,351 @@ +# ------------------------------------ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. +# ------------------------------------ +# cspell:ignore cafile +from json import loads +import hashlib +import http.client +import os +import ssl +import urllib.parse +from typing import Any, Iterator, MutableMapping, Optional + +from azure.core.configuration import ConnectionConfiguration +from azure.core.exceptions import ServiceRequestError, ServiceResponseError +from azure.core.rest import HttpRequest, HttpResponse +from azure.core.pipeline.transport import HttpTransport + + +class HttpClientTransportResponse(HttpResponse): + """Create a HttpResponse from an http.client response. + + :param HttpRequest request: The request. + :type request: ~azure.core.pipeline.transport.HttpRequest + :param httpclient_response: The response from http.client + :type httpclient_response: http.client.HTTPResponse + :param block_size: The block size to use for downloading the response content. + :type block_size: int + """ + + def __init__( + self, request: HttpRequest, httpclient_response: http.client.HTTPResponse, block_size: Optional[int] = None + ) -> None: + self._request = request + self._httpclient_response = httpclient_response + self._block_size = block_size or 4096 + self._data: Optional[bytes] = None + self._closed = False + self._headers = {k.lower(): v for k, v in httpclient_response.getheaders()} + self._content_type = self._headers.get("content-type") + self._encoding: Optional[str] = None + + def __enter__(self) -> "HttpClientTransportResponse": + return self + + def __exit__(self, *args: Any) -> None: + self.close() + + def close(self) -> None: + if not self._closed: + self._httpclient_response.close() + self._closed = True + + def read(self) -> bytes: + if self._data is None: + self._data = self._httpclient_response.read() + return self._data + + def iter_raw(self, **kwargs: Any) -> Iterator[bytes]: + if self._data: + yield self._data + else: + chunk = self._httpclient_response.read(self._block_size) + while chunk: + yield chunk + chunk = self._httpclient_response.read(self._block_size) + + def iter_bytes(self, **kwargs: Any) -> Iterator[bytes]: + # http.client doesn't support compressed encoding automatically, + # so the decompression is already done at read time. + # Just use iter_raw here + yield from self.iter_raw(**kwargs) + + @property + def request(self) -> HttpRequest: + return self._request + + @property + def status_code(self) -> int: + return self._httpclient_response.status + + @property + def headers(self) -> MutableMapping[str, str]: + return self._headers + + @property + def reason(self) -> str: + return self._httpclient_response.reason + + @property + def content_type(self) -> Optional[str]: + return self._content_type + + @property + def url(self) -> str: + return self._request.url + + @property + def is_closed(self) -> bool: + return self._closed + + @property + def is_stream_consumed(self) -> bool: + return self._data is not None + + @property + def encoding(self) -> Optional[str]: + return self._encoding + + @encoding.setter + def encoding(self, value: Optional[str]) -> None: + self._encoding = value + + @property + def content(self) -> bytes: + return self.read() + + def text(self, encoding: Optional[str] = None) -> str: + if encoding is None: + encoding = self.encoding or "utf-8" + return self.content.decode(encoding) + + def raise_for_status(self) -> None: + if self.status_code >= 400: + from azure.core.exceptions import HttpResponseError + + raise HttpResponseError(response=self) + + def json(self) -> Any: + return loads(self.text()) + + +class SniSSLContext(ssl.SSLContext): + def __new__(cls, sni_hostname: str, protocol=None): + instance = super().__new__(cls, protocol=protocol) + instance.sni_hostname = sni_hostname # type: ignore + return instance + + def wrap_socket(self, *args, **kwargs): + kwargs["server_hostname"] = self.sni_hostname # type: ignore + return super().wrap_socket(*args, **kwargs) + + +class HttpClientTransport(HttpTransport): + """Implements an HTTP sender using Python's built-in http.client library.""" + + def __init__(self, **kwargs) -> None: + self.connection_config = ConnectionConfiguration(**kwargs) + self._ca_data = kwargs.pop("ca_data", None) + self._ca_file = kwargs.pop("ca_file", None) + + if self._ca_file and self._ca_data: + raise ValueError("Both ca_file and ca_data are set. Only one should be set") + + self._sni = kwargs.pop("sni", None) + self._proxy_endpoint = kwargs.pop("proxy_endpoint", None) + if self._proxy_endpoint: + self._validate_url(self._proxy_endpoint) + + self._connection: Optional[http.client.HTTPSConnection] = None + self._ca_file_hash: Optional[str] = None + self._ca_file_mtime: Optional[float] = None + + # Initialize CA file tracking if a CA file is specified + if self._ca_file: + self._update_ca_file_tracking() + + def __enter__(self) -> "HttpClientTransport": + self.open() + return self + + def __exit__(self, *args): + self.close() + + def open(self) -> None: + pass # We create connections as needed + + def close(self) -> None: + if self._connection: + self._connection.close() + self._connection = None + + def _validate_url(self, url: str) -> None: + parsed_url = urllib.parse.urlparse(url) + if parsed_url.scheme != "https": + raise ValueError(f"Endpoint URL ({url}) must use the 'https' scheme. Got '{parsed_url.scheme}' instead.") + if parsed_url.username or parsed_url.password: + raise ValueError(f"Endpoint URL ({url}) must not contain username or password.") + if parsed_url.fragment: + raise ValueError(f"Endpoint URL ({url}) must not contain a fragment.") + if parsed_url.query: + raise ValueError(f"Endpoint URL ({url}) must not contain query parameters.") + + def _update_ca_file_tracking(self) -> None: + """Update the CA file hash and modification time for change detection.""" + if not self._ca_file or not os.path.exists(self._ca_file): + self._ca_file_hash = None + self._ca_file_mtime = None + return + + try: + # Read file content first to check if empty + with open(self._ca_file, "rb") as f: + content = f.read() + + # Check if the file is empty + if not content: + # If no prior tracking state exists (first read), fail + if self._ca_file_hash is None: + raise ValueError(f"CA file ({self._ca_file}) is empty. Cannot establish secure connection.") + return + + # File has content, update tracking + self._ca_file_mtime = os.path.getmtime(self._ca_file) + self._ca_file_hash = hashlib.sha256(content).hexdigest() + except (OSError, IOError): + # If we can't read the file, clear the tracking + self._ca_file_hash = None + self._ca_file_mtime = None + + def _has_ca_file_changed(self) -> bool: + """Check if the CA file has changed since last tracking update. + + :return: True if the CA file has changed, False otherwise. + :rtype: bool + """ + if not self._ca_file: + return False + + if not os.path.exists(self._ca_file): + # File was deleted, consider this a change + return self._ca_file_hash is not None or self._ca_file_mtime is not None + + try: + # Check modification time first (faster) + current_mtime = os.path.getmtime(self._ca_file) + if self._ca_file_mtime != current_mtime: + return True + + # If mtime is the same, check content hash to be sure + with open(self._ca_file, "rb") as f: + content = f.read() + current_hash = hashlib.sha256(content).hexdigest() + return self._ca_file_hash != current_hash + + except (OSError, IOError): + # If we can't read the file, assume it changed + return True + + def _get_connection(self, host: str) -> http.client.HTTPSConnection: + + # Check if CA file has changed and invalidate connections if needed + if self._ca_file and self._has_ca_file_changed(): + # CA file changed, close all existing connections and clear cache + if self._connection: + self._connection.close() + self._connection = None + # Update tracking with new CA file state + self._update_ca_file_tracking() + + # Use existing connection if available + if self._connection: + return self._connection + + # Create HTTPS connection + ssl_context: ssl.SSLContext + if self._sni: + ssl_context = SniSSLContext(self._sni, ssl.PROTOCOL_TLS_CLIENT) + ssl_context.verify_mode = ssl.CERT_REQUIRED + ssl_context.check_hostname = True + + if self._ca_data or self._ca_file: + ssl_context.load_verify_locations(cafile=self._ca_file, cadata=self._ca_data) + else: + ssl_context.load_default_certs() + else: + ssl_context = ssl.create_default_context(cafile=self._ca_file, cadata=self._ca_data) + + if not self.connection_config.verify: + ssl_context.check_hostname = False + ssl_context.verify_mode = ssl.CERT_NONE + + connection = http.client.HTTPSConnection( + host, + timeout=self.connection_config.timeout, + context=ssl_context, + ) + + self._connection = connection + return connection + + def _update_request_url(self, request: HttpRequest) -> None: + parsed_request_url = urllib.parse.urlparse(request.url) + if self._proxy_endpoint: + parsed_proxy_url = urllib.parse.urlparse(self._proxy_endpoint) + combined_path = parsed_proxy_url.path.rstrip("/") + "/" + parsed_request_url.path.lstrip("/") + new_url = urllib.parse.urlunparse( + ( + parsed_proxy_url.scheme, + parsed_proxy_url.netloc, + combined_path, + parsed_request_url.params, + parsed_request_url.query, + parsed_request_url.fragment, + ) + ) + request.url = new_url + + def send(self, request: HttpRequest, **kwargs) -> HttpResponse: + """Send request object according to configuration. + + :param request: The HTTP request object. + :type request: ~azure.core.rest.HttpRequest + :return: The HTTP response object. + :rtype: ~azure.core.rest.HttpResponse + """ + + # Get or create a connection for the URL + self._update_request_url(request) + parsed_url = urllib.parse.urlparse(request.url) + full_path = urllib.parse.urlunparse( + ("", "", parsed_url.path, parsed_url.params, parsed_url.query, parsed_url.fragment) + ) + + # Get connection timeout + connection_timeout = kwargs.pop("connection_timeout", self.connection_config.timeout) + + try: + # Get connection + connection = self._get_connection(parsed_url.netloc) + + if connection_timeout is not None: + connection.timeout = connection_timeout + + connection.request(request.method, full_path, body=request.data, headers=request.headers) + response = connection.getresponse() + return HttpClientTransportResponse( + request=request, + httpclient_response=response, + block_size=self.connection_config.data_block_size, + ) + + except http.client.HTTPException as err: + raise ServiceRequestError(err) from err + except ssl.SSLError as err: + raise ServiceResponseError(err) from err + except Exception as err: + raise ServiceRequestError(err) from err + + def __repr__(self) -> str: + return f"<{type(self).__name__}>" diff --git a/sdk/identity/azure-identity/azure/identity/aio/_credentials/workload_identity.py b/sdk/identity/azure-identity/azure/identity/aio/_credentials/workload_identity.py index 8c44369da6ff..e9a7cdd4e83c 100644 --- a/sdk/identity/azure-identity/azure/identity/aio/_credentials/workload_identity.py +++ b/sdk/identity/azure-identity/azure/identity/aio/_credentials/workload_identity.py @@ -7,6 +7,7 @@ from .client_assertion import ClientAssertionCredential from ..._credentials.workload_identity import TokenFileMixin, WORKLOAD_CONFIG_ERROR from ..._constants import EnvironmentVariables +from ..._internal.http_client_transport import HttpClientTransport class WorkloadIdentityCredential(ClientAssertionCredential, TokenFileMixin): @@ -30,6 +31,8 @@ class WorkloadIdentityCredential(ClientAssertionCredential, TokenFileMixin): :keyword str client_id: The client ID of a Microsoft Entra app registration. :keyword str token_file_path: The path to a file containing a Kubernetes service account token that authenticates the identity. + :keyword str use_token_proxy: Whether or not to to read token proxy configuration from environment variables and use + a token proxy to acquire tokens. Defaults to False. .. admonition:: Example: @@ -47,6 +50,7 @@ def __init__( tenant_id: Optional[str] = None, client_id: Optional[str] = None, token_file_path: Optional[str] = None, + use_token_proxy: bool = False, **kwargs: Any, ) -> None: tenant_id = tenant_id or os.environ.get(EnvironmentVariables.AZURE_TENANT_ID) @@ -72,6 +76,31 @@ def __init__( assert token_file_path is not None self._token_file_path = token_file_path + + if use_token_proxy: + token_proxy_endpoint = os.environ.get(EnvironmentVariables.AZURE_KUBERNETES_TOKEN_PROXY) + if not token_proxy_endpoint: + raise ValueError( + "use_token_proxy is True, but no token proxy endpoint was found. " + f"Ensure the {EnvironmentVariables.AZURE_KUBERNETES_TOKEN_PROXY} environment variable is set." + ) + + sni = os.environ.get(EnvironmentVariables.AZURE_KUBERNETES_SNI_NAME) + ca_file = os.environ.get(EnvironmentVariables.AZURE_KUBERNETES_CA_FILE) + ca_data = os.environ.get(EnvironmentVariables.AZURE_KUBERNETES_CA_DATA) + + if ca_file and ca_data: + raise ValueError( + "Both AZURE_KUBERNETES_CA_FILE and AZURE_KUBERNETES_CA_DATA are set. Only one should be set" + ) + + kwargs["transport"] = HttpClientTransport( + sni=sni, + proxy_endpoint=token_proxy_endpoint, + ca_file=ca_file, + ca_data=ca_data, + ) + super().__init__( tenant_id=tenant_id, client_id=client_id, diff --git a/sdk/identity/azure-identity/tests/proxy_server.py b/sdk/identity/azure-identity/tests/proxy_server.py new file mode 100644 index 000000000000..352ab8f1ba2a --- /dev/null +++ b/sdk/identity/azure-identity/tests/proxy_server.py @@ -0,0 +1,302 @@ +# ------------------------------------ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. +# ------------------------------------ +# cspell:ignore ests +""" +Local test server for HttpClientTransport testing. + +This server simulates a token proxy that can: +1. Accept HTTPS requests with custom SNI and CA certificates +2. Route requests to downstream services +3. Handle various error scenarios for testing +4. Support certificate rotation scenarios +""" + +import argparse +import ipaddress +import json +import logging +import os +import ssl +import tempfile +import threading +import time +import datetime +from http.server import HTTPServer, BaseHTTPRequestHandler +from socketserver import ThreadingMixIn +import uuid + +from cryptography import x509 +from cryptography.x509.oid import NameOID +from cryptography.hazmat.primitives import hashes, serialization +from cryptography.hazmat.primitives.asymmetric import rsa + + +class TokenProxyHandler(BaseHTTPRequestHandler): + """HTTP request handler that simulates a token proxy.""" + + def log_message(self, format, *args): + """Override to use proper logging instead of stderr.""" + logging.info(f"{self.address_string()} - {format % args}") + + def do_GET(self): + """Handle GET requests.""" + self._handle_request() + + def do_POST(self): + """Handle POST requests.""" + self._handle_request() + + def do_PUT(self): + """Handle PUT requests.""" + self._handle_request() + + def do_DELETE(self): + """Handle DELETE requests.""" + self._handle_request() + + def do_PATCH(self): + """Handle PATCH requests.""" + self._handle_request() + + def _handle_request(self): + """Common request handling logic.""" + path = self.path + headers = dict(self.headers) + + # Read request body if present + content_length = int(headers.get("content-length", 0)) + body = self.rfile.read(content_length) if content_length > 0 else b"" + + logging.info(f"Received {self.command} {path}") + logging.info(f"Headers: {headers}") + + # Simulate different responses based on path + if path == "/health": + self._send_health_response() + elif path.startswith("/oauth2/v2.0/token"): + self._send_token_response(body) + elif path == "/error/500": + self._send_error_response(500, "Internal Server Error") + elif path == "/error/ssl": + # Simulate SSL error by closing connection + self.wfile.close() + return + elif path == "/slow": + # Simulate slow response using threading.Event instead of time.sleep + # to avoid being mocked by conftest.py + event = threading.Event() + event.wait(timeout=2) + self._send_json_response({"message": "slow response"}) + else: + self._send_proxy_response(path, body, headers) + + def _send_health_response(self): + """Send a health check response.""" + response = {"status": "healthy", "timestamp": time.time(), "server": "token-proxy-test-server"} + self._send_json_response(response) + + def _send_token_response(self, body): + """Send a mock OAuth token response.""" + # Parse the request body to extract grant type, etc. + try: + if body: + body_str = body.decode("utf-8") + logging.info(f"Token request body: {body_str}") + except Exception as e: + logging.warning(f"Could not decode request body: {e}") + + # Mock token response + response = { + "access_token": f"mock_token_{uuid.uuid4().hex[:16]}", + "token_type": "Bearer", + "expires_in": 3600, + "scope": "https://graph.microsoft.com/.default", + } + self._send_json_response(response) + + def _send_proxy_response(self, path, body, headers): + """Send a generic proxy response.""" + response = { + "proxied_path": path, + "method": self.command, + "headers_received": dict(headers), + "body_length": len(body), + "proxy_server": "token-proxy-test-server", + } + self._send_json_response(response) + + def _send_json_response(self, data, status_code=200): + """Send a JSON response.""" + response_json = json.dumps(data, indent=2) + + self.send_response(status_code) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(response_json))) + self.send_header("Server", "token-proxy-test-server") + self.end_headers() + + self.wfile.write(response_json.encode("utf-8")) + + def _send_error_response(self, status_code, message): + """Send an error response.""" + self.send_response(status_code) + self.send_header("Content-Type", "text/plain") + self.send_header("Content-Length", str(len(message))) + self.end_headers() + + self.wfile.write(message.encode("utf-8")) + + +class ThreadedHTTPServer(ThreadingMixIn, HTTPServer): + """Threaded HTTP server for handling multiple concurrent requests.""" + + allow_reuse_address = True + daemon_threads = True + + +class TokenProxyTestServer: + """Test server that can be configured with SSL/TLS and custom certificates.""" + + def __init__(self, host="localhost", port=0, use_ssl=True): + self.host = host + self.port = port + self.use_ssl = use_ssl + self.server = None + self.server_thread = None + self.cert_file = None + self.key_file = None + self.ca_file = None + self._temp_files = [] + + def generate_test_certificates(self): + """Generate self-signed certificates for testing.""" + + # Generate private key + private_key = rsa.generate_private_key( + public_exponent=65537, + key_size=2048, + ) + + # Create certificate + subject = issuer = x509.Name( + [ + x509.NameAttribute(NameOID.COUNTRY_NAME, "US"), + x509.NameAttribute(NameOID.STATE_OR_PROVINCE_NAME, "Test"), + x509.NameAttribute(NameOID.LOCALITY_NAME, "Test"), + x509.NameAttribute(NameOID.ORGANIZATION_NAME, "Test Proxy Server"), + x509.NameAttribute(NameOID.COMMON_NAME, self.host), + ] + ) + + cert = ( + x509.CertificateBuilder() + .subject_name(subject) + .issuer_name(issuer) + .public_key(private_key.public_key()) + .serial_number(x509.random_serial_number()) + .not_valid_before(datetime.datetime.now(datetime.timezone.utc)) + .not_valid_after(datetime.datetime.now(datetime.timezone.utc) + datetime.timedelta(days=365)) + .add_extension( + x509.SubjectAlternativeName( + [ + x509.DNSName(self.host), + x509.DNSName("localhost"), + x509.DNSName("1234.ests.aks"), + x509.IPAddress(ipaddress.IPv4Address("127.0.0.1")), + ] + ), + critical=False, + ) + .sign(private_key, hashes.SHA256()) + ) + + # Write certificate to temp file + cert_fd, self.cert_file = tempfile.mkstemp(suffix=".pem", prefix="test_cert_") + with os.fdopen(cert_fd, "wb") as f: + f.write(cert.public_bytes(serialization.Encoding.PEM)) + self._temp_files.append(self.cert_file) + + # Write private key to temp file + key_fd, self.key_file = tempfile.mkstemp(suffix=".pem", prefix="test_key_") + with os.fdopen(key_fd, "wb") as f: + f.write( + private_key.private_bytes( + encoding=serialization.Encoding.PEM, + format=serialization.PrivateFormat.PKCS8, + encryption_algorithm=serialization.NoEncryption(), + ) + ) + self._temp_files.append(self.key_file) + + # Use the same cert as CA for simplicity + self.ca_file = self.cert_file + + logging.info(f"Generated test certificate: {self.cert_file}") + logging.info(f"Generated test key: {self.key_file}") + + def start(self): + """Start the test server.""" + if self.use_ssl: + self.generate_test_certificates() + + # Create server + self.server = ThreadedHTTPServer((self.host, self.port), TokenProxyHandler) + + if self.use_ssl: + # Configure SSL context + context = ssl.create_default_context(ssl.Purpose.CLIENT_AUTH) + if self.cert_file and self.key_file: + context.load_cert_chain(self.cert_file, self.key_file) + # Disable certificate verification for testing + context.check_hostname = False + context.verify_mode = ssl.CERT_NONE + + self.server.socket = context.wrap_socket(self.server.socket, server_side=True) + + # Update port if it was 0 (auto-assigned) + self.port = self.server.server_address[1] + + # Start server in background thread + self.server_thread = threading.Thread(target=self.server.serve_forever, daemon=True) + self.server_thread.start() + + scheme = "https" if self.use_ssl else "http" + logging.info(f"Test server started at {scheme}://{self.host}:{self.port}") + + return f"{scheme}://{self.host}:{self.port}" + + def stop(self): + """Stop the test server and clean up.""" + if self.server: + self.server.shutdown() + self.server.server_close() + + if self.server_thread: + self.server_thread.join(timeout=5) + + # Clean up temporary files + for temp_file in self._temp_files: + try: + os.unlink(temp_file) + except OSError: + pass + self._temp_files.clear() # Clear the list after cleanup + + logging.info("Test server stopped and cleaned up") + + def __enter__(self): + """Context manager entry.""" + self.start() + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + """Context manager exit.""" + self.stop() + + @property + def base_url(self): + """Get the base URL of the server.""" + scheme = "https" if self.use_ssl else "http" + return f"{scheme}://{self.host}:{self.port}" diff --git a/sdk/identity/azure-identity/tests/test_http_client_transport.py b/sdk/identity/azure-identity/tests/test_http_client_transport.py new file mode 100644 index 000000000000..7282b7c7fe8d --- /dev/null +++ b/sdk/identity/azure-identity/tests/test_http_client_transport.py @@ -0,0 +1,606 @@ +# ------------------------------------ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. +# ------------------------------------ +import hashlib +import http.client +import json +import os +import ssl +import tempfile +import time +from unittest import mock +from urllib.parse import urlparse + +import pytest + +from azure.core.rest import HttpRequest +from azure.core.exceptions import ServiceRequestError, ServiceResponseError +from azure.identity._internal.http_client_transport import HttpClientTransport, SniSSLContext + + +class TestHttpClientTransport: + """Test cases for HttpClientTransport class.""" + + def test_init_basic(self): + """Test basic initialization of HttpClientTransport.""" + transport = HttpClientTransport() + assert transport._ca_data is None + assert transport._ca_file is None + assert transport._sni is None + assert transport._proxy_endpoint is None + assert transport._connection is None + assert transport._ca_file_hash is None + assert transport._ca_file_mtime is None + + def test_init_with_ca_data(self): + """Test initialization with CA data.""" + ca_data = "-----BEGIN CERTIFICATE-----\nSome fake cert\n-----END CERTIFICATE-----" + transport = HttpClientTransport(ca_data=ca_data) + assert transport._ca_data == ca_data + assert transport._ca_file is None + + def test_init_with_ca_file(self): + """Test initialization with CA file.""" + # Use simple but valid-looking PEM content + ca_content = "-----BEGIN CERTIFICATE-----\nMIIDummy certificate content here\n-----END CERTIFICATE-----" + + with tempfile.NamedTemporaryFile(mode="w", delete=False) as temp_file: + temp_file.write(ca_content) + temp_file_path = temp_file.name + + try: + transport = HttpClientTransport(ca_file=temp_file_path) + assert transport._ca_file == temp_file_path + assert transport._ca_data is None + assert transport._ca_file_hash is not None + assert transport._ca_file_mtime is not None + finally: + os.unlink(temp_file_path) + + def test_init_with_both_ca_file_and_data_raises_error(self): + """Test that providing both CA file and data raises an error.""" + ca_data = "-----BEGIN CERTIFICATE-----\nMIIDummy certificate content\n-----END CERTIFICATE-----" + + with tempfile.NamedTemporaryFile(mode="w", delete=False) as temp_file: + temp_file.write(ca_data) + temp_file_path = temp_file.name + + try: + with pytest.raises(ValueError, match="Both ca_file and ca_data are set"): + HttpClientTransport(ca_file=temp_file_path, ca_data=ca_data) + finally: + os.unlink(temp_file_path) + + def test_init_with_empty_ca_file_raises_error(self): + """Test that empty CA file raises an error.""" + with tempfile.NamedTemporaryFile(mode="w", delete=False) as temp_file: + temp_file_path = temp_file.name + + try: + with pytest.raises(ValueError, match="CA file .* is empty"): + HttpClientTransport(ca_file=temp_file_path) + finally: + os.unlink(temp_file_path) + + def test_init_with_sni(self): + """Test initialization with SNI hostname.""" + sni_hostname = "example.com" + transport = HttpClientTransport(sni=sni_hostname) + assert transport._sni == sni_hostname + + def test_init_with_proxy_endpoint(self): + """Test initialization with proxy endpoint.""" + proxy_endpoint = "https://proxy.example.com:8080" + transport = HttpClientTransport(proxy_endpoint=proxy_endpoint) + assert transport._proxy_endpoint == proxy_endpoint + + def test_validate_url_valid_https(self): + """Test URL validation with valid HTTPS URL.""" + transport = HttpClientTransport() + # Should not raise any exception + transport._validate_url("https://example.com/path") + + def test_validate_url_non_https_scheme(self): + """Test URL validation rejects non-HTTPS schemes.""" + transport = HttpClientTransport() + with pytest.raises(ValueError, match="must use the 'https' scheme"): + transport._validate_url("http://example.com") + + def test_validate_url_with_user_info(self): + """Test URL validation rejects URLs with user info.""" + transport = HttpClientTransport() + with pytest.raises(ValueError, match="must not contain username or password"): + transport._validate_url("https://user:pass@example.com") + + def test_validate_url_with_fragment(self): + """Test URL validation rejects URLs with fragments.""" + transport = HttpClientTransport() + with pytest.raises(ValueError, match="must not contain a fragment"): + transport._validate_url("https://example.com#fragment") + + def test_validate_url_with_query(self): + """Test URL validation rejects URLs with query parameters.""" + transport = HttpClientTransport() + with pytest.raises(ValueError, match="must not contain query parameters"): + transport._validate_url("https://example.com?query=value") + + def test_ca_file_tracking_nonexistent_file(self): + """Test CA file tracking with non-existent file.""" + transport = HttpClientTransport() + transport._ca_file = "/nonexistent/file.pem" + transport._update_ca_file_tracking() + assert transport._ca_file_hash is None + assert transport._ca_file_mtime is None + + def test_ca_file_tracking_updates_hash_and_mtime(self): + """Test CA file tracking updates hash and modification time.""" + content = "-----BEGIN CERTIFICATE-----\nMIIDummy certificate content\n-----END CERTIFICATE-----" + + with tempfile.NamedTemporaryFile(mode="w", delete=False) as temp_file: + temp_file.write(content) + temp_file_path = temp_file.name + + try: + transport = HttpClientTransport(ca_file=temp_file_path) + expected_hash = hashlib.sha256(content.encode()).hexdigest() + assert transport._ca_file_hash == expected_hash + assert transport._ca_file_mtime == os.path.getmtime(temp_file_path) + finally: + os.unlink(temp_file_path) + + def test_ca_file_change_detection_no_change(self): + """Test CA file change detection when file hasn't changed.""" + content = "-----BEGIN CERTIFICATE-----\nMIIDummy certificate content\n-----END CERTIFICATE-----" + + with tempfile.NamedTemporaryFile(mode="w", delete=False) as temp_file: + temp_file.write(content) + temp_file_path = temp_file.name + + try: + transport = HttpClientTransport(ca_file=temp_file_path) + # File hasn't changed + assert not transport._has_ca_file_changed() + finally: + os.unlink(temp_file_path) + + def test_ca_file_change_detection_content_changed(self): + """Test CA file change detection when file content has changed.""" + original_content = "-----BEGIN CERTIFICATE-----\nMIIDOriginal certificate\n-----END CERTIFICATE-----" + + with tempfile.NamedTemporaryFile(mode="w", delete=False) as temp_file: + temp_file.write(original_content) + temp_file_path = temp_file.name + + try: + transport = HttpClientTransport(ca_file=temp_file_path) + + # Modify the file + time.sleep(0.1) # Ensure mtime changes + modified_content = "-----BEGIN CERTIFICATE-----\nMIIDModified certificate\n-----END CERTIFICATE-----" + with open(temp_file_path, "w") as f: + f.write(modified_content) + + # File should be detected as changed + assert transport._has_ca_file_changed() + finally: + os.unlink(temp_file_path) + + def test_ca_file_change_detection_file_deleted(self): + """Test CA file change detection when file is deleted.""" + with tempfile.NamedTemporaryFile(mode="w", delete=False) as temp_file: + temp_file.write("-----BEGIN CERTIFICATE-----\nSome cert\n-----END CERTIFICATE-----") + temp_file_path = temp_file.name + + transport = HttpClientTransport(ca_file=temp_file_path) + + # Delete the file + os.unlink(temp_file_path) + + # File deletion should be detected as a change + assert transport._has_ca_file_changed() + + def test_ca_file_empty_during_rotation(self): + """Test CA file becoming empty during rotation with existing connection.""" + with tempfile.NamedTemporaryFile(mode="w", delete=False) as temp_file: + temp_file.write("-----BEGIN CERTIFICATE-----\nSome cert\n-----END CERTIFICATE-----") + temp_file_path = temp_file.name + + try: + transport = HttpClientTransport(ca_file=temp_file_path) + original_hash = transport._ca_file_hash + + # Simulate having an existing connection + transport._connection = mock.Mock() + + # Make file empty + with open(temp_file_path, "w") as f: + f.write("") + + # Should not raise error and should preserve old hash + transport._update_ca_file_tracking() + assert transport._ca_file_hash == original_hash + finally: + os.unlink(temp_file_path) + + def test_update_request_url_no_proxy(self): + """Test request URL update with no proxy.""" + transport = HttpClientTransport() + request = HttpRequest("GET", "https://example.com/path?query=value") + original_url = request.url + + transport._update_request_url(request) + + # URL should remain unchanged + assert request.url == original_url + + def test_update_request_url_with_proxy(self): + """Test request URL update with proxy.""" + proxy_endpoint = "https://proxy.example.com:8080/proxy" + transport = HttpClientTransport(proxy_endpoint=proxy_endpoint) + request = HttpRequest("GET", "https://original.com/api/endpoint?query=value") + + transport._update_request_url(request) + + # URL should be updated to use proxy + expected_url = "https://proxy.example.com:8080/proxy/api/endpoint?query=value" + assert request.url == expected_url + + def test_update_request_url_proxy_path_combination(self): + """Test request URL update with proxy that has a path.""" + proxy_endpoint = "https://proxy.example.com/service" + transport = HttpClientTransport(proxy_endpoint=proxy_endpoint) + request = HttpRequest("GET", "https://original.com/oauth2/v2.0/token") + + transport._update_request_url(request) + + # Paths should be combined correctly + expected_url = "https://proxy.example.com/service/oauth2/v2.0/token" + assert request.url == expected_url + + def test_context_manager(self): + """Test HttpClientTransport as context manager.""" + transport = HttpClientTransport() + + with transport as t: + assert t is transport + + # Should be able to use after context manager + assert transport is not None + + def test_close_with_connection(self): + """Test closing transport with active connection.""" + transport = HttpClientTransport() + mock_connection = mock.Mock() + transport._connection = mock_connection + + transport.close() + + mock_connection.close.assert_called_once() + assert transport._connection is None + + def test_close_without_connection(self): + """Test closing transport without active connection.""" + transport = HttpClientTransport() + + # Should not raise any exception + transport.close() + assert transport._connection is None + + def test_repr(self): + """Test string representation of HttpClientTransport.""" + transport = HttpClientTransport() + repr_str = repr(transport) + assert "HttpClientTransport" in repr_str + + +class TestSniSSLContext: + """Test cases for SniSSLContext class.""" + + def test_init(self): + """Test SNI SSL context initialization.""" + hostname = "example.com" + context = SniSSLContext(hostname, ssl.PROTOCOL_TLS_CLIENT) + assert context.sni_hostname == hostname + + def test_wrap_socket_adds_server_hostname(self): + """Test that wrap_socket adds server_hostname parameter.""" + hostname = "example.com" + context = SniSSLContext(hostname, ssl.PROTOCOL_TLS_CLIENT) + + # Mock the parent wrap_socket method + with mock.patch.object(ssl.SSLContext, "wrap_socket") as mock_wrap: + mock_sock = mock.Mock() + context.wrap_socket(mock_sock) + + # Verify server_hostname was added to kwargs + mock_wrap.assert_called_once_with(mock_sock, server_hostname=hostname) + + +class TestHttpClientTransportIntegration: + """Integration tests for HttpClientTransport with mocked HTTP connections.""" + + def test_send_request_success(self): + """Test successful request sending.""" + transport = HttpClientTransport() + request = HttpRequest("GET", "https://example.com/test") + + # Mock the HTTP connection and response + with mock.patch("http.client.HTTPSConnection") as mock_conn_class: + mock_conn = mock.Mock() + mock_conn_class.return_value = mock_conn + + mock_response = mock.Mock() + mock_response.status = 200 + mock_response.reason = "OK" + mock_response.getheaders.return_value = [("content-type", "application/json")] + mock_response.read.return_value = b'{"success": true}' + mock_conn.getresponse.return_value = mock_response + + response = transport.send(request) + + # Verify connection was created and request was sent + mock_conn_class.assert_called_once() + mock_conn.request.assert_called_once_with("GET", "/test", body=None, headers=request.headers) + + # Verify response properties + assert response.status_code == 200 + assert response.reason == "OK" + assert response.headers["content-type"] == "application/json" + + def test_send_request_http_exception(self): + """Test request sending with HTTP exception.""" + transport = HttpClientTransport() + request = HttpRequest("GET", "https://example.com/test") + + with mock.patch("http.client.HTTPSConnection") as mock_conn_class: + mock_conn = mock.Mock() + mock_conn_class.return_value = mock_conn + mock_conn.request.side_effect = http.client.HTTPException("Connection failed") + + with pytest.raises(ServiceRequestError): + transport.send(request) + + def test_send_request_ssl_error(self): + """Test request sending with SSL error.""" + transport = HttpClientTransport() + request = HttpRequest("GET", "https://example.com/test") + + with mock.patch("http.client.HTTPSConnection") as mock_conn_class: + mock_conn = mock.Mock() + mock_conn_class.return_value = mock_conn + mock_conn.request.side_effect = ssl.SSLError("SSL handshake failed") + + with pytest.raises(ServiceResponseError): + transport.send(request) + + def test_send_request_with_proxy(self): + """Test request sending through proxy.""" + proxy_endpoint = "https://proxy.example.com" + transport = HttpClientTransport(proxy_endpoint=proxy_endpoint) + request = HttpRequest("GET", "https://original.com/api/test") + + with mock.patch("http.client.HTTPSConnection") as mock_conn_class: + mock_conn = mock.Mock() + mock_conn_class.return_value = mock_conn + + mock_response = mock.Mock() + mock_response.status = 200 + mock_response.getheaders.return_value = [] + mock_response.read.return_value = b"" + mock_conn.getresponse.return_value = mock_response + + transport.send(request) + + # Verify connection was made to proxy host + args, kwargs = mock_conn_class.call_args + assert args[0] == "proxy.example.com" + + # Verify request was sent to proxy path + mock_conn.request.assert_called_once() + call_args = mock_conn.request.call_args + method = call_args[0][0] # First positional arg + path = call_args[0][1] # Second positional arg + assert method == "GET" + assert path == "/api/test" # Original path + + def test_connection_reuse(self): + """Test that connections are reused for multiple requests.""" + transport = HttpClientTransport() + request1 = HttpRequest("GET", "https://example.com/test1") + request2 = HttpRequest("GET", "https://example.com/test2") + + with mock.patch("http.client.HTTPSConnection") as mock_conn_class: + mock_conn = mock.Mock() + mock_conn_class.return_value = mock_conn + + mock_response = mock.Mock() + mock_response.status = 200 + mock_response.getheaders.return_value = [] + mock_response.read.return_value = b"" + mock_conn.getresponse.return_value = mock_response + + # Send two requests + transport.send(request1) + transport.send(request2) + + # Connection should be created only once + assert mock_conn_class.call_count == 1 + # But request should be called twice + assert mock_conn.request.call_count == 2 + + def test_ca_file_change_invalidates_connection(self): + """Test that CA file changes invalidate existing connections.""" + # Create test content that looks like PEM format but is just test data + ca_content = "-----BEGIN CERTIFICATE-----\nMIIDummy certificate content here\n-----END CERTIFICATE-----" + + with tempfile.NamedTemporaryFile(mode="w", delete=False) as temp_file: + temp_file.write(ca_content) + temp_file_path = temp_file.name + + try: + # Mock the SSL context creation to avoid actual certificate validation + with mock.patch("ssl.create_default_context") as mock_ssl_context: + mock_ssl_context.return_value = mock.Mock() + + transport = HttpClientTransport(ca_file=temp_file_path) + request = HttpRequest("GET", "https://example.com/test") + + with mock.patch("http.client.HTTPSConnection") as mock_conn_class: + mock_conn1 = mock.Mock() + mock_conn2 = mock.Mock() + mock_conn_class.side_effect = [mock_conn1, mock_conn2] + + mock_response = mock.Mock() + mock_response.status = 200 + mock_response.getheaders.return_value = [] + mock_response.read.return_value = b"" + mock_conn1.getresponse.return_value = mock_response + mock_conn2.getresponse.return_value = mock_response + + # First request + transport.send(request) + assert mock_conn_class.call_count == 1 + + # Modify CA file + time.sleep(0.1) # Ensure mtime changes + modified_content = ca_content.replace("Dummy", "Foo") + with open(temp_file_path, "w") as f: + f.write(modified_content) + + # Second request should create new connection + transport.send(request) + assert mock_conn_class.call_count == 2 + + # First connection should be closed + mock_conn1.close.assert_called_once() + + finally: + os.unlink(temp_file_path) + + +class TestHttpClientTransportResponse: + """Test cases for HttpClientTransportResponse class.""" + + def test_response_properties(self): + """Test basic response properties.""" + request = HttpRequest("GET", "https://example.com/test") + + mock_http_response = mock.Mock() + mock_http_response.status = 200 + mock_http_response.reason = "OK" + mock_http_response.getheaders.return_value = [("Content-Type", "application/json"), ("Content-Length", "100")] + mock_http_response.read.return_value = b'{"test": "data"}' + + response = transport.HttpClientTransportResponse(request, mock_http_response) + + assert response.status_code == 200 + assert response.reason == "OK" + assert response.headers["content-type"] == "application/json" + assert response.headers["content-length"] == "100" + assert response.content == b'{"test": "data"}' + assert response.text() == '{"test": "data"}' + assert response.json() == {"test": "data"} + assert response.url == "https://example.com/test" + assert not response.is_closed + # Stream is consumed after calling .content, .text(), or .json() + assert response.is_stream_consumed + + def test_response_context_manager(self): + """Test response as context manager.""" + request = HttpRequest("GET", "https://example.com/test") + mock_http_response = mock.Mock() + mock_http_response.status = 200 + mock_http_response.getheaders.return_value = [] + + response = transport.HttpClientTransportResponse(request, mock_http_response) + + with response as r: + assert r is response + assert not r.is_closed + + assert response.is_closed + mock_http_response.close.assert_called_once() + + def test_response_raise_for_status_success(self): + """Test raise_for_status with successful response.""" + request = HttpRequest("GET", "https://example.com/test") + mock_http_response = mock.Mock() + mock_http_response.status = 200 + mock_http_response.getheaders.return_value = [] + + response = transport.HttpClientTransportResponse(request, mock_http_response) + + # Should not raise any exception + response.raise_for_status() + + def test_response_raise_for_status_error(self): + """Test raise_for_status with error response.""" + from azure.core.exceptions import HttpResponseError + + request = HttpRequest("GET", "https://example.com/test") + mock_http_response = mock.Mock() + mock_http_response.status = 404 + mock_http_response.getheaders.return_value = [] + + response = transport.HttpClientTransportResponse(request, mock_http_response) + + with pytest.raises(HttpResponseError): + response.raise_for_status() + + def test_response_iter_raw(self): + """Test response iter_raw method.""" + request = HttpRequest("GET", "https://example.com/test") + mock_http_response = mock.Mock() + mock_http_response.status = 200 + mock_http_response.getheaders.return_value = [] + mock_http_response.read.side_effect = [b"chunk1", b"chunk2", b""] + + response = transport.HttpClientTransportResponse(request, mock_http_response, block_size=6) + + chunks = list(response.iter_raw()) + assert chunks == [b"chunk1", b"chunk2"] + + def test_response_iter_bytes(self): + """Test response iter_bytes method.""" + request = HttpRequest("GET", "https://example.com/test") + mock_http_response = mock.Mock() + mock_http_response.status = 200 + mock_http_response.getheaders.return_value = [] + mock_http_response.read.side_effect = [b"chunk1", b"chunk2", b""] + + response = transport.HttpClientTransportResponse(request, mock_http_response, block_size=6) + + chunks = list(response.iter_bytes()) + assert chunks == [b"chunk1", b"chunk2"] + + def test_response_encoding(self): + """Test response encoding property.""" + request = HttpRequest("GET", "https://example.com/test") + mock_http_response = mock.Mock() + mock_http_response.status = 200 + mock_http_response.getheaders.return_value = [] + mock_http_response.read.return_value = b"\xc3\xa9" # é in UTF-8 + + response = transport.HttpClientTransportResponse(request, mock_http_response) + + # Default encoding (UTF-8) + assert response.text() == "é" + + # Create a new response for testing latin-1 encoding + mock_http_response2 = mock.Mock() + mock_http_response2.status = 200 + mock_http_response2.getheaders.return_value = [] + mock_http_response2.read.return_value = b"\xe9" # é in Latin-1 + + response2 = transport.HttpClientTransportResponse(request, mock_http_response2) + response2.encoding = "latin-1" + assert response2.text() == "é" + + # Test explicit encoding parameter overrides response encoding + response2.encoding = "utf-8" + assert response2.text(encoding="latin-1") == "é" + + +# Import the actual class for the response tests +from azure.identity._internal import http_client_transport as transport diff --git a/sdk/identity/azure-identity/tests/test_workload_identity_credential.py b/sdk/identity/azure-identity/tests/test_workload_identity_credential.py index 1db0874a77ce..8a24562f0d09 100644 --- a/sdk/identity/azure-identity/tests/test_workload_identity_credential.py +++ b/sdk/identity/azure-identity/tests/test_workload_identity_credential.py @@ -2,10 +2,13 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT License. # ------------------------------------ +import os +import tempfile from unittest.mock import mock_open, MagicMock, patch import pytest from azure.identity import WorkloadIdentityCredential +from azure.identity._internal.http_client_transport import HttpClientTransport from helpers import mock_response, build_aad_response, GET_TOKEN_METHODS @@ -44,3 +47,207 @@ def send(request, **kwargs): assert token.token == access_token open_mock.assert_called_once_with(token_file_path, encoding="utf-8") + + +class TestWorkloadIdentityCredentialTokenProxy: + """Test cases for WorkloadIdentityCredential with use_token_proxy=True.""" + + def test_use_token_proxy_creates_http_client_transport(self): + """Test that use_token_proxy=True creates HttpClientTransport with correct parameters.""" + tenant_id = "tenant-id" + client_id = "client-id" + token_file_path = "foo-path" + proxy_endpoint = "https://proxy.example.com:8080" + sni_hostname = "sni.example.com" + ca_file_path = "/path/to/ca.pem" + + env_vars = { + "AZURE_KUBERNETES_TOKEN_PROXY": proxy_endpoint, + "AZURE_KUBERNETES_SNI_NAME": sni_hostname, + "AZURE_KUBERNETES_CA_FILE": ca_file_path, + } + + with patch.dict(os.environ, env_vars, clear=False): + with patch("azure.identity._credentials.workload_identity.HttpClientTransport") as mock_transport_class: + mock_transport_instance = MagicMock() + mock_transport_class.return_value = mock_transport_instance + + credential = WorkloadIdentityCredential( + tenant_id=tenant_id, + client_id=client_id, + token_file_path=token_file_path, + use_token_proxy=True, + ) + + # Verify HttpClientTransport was called with correct parameters + mock_transport_class.assert_called_once_with( + sni=sni_hostname, + proxy_endpoint=proxy_endpoint, + ca_file=ca_file_path, + ca_data=None, + ) + + def test_use_token_proxy_with_ca_data(self): + """Test use_token_proxy with CA data instead of CA file.""" + tenant_id = "tenant-id" + client_id = "client-id" + token_file_path = "foo-path" + proxy_endpoint = "https://proxy.example.com:8080" + ca_data = "-----BEGIN CERTIFICATE-----\nTest CA data\n-----END CERTIFICATE-----" + + env_vars = { + "AZURE_KUBERNETES_TOKEN_PROXY": proxy_endpoint, + "AZURE_KUBERNETES_CA_DATA": ca_data, + } + + with patch.dict(os.environ, env_vars, clear=False): + with patch("azure.identity._credentials.workload_identity.HttpClientTransport") as mock_transport_class: + mock_transport_instance = MagicMock() + mock_transport_class.return_value = mock_transport_instance + + credential = WorkloadIdentityCredential( + tenant_id=tenant_id, + client_id=client_id, + token_file_path=token_file_path, + use_token_proxy=True, + ) + + # Verify HttpClientTransport was called with CA data + mock_transport_class.assert_called_once_with( + sni=None, + proxy_endpoint=proxy_endpoint, + ca_file=None, + ca_data=ca_data, + ) + + def test_use_token_proxy_minimal_config(self): + """Test use_token_proxy with minimal configuration (only proxy endpoint).""" + tenant_id = "tenant-id" + client_id = "client-id" + token_file_path = "foo-path" + proxy_endpoint = "https://proxy.example.com:8080" + + env_vars = { + "AZURE_KUBERNETES_TOKEN_PROXY": proxy_endpoint, + } + + with patch.dict(os.environ, env_vars, clear=False): + with patch("azure.identity._credentials.workload_identity.HttpClientTransport") as mock_transport_class: + mock_transport_instance = MagicMock() + mock_transport_class.return_value = mock_transport_instance + + credential = WorkloadIdentityCredential( + tenant_id=tenant_id, + client_id=client_id, + token_file_path=token_file_path, + use_token_proxy=True, + ) + + # Verify HttpClientTransport was called with minimal config + mock_transport_class.assert_called_once_with( + sni=None, + proxy_endpoint=proxy_endpoint, + ca_file=None, + ca_data=None, + ) + + def test_use_token_proxy_missing_proxy_endpoint_raises_error(self): + """Test that use_token_proxy=True without proxy endpoint raises ValueError.""" + tenant_id = "tenant-id" + client_id = "client-id" + token_file_path = "foo-path" + + # Ensure proxy endpoint env var is not set + with patch.dict(os.environ, {}, clear=False): + if "AZURE_KUBERNETES_TOKEN_PROXY" in os.environ: + del os.environ["AZURE_KUBERNETES_TOKEN_PROXY"] + + with pytest.raises(ValueError, match="use_token_proxy is True, but no token proxy endpoint was found"): + WorkloadIdentityCredential( + tenant_id=tenant_id, + client_id=client_id, + token_file_path=token_file_path, + use_token_proxy=True, + ) + + def test_use_token_proxy_both_ca_file_and_data_raises_error(self): + """Test that setting both CA file and CA data raises ValueError.""" + tenant_id = "tenant-id" + client_id = "client-id" + token_file_path = "foo-path" + proxy_endpoint = "https://proxy.example.com:8080" + ca_file_path = "/path/to/ca.pem" + ca_data = "-----BEGIN CERTIFICATE-----\nTest CA data\n-----END CERTIFICATE-----" + + env_vars = { + "AZURE_KUBERNETES_TOKEN_PROXY": proxy_endpoint, + "AZURE_KUBERNETES_CA_FILE": ca_file_path, + "AZURE_KUBERNETES_CA_DATA": ca_data, + } + + with patch.dict(os.environ, env_vars, clear=False): + with pytest.raises(ValueError, match="Both AZURE_KUBERNETES_CA_FILE and AZURE_KUBERNETES_CA_DATA are set"): + WorkloadIdentityCredential( + tenant_id=tenant_id, + client_id=client_id, + token_file_path=token_file_path, + use_token_proxy=True, + ) + + @pytest.mark.parametrize("get_token_method", GET_TOKEN_METHODS) + def test_use_token_proxy_get_token_success(self, get_token_method): + """Test successful token acquisition when using token proxy.""" + tenant_id = "tenant-id" + client_id = "client-id" + access_token = "foo-access-token" + token_file_path = "foo-path" + assertion = "foo-assertion" + proxy_endpoint = "https://proxy.example.com:8080" + + def send(request, **kwargs): + assert "claims" not in kwargs + assert "tenant_id" not in kwargs + assert request.data.get("client_assertion") == assertion + return mock_response(json_payload=build_aad_response(access_token=access_token)) + + # Mock the transport that would be created by HttpClientTransport + mock_transport_instance = MagicMock(send=send) + + env_vars = { + "AZURE_KUBERNETES_TOKEN_PROXY": proxy_endpoint, + } + + with patch.dict(os.environ, env_vars, clear=False): + with patch("azure.identity._credentials.workload_identity.HttpClientTransport") as mock_transport_class: + mock_transport_class.return_value = mock_transport_instance + + credential = WorkloadIdentityCredential( + tenant_id=tenant_id, + client_id=client_id, + token_file_path=token_file_path, + use_token_proxy=True, + ) + + open_mock = mock_open(read_data=assertion) + with patch("builtins.open", open_mock): + token = getattr(credential, get_token_method)("scope") + assert token.token == access_token + + open_mock.assert_called_once_with(token_file_path, encoding="utf-8") + + def test_use_token_proxy_false_does_not_create_transport(self): + """Test that use_token_proxy=False (default) does not create HttpClientTransport.""" + tenant_id = "tenant-id" + client_id = "client-id" + token_file_path = "foo-path" + + with patch("azure.identity._credentials.workload_identity.HttpClientTransport") as mock_transport_class: + credential = WorkloadIdentityCredential( + tenant_id=tenant_id, + client_id=client_id, + token_file_path=token_file_path, + use_token_proxy=False, + ) + + # Verify HttpClientTransport was NOT called + mock_transport_class.assert_not_called() diff --git a/sdk/identity/azure-identity/tests/test_workload_identity_credential_async.py b/sdk/identity/azure-identity/tests/test_workload_identity_credential_async.py index fc6f3e8c6cb5..3a40c2b5f18c 100644 --- a/sdk/identity/azure-identity/tests/test_workload_identity_credential_async.py +++ b/sdk/identity/azure-identity/tests/test_workload_identity_credential_async.py @@ -2,10 +2,13 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT License. # ------------------------------------ +import os +import tempfile from unittest.mock import mock_open, patch, MagicMock import pytest from azure.identity.aio import WorkloadIdentityCredential +from azure.identity._internal.http_client_transport import HttpClientTransport from helpers import mock_response, build_aad_response, GET_TOKEN_METHODS @@ -45,3 +48,208 @@ async def send(request, **kwargs): assert token.token == access_token open_mock.assert_called_once_with(token_file_path, encoding="utf-8") + + +class TestWorkloadIdentityCredentialTokenProxyAsync: + """Async test cases for WorkloadIdentityCredential with use_token_proxy=True.""" + + def test_use_token_proxy_creates_http_client_transport(self): + """Test that use_token_proxy=True creates HttpClientTransport with correct parameters.""" + tenant_id = "tenant-id" + client_id = "client-id" + token_file_path = "foo-path" + proxy_endpoint = "https://proxy.example.com:8080" + sni_hostname = "sni.example.com" + ca_file_path = "/path/to/ca.pem" + + env_vars = { + "AZURE_KUBERNETES_TOKEN_PROXY": proxy_endpoint, + "AZURE_KUBERNETES_SNI_NAME": sni_hostname, + "AZURE_KUBERNETES_CA_FILE": ca_file_path, + } + + with patch.dict(os.environ, env_vars, clear=False): + with patch("azure.identity.aio._credentials.workload_identity.HttpClientTransport") as mock_transport_class: + mock_transport_instance = MagicMock() + mock_transport_class.return_value = mock_transport_instance + + credential = WorkloadIdentityCredential( + tenant_id=tenant_id, + client_id=client_id, + token_file_path=token_file_path, + use_token_proxy=True, + ) + + # Verify HttpClientTransport was called with correct parameters + mock_transport_class.assert_called_once_with( + sni=sni_hostname, + proxy_endpoint=proxy_endpoint, + ca_file=ca_file_path, + ca_data=None, + ) + + def test_use_token_proxy_with_ca_data(self): + """Test use_token_proxy with CA data instead of CA file.""" + tenant_id = "tenant-id" + client_id = "client-id" + token_file_path = "foo-path" + proxy_endpoint = "https://proxy.example.com:8080" + ca_data = "-----BEGIN CERTIFICATE-----\nTest CA data\n-----END CERTIFICATE-----" + + env_vars = { + "AZURE_KUBERNETES_TOKEN_PROXY": proxy_endpoint, + "AZURE_KUBERNETES_CA_DATA": ca_data, + } + + with patch.dict(os.environ, env_vars, clear=False): + with patch("azure.identity.aio._credentials.workload_identity.HttpClientTransport") as mock_transport_class: + mock_transport_instance = MagicMock() + mock_transport_class.return_value = mock_transport_instance + + credential = WorkloadIdentityCredential( + tenant_id=tenant_id, + client_id=client_id, + token_file_path=token_file_path, + use_token_proxy=True, + ) + + # Verify HttpClientTransport was called with CA data + mock_transport_class.assert_called_once_with( + sni=None, + proxy_endpoint=proxy_endpoint, + ca_file=None, + ca_data=ca_data, + ) + + def test_use_token_proxy_minimal_config(self): + """Test use_token_proxy with minimal configuration (only proxy endpoint).""" + tenant_id = "tenant-id" + client_id = "client-id" + token_file_path = "foo-path" + proxy_endpoint = "https://proxy.example.com:8080" + + env_vars = { + "AZURE_KUBERNETES_TOKEN_PROXY": proxy_endpoint, + } + + with patch.dict(os.environ, env_vars, clear=False): + with patch("azure.identity.aio._credentials.workload_identity.HttpClientTransport") as mock_transport_class: + mock_transport_instance = MagicMock() + mock_transport_class.return_value = mock_transport_instance + + credential = WorkloadIdentityCredential( + tenant_id=tenant_id, + client_id=client_id, + token_file_path=token_file_path, + use_token_proxy=True, + ) + + # Verify HttpClientTransport was called with minimal config + mock_transport_class.assert_called_once_with( + sni=None, + proxy_endpoint=proxy_endpoint, + ca_file=None, + ca_data=None, + ) + + def test_use_token_proxy_missing_proxy_endpoint_raises_error(self): + """Test that use_token_proxy=True without proxy endpoint raises ValueError.""" + tenant_id = "tenant-id" + client_id = "client-id" + token_file_path = "foo-path" + + # Ensure proxy endpoint env var is not set + with patch.dict(os.environ, {}, clear=False): + if "AZURE_KUBERNETES_TOKEN_PROXY" in os.environ: + del os.environ["AZURE_KUBERNETES_TOKEN_PROXY"] + + with pytest.raises(ValueError, match="use_token_proxy is True, but no token proxy endpoint was found"): + WorkloadIdentityCredential( + tenant_id=tenant_id, + client_id=client_id, + token_file_path=token_file_path, + use_token_proxy=True, + ) + + def test_use_token_proxy_both_ca_file_and_data_raises_error(self): + """Test that setting both CA file and CA data raises ValueError.""" + tenant_id = "tenant-id" + client_id = "client-id" + token_file_path = "foo-path" + proxy_endpoint = "https://proxy.example.com:8080" + ca_file_path = "/path/to/ca.pem" + ca_data = "-----BEGIN CERTIFICATE-----\nTest CA data\n-----END CERTIFICATE-----" + + env_vars = { + "AZURE_KUBERNETES_TOKEN_PROXY": proxy_endpoint, + "AZURE_KUBERNETES_CA_FILE": ca_file_path, + "AZURE_KUBERNETES_CA_DATA": ca_data, + } + + with patch.dict(os.environ, env_vars, clear=False): + with pytest.raises(ValueError, match="Both AZURE_KUBERNETES_CA_FILE and AZURE_KUBERNETES_CA_DATA are set"): + WorkloadIdentityCredential( + tenant_id=tenant_id, + client_id=client_id, + token_file_path=token_file_path, + use_token_proxy=True, + ) + + @pytest.mark.asyncio + @pytest.mark.parametrize("get_token_method", GET_TOKEN_METHODS) + async def test_use_token_proxy_get_token_success(self, get_token_method): + """Test successful token acquisition when using token proxy.""" + tenant_id = "tenant-id" + client_id = "client-id" + access_token = "foo-access-token" + token_file_path = "foo-path" + assertion = "foo-assertion" + proxy_endpoint = "https://proxy.example.com:8080" + + async def send(request, **kwargs): + assert "claims" not in kwargs + assert "tenant_id" not in kwargs + assert request.data.get("client_assertion") == assertion + return mock_response(json_payload=build_aad_response(access_token=access_token)) + + # Mock the transport that would be created by HttpClientTransport + mock_transport_instance = MagicMock(send=send) + + env_vars = { + "AZURE_KUBERNETES_TOKEN_PROXY": proxy_endpoint, + } + + with patch.dict(os.environ, env_vars, clear=False): + with patch("azure.identity.aio._credentials.workload_identity.HttpClientTransport") as mock_transport_class: + mock_transport_class.return_value = mock_transport_instance + + credential = WorkloadIdentityCredential( + tenant_id=tenant_id, + client_id=client_id, + token_file_path=token_file_path, + use_token_proxy=True, + ) + + open_mock = mock_open(read_data=assertion) + with patch("builtins.open", open_mock): + token = await getattr(credential, get_token_method)("scope") + assert token.token == access_token + + open_mock.assert_called_once_with(token_file_path, encoding="utf-8") + + def test_use_token_proxy_false_does_not_create_transport(self): + """Test that use_token_proxy=False (default) does not create HttpClientTransport.""" + tenant_id = "tenant-id" + client_id = "client-id" + token_file_path = "foo-path" + + with patch("azure.identity.aio._credentials.workload_identity.HttpClientTransport") as mock_transport_class: + credential = WorkloadIdentityCredential( + tenant_id=tenant_id, + client_id=client_id, + token_file_path=token_file_path, + use_token_proxy=False, + ) + + # Verify HttpClientTransport was NOT called + mock_transport_class.assert_not_called() From 747f53dbb5ffcc1916b230e40c5ffce61fac0652 Mon Sep 17 00:00:00 2001 From: Paul Van Eck Date: Tue, 28 Oct 2025 09:53:40 +0000 Subject: [PATCH 02/12] Refactor http.client transport and add aio transport Signed-off-by: Paul Van Eck --- .../_internal/http_client_transport.py | 176 ++--- .../token_binding_transport_mixin.py | 123 ++++ .../aio/_credentials/workload_identity.py | 73 +- .../azure-identity/tests/proxy_server.py | 39 +- .../tests/test_http_client_transport.py | 666 +++++++++++------- .../test_workload_identity_credential.py | 6 +- ...test_workload_identity_credential_async.py | 660 ++++++++++++++++- 7 files changed, 1300 insertions(+), 443 deletions(-) create mode 100644 sdk/identity/azure-identity/azure/identity/_internal/token_binding_transport_mixin.py diff --git a/sdk/identity/azure-identity/azure/identity/_internal/http_client_transport.py b/sdk/identity/azure-identity/azure/identity/_internal/http_client_transport.py index 4677deb28bae..8dd729899c70 100644 --- a/sdk/identity/azure-identity/azure/identity/_internal/http_client_transport.py +++ b/sdk/identity/azure-identity/azure/identity/_internal/http_client_transport.py @@ -3,19 +3,19 @@ # Licensed under the MIT License. # ------------------------------------ # cspell:ignore cafile -from json import loads -import hashlib +from json import loads, dumps import http.client -import os import ssl import urllib.parse from typing import Any, Iterator, MutableMapping, Optional from azure.core.configuration import ConnectionConfiguration -from azure.core.exceptions import ServiceRequestError, ServiceResponseError +from azure.core.exceptions import ServiceRequestError from azure.core.rest import HttpRequest, HttpResponse from azure.core.pipeline.transport import HttpTransport +from .token_binding_transport_mixin import TokenBindingTransportMixin + class HttpClientTransportResponse(HttpResponse): """Create a HttpResponse from an http.client response. @@ -132,6 +132,8 @@ def json(self) -> Any: class SniSSLContext(ssl.SSLContext): def __new__(cls, sni_hostname: str, protocol=None): + if protocol is None: + protocol = ssl.PROTOCOL_TLS_CLIENT instance = super().__new__(cls, protocol=protocol) instance.sni_hostname = sni_hostname # type: ignore return instance @@ -141,29 +143,16 @@ def wrap_socket(self, *args, **kwargs): return super().wrap_socket(*args, **kwargs) -class HttpClientTransport(HttpTransport): +class HttpClientTransport(TokenBindingTransportMixin, HttpTransport): """Implements an HTTP sender using Python's built-in http.client library.""" def __init__(self, **kwargs) -> None: self.connection_config = ConnectionConfiguration(**kwargs) - self._ca_data = kwargs.pop("ca_data", None) - self._ca_file = kwargs.pop("ca_file", None) - - if self._ca_file and self._ca_data: - raise ValueError("Both ca_file and ca_data are set. Only one should be set") - - self._sni = kwargs.pop("sni", None) - self._proxy_endpoint = kwargs.pop("proxy_endpoint", None) - if self._proxy_endpoint: - self._validate_url(self._proxy_endpoint) - - self._connection: Optional[http.client.HTTPSConnection] = None - self._ca_file_hash: Optional[str] = None - self._ca_file_mtime: Optional[float] = None - - # Initialize CA file tracking if a CA file is specified - if self._ca_file: - self._update_ca_file_tracking() + self._ssl_context: Optional[ssl.SSLContext] = None + super().__init__(**kwargs) + # Initialize SSL context if we have CA data + if self._ca_data: + self._ssl_context = self._create_ssl_context() def __enter__(self) -> "HttpClientTransport": self.open() @@ -176,136 +165,48 @@ def open(self) -> None: pass # We create connections as needed def close(self) -> None: - if self._connection: - self._connection.close() - self._connection = None - - def _validate_url(self, url: str) -> None: - parsed_url = urllib.parse.urlparse(url) - if parsed_url.scheme != "https": - raise ValueError(f"Endpoint URL ({url}) must use the 'https' scheme. Got '{parsed_url.scheme}' instead.") - if parsed_url.username or parsed_url.password: - raise ValueError(f"Endpoint URL ({url}) must not contain username or password.") - if parsed_url.fragment: - raise ValueError(f"Endpoint URL ({url}) must not contain a fragment.") - if parsed_url.query: - raise ValueError(f"Endpoint URL ({url}) must not contain query parameters.") - - def _update_ca_file_tracking(self) -> None: - """Update the CA file hash and modification time for change detection.""" - if not self._ca_file or not os.path.exists(self._ca_file): - self._ca_file_hash = None - self._ca_file_mtime = None - return - - try: - # Read file content first to check if empty - with open(self._ca_file, "rb") as f: - content = f.read() - - # Check if the file is empty - if not content: - # If no prior tracking state exists (first read), fail - if self._ca_file_hash is None: - raise ValueError(f"CA file ({self._ca_file}) is empty. Cannot establish secure connection.") - return - - # File has content, update tracking - self._ca_file_mtime = os.path.getmtime(self._ca_file) - self._ca_file_hash = hashlib.sha256(content).hexdigest() - except (OSError, IOError): - # If we can't read the file, clear the tracking - self._ca_file_hash = None - self._ca_file_mtime = None - - def _has_ca_file_changed(self) -> bool: - """Check if the CA file has changed since last tracking update. - - :return: True if the CA file has changed, False otherwise. - :rtype: bool - """ - if not self._ca_file: - return False - - if not os.path.exists(self._ca_file): - # File was deleted, consider this a change - return self._ca_file_hash is not None or self._ca_file_mtime is not None - - try: - # Check modification time first (faster) - current_mtime = os.path.getmtime(self._ca_file) - if self._ca_file_mtime != current_mtime: - return True + pass # No persistent connections to close - # If mtime is the same, check content hash to be sure - with open(self._ca_file, "rb") as f: - content = f.read() - current_hash = hashlib.sha256(content).hexdigest() - return self._ca_file_hash != current_hash - - except (OSError, IOError): - # If we can't read the file, assume it changed - return True - - def _get_connection(self, host: str) -> http.client.HTTPSConnection: - - # Check if CA file has changed and invalidate connections if needed - if self._ca_file and self._has_ca_file_changed(): - # CA file changed, close all existing connections and clear cache - if self._connection: - self._connection.close() - self._connection = None - # Update tracking with new CA file state - self._update_ca_file_tracking() - - # Use existing connection if available - if self._connection: - return self._connection - - # Create HTTPS connection + def _create_ssl_context(self) -> ssl.SSLContext: + # Create SSL context using current CA data or file. ssl_context: ssl.SSLContext if self._sni: ssl_context = SniSSLContext(self._sni, ssl.PROTOCOL_TLS_CLIENT) ssl_context.verify_mode = ssl.CERT_REQUIRED ssl_context.check_hostname = True - if self._ca_data or self._ca_file: - ssl_context.load_verify_locations(cafile=self._ca_file, cadata=self._ca_data) + if self._ca_data: + ssl_context.load_verify_locations(cadata=self._ca_data) else: ssl_context.load_default_certs() else: - ssl_context = ssl.create_default_context(cafile=self._ca_file, cadata=self._ca_data) + ssl_context = ssl.create_default_context(cadata=self._ca_data) if not self.connection_config.verify: ssl_context.check_hostname = False ssl_context.verify_mode = ssl.CERT_NONE + return ssl_context + + def _get_connection(self, host: str) -> http.client.HTTPSConnection: + # Check if CA file has changed and reload ca_data if needed + if self._ca_file and self._has_ca_file_changed(): + self._load_ca_file_to_data() + # If ca_data was updated, recreate SSL context with the new data + if self._ca_data: + self._ssl_context = self._create_ssl_context() + + # Use cached SSL context or create a new one + ssl_context = self._ssl_context or self._create_ssl_context() + connection = http.client.HTTPSConnection( host, timeout=self.connection_config.timeout, context=ssl_context, ) - self._connection = connection return connection - def _update_request_url(self, request: HttpRequest) -> None: - parsed_request_url = urllib.parse.urlparse(request.url) - if self._proxy_endpoint: - parsed_proxy_url = urllib.parse.urlparse(self._proxy_endpoint) - combined_path = parsed_proxy_url.path.rstrip("/") + "/" + parsed_request_url.path.lstrip("/") - new_url = urllib.parse.urlunparse( - ( - parsed_proxy_url.scheme, - parsed_proxy_url.netloc, - combined_path, - parsed_request_url.params, - parsed_request_url.query, - parsed_request_url.fragment, - ) - ) - request.url = new_url - def send(self, request: HttpRequest, **kwargs) -> HttpResponse: """Send request object according to configuration. @@ -331,19 +232,26 @@ def send(self, request: HttpRequest, **kwargs) -> HttpResponse: if connection_timeout is not None: connection.timeout = connection_timeout - - connection.request(request.method, full_path, body=request.data, headers=request.headers) + connection.request( + request.method, + full_path, + body=dumps(request.data) if isinstance(request.data, (dict, list)) else request.data, + headers=request.headers, + ) response = connection.getresponse() - return HttpClientTransportResponse( + transport_response = HttpClientTransportResponse( request=request, httpclient_response=response, block_size=self.connection_config.data_block_size, ) + connection.close() + return transport_response + except http.client.HTTPException as err: raise ServiceRequestError(err) from err except ssl.SSLError as err: - raise ServiceResponseError(err) from err + raise ServiceRequestError(err) from err except Exception as err: raise ServiceRequestError(err) from err diff --git a/sdk/identity/azure-identity/azure/identity/_internal/token_binding_transport_mixin.py b/sdk/identity/azure-identity/azure/identity/_internal/token_binding_transport_mixin.py new file mode 100644 index 000000000000..6b600612068d --- /dev/null +++ b/sdk/identity/azure-identity/azure/identity/_internal/token_binding_transport_mixin.py @@ -0,0 +1,123 @@ +# ------------------------------------ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. +# ------------------------------------ +# cspell:ignore cafile +import os +import urllib.parse +from typing import Optional + +from azure.core.rest import HttpRequest + + +class TokenBindingTransportMixin: + """Mixin class providing URL validation, CA file tracking, and proxy URL functionality for transport classes.""" + + def __init__(self, **kwargs): + """Initialize CA file tracking and proxy attributes.""" + self._ca_file = kwargs.pop("ca_file", None) + self._ca_data = kwargs.pop("ca_data", None) + self._proxy_endpoint = kwargs.pop("proxy_endpoint", None) + self._sni = kwargs.pop("sni", None) + + self._ca_file_mtime: Optional[float] = None + + if self._ca_file and self._ca_data: + raise ValueError("Both ca_file and ca_data are set. Only one should be set") + + if self._proxy_endpoint: + self._validate_url(self._proxy_endpoint) + + # If we have a ca_file, read it once and store as ca_data + if self._ca_file: + self._load_ca_file_to_data() + + super().__init__() + + def _validate_url(self, url: str) -> None: + """Validate that a URL meets security requirements for HTTPS connections. + + :param url: The URL to validate. + :type url: str + :raises ValueError: If the URL does not meet security requirements. + """ + parsed_url = urllib.parse.urlparse(url) + if parsed_url.scheme != "https": + raise ValueError(f"Endpoint URL ({url}) must use the 'https' scheme. Got '{parsed_url.scheme}' instead.") + if parsed_url.username or parsed_url.password: + raise ValueError(f"Endpoint URL ({url}) must not contain username or password.") + if parsed_url.fragment: + raise ValueError(f"Endpoint URL ({url}) must not contain a fragment.") + if parsed_url.query: + raise ValueError(f"Endpoint URL ({url}) must not contain query parameters.") + + def _load_ca_file_to_data(self) -> None: + """Load CA file content into ca_data and track modification time. + + :raises ValueError: If the CA file is empty on first read. + """ + try: + with open(self._ca_file, "r", encoding="utf-8") as f: + content = f.read() + + # Check if the file is empty + if not content: + # If no prior ca_data exists (first read), fail + if self._ca_data is None: + raise ValueError(f"CA file ({self._ca_file}) is empty. Cannot establish secure connection.") + # If we had prior ca_data, keep it (mid-rotation scenario) + return + + # File has content, update ca_data and tracking + self._ca_data = content + self._ca_file_mtime = os.path.getmtime(self._ca_file) + except (OSError, IOError) as e: + # If no prior ca_data exists (first read), fail + if self._ca_data is None: + raise ValueError(f"Failed to read CA file ({self._ca_file}): {e}") from e + # If we can't read the file, keep existing ca_data but clear mtime + # so we'll try to reload on the next change check + self._ca_file_mtime = None + + def _has_ca_file_changed(self) -> bool: + """Check if the CA file has changed since last load. + + :return: True if the CA file has changed, False otherwise. + :rtype: bool + """ + if not self._ca_file: + return False + + if not os.path.exists(self._ca_file): + # File was deleted, consider this a change if we had data before + return self._ca_data is not None or self._ca_file_mtime is not None + + try: + # Check modification time + current_mtime = os.path.getmtime(self._ca_file) + return self._ca_file_mtime != current_mtime + except (OSError, IOError): + # If we can't read the file stats, assume it changed + return True + + def _update_request_url(self, request: HttpRequest) -> None: + """Update the request URL to use proxy endpoint if configured. + + :param request: The HTTP request object to update. + :type request: ~azure.core.rest.HttpRequest + """ + parsed_request_url = urllib.parse.urlparse(request.url) + if self._proxy_endpoint: + parsed_proxy_url = urllib.parse.urlparse(self._proxy_endpoint) + combined_path = parsed_proxy_url.path.rstrip("/") + "/" + parsed_request_url.path.lstrip("/") + new_url = urllib.parse.urlunparse( + ( + parsed_proxy_url.scheme, + parsed_proxy_url.netloc, + combined_path, + parsed_request_url.params, + parsed_request_url.query, + parsed_request_url.fragment, + ) + ) + request.url = new_url diff --git a/sdk/identity/azure-identity/azure/identity/aio/_credentials/workload_identity.py b/sdk/identity/azure-identity/azure/identity/aio/_credentials/workload_identity.py index e9a7cdd4e83c..70fa6aef4507 100644 --- a/sdk/identity/azure-identity/azure/identity/aio/_credentials/workload_identity.py +++ b/sdk/identity/azure-identity/azure/identity/aio/_credentials/workload_identity.py @@ -2,12 +2,19 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT License. # ------------------------------------ +# cspell:ignore cafile import os +import ssl +import logging from typing import Any, Optional + from .client_assertion import ClientAssertionCredential from ..._credentials.workload_identity import TokenFileMixin, WORKLOAD_CONFIG_ERROR from ..._constants import EnvironmentVariables -from ..._internal.http_client_transport import HttpClientTransport +from ..._internal.token_binding_transport_mixin import TokenBindingTransportMixin + + +_LOGGER = logging.getLogger(__name__) class WorkloadIdentityCredential(ClientAssertionCredential, TokenFileMixin): @@ -78,14 +85,14 @@ def __init__( self._token_file_path = token_file_path if use_token_proxy: - token_proxy_endpoint = os.environ.get(EnvironmentVariables.AZURE_KUBERNETES_TOKEN_PROXY) - if not token_proxy_endpoint: + self._token_proxy_endpoint = os.environ.get(EnvironmentVariables.AZURE_KUBERNETES_TOKEN_PROXY) + if not self._token_proxy_endpoint: raise ValueError( "use_token_proxy is True, but no token proxy endpoint was found. " f"Ensure the {EnvironmentVariables.AZURE_KUBERNETES_TOKEN_PROXY} environment variable is set." ) - sni = os.environ.get(EnvironmentVariables.AZURE_KUBERNETES_SNI_NAME) + self._sni = os.environ.get(EnvironmentVariables.AZURE_KUBERNETES_SNI_NAME) ca_file = os.environ.get(EnvironmentVariables.AZURE_KUBERNETES_CA_FILE) ca_data = os.environ.get(EnvironmentVariables.AZURE_KUBERNETES_CA_DATA) @@ -94,13 +101,20 @@ def __init__( "Both AZURE_KUBERNETES_CA_FILE and AZURE_KUBERNETES_CA_DATA are set. Only one should be set" ) - kwargs["transport"] = HttpClientTransport( - sni=sni, - proxy_endpoint=token_proxy_endpoint, + transport = _get_transport( + sni=self._sni, + token_proxy_endpoint=self._token_proxy_endpoint, ca_file=ca_file, ca_data=ca_data, ) + if transport: + kwargs["transport"] = transport + else: + raise ValueError( + "Async transport creation failed. Ensure that the aiohttp package is installed to enable token proxy usage." + ) + super().__init__( tenant_id=tenant_id, client_id=client_id, @@ -108,3 +122,48 @@ def __init__( token_file_path=token_file_path, **kwargs, ) + + +def _get_transport(sni, token_proxy_endpoint, ca_file, ca_data): + try: + from azure.core.pipeline.transport import ( # pylint: disable=non-abstract-transport-import, no-name-in-module + AioHttpTransport, + ) + + class WorkloadIdentityAioHttpTransport(TokenBindingTransportMixin, AioHttpTransport): + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self._ssl_context = ssl.create_default_context(cadata=self._ca_data) + + async def send(self, request, **kwargs): + self._update_request_url(request) + kwargs.setdefault("server_hostname", self._sni) + + # Check if CA file has changed and reload ca_data if needed + if self._ca_file and self._has_ca_file_changed(): + self._load_ca_file_to_data() + # If ca_data was updated, recreate SSL context with the new data + if self._ca_data: + self._ssl_context = ssl.create_default_context(cadata=self._ca_data) + + if self._ssl_context: + kwargs.setdefault("ssl", self._ssl_context) + return await super().send(request, **kwargs) + + async def __aenter__(self): + await super().__aenter__() + return self + + async def __aexit__(self, *args): + await super().__aexit__(*args) + + transport = WorkloadIdentityAioHttpTransport( + sni=sni, + proxy_endpoint=token_proxy_endpoint, + ca_file=ca_file, + ca_data=ca_data, + ) + except ImportError: + transport = None + return transport diff --git a/sdk/identity/azure-identity/tests/proxy_server.py b/sdk/identity/azure-identity/tests/proxy_server.py index 352ab8f1ba2a..918d0a674c8b 100644 --- a/sdk/identity/azure-identity/tests/proxy_server.py +++ b/sdk/identity/azure-identity/tests/proxy_server.py @@ -75,7 +75,7 @@ def _handle_request(self): # Simulate different responses based on path if path == "/health": self._send_health_response() - elif path.startswith("/oauth2/v2.0/token"): + elif path.endswith("/oauth2/v2.0/token"): self._send_token_response(body) elif path == "/error/500": self._send_error_response(500, "Internal Server Error") @@ -300,3 +300,40 @@ def base_url(self): """Get the base URL of the server.""" scheme = "https" if self.use_ssl else "http" return f"{scheme}://{self.host}:{self.port}" + + +def main(): + """Run the test server standalone.""" + parser = argparse.ArgumentParser(description="Token Proxy Test Server") + parser.add_argument("--host", default="localhost", help="Host to bind to") + parser.add_argument("--port", type=int, default=8443, help="Port to bind to") + parser.add_argument("--no-ssl", action="store_true", help="Disable SSL/TLS") + parser.add_argument("--verbose", "-v", action="store_true", help="Enable verbose logging") + + args = parser.parse_args() + + # Configure logging + log_level = logging.DEBUG if args.verbose else logging.INFO + logging.basicConfig(level=log_level, format="%(asctime)s - %(levelname)s - %(message)s") + + # Start server + with TokenProxyTestServer(host=args.host, port=args.port, use_ssl=not args.no_ssl) as server: + print(f"Server running at {server.base_url}") + print("Available endpoints:") + print(f" {server.base_url}/health - Health check") + print(f" {server.base_url}/oauth2/v2.0/token - Mock OAuth token endpoint") + print(f" {server.base_url}/error/500 - Simulate server error") + print(f" {server.base_url}/error/ssl - Simulate SSL error") + print(f" {server.base_url}/slow - Simulate slow response") + print(f" {server.base_url}/ - Generic proxy response") + print("\nPress Ctrl+C to stop") + + try: + while True: + time.sleep(1) + except KeyboardInterrupt: + print("\nShutting down...") + + +if __name__ == "__main__": + main() diff --git a/sdk/identity/azure-identity/tests/test_http_client_transport.py b/sdk/identity/azure-identity/tests/test_http_client_transport.py index 7282b7c7fe8d..95f40235defc 100644 --- a/sdk/identity/azure-identity/tests/test_http_client_transport.py +++ b/sdk/identity/azure-identity/tests/test_http_client_transport.py @@ -2,22 +2,32 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT License. # ------------------------------------ -import hashlib -import http.client -import json import os import ssl import tempfile import time +from time import sleep as real_sleep from unittest import mock -from urllib.parse import urlparse import pytest from azure.core.rest import HttpRequest from azure.core.exceptions import ServiceRequestError, ServiceResponseError +from azure.identity._internal import http_client_transport as transport from azure.identity._internal.http_client_transport import HttpClientTransport, SniSSLContext +from proxy_server import TokenProxyTestServer + + +PEM_CERT_PATH = os.path.join(os.path.dirname(__file__), "certificate.pem") + + +@pytest.fixture(scope="module") +def ca_data() -> str: + """Read CA certificate data from a PEM file for testing.""" + with open(PEM_CERT_PATH, "r", encoding="utf-8") as f: + return f.read() + class TestHttpClientTransport: """Test cases for HttpClientTransport class.""" @@ -29,48 +39,26 @@ def test_init_basic(self): assert transport._ca_file is None assert transport._sni is None assert transport._proxy_endpoint is None - assert transport._connection is None - assert transport._ca_file_hash is None assert transport._ca_file_mtime is None - def test_init_with_ca_data(self): + def test_init_with_ca_data(self, ca_data): """Test initialization with CA data.""" - ca_data = "-----BEGIN CERTIFICATE-----\nSome fake cert\n-----END CERTIFICATE-----" transport = HttpClientTransport(ca_data=ca_data) assert transport._ca_data == ca_data assert transport._ca_file is None - def test_init_with_ca_file(self): + def test_init_with_ca_file(self, ca_data): """Test initialization with CA file.""" - # Use simple but valid-looking PEM content - ca_content = "-----BEGIN CERTIFICATE-----\nMIIDummy certificate content here\n-----END CERTIFICATE-----" - with tempfile.NamedTemporaryFile(mode="w", delete=False) as temp_file: - temp_file.write(ca_content) - temp_file_path = temp_file.name - - try: - transport = HttpClientTransport(ca_file=temp_file_path) - assert transport._ca_file == temp_file_path - assert transport._ca_data is None - assert transport._ca_file_hash is not None - assert transport._ca_file_mtime is not None - finally: - os.unlink(temp_file_path) + transport = HttpClientTransport(ca_file=PEM_CERT_PATH) + assert transport._ca_file == PEM_CERT_PATH + assert transport._ca_data == ca_data + assert transport._ca_file_mtime is not None - def test_init_with_both_ca_file_and_data_raises_error(self): + def test_init_with_both_ca_file_and_data_raises_error(self, ca_data): """Test that providing both CA file and data raises an error.""" - ca_data = "-----BEGIN CERTIFICATE-----\nMIIDummy certificate content\n-----END CERTIFICATE-----" - - with tempfile.NamedTemporaryFile(mode="w", delete=False) as temp_file: - temp_file.write(ca_data) - temp_file_path = temp_file.name - - try: - with pytest.raises(ValueError, match="Both ca_file and ca_data are set"): - HttpClientTransport(ca_file=temp_file_path, ca_data=ca_data) - finally: - os.unlink(temp_file_path) + with pytest.raises(ValueError, match="Both ca_file and ca_data are set"): + HttpClientTransport(ca_file=PEM_CERT_PATH, ca_data=ca_data) def test_init_with_empty_ca_file_raises_error(self): """Test that empty CA file raises an error.""" @@ -125,71 +113,41 @@ def test_validate_url_with_query(self): with pytest.raises(ValueError, match="must not contain query parameters"): transport._validate_url("https://example.com?query=value") - def test_ca_file_tracking_nonexistent_file(self): - """Test CA file tracking with non-existent file.""" - transport = HttpClientTransport() - transport._ca_file = "/nonexistent/file.pem" - transport._update_ca_file_tracking() - assert transport._ca_file_hash is None - assert transport._ca_file_mtime is None - - def test_ca_file_tracking_updates_hash_and_mtime(self): - """Test CA file tracking updates hash and modification time.""" - content = "-----BEGIN CERTIFICATE-----\nMIIDummy certificate content\n-----END CERTIFICATE-----" - - with tempfile.NamedTemporaryFile(mode="w", delete=False) as temp_file: - temp_file.write(content) - temp_file_path = temp_file.name - - try: - transport = HttpClientTransport(ca_file=temp_file_path) - expected_hash = hashlib.sha256(content.encode()).hexdigest() - assert transport._ca_file_hash == expected_hash - assert transport._ca_file_mtime == os.path.getmtime(temp_file_path) - finally: - os.unlink(temp_file_path) + def test_ca_file_tracking_updates_mtime(self, ca_data): + """Test CA file tracking updates modification time.""" + transport = HttpClientTransport(ca_file=PEM_CERT_PATH) + assert transport._ca_file_mtime == os.path.getmtime(PEM_CERT_PATH) def test_ca_file_change_detection_no_change(self): """Test CA file change detection when file hasn't changed.""" - content = "-----BEGIN CERTIFICATE-----\nMIIDummy certificate content\n-----END CERTIFICATE-----" + transport = HttpClientTransport(ca_file=PEM_CERT_PATH) + assert not transport._has_ca_file_changed() - with tempfile.NamedTemporaryFile(mode="w", delete=False) as temp_file: - temp_file.write(content) - temp_file_path = temp_file.name - - try: - transport = HttpClientTransport(ca_file=temp_file_path) - # File hasn't changed - assert not transport._has_ca_file_changed() - finally: - os.unlink(temp_file_path) - - def test_ca_file_change_detection_content_changed(self): + def test_ca_file_change_detection_content_changed(self, ca_data): """Test CA file change detection when file content has changed.""" - original_content = "-----BEGIN CERTIFICATE-----\nMIIDOriginal certificate\n-----END CERTIFICATE-----" with tempfile.NamedTemporaryFile(mode="w", delete=False) as temp_file: - temp_file.write(original_content) + temp_file.write(ca_data) temp_file_path = temp_file.name try: transport = HttpClientTransport(ca_file=temp_file_path) # Modify the file - time.sleep(0.1) # Ensure mtime changes - modified_content = "-----BEGIN CERTIFICATE-----\nMIIDModified certificate\n-----END CERTIFICATE-----" - with open(temp_file_path, "w") as f: - f.write(modified_content) + real_sleep(0.1) # Ensure mtime changes + with open(temp_file_path, "a") as f: + f.write("\n") # File should be detected as changed assert transport._has_ca_file_changed() finally: os.unlink(temp_file_path) - def test_ca_file_change_detection_file_deleted(self): + def test_ca_file_change_detection_file_deleted(self, ca_data): """Test CA file change detection when file is deleted.""" + with tempfile.NamedTemporaryFile(mode="w", delete=False) as temp_file: - temp_file.write("-----BEGIN CERTIFICATE-----\nSome cert\n-----END CERTIFICATE-----") + temp_file.write(ca_data) temp_file_path = temp_file.name transport = HttpClientTransport(ca_file=temp_file_path) @@ -200,26 +158,24 @@ def test_ca_file_change_detection_file_deleted(self): # File deletion should be detected as a change assert transport._has_ca_file_changed() - def test_ca_file_empty_during_rotation(self): + def test_ca_file_empty_during_rotation(self, ca_data): """Test CA file becoming empty during rotation with existing connection.""" + with tempfile.NamedTemporaryFile(mode="w", delete=False) as temp_file: - temp_file.write("-----BEGIN CERTIFICATE-----\nSome cert\n-----END CERTIFICATE-----") + temp_file.write(ca_data) temp_file_path = temp_file.name try: transport = HttpClientTransport(ca_file=temp_file_path) - original_hash = transport._ca_file_hash - - # Simulate having an existing connection - transport._connection = mock.Mock() + original_mtime = transport._ca_file_mtime # Make file empty with open(temp_file_path, "w") as f: f.write("") - # Should not raise error and should preserve old hash - transport._update_ca_file_tracking() - assert transport._ca_file_hash == original_hash + # Should not raise error and should preserve old data + assert transport._ca_data == ca_data + assert transport._ca_file_mtime == original_mtime finally: os.unlink(temp_file_path) @@ -268,25 +224,6 @@ def test_context_manager(self): # Should be able to use after context manager assert transport is not None - def test_close_with_connection(self): - """Test closing transport with active connection.""" - transport = HttpClientTransport() - mock_connection = mock.Mock() - transport._connection = mock_connection - - transport.close() - - mock_connection.close.assert_called_once() - assert transport._connection is None - - def test_close_without_connection(self): - """Test closing transport without active connection.""" - transport = HttpClientTransport() - - # Should not raise any exception - transport.close() - assert transport._connection is None - def test_repr(self): """Test string representation of HttpClientTransport.""" transport = HttpClientTransport() @@ -317,168 +254,6 @@ def test_wrap_socket_adds_server_hostname(self): mock_wrap.assert_called_once_with(mock_sock, server_hostname=hostname) -class TestHttpClientTransportIntegration: - """Integration tests for HttpClientTransport with mocked HTTP connections.""" - - def test_send_request_success(self): - """Test successful request sending.""" - transport = HttpClientTransport() - request = HttpRequest("GET", "https://example.com/test") - - # Mock the HTTP connection and response - with mock.patch("http.client.HTTPSConnection") as mock_conn_class: - mock_conn = mock.Mock() - mock_conn_class.return_value = mock_conn - - mock_response = mock.Mock() - mock_response.status = 200 - mock_response.reason = "OK" - mock_response.getheaders.return_value = [("content-type", "application/json")] - mock_response.read.return_value = b'{"success": true}' - mock_conn.getresponse.return_value = mock_response - - response = transport.send(request) - - # Verify connection was created and request was sent - mock_conn_class.assert_called_once() - mock_conn.request.assert_called_once_with("GET", "/test", body=None, headers=request.headers) - - # Verify response properties - assert response.status_code == 200 - assert response.reason == "OK" - assert response.headers["content-type"] == "application/json" - - def test_send_request_http_exception(self): - """Test request sending with HTTP exception.""" - transport = HttpClientTransport() - request = HttpRequest("GET", "https://example.com/test") - - with mock.patch("http.client.HTTPSConnection") as mock_conn_class: - mock_conn = mock.Mock() - mock_conn_class.return_value = mock_conn - mock_conn.request.side_effect = http.client.HTTPException("Connection failed") - - with pytest.raises(ServiceRequestError): - transport.send(request) - - def test_send_request_ssl_error(self): - """Test request sending with SSL error.""" - transport = HttpClientTransport() - request = HttpRequest("GET", "https://example.com/test") - - with mock.patch("http.client.HTTPSConnection") as mock_conn_class: - mock_conn = mock.Mock() - mock_conn_class.return_value = mock_conn - mock_conn.request.side_effect = ssl.SSLError("SSL handshake failed") - - with pytest.raises(ServiceResponseError): - transport.send(request) - - def test_send_request_with_proxy(self): - """Test request sending through proxy.""" - proxy_endpoint = "https://proxy.example.com" - transport = HttpClientTransport(proxy_endpoint=proxy_endpoint) - request = HttpRequest("GET", "https://original.com/api/test") - - with mock.patch("http.client.HTTPSConnection") as mock_conn_class: - mock_conn = mock.Mock() - mock_conn_class.return_value = mock_conn - - mock_response = mock.Mock() - mock_response.status = 200 - mock_response.getheaders.return_value = [] - mock_response.read.return_value = b"" - mock_conn.getresponse.return_value = mock_response - - transport.send(request) - - # Verify connection was made to proxy host - args, kwargs = mock_conn_class.call_args - assert args[0] == "proxy.example.com" - - # Verify request was sent to proxy path - mock_conn.request.assert_called_once() - call_args = mock_conn.request.call_args - method = call_args[0][0] # First positional arg - path = call_args[0][1] # Second positional arg - assert method == "GET" - assert path == "/api/test" # Original path - - def test_connection_reuse(self): - """Test that connections are reused for multiple requests.""" - transport = HttpClientTransport() - request1 = HttpRequest("GET", "https://example.com/test1") - request2 = HttpRequest("GET", "https://example.com/test2") - - with mock.patch("http.client.HTTPSConnection") as mock_conn_class: - mock_conn = mock.Mock() - mock_conn_class.return_value = mock_conn - - mock_response = mock.Mock() - mock_response.status = 200 - mock_response.getheaders.return_value = [] - mock_response.read.return_value = b"" - mock_conn.getresponse.return_value = mock_response - - # Send two requests - transport.send(request1) - transport.send(request2) - - # Connection should be created only once - assert mock_conn_class.call_count == 1 - # But request should be called twice - assert mock_conn.request.call_count == 2 - - def test_ca_file_change_invalidates_connection(self): - """Test that CA file changes invalidate existing connections.""" - # Create test content that looks like PEM format but is just test data - ca_content = "-----BEGIN CERTIFICATE-----\nMIIDummy certificate content here\n-----END CERTIFICATE-----" - - with tempfile.NamedTemporaryFile(mode="w", delete=False) as temp_file: - temp_file.write(ca_content) - temp_file_path = temp_file.name - - try: - # Mock the SSL context creation to avoid actual certificate validation - with mock.patch("ssl.create_default_context") as mock_ssl_context: - mock_ssl_context.return_value = mock.Mock() - - transport = HttpClientTransport(ca_file=temp_file_path) - request = HttpRequest("GET", "https://example.com/test") - - with mock.patch("http.client.HTTPSConnection") as mock_conn_class: - mock_conn1 = mock.Mock() - mock_conn2 = mock.Mock() - mock_conn_class.side_effect = [mock_conn1, mock_conn2] - - mock_response = mock.Mock() - mock_response.status = 200 - mock_response.getheaders.return_value = [] - mock_response.read.return_value = b"" - mock_conn1.getresponse.return_value = mock_response - mock_conn2.getresponse.return_value = mock_response - - # First request - transport.send(request) - assert mock_conn_class.call_count == 1 - - # Modify CA file - time.sleep(0.1) # Ensure mtime changes - modified_content = ca_content.replace("Dummy", "Foo") - with open(temp_file_path, "w") as f: - f.write(modified_content) - - # Second request should create new connection - transport.send(request) - assert mock_conn_class.call_count == 2 - - # First connection should be closed - mock_conn1.close.assert_called_once() - - finally: - os.unlink(temp_file_path) - - class TestHttpClientTransportResponse: """Test cases for HttpClientTransportResponse class.""" @@ -602,5 +377,354 @@ def test_response_encoding(self): assert response2.text(encoding="latin-1") == "é" -# Import the actual class for the response tests -from azure.identity._internal import http_client_transport as transport +class TestHttpClientTransportWithLocalServer: + """Integration tests using a local test server.""" + + def test_basic_https_request(self): + """Test basic HTTPS request to test server.""" + with TokenProxyTestServer(use_ssl=True) as server: + # Create transport with server's CA certificate + transport = HttpClientTransport(ca_file=server.ca_file) + request = HttpRequest("GET", f"{server.base_url}/health") + + response = transport.send(request) + + assert response.status_code == 200 + data = response.json() + assert data["status"] == "healthy" + assert "timestamp" in data + + def test_post_request_with_body(self): + """Test POST request with request body.""" + with TokenProxyTestServer(use_ssl=True) as server: + transport = HttpClientTransport(ca_file=server.ca_file) + + # Prepare OAuth-like request + body = "grant_type=client_credentials&scope=https://graph.microsoft.com/.default" + request = HttpRequest( + "POST", + f"{server.base_url}/tenant/oauth2/v2.0/token", + headers={"Content-Type": "application/x-www-form-urlencoded"}, + data=body.encode("utf-8"), + ) + + response = transport.send(request) + + assert response.status_code == 200 + data = response.json() + assert "access_token" in data + assert data["token_type"] == "Bearer" + assert data["expires_in"] == 3600 + + def test_proxy_endpoint_comprehensive(self): + """Test comprehensive proxy endpoint functionality with various HTTP methods and scenarios.""" + with TokenProxyTestServer(use_ssl=True) as server: + # Configure transport with proxy endpoint + transport = HttpClientTransport(proxy_endpoint=server.base_url, ca_file=server.ca_file) + + # Test 1: POST request with JSON body through proxy + post_data = {"grant_type": "client_credentials", "scope": "https://graph.microsoft.com/.default"} + post_request = HttpRequest( + "POST", + "https://login.microsoftonline.com/tenant/oauth2/v2.0/token2", + headers={"Content-Type": "application/json"}, + json=post_data, + ) + + post_response = transport.send(post_request) + assert post_response.status_code == 200 + post_data_response = post_response.json() + assert post_data_response["method"] == "POST" + assert post_data_response["proxied_path"] == "/tenant/oauth2/v2.0/token2" + + # Test 2: PUT request through proxy + put_request = HttpRequest( + "PUT", + "https://graph.microsoft.com/v1.0/me/profile", + headers={"Content-Type": "application/json"}, + json={"displayName": "Test User"}, + ) + + put_response = transport.send(put_request) + assert put_response.status_code == 200 + put_data_response = put_response.json() + assert put_data_response["method"] == "PUT" + assert put_data_response["proxied_path"] == "/v1.0/me/profile" + + # Test 3: DELETE request through proxy + delete_request = HttpRequest("DELETE", "https://graph.microsoft.com/v1.0/applications/app-id") + + delete_response = transport.send(delete_request) + assert delete_response.status_code == 200 + delete_data_response = delete_response.json() + assert delete_data_response["method"] == "DELETE" + assert delete_data_response["proxied_path"] == "/v1.0/applications/app-id" + + # Test 4: Complex URL with multiple path segments and query parameters + complex_url = "https://management.azure.com/subscriptions/sub-id/resourceGroups/rg/providers/Microsoft.Storage/storageAccounts/account?api-version=2021-04-01&expand=properties" + complex_request = HttpRequest("GET", complex_url) + + complex_response = transport.send(complex_request) + assert complex_response.status_code == 200 + complex_data_response = complex_response.json() + assert complex_data_response["method"] == "GET" + expected_path = "/subscriptions/sub-id/resourceGroups/rg/providers/Microsoft.Storage/storageAccounts/account?api-version=2021-04-01&expand=properties" + assert complex_data_response["proxied_path"] == expected_path + + # Test 5: Request with custom headers through proxy + headers_request = HttpRequest( + "GET", + "https://vault.azure.net/secrets/test-secret?api-version=7.3", + headers={ + "Authorization": "Bearer test-token", + "X-Custom-Header": "proxy-test-value", + "User-Agent": "Azure-SDK-For-Python", + }, + ) + + headers_response = transport.send(headers_request) + assert headers_response.status_code == 200 + headers_data_response = headers_response.json() + assert headers_data_response["method"] == "GET" + assert headers_data_response["proxied_path"] == "/secrets/test-secret?api-version=7.3" + + # Verify headers were forwarded through proxy + received_headers = headers_data_response["headers_received"] + assert "Authorization" in received_headers + assert "X-Custom-Header" in received_headers + assert received_headers["Authorization"] == "Bearer test-token" + assert received_headers["X-Custom-Header"] == "proxy-test-value" + + def test_sni_with_custom_hostname(self): + """Test SNI (Server Name Indication) with custom hostname.""" + with TokenProxyTestServer(use_ssl=True) as server: + # Use SNI with a different hostname than the server + transport = HttpClientTransport(sni="1234.ests.aks", ca_file=server.ca_file) + + request = HttpRequest("GET", f"{server.base_url}/health") + + response = transport.send(request) + + assert response.status_code == 200 + data = response.json() + assert data["status"] == "healthy" + + def test_connection_reuse(self): + """Test that connections are reused for multiple requests.""" + with TokenProxyTestServer(use_ssl=True) as server: + transport = HttpClientTransport(ca_file=server.ca_file) + + # Make multiple requests + requests_data = [] + for i in range(3): + request = HttpRequest("GET", f"{server.base_url}/health") + response = transport.send(request) + assert response.status_code == 200 + requests_data.append(response.json()) + + # All requests should succeed + assert len(requests_data) == 3 + for data in requests_data: + assert data["status"] == "healthy" + + def test_ca_file_change_detection(self): + """Test CA file change detection with real certificates.""" + with TokenProxyTestServer(use_ssl=True) as server: + # Create a copy of the CA file that we can modify + ca_file = server.ca_file + if ca_file is None: + pytest.skip("CA file not available") + + assert ca_file is not None # Type hint for mypy + with tempfile.NamedTemporaryFile(mode="w", delete=False, suffix=".pem") as temp_ca: + with open(ca_file, "r") as original: + original_content = original.read() + temp_ca.write(original_content) + temp_ca_path = temp_ca.name + + try: + transport = HttpClientTransport(ca_file=temp_ca_path) + + # First request should work + request = HttpRequest("GET", f"{server.base_url}/health") + response1 = transport.send(request) + assert response1.status_code == 200 + + # Modify the CA file (add some content) + real_sleep(0.1) + with open(temp_ca_path, "a") as f: + f.write("\n# Modified for testing\n") + + # Second request should still work (using the same cert content) + response2 = transport.send(request) + assert response2.status_code == 200 + + finally: + os.unlink(temp_ca_path) + + def test_ssl_error_handling(self): + """Test SSL error handling.""" + with TokenProxyTestServer(use_ssl=True) as server: + # Create transport without proper CA file (will cause SSL error) + transport = HttpClientTransport() # No CA file provided + request = HttpRequest("GET", f"{server.base_url}/health") + + # Should raise SSL-related error + with pytest.raises((ServiceRequestError, ServiceResponseError)): + transport.send(request) + + def test_server_error_response(self): + """Test handling of server error responses.""" + with TokenProxyTestServer(use_ssl=True) as server: + transport = HttpClientTransport(ca_file=server.ca_file) + request = HttpRequest("GET", f"{server.base_url}/error/500") + + response = transport.send(request) + + assert response.status_code == 500 + # Should not raise exception, just return error response + + def test_slow_server_response(self): + """Test handling of slow server responses.""" + with TokenProxyTestServer(use_ssl=True) as server: + # Set a longer timeout for this test + transport = HttpClientTransport(ca_file=server.ca_file, timeout=5) + # # Ensure fresh connection by closing any existing ones + # transport.close() + request = HttpRequest("GET", f"{server.base_url}/slow") + + start_time = time.time() + response = transport.send(request) + elapsed_time = time.time() - start_time + + assert response.status_code == 200 + # Should take at least 2 seconds (server waits for 2s) + assert elapsed_time >= 2.0 + data = response.json() + assert data["message"] == "slow response" + + def test_custom_headers_preserved(self): + """Test that custom headers are preserved and sent to server.""" + with TokenProxyTestServer(use_ssl=True) as server: + transport = HttpClientTransport(ca_file=server.ca_file) + + custom_headers = { + "Authorization": "Bearer test-token", + "User-Agent": "HttpClientTransport/1.0", + "X-Custom-Header": "test-value", + } + + request = HttpRequest("GET", f"{server.base_url}/proxy/test", headers=custom_headers) + + response = transport.send(request) + + assert response.status_code == 200 + data = response.json() + + # Server echoes back the headers it received + received_headers = data["headers_received"] + assert "Authorization" in received_headers + assert "User-Agent" in received_headers + assert "X-Custom-Header" in received_headers + assert received_headers["Authorization"] == "Bearer test-token" + + def test_query_parameters_preserved(self): + """Test that query parameters are preserved in proxy requests.""" + with TokenProxyTestServer(use_ssl=True) as server: + transport = HttpClientTransport(proxy_endpoint=server.base_url, ca_file=server.ca_file) + + # Request with query parameters + original_url = "https://example.com/api/data?scope=read&limit=10&format=json" + request = HttpRequest("GET", original_url) + + response = transport.send(request) + + assert response.status_code == 200 + data = response.json() + + # The path should include the query parameters + expected_path = "/api/data?scope=read&limit=10&format=json" + assert data["proxied_path"] == expected_path + + def test_concurrent_requests(self): + """Test handling multiple concurrent requests.""" + import threading + import concurrent.futures + + with TokenProxyTestServer(use_ssl=True) as server: + + transport = HttpClientTransport(ca_file=server.ca_file) + + def make_request(request_id): + request = HttpRequest("GET", f"{server.base_url}/health") + response = transport.send(request) + return request_id, response.status_code, response.json() + + # Make 5 concurrent requests + with concurrent.futures.ThreadPoolExecutor(max_workers=5) as executor: + futures = [executor.submit(make_request, i) for i in range(5)] + results = [future.result() for future in concurrent.futures.as_completed(futures)] + + # All requests should succeed + assert len(results) == 5 + for request_id, status_code, data in results: + assert status_code == 200 + assert data["status"] == "healthy" + + +class TestTokenProxyTestServer: + """Tests for the test server itself.""" + + def test_server_startup_and_shutdown(self): + """Test that server starts and stops properly.""" + server = TokenProxyTestServer(use_ssl=True) + + # Server should not be running initially + assert server.server is None + + # Start server + base_url = server.start() + assert server.server is not None + assert base_url.startswith("https://") + assert str(server.port) in base_url + + # Stop server + server.stop() + + # Should clean up properly + assert len(server._temp_files) == 0 # Files should be cleaned up + + def test_context_manager(self): + """Test using server as context manager.""" + with TokenProxyTestServer(use_ssl=False) as server: + assert server.server is not None + assert server.base_url.startswith("http://") + + # Server should be stopped after context exit + # Note: We can't easily test this without making a request + + def test_certificate_generation(self): + """Test certificate generation.""" + server = TokenProxyTestServer(use_ssl=True) + + try: + server.generate_test_certificates() + + # Should have created certificate files + assert server.cert_file is not None + assert server.key_file is not None + assert server.ca_file is not None + + # Files should exist + assert os.path.exists(server.cert_file) + assert os.path.exists(server.key_file) + assert os.path.exists(server.ca_file) + + # Files should contain certificate data + with open(server.cert_file, "r") as f: + cert_content = f.read() + assert "-----BEGIN CERTIFICATE-----" in cert_content + assert "-----END CERTIFICATE-----" in cert_content + + finally: + server.stop() # Clean up temp files diff --git a/sdk/identity/azure-identity/tests/test_workload_identity_credential.py b/sdk/identity/azure-identity/tests/test_workload_identity_credential.py index 8a24562f0d09..20470b38a536 100644 --- a/sdk/identity/azure-identity/tests/test_workload_identity_credential.py +++ b/sdk/identity/azure-identity/tests/test_workload_identity_credential.py @@ -105,7 +105,7 @@ def test_use_token_proxy_with_ca_data(self): mock_transport_instance = MagicMock() mock_transport_class.return_value = mock_transport_instance - credential = WorkloadIdentityCredential( + WorkloadIdentityCredential( tenant_id=tenant_id, client_id=client_id, token_file_path=token_file_path, @@ -136,7 +136,7 @@ def test_use_token_proxy_minimal_config(self): mock_transport_instance = MagicMock() mock_transport_class.return_value = mock_transport_instance - credential = WorkloadIdentityCredential( + WorkloadIdentityCredential( tenant_id=tenant_id, client_id=client_id, token_file_path=token_file_path, @@ -242,7 +242,7 @@ def test_use_token_proxy_false_does_not_create_transport(self): token_file_path = "foo-path" with patch("azure.identity._credentials.workload_identity.HttpClientTransport") as mock_transport_class: - credential = WorkloadIdentityCredential( + WorkloadIdentityCredential( tenant_id=tenant_id, client_id=client_id, token_file_path=token_file_path, diff --git a/sdk/identity/azure-identity/tests/test_workload_identity_credential_async.py b/sdk/identity/azure-identity/tests/test_workload_identity_credential_async.py index 3a40c2b5f18c..42913b5c9f76 100644 --- a/sdk/identity/azure-identity/tests/test_workload_identity_credential_async.py +++ b/sdk/identity/azure-identity/tests/test_workload_identity_credential_async.py @@ -2,15 +2,29 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT License. # ------------------------------------ +# cspell:ignore cafile import os import tempfile +from time import sleep as real_sleep from unittest.mock import mock_open, patch, MagicMock import pytest +from azure.core.rest import HttpRequest from azure.identity.aio import WorkloadIdentityCredential -from azure.identity._internal.http_client_transport import HttpClientTransport +from azure.identity.aio._credentials.workload_identity import _get_transport from helpers import mock_response, build_aad_response, GET_TOKEN_METHODS +from proxy_server import TokenProxyTestServer + + +PEM_CERT_PATH = os.path.join(os.path.dirname(__file__), "certificate.pem") + + +@pytest.fixture(scope="module") +def ca_data() -> str: + """Read CA certificate data from a PEM file for testing.""" + with open(PEM_CERT_PATH, "r", encoding="utf-8") as f: + return f.read() def test_workload_identity_credential_initialize(): @@ -53,8 +67,8 @@ async def send(request, **kwargs): class TestWorkloadIdentityCredentialTokenProxyAsync: """Async test cases for WorkloadIdentityCredential with use_token_proxy=True.""" - def test_use_token_proxy_creates_http_client_transport(self): - """Test that use_token_proxy=True creates HttpClientTransport with correct parameters.""" + def test_use_token_proxy_creates_custom_aiohttp_transport(self): + """Test that use_token_proxy=True creates a custom aiohttp transport with correct parameters.""" tenant_id = "tenant-id" client_id = "client-id" token_file_path = "foo-path" @@ -69,21 +83,20 @@ def test_use_token_proxy_creates_http_client_transport(self): } with patch.dict(os.environ, env_vars, clear=False): - with patch("azure.identity.aio._credentials.workload_identity.HttpClientTransport") as mock_transport_class: + with patch("azure.identity.aio._credentials.workload_identity._get_transport") as mock_get_transport: mock_transport_instance = MagicMock() - mock_transport_class.return_value = mock_transport_instance + mock_get_transport.return_value = mock_transport_instance - credential = WorkloadIdentityCredential( + WorkloadIdentityCredential( tenant_id=tenant_id, client_id=client_id, token_file_path=token_file_path, use_token_proxy=True, ) - # Verify HttpClientTransport was called with correct parameters - mock_transport_class.assert_called_once_with( + mock_get_transport.assert_called_once_with( sni=sni_hostname, - proxy_endpoint=proxy_endpoint, + token_proxy_endpoint=proxy_endpoint, ca_file=ca_file_path, ca_data=None, ) @@ -102,21 +115,19 @@ def test_use_token_proxy_with_ca_data(self): } with patch.dict(os.environ, env_vars, clear=False): - with patch("azure.identity.aio._credentials.workload_identity.HttpClientTransport") as mock_transport_class: + with patch("azure.identity.aio._credentials.workload_identity._get_transport") as mock_get_transport: mock_transport_instance = MagicMock() - mock_transport_class.return_value = mock_transport_instance + mock_get_transport.return_value = mock_transport_instance - credential = WorkloadIdentityCredential( + WorkloadIdentityCredential( tenant_id=tenant_id, client_id=client_id, token_file_path=token_file_path, use_token_proxy=True, ) - - # Verify HttpClientTransport was called with CA data - mock_transport_class.assert_called_once_with( + mock_get_transport.assert_called_once_with( sni=None, - proxy_endpoint=proxy_endpoint, + token_proxy_endpoint=proxy_endpoint, ca_file=None, ca_data=ca_data, ) @@ -133,11 +144,11 @@ def test_use_token_proxy_minimal_config(self): } with patch.dict(os.environ, env_vars, clear=False): - with patch("azure.identity.aio._credentials.workload_identity.HttpClientTransport") as mock_transport_class: + with patch("azure.identity.aio._credentials.workload_identity._get_transport") as mock_get_transport: mock_transport_instance = MagicMock() - mock_transport_class.return_value = mock_transport_instance + mock_get_transport.return_value = mock_transport_instance - credential = WorkloadIdentityCredential( + WorkloadIdentityCredential( tenant_id=tenant_id, client_id=client_id, token_file_path=token_file_path, @@ -145,9 +156,9 @@ def test_use_token_proxy_minimal_config(self): ) # Verify HttpClientTransport was called with minimal config - mock_transport_class.assert_called_once_with( + mock_get_transport.assert_called_once_with( sni=None, - proxy_endpoint=proxy_endpoint, + token_proxy_endpoint=proxy_endpoint, ca_file=None, ca_data=None, ) @@ -220,8 +231,8 @@ async def send(request, **kwargs): } with patch.dict(os.environ, env_vars, clear=False): - with patch("azure.identity.aio._credentials.workload_identity.HttpClientTransport") as mock_transport_class: - mock_transport_class.return_value = mock_transport_instance + with patch("azure.identity.aio._credentials.workload_identity._get_transport") as mock_get_transport: + mock_get_transport.return_value = mock_transport_instance credential = WorkloadIdentityCredential( tenant_id=tenant_id, @@ -243,13 +254,608 @@ def test_use_token_proxy_false_does_not_create_transport(self): client_id = "client-id" token_file_path = "foo-path" - with patch("azure.identity.aio._credentials.workload_identity.HttpClientTransport") as mock_transport_class: - credential = WorkloadIdentityCredential( + with patch("azure.identity.aio._credentials.workload_identity._get_transport") as mock_get_transport: + WorkloadIdentityCredential( tenant_id=tenant_id, client_id=client_id, token_file_path=token_file_path, use_token_proxy=False, ) + mock_get_transport.assert_not_called() + + +class TestCustomAioHttpTransport: + """Test cases for the custom AioHttpTransport used by WorkloadIdentityCredential.""" + + def test_get_transport_creates_workload_identity_aiohttp_transport(self, ca_data): + """Test that _get_transport creates WorkloadIdentityAioHttpTransport with correct parameters.""" + sni = "test.sni.com" + proxy_endpoint = "https://proxy.example.com:8080" + ca_file = PEM_CERT_PATH + + transport = _get_transport( + sni=sni, + token_proxy_endpoint=proxy_endpoint, + ca_file=ca_file, + ca_data=None, + ) + + assert transport is not None + assert hasattr(transport, "_sni") + assert hasattr(transport, "_proxy_endpoint") + assert hasattr(transport, "_ca_file") + assert hasattr(transport, "_ca_data") + assert transport._sni == sni + assert transport._proxy_endpoint == proxy_endpoint + assert transport._ca_file == ca_file + assert transport._ca_data == ca_data + + def test_get_transport_with_minimal_config(self): + """Test _get_transport with minimal configuration.""" + proxy_endpoint = "https://proxy.example.com:8080" + + transport = _get_transport( + sni=None, + token_proxy_endpoint=proxy_endpoint, + ca_file=None, + ca_data=None, + ) + + assert transport is not None + assert transport._sni is None + assert transport._proxy_endpoint == proxy_endpoint + assert transport._ca_file is None + assert transport._ca_data is None + + @pytest.mark.asyncio + async def test_workload_identity_aiohttp_transport_send_with_sni(self): + """Test that WorkloadIdentityAioHttpTransport.send sets server_hostname correctly.""" + sni = "test.sni.com" + proxy_endpoint = "https://proxy.example.com:8080" + + transport = _get_transport( + sni=sni, + token_proxy_endpoint=proxy_endpoint, + ca_file=None, + ca_data=None, + ) + assert transport is not None + + # Mock the parent send method + mock_request = MagicMock() + mock_request.url = "https://login.microsoftonline.com/tenant/oauth2/v2.0/token" + + with patch.object(transport.__class__.__bases__[1], "send") as mock_parent_send: + mock_parent_send.return_value = MagicMock() + + await transport.send(mock_request) + + # Verify parent send was called with server_hostname set + mock_parent_send.assert_called_once() + call_args = mock_parent_send.call_args + assert call_args[1]["server_hostname"] == sni + + @pytest.mark.asyncio + async def test_workload_identity_aiohttp_transport_send_updates_url(self): + """Test that WorkloadIdentityAioHttpTransport.send updates request URL with proxy endpoint.""" + proxy_endpoint = "https://proxy.example.com:8080/path" + original_url = "https://login.microsoftonline.com/tenant/oauth2/v2.0/token" + expected_url = "https://proxy.example.com:8080/path/tenant/oauth2/v2.0/token" + + transport = _get_transport( + sni=None, + token_proxy_endpoint=proxy_endpoint, + ca_file=None, + ca_data=None, + ) + assert transport is not None + + mock_request = MagicMock() + mock_request.url = original_url + + with patch.object(transport.__class__.__bases__[1], "send") as mock_parent_send: + mock_parent_send.return_value = MagicMock() + + await transport.send(mock_request) + + # Verify URL was updated to use proxy endpoint + assert mock_request.url == expected_url + + @pytest.mark.asyncio + async def test_workload_identity_aiohttp_transport_send_with_ca_data(self, ca_data): + """Test that WorkloadIdentityAioHttpTransport.send creates SSL context from CA data.""" + proxy_endpoint = "https://proxy.example.com:8080" + + transport = _get_transport( + sni=None, + token_proxy_endpoint=proxy_endpoint, + ca_file=None, + ca_data=ca_data, + ) + assert transport is not None + + mock_request = MagicMock() + mock_request.url = "https://login.microsoftonline.com/tenant/oauth2/v2.0/token" + + with patch.object(transport.__class__.__bases__[1], "send") as mock_parent_send: + mock_parent_send.return_value = MagicMock() + + await transport.send(mock_request) + + # Verify SSL context was set + mock_parent_send.assert_called_once() + call_args = mock_parent_send.call_args + assert "ssl" in call_args[1] + assert call_args[1]["ssl"] is not None + + @pytest.mark.asyncio + async def test_workload_identity_aiohttp_transport_send_with_ca_file_reload(self, ca_data): + """Test that WorkloadIdentityAioHttpTransport.send reloads CA file when changed.""" + + # Create a temporary CA file + with tempfile.NamedTemporaryFile(mode="w", delete=False, suffix=".pem") as ca_file: + ca_file.write(ca_data) + ca_file_path = ca_file.name + + try: + proxy_endpoint = "https://proxy.example.com:8080" + + transport = _get_transport( + sni=None, + token_proxy_endpoint=proxy_endpoint, + ca_file=ca_file_path, + ca_data=None, + ) + + assert transport is not None + # Store original CA data + original_ca_data = transport._ca_data + + # Simulate file change by modifying mtime tracking + real_sleep(0.1) # Ensure different mtime + with open(ca_file_path, "a") as f: + f.write("\n") + + mock_request = MagicMock() + mock_request.url = "https://login.microsoftonline.com/tenant/oauth2/v2.0/token" + + with patch.object(transport.__class__.__bases__[1], "send") as mock_parent_send: + mock_parent_send.return_value = MagicMock() + + await transport.send(mock_request) + + # Verify CA data was reloaded + assert transport._ca_data != original_ca_data + + finally: + # Clean up temporary file + os.unlink(ca_file_path) + + @pytest.mark.asyncio + async def test_workload_identity_aiohttp_transport_context_manager(self): + """Test that WorkloadIdentityAioHttpTransport works as async context manager.""" + transport = _get_transport( + sni=None, + token_proxy_endpoint="https://proxy.example.com:8080", + ca_file=None, + ca_data=None, + ) + assert transport is not None + + # Mock the parent context manager methods + with patch.object(transport.__class__.__bases__[1], "__aenter__") as mock_aenter, patch.object( + transport.__class__.__bases__[1], "__aexit__" + ) as mock_aexit: + + mock_aenter.return_value = transport + mock_aexit.return_value = None + + async with transport as ctx_transport: + assert ctx_transport == transport + + mock_aenter.assert_called_once() + mock_aexit.assert_called_once() + + def test_workload_identity_aiohttp_transport_initialization_with_ca_data(self, ca_data): + """Test WorkloadIdentityAioHttpTransport initialization with CA data creates SSL context.""" + transport = _get_transport( + sni=None, + token_proxy_endpoint="https://proxy.example.com:8080", + ca_file=None, + ca_data=ca_data, + ) + assert transport is not None + + # Verify SSL context was created during initialization + assert hasattr(transport, "_ssl_context") + assert transport._ssl_context is not None + + def test_workload_identity_aiohttp_transport_initialization_without_ca_data(self): + """Test WorkloadIdentityAioHttpTransport initialization without CA data.""" + transport = _get_transport( + sni=None, + token_proxy_endpoint="https://proxy.example.com:8080", + ca_file=None, + ca_data=None, + ) + assert transport is not None + + # Verify SSL context is created with None ca_data (creates default context) + assert hasattr(transport, "_ssl_context") + # SSL context should still be created even with None ca_data + + @pytest.mark.asyncio + async def test_workload_identity_aiohttp_transport_send_no_ssl_context_when_no_ca_data(self): + """Test that no SSL context is passed when ca_data is None and SSL context creation fails.""" + transport = _get_transport( + sni=None, + token_proxy_endpoint="https://proxy.example.com:8080", + ca_file=None, + ca_data=None, + ) + assert transport is not None + + # Mock SSL context creation to return None + with patch("ssl.create_default_context", return_value=None): + # Manually set ssl_context to None to test the conditional logic + with patch.object(transport, "_ssl_context", None): + mock_request = MagicMock() + mock_request.url = "https://login.microsoftonline.com/tenant/oauth2/v2.0/token" + + with patch.object(transport.__class__.__bases__[1], "send") as mock_parent_send: + mock_parent_send.return_value = MagicMock() + + await transport.send(mock_request) + + # Verify SSL context was not set when None + mock_parent_send.assert_called_once() + call_args = mock_parent_send.call_args + assert "ssl" not in call_args[1] or call_args[1].get("ssl") is None + + def test_workload_identity_aiohttp_transport_inherits_from_token_binding_mixin(self): + """Test that WorkloadIdentityAioHttpTransport inherits from TokenBindingTransportMixin.""" + transport = _get_transport( + sni="test.sni.com", + token_proxy_endpoint="https://proxy.example.com:8080", + ca_file=None, + ca_data=None, + ) + + assert transport is not None + + # Verify inheritance from TokenBindingTransportMixin + from azure.identity._internal.token_binding_transport_mixin import TokenBindingTransportMixin + + assert isinstance(transport, TokenBindingTransportMixin) + + # Verify TokenBindingTransportMixin methods are available + assert hasattr(transport, "_update_request_url") + assert hasattr(transport, "_has_ca_file_changed") + assert hasattr(transport, "_load_ca_file_to_data") + assert hasattr(transport, "_validate_url") + + +class TestCustomAioHttpTransportWithLocalServer: + """Integration tests using a local test server for WorkloadIdentityAioHttpTransport.""" + + @pytest.mark.asyncio + async def test_basic_https_request(self): + """Test basic HTTPS request to test server.""" + with TokenProxyTestServer(use_ssl=True) as server: + # Create transport with server's CA certificate + transport = _get_transport(sni=None, token_proxy_endpoint=None, ca_file=server.ca_file, ca_data=None) + assert transport is not None + request = HttpRequest("GET", f"{server.base_url}/health") + + response = await transport.send(request) - # Verify HttpClientTransport was NOT called - mock_transport_class.assert_not_called() + assert response.status_code == 200 + data = response.json() + assert data["status"] == "healthy" + assert "timestamp" in data + + @pytest.mark.asyncio + async def test_post_request_with_body(self): + """Test POST request with request body.""" + with TokenProxyTestServer(use_ssl=True) as server: + transport = _get_transport(sni=None, token_proxy_endpoint=None, ca_file=server.ca_file, ca_data=None) + assert transport is not None + + # Prepare OAuth-like request + body = "grant_type=client_credentials&scope=https://graph.microsoft.com/.default" + request = HttpRequest( + "POST", + f"{server.base_url}/tenant/oauth2/v2.0/token", + headers={"Content-Type": "application/x-www-form-urlencoded"}, + data=body.encode("utf-8"), + ) + + response = await transport.send(request) + + assert response.status_code == 200 + data = response.json() + assert "access_token" in data + assert data["token_type"] == "Bearer" + assert data["expires_in"] == 3600 + + @pytest.mark.asyncio + async def test_proxy_endpoint_comprehensive(self): + """Test comprehensive proxy endpoint functionality with various HTTP methods and scenarios.""" + with TokenProxyTestServer(use_ssl=True) as server: + # Configure transport with proxy endpoint + transport = _get_transport( + sni=None, token_proxy_endpoint=server.base_url, ca_file=server.ca_file, ca_data=None + ) + assert transport is not None + + # Test 1: POST request with JSON body through proxy + post_data = {"grant_type": "client_credentials", "scope": "https://graph.microsoft.com/.default"} + post_request = HttpRequest( + "POST", + "https://login.microsoftonline.com/tenant/oauth2/v2.0/token2", + headers={"Content-Type": "application/json"}, + json=post_data, + ) + + post_response = await transport.send(post_request) + assert post_response.status_code == 200 + post_data_response = post_response.json() + assert post_data_response["method"] == "POST" + assert post_data_response["proxied_path"] == "/tenant/oauth2/v2.0/token2" + + # Test 2: PUT request through proxy + put_request = HttpRequest( + "PUT", + "https://graph.microsoft.com/v1.0/me/profile", + headers={"Content-Type": "application/json"}, + json={"displayName": "Test User"}, + ) + + put_response = await transport.send(put_request) + assert put_response.status_code == 200 + put_data_response = put_response.json() + assert put_data_response["method"] == "PUT" + assert put_data_response["proxied_path"] == "/v1.0/me/profile" + + # Test 3: DELETE request through proxy + delete_request = HttpRequest("DELETE", "https://graph.microsoft.com/v1.0/applications/app-id") + + delete_response = await transport.send(delete_request) + assert delete_response.status_code == 200 + delete_data_response = delete_response.json() + assert delete_data_response["method"] == "DELETE" + assert delete_data_response["proxied_path"] == "/v1.0/applications/app-id" + + # Test 4: Complex URL with multiple path segments and query parameters + complex_url = "https://management.azure.com/subscriptions/sub-id/resourceGroups/rg/providers/Microsoft.Storage/storageAccounts/account?api-version=2021-04-01&expand=properties" + complex_request = HttpRequest("GET", complex_url) + + complex_response = await transport.send(complex_request) + assert complex_response.status_code == 200 + complex_data_response = complex_response.json() + assert complex_data_response["method"] == "GET" + expected_path = "/subscriptions/sub-id/resourceGroups/rg/providers/Microsoft.Storage/storageAccounts/account?api-version=2021-04-01&expand=properties" + assert complex_data_response["proxied_path"] == expected_path + + # Test 5: Request with custom headers through proxy + headers_request = HttpRequest( + "GET", + "https://vault.azure.net/secrets/test-secret?api-version=7.3", + headers={ + "Authorization": "Bearer test-token", + "X-Custom-Header": "proxy-test-value", + "User-Agent": "Azure-SDK-For-Python", + }, + ) + + headers_response = await transport.send(headers_request) + assert headers_response.status_code == 200 + headers_data_response = headers_response.json() + assert headers_data_response["method"] == "GET" + assert headers_data_response["proxied_path"] == "/secrets/test-secret?api-version=7.3" + + # Verify headers were forwarded through proxy + received_headers = headers_data_response["headers_received"] + assert "Authorization" in received_headers + assert "X-Custom-Header" in received_headers + assert received_headers["Authorization"] == "Bearer test-token" + assert received_headers["X-Custom-Header"] == "proxy-test-value" + + @pytest.mark.asyncio + async def test_sni_with_custom_hostname(self): + """Test SNI (Server Name Indication) with custom hostname.""" + with TokenProxyTestServer(use_ssl=True) as server: + # Use SNI with a different hostname than the server + transport = _get_transport( + sni="1234.ests.aks", token_proxy_endpoint=None, ca_file=server.ca_file, ca_data=None + ) + assert transport is not None + + request = HttpRequest("GET", f"{server.base_url}/health") + + response = await transport.send(request) + + assert response.status_code == 200 + data = response.json() + assert data["status"] == "healthy" + + @pytest.mark.asyncio + async def test_connection_reuse(self): + """Test that connections are reused for multiple requests.""" + with TokenProxyTestServer(use_ssl=True) as server: + transport = _get_transport(sni=None, token_proxy_endpoint=None, ca_file=server.ca_file, ca_data=None) + assert transport is not None + + # Make multiple requests + requests_data = [] + for i in range(3): + request = HttpRequest("GET", f"{server.base_url}/health") + response = await transport.send(request) + assert response.status_code == 200 + requests_data.append(response.json()) + + # All requests should succeed + assert len(requests_data) == 3 + for data in requests_data: + assert data["status"] == "healthy" + + @pytest.mark.asyncio + async def test_ca_file_change_detection(self): + """Test CA file change detection with real certificates.""" + with TokenProxyTestServer(use_ssl=True) as server: + # Create a copy of the CA file that we can modify + ca_file = server.ca_file + if ca_file is None: + pytest.skip("CA file not available") + + assert ca_file is not None # Type hint for mypy + with tempfile.NamedTemporaryFile(mode="w", delete=False, suffix=".pem") as temp_ca: + with open(ca_file, "r") as original: + original_content = original.read() + temp_ca.write(original_content) + temp_ca_path = temp_ca.name + + try: + transport = _get_transport(sni=None, token_proxy_endpoint=None, ca_file=temp_ca_path, ca_data=None) + assert transport is not None + + # First request should work + request = HttpRequest("GET", f"{server.base_url}/health") + response1 = await transport.send(request) + assert response1.status_code == 200 + + # Modify the CA file (add some content) + real_sleep(0.1) + with open(temp_ca_path, "a") as f: + f.write("\n# Modified for testing\n") + + # Second request should still work (using the same cert content) + response2 = await transport.send(request) + assert response2.status_code == 200 + + finally: + os.unlink(temp_ca_path) + + @pytest.mark.asyncio + async def test_ssl_error_handling(self): + """Test SSL error handling.""" + from azure.core.exceptions import ServiceRequestError, ServiceResponseError + + with TokenProxyTestServer(use_ssl=True) as server: + # Create transport without proper CA file (will cause SSL error) + transport = _get_transport(sni=None, token_proxy_endpoint=None, ca_file=None, ca_data=None) + assert transport is not None + + request = HttpRequest("GET", f"{server.base_url}/health") + + # Should raise SSL-related error + with pytest.raises((ServiceRequestError, ServiceResponseError)): + await transport.send(request) + + @pytest.mark.asyncio + async def test_server_error_response(self): + """Test handling of server error responses.""" + with TokenProxyTestServer(use_ssl=True) as server: + transport = _get_transport(sni=None, token_proxy_endpoint=None, ca_file=server.ca_file, ca_data=None) + assert transport is not None + + request = HttpRequest("GET", f"{server.base_url}/error/500") + + response = await transport.send(request) + + assert response.status_code == 500 + # Should not raise exception, just return error response + + @pytest.mark.asyncio + async def test_slow_server_response(self): + """Test handling of slow server responses.""" + import time + + with TokenProxyTestServer(use_ssl=True) as server: + transport = _get_transport(sni=None, token_proxy_endpoint=None, ca_file=server.ca_file, ca_data=None) + assert transport is not None + + request = HttpRequest("GET", f"{server.base_url}/slow") + + start_time = time.time() + response = await transport.send(request) + elapsed_time = time.time() - start_time + + assert response.status_code == 200 + # Should take at least 2 seconds (server waits for 2s) + assert elapsed_time >= 2.0 + data = response.json() + assert data["message"] == "slow response" + + @pytest.mark.asyncio + async def test_custom_headers_preserved(self): + """Test that custom headers are preserved and sent to server.""" + with TokenProxyTestServer(use_ssl=True) as server: + transport = _get_transport(sni=None, token_proxy_endpoint=None, ca_file=server.ca_file, ca_data=None) + assert transport is not None + + custom_headers = { + "Authorization": "Bearer test-token", + "User-Agent": "WorkloadIdentityAioHttpTransport/1.0", + "X-Custom-Header": "test-value", + } + + request = HttpRequest("GET", f"{server.base_url}/proxy/test", headers=custom_headers) + + response = await transport.send(request) + + assert response.status_code == 200 + data = response.json() + + # Server echoes back the headers it received + received_headers = data["headers_received"] + assert "Authorization" in received_headers + assert "User-Agent" in received_headers + assert "X-Custom-Header" in received_headers + assert received_headers["Authorization"] == "Bearer test-token" + + @pytest.mark.asyncio + async def test_query_parameters_preserved(self): + """Test that query parameters are preserved in proxy requests.""" + with TokenProxyTestServer(use_ssl=True) as server: + transport = _get_transport( + sni=None, token_proxy_endpoint=server.base_url, ca_file=server.ca_file, ca_data=None + ) + assert transport is not None + + # Request with query parameters + original_url = "https://example.com/api/data?scope=read&limit=10&format=json" + request = HttpRequest("GET", original_url) + + response = await transport.send(request) + + assert response.status_code == 200 + data = response.json() + + # The path should include the query parameters + expected_path = "/api/data?scope=read&limit=10&format=json" + assert data["proxied_path"] == expected_path + + @pytest.mark.asyncio + async def test_concurrent_requests(self): + """Test handling multiple concurrent requests.""" + import asyncio + + with TokenProxyTestServer(use_ssl=True) as server: + transport = _get_transport(sni=None, token_proxy_endpoint=None, ca_file=server.ca_file, ca_data=None) + assert transport is not None + + async def make_request(request_id): + request = HttpRequest("GET", f"{server.base_url}/health") + response = await transport.send(request) + return request_id, response.status_code, response.json() + + # Make 5 concurrent requests + tasks = [make_request(i) for i in range(5)] + results = await asyncio.gather(*tasks) + + # All requests should succeed + assert len(results) == 5 + for request_id, status_code, data in results: + assert status_code == 200 + assert data["status"] == "healthy" From c62e4b801009ec9e3610a2b6e11fd99ddbb6f14a Mon Sep 17 00:00:00 2001 From: Paul Van Eck Date: Wed, 29 Oct 2025 21:04:00 +0000 Subject: [PATCH 03/12] Updates Signed-off-by: Paul Van Eck --- .../azure/identity/_credentials/workload_identity.py | 4 ++-- .../azure/identity/_internal/http_client_transport.py | 2 +- .../identity/_internal/token_binding_transport_mixin.py | 2 +- .../azure/identity/aio/_credentials/workload_identity.py | 9 +++++---- sdk/identity/azure-identity/tests/proxy_server.py | 1 - .../azure-identity/tests/test_http_client_transport.py | 1 + .../tests/test_workload_identity_credential_async.py | 2 +- 7 files changed, 11 insertions(+), 10 deletions(-) diff --git a/sdk/identity/azure-identity/azure/identity/_credentials/workload_identity.py b/sdk/identity/azure-identity/azure/identity/_credentials/workload_identity.py index 0dc44294a8cc..e4b59b0e8c26 100644 --- a/sdk/identity/azure-identity/azure/identity/_credentials/workload_identity.py +++ b/sdk/identity/azure-identity/azure/identity/_credentials/workload_identity.py @@ -58,7 +58,7 @@ class WorkloadIdentityCredential(ClientAssertionCredential, TokenFileMixin): :keyword str client_id: The client ID of a Microsoft Entra app registration. :keyword str token_file_path: The path to a file containing a Kubernetes service account token that authenticates the identity. - :keyword str use_token_proxy: Whether or not to to read token proxy configuration from environment variables and use + :keyword bool use_token_proxy: Whether or not to read token proxy configuration from environment variables and use a token proxy to acquire tokens. Defaults to False. .. admonition:: Example: @@ -118,7 +118,7 @@ def __init__( if ca_file and ca_data: raise ValueError( - "Both AZURE_KUBERNETES_CA_FILE and AZURE_KUBERNETES_CA_DATA are set. Only one should be set" + "Both AZURE_KUBERNETES_CA_FILE and AZURE_KUBERNETES_CA_DATA are set. Only one should be set." ) kwargs["transport"] = HttpClientTransport( diff --git a/sdk/identity/azure-identity/azure/identity/_internal/http_client_transport.py b/sdk/identity/azure-identity/azure/identity/_internal/http_client_transport.py index 8dd729899c70..1737820b85e8 100644 --- a/sdk/identity/azure-identity/azure/identity/_internal/http_client_transport.py +++ b/sdk/identity/azure-identity/azure/identity/_internal/http_client_transport.py @@ -199,7 +199,7 @@ def _get_connection(self, host: str) -> http.client.HTTPSConnection: # Use cached SSL context or create a new one ssl_context = self._ssl_context or self._create_ssl_context() - connection = http.client.HTTPSConnection( + connection = http.client.HTTPSConnection( # nosec host, timeout=self.connection_config.timeout, context=ssl_context, diff --git a/sdk/identity/azure-identity/azure/identity/_internal/token_binding_transport_mixin.py b/sdk/identity/azure-identity/azure/identity/_internal/token_binding_transport_mixin.py index 6b600612068d..44dbf426e98f 100644 --- a/sdk/identity/azure-identity/azure/identity/_internal/token_binding_transport_mixin.py +++ b/sdk/identity/azure-identity/azure/identity/_internal/token_binding_transport_mixin.py @@ -106,8 +106,8 @@ def _update_request_url(self, request: HttpRequest) -> None: :param request: The HTTP request object to update. :type request: ~azure.core.rest.HttpRequest """ - parsed_request_url = urllib.parse.urlparse(request.url) if self._proxy_endpoint: + parsed_request_url = urllib.parse.urlparse(request.url) parsed_proxy_url = urllib.parse.urlparse(self._proxy_endpoint) combined_path = parsed_proxy_url.path.rstrip("/") + "/" + parsed_request_url.path.lstrip("/") new_url = urllib.parse.urlunparse( diff --git a/sdk/identity/azure-identity/azure/identity/aio/_credentials/workload_identity.py b/sdk/identity/azure-identity/azure/identity/aio/_credentials/workload_identity.py index 70fa6aef4507..1801d4cf8d2b 100644 --- a/sdk/identity/azure-identity/azure/identity/aio/_credentials/workload_identity.py +++ b/sdk/identity/azure-identity/azure/identity/aio/_credentials/workload_identity.py @@ -38,7 +38,7 @@ class WorkloadIdentityCredential(ClientAssertionCredential, TokenFileMixin): :keyword str client_id: The client ID of a Microsoft Entra app registration. :keyword str token_file_path: The path to a file containing a Kubernetes service account token that authenticates the identity. - :keyword str use_token_proxy: Whether or not to to read token proxy configuration from environment variables and use + :keyword bool use_token_proxy: Whether or not to read token proxy configuration from environment variables and use a token proxy to acquire tokens. Defaults to False. .. admonition:: Example: @@ -89,7 +89,7 @@ def __init__( if not self._token_proxy_endpoint: raise ValueError( "use_token_proxy is True, but no token proxy endpoint was found. " - f"Ensure the {EnvironmentVariables.AZURE_KUBERNETES_TOKEN_PROXY} environment variable is set." + f"Ensure that the {EnvironmentVariables.AZURE_KUBERNETES_TOKEN_PROXY} environment variable is set." ) self._sni = os.environ.get(EnvironmentVariables.AZURE_KUBERNETES_SNI_NAME) @@ -98,7 +98,7 @@ def __init__( if ca_file and ca_data: raise ValueError( - "Both AZURE_KUBERNETES_CA_FILE and AZURE_KUBERNETES_CA_DATA are set. Only one should be set" + "Both AZURE_KUBERNETES_CA_FILE and AZURE_KUBERNETES_CA_DATA are set. Only one should be set." ) transport = _get_transport( @@ -112,7 +112,8 @@ def __init__( kwargs["transport"] = transport else: raise ValueError( - "Async transport creation failed. Ensure that the aiohttp package is installed to enable token proxy usage." + "Async transport creation failed. Ensure that the aiohttp package is installed to enable token " + "proxy usage in this credential." ) super().__init__( diff --git a/sdk/identity/azure-identity/tests/proxy_server.py b/sdk/identity/azure-identity/tests/proxy_server.py index 918d0a674c8b..fb21e04c1bab 100644 --- a/sdk/identity/azure-identity/tests/proxy_server.py +++ b/sdk/identity/azure-identity/tests/proxy_server.py @@ -134,7 +134,6 @@ def _send_json_response(self, data, status_code=200): self.send_response(status_code) self.send_header("Content-Type", "application/json") self.send_header("Content-Length", str(len(response_json))) - self.send_header("Server", "token-proxy-test-server") self.end_headers() self.wfile.write(response_json.encode("utf-8")) diff --git a/sdk/identity/azure-identity/tests/test_http_client_transport.py b/sdk/identity/azure-identity/tests/test_http_client_transport.py index 95f40235defc..f96b9fa71de7 100644 --- a/sdk/identity/azure-identity/tests/test_http_client_transport.py +++ b/sdk/identity/azure-identity/tests/test_http_client_transport.py @@ -2,6 +2,7 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT License. # ------------------------------------ +# cspell:ignore ests import os import ssl import tempfile diff --git a/sdk/identity/azure-identity/tests/test_workload_identity_credential_async.py b/sdk/identity/azure-identity/tests/test_workload_identity_credential_async.py index 42913b5c9f76..7a33656d617a 100644 --- a/sdk/identity/azure-identity/tests/test_workload_identity_credential_async.py +++ b/sdk/identity/azure-identity/tests/test_workload_identity_credential_async.py @@ -2,7 +2,7 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT License. # ------------------------------------ -# cspell:ignore cafile +# cspell:ignore cafile aexit ests import os import tempfile from time import sleep as real_sleep From 0e08bbf326f7fb7a33db76de2f9f437eab983922 Mon Sep 17 00:00:00 2001 From: Paul Van Eck Date: Wed, 29 Oct 2025 22:24:15 +0000 Subject: [PATCH 04/12] Update changelog Signed-off-by: Paul Van Eck --- sdk/identity/azure-identity/CHANGELOG.md | 4 +++- sdk/identity/azure-identity/azure/identity/_version.py | 2 +- sdk/identity/azure-identity/pyproject.toml | 4 ++-- 3 files changed, 6 insertions(+), 4 deletions(-) diff --git a/sdk/identity/azure-identity/CHANGELOG.md b/sdk/identity/azure-identity/CHANGELOG.md index 4928a5f84707..c0b31ed4f9b0 100644 --- a/sdk/identity/azure-identity/CHANGELOG.md +++ b/sdk/identity/azure-identity/CHANGELOG.md @@ -1,9 +1,11 @@ # Release History -## 1.25.2 (Unreleased) +## 1.26.0b1 (Unreleased) ### Features Added +- Added support for WorkloadIdentityCredential identity binding mode in AKS. ([#43287](https://github.com/Azure/azure-sdk-for-python/pull/43287)) + ### Breaking Changes ### Bugs Fixed diff --git a/sdk/identity/azure-identity/azure/identity/_version.py b/sdk/identity/azure-identity/azure/identity/_version.py index f80f2f3644d7..92c78eb1a721 100644 --- a/sdk/identity/azure-identity/azure/identity/_version.py +++ b/sdk/identity/azure-identity/azure/identity/_version.py @@ -2,4 +2,4 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT License. # ------------------------------------ -VERSION = "1.25.2" +VERSION = "1.26.0b1" diff --git a/sdk/identity/azure-identity/pyproject.toml b/sdk/identity/azure-identity/pyproject.toml index 5d61880b9f88..c9fa89cb338a 100644 --- a/sdk/identity/azure-identity/pyproject.toml +++ b/sdk/identity/azure-identity/pyproject.toml @@ -12,7 +12,7 @@ keywords = ["azure", "azure sdk"] requires-python = ">=3.9" license = "MIT" classifiers = [ - "Development Status :: 5 - Production/Stable", + "Development Status :: 4 - Beta", "Programming Language :: Python", "Programming Language :: Python :: 3 :: Only", "Programming Language :: Python :: 3", @@ -47,4 +47,4 @@ pytyped = ["py.typed"] [tool.azure-sdk-build] pyright = false verifytypes = true -black = true \ No newline at end of file +black = true From 321d54568eeb3ae4c80bfd567c1b0680f252adc0 Mon Sep 17 00:00:00 2001 From: Paul Van Eck Date: Thu, 30 Oct 2025 02:27:55 +0000 Subject: [PATCH 05/12] Use RequestsTransports instead Signed-off-by: Paul Van Eck --- .../pipeline/transport/_requests_basic.py | 1 + .../_credentials/workload_identity.py | 74 +- .../_internal/http_client_transport.py | 259 ------- .../tests/test_http_client_transport.py | 731 ------------------ .../test_workload_identity_credential.py | 543 ++++++++++++- ...test_workload_identity_credential_async.py | 24 +- 6 files changed, 590 insertions(+), 1042 deletions(-) delete mode 100644 sdk/identity/azure-identity/azure/identity/_internal/http_client_transport.py delete mode 100644 sdk/identity/azure-identity/tests/test_http_client_transport.py diff --git a/sdk/core/azure-core/azure/core/pipeline/transport/_requests_basic.py b/sdk/core/azure-core/azure/core/pipeline/transport/_requests_basic.py index 9f102f4b0b20..44e150eddf7a 100644 --- a/sdk/core/azure-core/azure/core/pipeline/transport/_requests_basic.py +++ b/sdk/core/azure-core/azure/core/pipeline/transport/_requests_basic.py @@ -270,6 +270,7 @@ def __init__(self, **kwargs) -> None: self._use_env_settings = kwargs.pop("use_env_settings", True) # See https://github.com/Azure/azure-sdk-for-python/issues/25640 to understand why we track this self._has_been_opened = False + print('---done init---') def __enter__(self) -> "RequestsTransport": self.open() diff --git a/sdk/identity/azure-identity/azure/identity/_credentials/workload_identity.py b/sdk/identity/azure-identity/azure/identity/_credentials/workload_identity.py index e4b59b0e8c26..6216f021457c 100644 --- a/sdk/identity/azure-identity/azure/identity/_credentials/workload_identity.py +++ b/sdk/identity/azure-identity/azure/identity/_credentials/workload_identity.py @@ -4,12 +4,13 @@ # ------------------------------------ import os import time +import ssl from typing import Any from typing import Optional from .client_assertion import ClientAssertionCredential from .._constants import EnvironmentVariables -from .._internal.http_client_transport import HttpClientTransport +from .._internal.token_binding_transport_mixin import TokenBindingTransportMixin WORKLOAD_CONFIG_ERROR = ( @@ -121,12 +122,21 @@ def __init__( "Both AZURE_KUBERNETES_CA_FILE and AZURE_KUBERNETES_CA_DATA are set. Only one should be set." ) - kwargs["transport"] = HttpClientTransport( + transport = _get_transport( sni=sni, - proxy_endpoint=token_proxy_endpoint, + token_proxy_endpoint=token_proxy_endpoint, ca_file=ca_file, ca_data=ca_data, ) + + if transport: + kwargs["transport"] = transport + else: + raise ValueError( + "Transport creation failed. Ensure that the requests package is installed to enable token " + "proxy usage in this credential." + ) + super(WorkloadIdentityCredential, self).__init__( tenant_id=tenant_id, client_id=client_id, @@ -134,3 +144,61 @@ def __init__( token_file_path=token_file_path, **kwargs, ) + + +def _get_transport(sni, token_proxy_endpoint, ca_file, ca_data): + try: + from azure.core.pipeline.transport import ( # pylint: disable=non-abstract-transport-import, no-name-in-module + RequestsTransport, + ) + from requests.adapters import HTTPAdapter + from requests import Session + + class SNIAdapter(HTTPAdapter): + """A custom HTTPAdapter that allows setting a custom SNI hostname.""" + + def __init__(self, server_hostname, ca_data, **kwargs): + self.server_hostname = server_hostname + self.ca_data = ca_data + super().__init__(**kwargs) + + def init_poolmanager(self, connections, maxsize, block=False, **pool_kwargs): + if self.server_hostname: + pool_kwargs["server_hostname"] = self.server_hostname + pool_kwargs["ssl_context"] = ssl.create_default_context(cadata=self.ca_data) + super().init_poolmanager(connections, maxsize, block, **pool_kwargs) + + class CustomRequestsTransport(TokenBindingTransportMixin, RequestsTransport): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self._create_session() + + def _create_session(self): + if self.session: # pylint: disable=access-member-before-definition + self.session.close() # pylint: disable=access-member-before-definition + + self.session = Session() + adapter = SNIAdapter(self._sni, self._ca_data) + self.session.mount("https://", adapter) + + def send(self, request, **kwargs): + self._update_request_url(request) + + # Check if CA file has changed and reload ca_data if needed + if self._ca_file and self._has_ca_file_changed(): + self._load_ca_file_to_data() + # If ca_data was updated, recreate SSL context with the new data + if self._ca_data: + self._create_session() + return super().send(request, **kwargs) + + transport = CustomRequestsTransport( + sni=sni, + proxy_endpoint=token_proxy_endpoint, + ca_file=ca_file, + ca_data=ca_data, + ) + + except ImportError: + return None + return transport diff --git a/sdk/identity/azure-identity/azure/identity/_internal/http_client_transport.py b/sdk/identity/azure-identity/azure/identity/_internal/http_client_transport.py deleted file mode 100644 index 1737820b85e8..000000000000 --- a/sdk/identity/azure-identity/azure/identity/_internal/http_client_transport.py +++ /dev/null @@ -1,259 +0,0 @@ -# ------------------------------------ -# Copyright (c) Microsoft Corporation. -# Licensed under the MIT License. -# ------------------------------------ -# cspell:ignore cafile -from json import loads, dumps -import http.client -import ssl -import urllib.parse -from typing import Any, Iterator, MutableMapping, Optional - -from azure.core.configuration import ConnectionConfiguration -from azure.core.exceptions import ServiceRequestError -from azure.core.rest import HttpRequest, HttpResponse -from azure.core.pipeline.transport import HttpTransport - -from .token_binding_transport_mixin import TokenBindingTransportMixin - - -class HttpClientTransportResponse(HttpResponse): - """Create a HttpResponse from an http.client response. - - :param HttpRequest request: The request. - :type request: ~azure.core.pipeline.transport.HttpRequest - :param httpclient_response: The response from http.client - :type httpclient_response: http.client.HTTPResponse - :param block_size: The block size to use for downloading the response content. - :type block_size: int - """ - - def __init__( - self, request: HttpRequest, httpclient_response: http.client.HTTPResponse, block_size: Optional[int] = None - ) -> None: - self._request = request - self._httpclient_response = httpclient_response - self._block_size = block_size or 4096 - self._data: Optional[bytes] = None - self._closed = False - self._headers = {k.lower(): v for k, v in httpclient_response.getheaders()} - self._content_type = self._headers.get("content-type") - self._encoding: Optional[str] = None - - def __enter__(self) -> "HttpClientTransportResponse": - return self - - def __exit__(self, *args: Any) -> None: - self.close() - - def close(self) -> None: - if not self._closed: - self._httpclient_response.close() - self._closed = True - - def read(self) -> bytes: - if self._data is None: - self._data = self._httpclient_response.read() - return self._data - - def iter_raw(self, **kwargs: Any) -> Iterator[bytes]: - if self._data: - yield self._data - else: - chunk = self._httpclient_response.read(self._block_size) - while chunk: - yield chunk - chunk = self._httpclient_response.read(self._block_size) - - def iter_bytes(self, **kwargs: Any) -> Iterator[bytes]: - # http.client doesn't support compressed encoding automatically, - # so the decompression is already done at read time. - # Just use iter_raw here - yield from self.iter_raw(**kwargs) - - @property - def request(self) -> HttpRequest: - return self._request - - @property - def status_code(self) -> int: - return self._httpclient_response.status - - @property - def headers(self) -> MutableMapping[str, str]: - return self._headers - - @property - def reason(self) -> str: - return self._httpclient_response.reason - - @property - def content_type(self) -> Optional[str]: - return self._content_type - - @property - def url(self) -> str: - return self._request.url - - @property - def is_closed(self) -> bool: - return self._closed - - @property - def is_stream_consumed(self) -> bool: - return self._data is not None - - @property - def encoding(self) -> Optional[str]: - return self._encoding - - @encoding.setter - def encoding(self, value: Optional[str]) -> None: - self._encoding = value - - @property - def content(self) -> bytes: - return self.read() - - def text(self, encoding: Optional[str] = None) -> str: - if encoding is None: - encoding = self.encoding or "utf-8" - return self.content.decode(encoding) - - def raise_for_status(self) -> None: - if self.status_code >= 400: - from azure.core.exceptions import HttpResponseError - - raise HttpResponseError(response=self) - - def json(self) -> Any: - return loads(self.text()) - - -class SniSSLContext(ssl.SSLContext): - def __new__(cls, sni_hostname: str, protocol=None): - if protocol is None: - protocol = ssl.PROTOCOL_TLS_CLIENT - instance = super().__new__(cls, protocol=protocol) - instance.sni_hostname = sni_hostname # type: ignore - return instance - - def wrap_socket(self, *args, **kwargs): - kwargs["server_hostname"] = self.sni_hostname # type: ignore - return super().wrap_socket(*args, **kwargs) - - -class HttpClientTransport(TokenBindingTransportMixin, HttpTransport): - """Implements an HTTP sender using Python's built-in http.client library.""" - - def __init__(self, **kwargs) -> None: - self.connection_config = ConnectionConfiguration(**kwargs) - self._ssl_context: Optional[ssl.SSLContext] = None - super().__init__(**kwargs) - # Initialize SSL context if we have CA data - if self._ca_data: - self._ssl_context = self._create_ssl_context() - - def __enter__(self) -> "HttpClientTransport": - self.open() - return self - - def __exit__(self, *args): - self.close() - - def open(self) -> None: - pass # We create connections as needed - - def close(self) -> None: - pass # No persistent connections to close - - def _create_ssl_context(self) -> ssl.SSLContext: - # Create SSL context using current CA data or file. - ssl_context: ssl.SSLContext - if self._sni: - ssl_context = SniSSLContext(self._sni, ssl.PROTOCOL_TLS_CLIENT) - ssl_context.verify_mode = ssl.CERT_REQUIRED - ssl_context.check_hostname = True - - if self._ca_data: - ssl_context.load_verify_locations(cadata=self._ca_data) - else: - ssl_context.load_default_certs() - else: - ssl_context = ssl.create_default_context(cadata=self._ca_data) - - if not self.connection_config.verify: - ssl_context.check_hostname = False - ssl_context.verify_mode = ssl.CERT_NONE - - return ssl_context - - def _get_connection(self, host: str) -> http.client.HTTPSConnection: - # Check if CA file has changed and reload ca_data if needed - if self._ca_file and self._has_ca_file_changed(): - self._load_ca_file_to_data() - # If ca_data was updated, recreate SSL context with the new data - if self._ca_data: - self._ssl_context = self._create_ssl_context() - - # Use cached SSL context or create a new one - ssl_context = self._ssl_context or self._create_ssl_context() - - connection = http.client.HTTPSConnection( # nosec - host, - timeout=self.connection_config.timeout, - context=ssl_context, - ) - - return connection - - def send(self, request: HttpRequest, **kwargs) -> HttpResponse: - """Send request object according to configuration. - - :param request: The HTTP request object. - :type request: ~azure.core.rest.HttpRequest - :return: The HTTP response object. - :rtype: ~azure.core.rest.HttpResponse - """ - - # Get or create a connection for the URL - self._update_request_url(request) - parsed_url = urllib.parse.urlparse(request.url) - full_path = urllib.parse.urlunparse( - ("", "", parsed_url.path, parsed_url.params, parsed_url.query, parsed_url.fragment) - ) - - # Get connection timeout - connection_timeout = kwargs.pop("connection_timeout", self.connection_config.timeout) - - try: - # Get connection - connection = self._get_connection(parsed_url.netloc) - - if connection_timeout is not None: - connection.timeout = connection_timeout - connection.request( - request.method, - full_path, - body=dumps(request.data) if isinstance(request.data, (dict, list)) else request.data, - headers=request.headers, - ) - response = connection.getresponse() - transport_response = HttpClientTransportResponse( - request=request, - httpclient_response=response, - block_size=self.connection_config.data_block_size, - ) - - connection.close() - return transport_response - - except http.client.HTTPException as err: - raise ServiceRequestError(err) from err - except ssl.SSLError as err: - raise ServiceRequestError(err) from err - except Exception as err: - raise ServiceRequestError(err) from err - - def __repr__(self) -> str: - return f"<{type(self).__name__}>" diff --git a/sdk/identity/azure-identity/tests/test_http_client_transport.py b/sdk/identity/azure-identity/tests/test_http_client_transport.py deleted file mode 100644 index f96b9fa71de7..000000000000 --- a/sdk/identity/azure-identity/tests/test_http_client_transport.py +++ /dev/null @@ -1,731 +0,0 @@ -# ------------------------------------ -# Copyright (c) Microsoft Corporation. -# Licensed under the MIT License. -# ------------------------------------ -# cspell:ignore ests -import os -import ssl -import tempfile -import time -from time import sleep as real_sleep -from unittest import mock - -import pytest - -from azure.core.rest import HttpRequest -from azure.core.exceptions import ServiceRequestError, ServiceResponseError -from azure.identity._internal import http_client_transport as transport -from azure.identity._internal.http_client_transport import HttpClientTransport, SniSSLContext - -from proxy_server import TokenProxyTestServer - - -PEM_CERT_PATH = os.path.join(os.path.dirname(__file__), "certificate.pem") - - -@pytest.fixture(scope="module") -def ca_data() -> str: - """Read CA certificate data from a PEM file for testing.""" - with open(PEM_CERT_PATH, "r", encoding="utf-8") as f: - return f.read() - - -class TestHttpClientTransport: - """Test cases for HttpClientTransport class.""" - - def test_init_basic(self): - """Test basic initialization of HttpClientTransport.""" - transport = HttpClientTransport() - assert transport._ca_data is None - assert transport._ca_file is None - assert transport._sni is None - assert transport._proxy_endpoint is None - assert transport._ca_file_mtime is None - - def test_init_with_ca_data(self, ca_data): - """Test initialization with CA data.""" - transport = HttpClientTransport(ca_data=ca_data) - assert transport._ca_data == ca_data - assert transport._ca_file is None - - def test_init_with_ca_file(self, ca_data): - """Test initialization with CA file.""" - - transport = HttpClientTransport(ca_file=PEM_CERT_PATH) - assert transport._ca_file == PEM_CERT_PATH - assert transport._ca_data == ca_data - assert transport._ca_file_mtime is not None - - def test_init_with_both_ca_file_and_data_raises_error(self, ca_data): - """Test that providing both CA file and data raises an error.""" - with pytest.raises(ValueError, match="Both ca_file and ca_data are set"): - HttpClientTransport(ca_file=PEM_CERT_PATH, ca_data=ca_data) - - def test_init_with_empty_ca_file_raises_error(self): - """Test that empty CA file raises an error.""" - with tempfile.NamedTemporaryFile(mode="w", delete=False) as temp_file: - temp_file_path = temp_file.name - - try: - with pytest.raises(ValueError, match="CA file .* is empty"): - HttpClientTransport(ca_file=temp_file_path) - finally: - os.unlink(temp_file_path) - - def test_init_with_sni(self): - """Test initialization with SNI hostname.""" - sni_hostname = "example.com" - transport = HttpClientTransport(sni=sni_hostname) - assert transport._sni == sni_hostname - - def test_init_with_proxy_endpoint(self): - """Test initialization with proxy endpoint.""" - proxy_endpoint = "https://proxy.example.com:8080" - transport = HttpClientTransport(proxy_endpoint=proxy_endpoint) - assert transport._proxy_endpoint == proxy_endpoint - - def test_validate_url_valid_https(self): - """Test URL validation with valid HTTPS URL.""" - transport = HttpClientTransport() - # Should not raise any exception - transport._validate_url("https://example.com/path") - - def test_validate_url_non_https_scheme(self): - """Test URL validation rejects non-HTTPS schemes.""" - transport = HttpClientTransport() - with pytest.raises(ValueError, match="must use the 'https' scheme"): - transport._validate_url("http://example.com") - - def test_validate_url_with_user_info(self): - """Test URL validation rejects URLs with user info.""" - transport = HttpClientTransport() - with pytest.raises(ValueError, match="must not contain username or password"): - transport._validate_url("https://user:pass@example.com") - - def test_validate_url_with_fragment(self): - """Test URL validation rejects URLs with fragments.""" - transport = HttpClientTransport() - with pytest.raises(ValueError, match="must not contain a fragment"): - transport._validate_url("https://example.com#fragment") - - def test_validate_url_with_query(self): - """Test URL validation rejects URLs with query parameters.""" - transport = HttpClientTransport() - with pytest.raises(ValueError, match="must not contain query parameters"): - transport._validate_url("https://example.com?query=value") - - def test_ca_file_tracking_updates_mtime(self, ca_data): - """Test CA file tracking updates modification time.""" - transport = HttpClientTransport(ca_file=PEM_CERT_PATH) - assert transport._ca_file_mtime == os.path.getmtime(PEM_CERT_PATH) - - def test_ca_file_change_detection_no_change(self): - """Test CA file change detection when file hasn't changed.""" - transport = HttpClientTransport(ca_file=PEM_CERT_PATH) - assert not transport._has_ca_file_changed() - - def test_ca_file_change_detection_content_changed(self, ca_data): - """Test CA file change detection when file content has changed.""" - - with tempfile.NamedTemporaryFile(mode="w", delete=False) as temp_file: - temp_file.write(ca_data) - temp_file_path = temp_file.name - - try: - transport = HttpClientTransport(ca_file=temp_file_path) - - # Modify the file - real_sleep(0.1) # Ensure mtime changes - with open(temp_file_path, "a") as f: - f.write("\n") - - # File should be detected as changed - assert transport._has_ca_file_changed() - finally: - os.unlink(temp_file_path) - - def test_ca_file_change_detection_file_deleted(self, ca_data): - """Test CA file change detection when file is deleted.""" - - with tempfile.NamedTemporaryFile(mode="w", delete=False) as temp_file: - temp_file.write(ca_data) - temp_file_path = temp_file.name - - transport = HttpClientTransport(ca_file=temp_file_path) - - # Delete the file - os.unlink(temp_file_path) - - # File deletion should be detected as a change - assert transport._has_ca_file_changed() - - def test_ca_file_empty_during_rotation(self, ca_data): - """Test CA file becoming empty during rotation with existing connection.""" - - with tempfile.NamedTemporaryFile(mode="w", delete=False) as temp_file: - temp_file.write(ca_data) - temp_file_path = temp_file.name - - try: - transport = HttpClientTransport(ca_file=temp_file_path) - original_mtime = transport._ca_file_mtime - - # Make file empty - with open(temp_file_path, "w") as f: - f.write("") - - # Should not raise error and should preserve old data - assert transport._ca_data == ca_data - assert transport._ca_file_mtime == original_mtime - finally: - os.unlink(temp_file_path) - - def test_update_request_url_no_proxy(self): - """Test request URL update with no proxy.""" - transport = HttpClientTransport() - request = HttpRequest("GET", "https://example.com/path?query=value") - original_url = request.url - - transport._update_request_url(request) - - # URL should remain unchanged - assert request.url == original_url - - def test_update_request_url_with_proxy(self): - """Test request URL update with proxy.""" - proxy_endpoint = "https://proxy.example.com:8080/proxy" - transport = HttpClientTransport(proxy_endpoint=proxy_endpoint) - request = HttpRequest("GET", "https://original.com/api/endpoint?query=value") - - transport._update_request_url(request) - - # URL should be updated to use proxy - expected_url = "https://proxy.example.com:8080/proxy/api/endpoint?query=value" - assert request.url == expected_url - - def test_update_request_url_proxy_path_combination(self): - """Test request URL update with proxy that has a path.""" - proxy_endpoint = "https://proxy.example.com/service" - transport = HttpClientTransport(proxy_endpoint=proxy_endpoint) - request = HttpRequest("GET", "https://original.com/oauth2/v2.0/token") - - transport._update_request_url(request) - - # Paths should be combined correctly - expected_url = "https://proxy.example.com/service/oauth2/v2.0/token" - assert request.url == expected_url - - def test_context_manager(self): - """Test HttpClientTransport as context manager.""" - transport = HttpClientTransport() - - with transport as t: - assert t is transport - - # Should be able to use after context manager - assert transport is not None - - def test_repr(self): - """Test string representation of HttpClientTransport.""" - transport = HttpClientTransport() - repr_str = repr(transport) - assert "HttpClientTransport" in repr_str - - -class TestSniSSLContext: - """Test cases for SniSSLContext class.""" - - def test_init(self): - """Test SNI SSL context initialization.""" - hostname = "example.com" - context = SniSSLContext(hostname, ssl.PROTOCOL_TLS_CLIENT) - assert context.sni_hostname == hostname - - def test_wrap_socket_adds_server_hostname(self): - """Test that wrap_socket adds server_hostname parameter.""" - hostname = "example.com" - context = SniSSLContext(hostname, ssl.PROTOCOL_TLS_CLIENT) - - # Mock the parent wrap_socket method - with mock.patch.object(ssl.SSLContext, "wrap_socket") as mock_wrap: - mock_sock = mock.Mock() - context.wrap_socket(mock_sock) - - # Verify server_hostname was added to kwargs - mock_wrap.assert_called_once_with(mock_sock, server_hostname=hostname) - - -class TestHttpClientTransportResponse: - """Test cases for HttpClientTransportResponse class.""" - - def test_response_properties(self): - """Test basic response properties.""" - request = HttpRequest("GET", "https://example.com/test") - - mock_http_response = mock.Mock() - mock_http_response.status = 200 - mock_http_response.reason = "OK" - mock_http_response.getheaders.return_value = [("Content-Type", "application/json"), ("Content-Length", "100")] - mock_http_response.read.return_value = b'{"test": "data"}' - - response = transport.HttpClientTransportResponse(request, mock_http_response) - - assert response.status_code == 200 - assert response.reason == "OK" - assert response.headers["content-type"] == "application/json" - assert response.headers["content-length"] == "100" - assert response.content == b'{"test": "data"}' - assert response.text() == '{"test": "data"}' - assert response.json() == {"test": "data"} - assert response.url == "https://example.com/test" - assert not response.is_closed - # Stream is consumed after calling .content, .text(), or .json() - assert response.is_stream_consumed - - def test_response_context_manager(self): - """Test response as context manager.""" - request = HttpRequest("GET", "https://example.com/test") - mock_http_response = mock.Mock() - mock_http_response.status = 200 - mock_http_response.getheaders.return_value = [] - - response = transport.HttpClientTransportResponse(request, mock_http_response) - - with response as r: - assert r is response - assert not r.is_closed - - assert response.is_closed - mock_http_response.close.assert_called_once() - - def test_response_raise_for_status_success(self): - """Test raise_for_status with successful response.""" - request = HttpRequest("GET", "https://example.com/test") - mock_http_response = mock.Mock() - mock_http_response.status = 200 - mock_http_response.getheaders.return_value = [] - - response = transport.HttpClientTransportResponse(request, mock_http_response) - - # Should not raise any exception - response.raise_for_status() - - def test_response_raise_for_status_error(self): - """Test raise_for_status with error response.""" - from azure.core.exceptions import HttpResponseError - - request = HttpRequest("GET", "https://example.com/test") - mock_http_response = mock.Mock() - mock_http_response.status = 404 - mock_http_response.getheaders.return_value = [] - - response = transport.HttpClientTransportResponse(request, mock_http_response) - - with pytest.raises(HttpResponseError): - response.raise_for_status() - - def test_response_iter_raw(self): - """Test response iter_raw method.""" - request = HttpRequest("GET", "https://example.com/test") - mock_http_response = mock.Mock() - mock_http_response.status = 200 - mock_http_response.getheaders.return_value = [] - mock_http_response.read.side_effect = [b"chunk1", b"chunk2", b""] - - response = transport.HttpClientTransportResponse(request, mock_http_response, block_size=6) - - chunks = list(response.iter_raw()) - assert chunks == [b"chunk1", b"chunk2"] - - def test_response_iter_bytes(self): - """Test response iter_bytes method.""" - request = HttpRequest("GET", "https://example.com/test") - mock_http_response = mock.Mock() - mock_http_response.status = 200 - mock_http_response.getheaders.return_value = [] - mock_http_response.read.side_effect = [b"chunk1", b"chunk2", b""] - - response = transport.HttpClientTransportResponse(request, mock_http_response, block_size=6) - - chunks = list(response.iter_bytes()) - assert chunks == [b"chunk1", b"chunk2"] - - def test_response_encoding(self): - """Test response encoding property.""" - request = HttpRequest("GET", "https://example.com/test") - mock_http_response = mock.Mock() - mock_http_response.status = 200 - mock_http_response.getheaders.return_value = [] - mock_http_response.read.return_value = b"\xc3\xa9" # é in UTF-8 - - response = transport.HttpClientTransportResponse(request, mock_http_response) - - # Default encoding (UTF-8) - assert response.text() == "é" - - # Create a new response for testing latin-1 encoding - mock_http_response2 = mock.Mock() - mock_http_response2.status = 200 - mock_http_response2.getheaders.return_value = [] - mock_http_response2.read.return_value = b"\xe9" # é in Latin-1 - - response2 = transport.HttpClientTransportResponse(request, mock_http_response2) - response2.encoding = "latin-1" - assert response2.text() == "é" - - # Test explicit encoding parameter overrides response encoding - response2.encoding = "utf-8" - assert response2.text(encoding="latin-1") == "é" - - -class TestHttpClientTransportWithLocalServer: - """Integration tests using a local test server.""" - - def test_basic_https_request(self): - """Test basic HTTPS request to test server.""" - with TokenProxyTestServer(use_ssl=True) as server: - # Create transport with server's CA certificate - transport = HttpClientTransport(ca_file=server.ca_file) - request = HttpRequest("GET", f"{server.base_url}/health") - - response = transport.send(request) - - assert response.status_code == 200 - data = response.json() - assert data["status"] == "healthy" - assert "timestamp" in data - - def test_post_request_with_body(self): - """Test POST request with request body.""" - with TokenProxyTestServer(use_ssl=True) as server: - transport = HttpClientTransport(ca_file=server.ca_file) - - # Prepare OAuth-like request - body = "grant_type=client_credentials&scope=https://graph.microsoft.com/.default" - request = HttpRequest( - "POST", - f"{server.base_url}/tenant/oauth2/v2.0/token", - headers={"Content-Type": "application/x-www-form-urlencoded"}, - data=body.encode("utf-8"), - ) - - response = transport.send(request) - - assert response.status_code == 200 - data = response.json() - assert "access_token" in data - assert data["token_type"] == "Bearer" - assert data["expires_in"] == 3600 - - def test_proxy_endpoint_comprehensive(self): - """Test comprehensive proxy endpoint functionality with various HTTP methods and scenarios.""" - with TokenProxyTestServer(use_ssl=True) as server: - # Configure transport with proxy endpoint - transport = HttpClientTransport(proxy_endpoint=server.base_url, ca_file=server.ca_file) - - # Test 1: POST request with JSON body through proxy - post_data = {"grant_type": "client_credentials", "scope": "https://graph.microsoft.com/.default"} - post_request = HttpRequest( - "POST", - "https://login.microsoftonline.com/tenant/oauth2/v2.0/token2", - headers={"Content-Type": "application/json"}, - json=post_data, - ) - - post_response = transport.send(post_request) - assert post_response.status_code == 200 - post_data_response = post_response.json() - assert post_data_response["method"] == "POST" - assert post_data_response["proxied_path"] == "/tenant/oauth2/v2.0/token2" - - # Test 2: PUT request through proxy - put_request = HttpRequest( - "PUT", - "https://graph.microsoft.com/v1.0/me/profile", - headers={"Content-Type": "application/json"}, - json={"displayName": "Test User"}, - ) - - put_response = transport.send(put_request) - assert put_response.status_code == 200 - put_data_response = put_response.json() - assert put_data_response["method"] == "PUT" - assert put_data_response["proxied_path"] == "/v1.0/me/profile" - - # Test 3: DELETE request through proxy - delete_request = HttpRequest("DELETE", "https://graph.microsoft.com/v1.0/applications/app-id") - - delete_response = transport.send(delete_request) - assert delete_response.status_code == 200 - delete_data_response = delete_response.json() - assert delete_data_response["method"] == "DELETE" - assert delete_data_response["proxied_path"] == "/v1.0/applications/app-id" - - # Test 4: Complex URL with multiple path segments and query parameters - complex_url = "https://management.azure.com/subscriptions/sub-id/resourceGroups/rg/providers/Microsoft.Storage/storageAccounts/account?api-version=2021-04-01&expand=properties" - complex_request = HttpRequest("GET", complex_url) - - complex_response = transport.send(complex_request) - assert complex_response.status_code == 200 - complex_data_response = complex_response.json() - assert complex_data_response["method"] == "GET" - expected_path = "/subscriptions/sub-id/resourceGroups/rg/providers/Microsoft.Storage/storageAccounts/account?api-version=2021-04-01&expand=properties" - assert complex_data_response["proxied_path"] == expected_path - - # Test 5: Request with custom headers through proxy - headers_request = HttpRequest( - "GET", - "https://vault.azure.net/secrets/test-secret?api-version=7.3", - headers={ - "Authorization": "Bearer test-token", - "X-Custom-Header": "proxy-test-value", - "User-Agent": "Azure-SDK-For-Python", - }, - ) - - headers_response = transport.send(headers_request) - assert headers_response.status_code == 200 - headers_data_response = headers_response.json() - assert headers_data_response["method"] == "GET" - assert headers_data_response["proxied_path"] == "/secrets/test-secret?api-version=7.3" - - # Verify headers were forwarded through proxy - received_headers = headers_data_response["headers_received"] - assert "Authorization" in received_headers - assert "X-Custom-Header" in received_headers - assert received_headers["Authorization"] == "Bearer test-token" - assert received_headers["X-Custom-Header"] == "proxy-test-value" - - def test_sni_with_custom_hostname(self): - """Test SNI (Server Name Indication) with custom hostname.""" - with TokenProxyTestServer(use_ssl=True) as server: - # Use SNI with a different hostname than the server - transport = HttpClientTransport(sni="1234.ests.aks", ca_file=server.ca_file) - - request = HttpRequest("GET", f"{server.base_url}/health") - - response = transport.send(request) - - assert response.status_code == 200 - data = response.json() - assert data["status"] == "healthy" - - def test_connection_reuse(self): - """Test that connections are reused for multiple requests.""" - with TokenProxyTestServer(use_ssl=True) as server: - transport = HttpClientTransport(ca_file=server.ca_file) - - # Make multiple requests - requests_data = [] - for i in range(3): - request = HttpRequest("GET", f"{server.base_url}/health") - response = transport.send(request) - assert response.status_code == 200 - requests_data.append(response.json()) - - # All requests should succeed - assert len(requests_data) == 3 - for data in requests_data: - assert data["status"] == "healthy" - - def test_ca_file_change_detection(self): - """Test CA file change detection with real certificates.""" - with TokenProxyTestServer(use_ssl=True) as server: - # Create a copy of the CA file that we can modify - ca_file = server.ca_file - if ca_file is None: - pytest.skip("CA file not available") - - assert ca_file is not None # Type hint for mypy - with tempfile.NamedTemporaryFile(mode="w", delete=False, suffix=".pem") as temp_ca: - with open(ca_file, "r") as original: - original_content = original.read() - temp_ca.write(original_content) - temp_ca_path = temp_ca.name - - try: - transport = HttpClientTransport(ca_file=temp_ca_path) - - # First request should work - request = HttpRequest("GET", f"{server.base_url}/health") - response1 = transport.send(request) - assert response1.status_code == 200 - - # Modify the CA file (add some content) - real_sleep(0.1) - with open(temp_ca_path, "a") as f: - f.write("\n# Modified for testing\n") - - # Second request should still work (using the same cert content) - response2 = transport.send(request) - assert response2.status_code == 200 - - finally: - os.unlink(temp_ca_path) - - def test_ssl_error_handling(self): - """Test SSL error handling.""" - with TokenProxyTestServer(use_ssl=True) as server: - # Create transport without proper CA file (will cause SSL error) - transport = HttpClientTransport() # No CA file provided - request = HttpRequest("GET", f"{server.base_url}/health") - - # Should raise SSL-related error - with pytest.raises((ServiceRequestError, ServiceResponseError)): - transport.send(request) - - def test_server_error_response(self): - """Test handling of server error responses.""" - with TokenProxyTestServer(use_ssl=True) as server: - transport = HttpClientTransport(ca_file=server.ca_file) - request = HttpRequest("GET", f"{server.base_url}/error/500") - - response = transport.send(request) - - assert response.status_code == 500 - # Should not raise exception, just return error response - - def test_slow_server_response(self): - """Test handling of slow server responses.""" - with TokenProxyTestServer(use_ssl=True) as server: - # Set a longer timeout for this test - transport = HttpClientTransport(ca_file=server.ca_file, timeout=5) - # # Ensure fresh connection by closing any existing ones - # transport.close() - request = HttpRequest("GET", f"{server.base_url}/slow") - - start_time = time.time() - response = transport.send(request) - elapsed_time = time.time() - start_time - - assert response.status_code == 200 - # Should take at least 2 seconds (server waits for 2s) - assert elapsed_time >= 2.0 - data = response.json() - assert data["message"] == "slow response" - - def test_custom_headers_preserved(self): - """Test that custom headers are preserved and sent to server.""" - with TokenProxyTestServer(use_ssl=True) as server: - transport = HttpClientTransport(ca_file=server.ca_file) - - custom_headers = { - "Authorization": "Bearer test-token", - "User-Agent": "HttpClientTransport/1.0", - "X-Custom-Header": "test-value", - } - - request = HttpRequest("GET", f"{server.base_url}/proxy/test", headers=custom_headers) - - response = transport.send(request) - - assert response.status_code == 200 - data = response.json() - - # Server echoes back the headers it received - received_headers = data["headers_received"] - assert "Authorization" in received_headers - assert "User-Agent" in received_headers - assert "X-Custom-Header" in received_headers - assert received_headers["Authorization"] == "Bearer test-token" - - def test_query_parameters_preserved(self): - """Test that query parameters are preserved in proxy requests.""" - with TokenProxyTestServer(use_ssl=True) as server: - transport = HttpClientTransport(proxy_endpoint=server.base_url, ca_file=server.ca_file) - - # Request with query parameters - original_url = "https://example.com/api/data?scope=read&limit=10&format=json" - request = HttpRequest("GET", original_url) - - response = transport.send(request) - - assert response.status_code == 200 - data = response.json() - - # The path should include the query parameters - expected_path = "/api/data?scope=read&limit=10&format=json" - assert data["proxied_path"] == expected_path - - def test_concurrent_requests(self): - """Test handling multiple concurrent requests.""" - import threading - import concurrent.futures - - with TokenProxyTestServer(use_ssl=True) as server: - - transport = HttpClientTransport(ca_file=server.ca_file) - - def make_request(request_id): - request = HttpRequest("GET", f"{server.base_url}/health") - response = transport.send(request) - return request_id, response.status_code, response.json() - - # Make 5 concurrent requests - with concurrent.futures.ThreadPoolExecutor(max_workers=5) as executor: - futures = [executor.submit(make_request, i) for i in range(5)] - results = [future.result() for future in concurrent.futures.as_completed(futures)] - - # All requests should succeed - assert len(results) == 5 - for request_id, status_code, data in results: - assert status_code == 200 - assert data["status"] == "healthy" - - -class TestTokenProxyTestServer: - """Tests for the test server itself.""" - - def test_server_startup_and_shutdown(self): - """Test that server starts and stops properly.""" - server = TokenProxyTestServer(use_ssl=True) - - # Server should not be running initially - assert server.server is None - - # Start server - base_url = server.start() - assert server.server is not None - assert base_url.startswith("https://") - assert str(server.port) in base_url - - # Stop server - server.stop() - - # Should clean up properly - assert len(server._temp_files) == 0 # Files should be cleaned up - - def test_context_manager(self): - """Test using server as context manager.""" - with TokenProxyTestServer(use_ssl=False) as server: - assert server.server is not None - assert server.base_url.startswith("http://") - - # Server should be stopped after context exit - # Note: We can't easily test this without making a request - - def test_certificate_generation(self): - """Test certificate generation.""" - server = TokenProxyTestServer(use_ssl=True) - - try: - server.generate_test_certificates() - - # Should have created certificate files - assert server.cert_file is not None - assert server.key_file is not None - assert server.ca_file is not None - - # Files should exist - assert os.path.exists(server.cert_file) - assert os.path.exists(server.key_file) - assert os.path.exists(server.ca_file) - - # Files should contain certificate data - with open(server.cert_file, "r") as f: - cert_content = f.read() - assert "-----BEGIN CERTIFICATE-----" in cert_content - assert "-----END CERTIFICATE-----" in cert_content - - finally: - server.stop() # Clean up temp files diff --git a/sdk/identity/azure-identity/tests/test_workload_identity_credential.py b/sdk/identity/azure-identity/tests/test_workload_identity_credential.py index 20470b38a536..06ad11f984f4 100644 --- a/sdk/identity/azure-identity/tests/test_workload_identity_credential.py +++ b/sdk/identity/azure-identity/tests/test_workload_identity_credential.py @@ -2,15 +2,29 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT License. # ------------------------------------ +# cspell:ignore cafile ests import os import tempfile +import time from unittest.mock import mock_open, MagicMock, patch import pytest +from azure.core.rest import HttpRequest from azure.identity import WorkloadIdentityCredential -from azure.identity._internal.http_client_transport import HttpClientTransport +from azure.identity._credentials.workload_identity import _get_transport from helpers import mock_response, build_aad_response, GET_TOKEN_METHODS +from proxy_server import TokenProxyTestServer + + +PEM_CERT_PATH = os.path.join(os.path.dirname(__file__), "certificate.pem") + + +@pytest.fixture(scope="module") +def ca_data() -> str: + """Read CA certificate data from a PEM file for testing.""" + with open(PEM_CERT_PATH, "r", encoding="utf-8") as f: + return f.read() def test_workload_identity_credential_initialize(): @@ -52,8 +66,8 @@ def send(request, **kwargs): class TestWorkloadIdentityCredentialTokenProxy: """Test cases for WorkloadIdentityCredential with use_token_proxy=True.""" - def test_use_token_proxy_creates_http_client_transport(self): - """Test that use_token_proxy=True creates HttpClientTransport with correct parameters.""" + def test_use_token_proxy_creates_custom_transport(self): + """Test that use_token_proxy=True creates a custom transport with correct parameters.""" tenant_id = "tenant-id" client_id = "client-id" token_file_path = "foo-path" @@ -68,21 +82,20 @@ def test_use_token_proxy_creates_http_client_transport(self): } with patch.dict(os.environ, env_vars, clear=False): - with patch("azure.identity._credentials.workload_identity.HttpClientTransport") as mock_transport_class: + with patch("azure.identity._credentials.workload_identity._get_transport") as mock_get_transport: mock_transport_instance = MagicMock() - mock_transport_class.return_value = mock_transport_instance + mock_get_transport.return_value = mock_transport_instance - credential = WorkloadIdentityCredential( + WorkloadIdentityCredential( tenant_id=tenant_id, client_id=client_id, token_file_path=token_file_path, use_token_proxy=True, ) - # Verify HttpClientTransport was called with correct parameters - mock_transport_class.assert_called_once_with( + mock_get_transport.assert_called_once_with( sni=sni_hostname, - proxy_endpoint=proxy_endpoint, + token_proxy_endpoint=proxy_endpoint, ca_file=ca_file_path, ca_data=None, ) @@ -101,9 +114,9 @@ def test_use_token_proxy_with_ca_data(self): } with patch.dict(os.environ, env_vars, clear=False): - with patch("azure.identity._credentials.workload_identity.HttpClientTransport") as mock_transport_class: + with patch("azure.identity._credentials.workload_identity._get_transport") as mock_get_transport: mock_transport_instance = MagicMock() - mock_transport_class.return_value = mock_transport_instance + mock_get_transport.return_value = mock_transport_instance WorkloadIdentityCredential( tenant_id=tenant_id, @@ -112,10 +125,9 @@ def test_use_token_proxy_with_ca_data(self): use_token_proxy=True, ) - # Verify HttpClientTransport was called with CA data - mock_transport_class.assert_called_once_with( + mock_get_transport.assert_called_once_with( sni=None, - proxy_endpoint=proxy_endpoint, + token_proxy_endpoint=proxy_endpoint, ca_file=None, ca_data=ca_data, ) @@ -132,9 +144,9 @@ def test_use_token_proxy_minimal_config(self): } with patch.dict(os.environ, env_vars, clear=False): - with patch("azure.identity._credentials.workload_identity.HttpClientTransport") as mock_transport_class: + with patch("azure.identity._credentials.workload_identity._get_transport") as mock_get_transport: mock_transport_instance = MagicMock() - mock_transport_class.return_value = mock_transport_instance + mock_get_transport.return_value = mock_transport_instance WorkloadIdentityCredential( tenant_id=tenant_id, @@ -143,10 +155,9 @@ def test_use_token_proxy_minimal_config(self): use_token_proxy=True, ) - # Verify HttpClientTransport was called with minimal config - mock_transport_class.assert_called_once_with( + mock_get_transport.assert_called_once_with( sni=None, - proxy_endpoint=proxy_endpoint, + token_proxy_endpoint=proxy_endpoint, ca_file=None, ca_data=None, ) @@ -210,7 +221,6 @@ def send(request, **kwargs): assert request.data.get("client_assertion") == assertion return mock_response(json_payload=build_aad_response(access_token=access_token)) - # Mock the transport that would be created by HttpClientTransport mock_transport_instance = MagicMock(send=send) env_vars = { @@ -218,8 +228,8 @@ def send(request, **kwargs): } with patch.dict(os.environ, env_vars, clear=False): - with patch("azure.identity._credentials.workload_identity.HttpClientTransport") as mock_transport_class: - mock_transport_class.return_value = mock_transport_instance + with patch("azure.identity._credentials.workload_identity._get_transport") as mock_get_transport: + mock_get_transport.return_value = mock_transport_instance credential = WorkloadIdentityCredential( tenant_id=tenant_id, @@ -236,18 +246,499 @@ def send(request, **kwargs): open_mock.assert_called_once_with(token_file_path, encoding="utf-8") def test_use_token_proxy_false_does_not_create_transport(self): - """Test that use_token_proxy=False (default) does not create HttpClientTransport.""" + """Test that use_token_proxy=False (default) does not create a custom transport.""" tenant_id = "tenant-id" client_id = "client-id" token_file_path = "foo-path" - with patch("azure.identity._credentials.workload_identity.HttpClientTransport") as mock_transport_class: + with patch("azure.identity._credentials.workload_identity._get_transport") as mock_get_transport: WorkloadIdentityCredential( tenant_id=tenant_id, client_id=client_id, token_file_path=token_file_path, use_token_proxy=False, ) + mock_get_transport.assert_not_called() + + +class TestCustomRequestsTransport: + """Test cases for the custom RequestsTransport used by WorkloadIdentityCredential.""" + + def test_get_transport_creates_custom_requests_transport(self, ca_data): + """Test that _get_transport creates CustomRequestsTransport with correct parameters.""" + sni = "test.sni.com" + proxy_endpoint = "https://proxy.example.com:8080" + ca_file = PEM_CERT_PATH + + transport = _get_transport( + sni=sni, + token_proxy_endpoint=proxy_endpoint, + ca_file=ca_file, + ca_data=None, + ) + + assert transport is not None + assert hasattr(transport, "_sni") + assert hasattr(transport, "_proxy_endpoint") + assert hasattr(transport, "_ca_file") + assert hasattr(transport, "_ca_data") + assert transport._sni == sni + assert transport._proxy_endpoint == proxy_endpoint + assert transport._ca_file == ca_file + assert transport._ca_data == ca_data + + def test_get_transport_with_minimal_config(self): + """Test _get_transport with minimal configuration.""" + proxy_endpoint = "https://proxy.example.com:8080" + + transport = _get_transport( + sni=None, + token_proxy_endpoint=proxy_endpoint, + ca_file=None, + ca_data=None, + ) + + assert transport is not None + assert transport._sni is None + assert transport._proxy_endpoint == proxy_endpoint + assert transport._ca_file is None + assert transport._ca_data is None + + def test_custom_requests_transport_send_with_sni(self): + """Test that CustomRequestsTransport.send works with SNI configuration.""" + sni = "1234.ests.aks" + proxy_endpoint = "https://proxy.example.com:8080" + + transport = _get_transport( + sni=sni, + token_proxy_endpoint=proxy_endpoint, + ca_file=None, + ca_data=None, + ) + assert transport is not None + + # Mock request + mock_request = HttpRequest("POST", "https://login.microsoftonline.com/tenant/oauth2/v2.0/token") + mock_request.data = {"grant_type": "client_credentials"} + + with patch.object(transport.__class__.__bases__[1], "send") as mock_parent_send: + mock_response_obj = MagicMock() + mock_parent_send.return_value = mock_response_obj + + response = transport.send(mock_request) + + assert response == mock_response_obj + mock_parent_send.assert_called_once_with(mock_request) + + def test_custom_requests_transport_send_updates_url(self): + """Test that CustomRequestsTransport.send updates request URL through proxy endpoint.""" + proxy_endpoint = "https://proxy.example.com:8080" + + transport = _get_transport( + sni=None, + token_proxy_endpoint=proxy_endpoint, + ca_file=None, + ca_data=None, + ) + assert transport is not None + + # Mock request with original URL + original_url = "https://login.microsoftonline.com/tenant/oauth2/v2.0/token" + mock_request = HttpRequest("POST", original_url) + + with patch.object(transport.__class__.__bases__[1], "send") as mock_parent_send: + mock_response_obj = MagicMock() + mock_parent_send.return_value = mock_response_obj + + transport.send(mock_request) + + # Verify that _update_request_url was called (URL should be modified) + # The exact URL modification is tested in the mixin tests + mock_parent_send.assert_called_once_with(mock_request) + + def test_custom_requests_transport_send_with_ca_data(self, ca_data): + """Test that CustomRequestsTransport.send works with CA data.""" + transport = _get_transport( + sni=None, + token_proxy_endpoint="https://proxy.example.com:8080", + ca_file=None, + ca_data=ca_data, + ) + assert transport is not None + + mock_request = HttpRequest("GET", "https://login.microsoftonline.com/tenant/oauth2/v2.0/token") + + with patch.object(transport.__class__.__bases__[1], "send") as mock_parent_send: + mock_response_obj = MagicMock() + mock_parent_send.return_value = mock_response_obj + + response = transport.send(mock_request) + + assert response == mock_response_obj + mock_parent_send.assert_called_once_with(mock_request) + + def test_custom_requests_transport_send_with_ca_file_reload(self, ca_data): + """Test that CustomRequestsTransport.send reloads CA file when changed.""" + # Create a temporary CA file + with tempfile.NamedTemporaryFile(mode="w", delete=False, suffix=".pem") as ca_file: + ca_file.write(ca_data) + ca_file_path = ca_file.name + + try: + transport = _get_transport( + sni=None, + token_proxy_endpoint="https://proxy.example.com:8080", + ca_file=ca_file_path, + ca_data=None, + ) + assert transport is not None + + mock_request = HttpRequest("GET", "https://login.microsoftonline.com/tenant/oauth2/v2.0/token") + + # Mock the parent send method and file modification time tracking + with patch.object(transport.__class__.__bases__[1], "send") as mock_parent_send: + mock_response_obj = MagicMock() + mock_parent_send.return_value = mock_response_obj + + # First request + transport.send(mock_request) + + # Simulate file change by modifying the file + time.sleep(0.1) # Ensure different modification time + with open(ca_file_path, "a") as f: + f.write("\n# Modified for testing\n") + + # Mock the CA file change detection + with patch.object(transport, "_has_ca_file_changed", return_value=True): + with patch.object(transport, "_load_ca_file_to_data") as mock_load_ca: + with patch.object(transport, "_create_session") as mock_create_session: + transport.send(mock_request) + + mock_load_ca.assert_called_once() + # _create_session should be called if _ca_data is truthy after loading + # This depends on the implementation details + + assert mock_parent_send.call_count == 2 + + finally: + os.unlink(ca_file_path) + + def test_custom_requests_transport_initialization_with_ca_data(self, ca_data): + """Test CustomRequestsTransport initialization with CA data creates SSL context.""" + transport = _get_transport( + sni=None, + token_proxy_endpoint="https://proxy.example.com:8080", + ca_file=None, + ca_data=ca_data, + ) + assert transport is not None + + # Verify transport was created properly + assert transport._ca_data == ca_data + + def test_custom_requests_transport_initialization_without_ca_data(self): + """Test CustomRequestsTransport initialization without CA data.""" + transport = _get_transport( + sni=None, + token_proxy_endpoint="https://proxy.example.com:8080", + ca_file=None, + ca_data=None, + ) + assert transport is not None + + # Verify transport was created properly + assert transport._ca_data is None + + def test_custom_requests_transport_inherits_from_token_binding_mixin(self): + """Test that CustomRequestsTransport inherits from TokenBindingTransportMixin.""" + transport = _get_transport( + sni="test.sni.com", + token_proxy_endpoint="https://proxy.example.com:8080", + ca_file=None, + ca_data=None, + ) + + assert transport is not None + + # Verify inheritance from TokenBindingTransportMixin + from azure.identity._internal.token_binding_transport_mixin import TokenBindingTransportMixin + + assert isinstance(transport, TokenBindingTransportMixin) + + # Verify TokenBindingTransportMixin methods are available + assert hasattr(transport, "_update_request_url") + assert hasattr(transport, "_has_ca_file_changed") + assert hasattr(transport, "_load_ca_file_to_data") + assert hasattr(transport, "_validate_url") + + +class TestCustomRequestsTransportWithLocalServer: + """Integration tests using a local test server for CustomRequestsTransport.""" + + def test_basic_https_request(self): + """Test basic HTTPS request to test server.""" + with TokenProxyTestServer(use_ssl=True) as server: + # Create transport with server's CA certificate + transport = _get_transport(sni=None, token_proxy_endpoint=None, ca_file=server.ca_file, ca_data=None) + assert transport is not None + request = HttpRequest("GET", f"{server.base_url}/health") + + response = transport.send(request) + + assert response.status_code == 200 + data = response.json() + assert data["status"] == "healthy" + assert "timestamp" in data + + def test_post_request_with_body(self): + """Test POST request with request body.""" + with TokenProxyTestServer(use_ssl=True) as server: + transport = _get_transport(sni=None, token_proxy_endpoint=None, ca_file=server.ca_file, ca_data=None) + assert transport is not None + + # Prepare OAuth-like request + body = "grant_type=client_credentials&scope=https://graph.microsoft.com/.default" + request = HttpRequest( + "POST", + f"{server.base_url}/tenant/oauth2/v2.0/token", + headers={"Content-Type": "application/x-www-form-urlencoded"}, + data=body.encode("utf-8"), + ) + + response = transport.send(request) + + assert response.status_code == 200 + data = response.json() + assert "access_token" in data + assert data["token_type"] == "Bearer" + assert data["expires_in"] == 3600 + + def test_proxy_endpoint_comprehensive(self): + """Test comprehensive proxy endpoint functionality with various HTTP methods and scenarios.""" + with TokenProxyTestServer(use_ssl=True) as server: + # Configure transport with proxy endpoint + transport = _get_transport( + sni=None, token_proxy_endpoint=server.base_url, ca_file=server.ca_file, ca_data=None + ) + assert transport is not None + + # Test 1: POST request with JSON body through proxy + post_data = {"grant_type": "client_credentials", "scope": "https://graph.microsoft.com/.default"} + post_request = HttpRequest( + "POST", + "https://login.microsoftonline.com/tenant/oauth2/v2.0/token2", + headers={"Content-Type": "application/json"}, + json=post_data, + ) + + post_response = transport.send(post_request) + assert post_response.status_code == 200 + post_data_response = post_response.json() + assert post_data_response["method"] == "POST" + assert post_data_response["proxied_path"] == "/tenant/oauth2/v2.0/token2" + + # Test 2: PUT request through proxy + put_request = HttpRequest( + "PUT", + "https://graph.microsoft.com/v1.0/me/profile", + headers={"Content-Type": "application/json"}, + json={"displayName": "Test User"}, + ) + + put_response = transport.send(put_request) + assert put_response.status_code == 200 + put_data_response = put_response.json() + assert put_data_response["method"] == "PUT" + assert put_data_response["proxied_path"] == "/v1.0/me/profile" + + # Test 3: DELETE request through proxy + delete_request = HttpRequest("DELETE", "https://graph.microsoft.com/v1.0/applications/app-id") + + delete_response = transport.send(delete_request) + assert delete_response.status_code == 200 + delete_data_response = delete_response.json() + assert delete_data_response["method"] == "DELETE" + assert delete_data_response["proxied_path"] == "/v1.0/applications/app-id" + + # Test 4: Complex URL with multiple path segments and query parameters + complex_url = "https://management.azure.com/subscriptions/sub-id/resourceGroups/rg/providers/Microsoft.Storage/storageAccounts/account?api-version=2021-04-01&expand=properties" + complex_request = HttpRequest("GET", complex_url) + + complex_response = transport.send(complex_request) + assert complex_response.status_code == 200 + complex_data_response = complex_response.json() + assert complex_data_response["method"] == "GET" + expected_path = "/subscriptions/sub-id/resourceGroups/rg/providers/Microsoft.Storage/storageAccounts/account?api-version=2021-04-01&expand=properties" + assert complex_data_response["proxied_path"] == expected_path + + # Test 5: Request with custom headers through proxy + headers_request = HttpRequest( + "GET", + "https://vault.azure.net/secrets/test-secret?api-version=7.3", + headers={ + "Authorization": "Bearer test-token", + "X-Custom-Header": "proxy-test-value", + "User-Agent": "Azure-SDK-For-Python", + }, + ) + + headers_response = transport.send(headers_request) + assert headers_response.status_code == 200 + headers_data_response = headers_response.json() + assert headers_data_response["method"] == "GET" + assert headers_data_response["proxied_path"] == "/secrets/test-secret?api-version=7.3" + + # Verify headers were forwarded through proxy + received_headers = headers_data_response["headers_received"] + assert "Authorization" in received_headers + assert "X-Custom-Header" in received_headers + assert received_headers["Authorization"] == "Bearer test-token" + assert received_headers["X-Custom-Header"] == "proxy-test-value" + + def test_sni_with_custom_hostname(self): + """Test SNI (Server Name Indication) with custom hostname.""" + with TokenProxyTestServer(use_ssl=True) as server: + # Use SNI with a different hostname than the server + transport = _get_transport( + sni="1234.ests.aks", token_proxy_endpoint=None, ca_file=server.ca_file, ca_data=None + ) + assert transport is not None + + request = HttpRequest("GET", f"{server.base_url}/health") + + response = transport.send(request) + + assert response.status_code == 200 + data = response.json() + assert data["status"] == "healthy" + + def test_ca_file_change_detection(self): + """Test CA file change detection with real certificates.""" + with TokenProxyTestServer(use_ssl=True) as server: + # Create a copy of the CA file that we can modify + ca_file = server.ca_file + if ca_file is None: + pytest.skip("CA file not available") + + assert ca_file is not None # Type hint for mypy + with tempfile.NamedTemporaryFile(mode="w", delete=False, suffix=".pem") as temp_ca: + with open(ca_file, "r") as original: + original_content = original.read() + temp_ca.write(original_content) + temp_ca_path = temp_ca.name + + try: + transport = _get_transport(sni=None, token_proxy_endpoint=None, ca_file=temp_ca_path, ca_data=None) + assert transport is not None + + # First request should work + request = HttpRequest("GET", f"{server.base_url}/health") + response1 = transport.send(request) + assert response1.status_code == 200 + + # Modify the CA file (add some content) + time.sleep(0.1) + with open(temp_ca_path, "a") as f: + f.write("\n# Modified for testing\n") + + # Second request should still work (using the same cert content) + response2 = transport.send(request) + assert response2.status_code == 200 + + finally: + os.unlink(temp_ca_path) + + def test_ssl_error_handling(self): + """Test SSL error handling.""" + from azure.core.exceptions import ServiceRequestError, ServiceResponseError + + with TokenProxyTestServer(use_ssl=True) as server: + # Create transport without proper CA file (will cause SSL error) + transport = _get_transport(sni=None, token_proxy_endpoint=None, ca_file=None, ca_data=None) + assert transport is not None + + request = HttpRequest("GET", f"{server.base_url}/health") + + # Should raise SSL-related error + with pytest.raises((ServiceRequestError, ServiceResponseError)): + transport.send(request) + + def test_server_error_response(self): + """Test handling of server error responses.""" + with TokenProxyTestServer(use_ssl=True) as server: + transport = _get_transport(sni=None, token_proxy_endpoint=None, ca_file=server.ca_file, ca_data=None) + assert transport is not None + + request = HttpRequest("GET", f"{server.base_url}/error/500") + + response = transport.send(request) + + assert response.status_code == 500 + # Should not raise exception, just return error response + + def test_slow_server_response(self): + """Test handling of slow server responses.""" + import time + + with TokenProxyTestServer(use_ssl=True) as server: + transport = _get_transport(sni=None, token_proxy_endpoint=None, ca_file=server.ca_file, ca_data=None) + assert transport is not None + + request = HttpRequest("GET", f"{server.base_url}/slow") + + start_time = time.time() + response = transport.send(request) + elapsed_time = time.time() - start_time + + assert response.status_code == 200 + # Should take at least 2 seconds (server waits for 2s) + assert elapsed_time >= 2.0 + data = response.json() + assert data["message"] == "slow response" + + def test_custom_headers_preserved(self): + """Test that custom headers are preserved and sent to server.""" + with TokenProxyTestServer(use_ssl=True) as server: + transport = _get_transport(sni=None, token_proxy_endpoint=None, ca_file=server.ca_file, ca_data=None) + assert transport is not None + + custom_headers = { + "Authorization": "Bearer test-token", + "User-Agent": "CustomRequestsTransport/1.0", + "X-Custom-Header": "test-value", + } + + request = HttpRequest("GET", f"{server.base_url}/proxy/test", headers=custom_headers) + + response = transport.send(request) + + assert response.status_code == 200 + data = response.json() + + # Server echoes back the headers it received + received_headers = data["headers_received"] + assert "Authorization" in received_headers + assert "User-Agent" in received_headers + assert "X-Custom-Header" in received_headers + assert received_headers["Authorization"] == "Bearer test-token" + + def test_query_parameters_preserved(self): + """Test that query parameters are preserved in proxy requests.""" + with TokenProxyTestServer(use_ssl=True) as server: + transport = _get_transport( + sni=None, token_proxy_endpoint=server.base_url, ca_file=server.ca_file, ca_data=None + ) + assert transport is not None + + # Request with query parameters + original_url = "https://example.com/api/data?scope=read&limit=10&format=json" + request = HttpRequest("GET", original_url) + + response = transport.send(request) + + assert response.status_code == 200 + data = response.json() - # Verify HttpClientTransport was NOT called - mock_transport_class.assert_not_called() + # The path should include the query parameters + expected_path = "/api/data?scope=read&limit=10&format=json" + assert data["proxied_path"] == expected_path diff --git a/sdk/identity/azure-identity/tests/test_workload_identity_credential_async.py b/sdk/identity/azure-identity/tests/test_workload_identity_credential_async.py index 7a33656d617a..7fdc7c5d973d 100644 --- a/sdk/identity/azure-identity/tests/test_workload_identity_credential_async.py +++ b/sdk/identity/azure-identity/tests/test_workload_identity_credential_async.py @@ -155,7 +155,6 @@ def test_use_token_proxy_minimal_config(self): use_token_proxy=True, ) - # Verify HttpClientTransport was called with minimal config mock_get_transport.assert_called_once_with( sni=None, token_proxy_endpoint=proxy_endpoint, @@ -223,7 +222,6 @@ async def send(request, **kwargs): assert request.data.get("client_assertion") == assertion return mock_response(json_payload=build_aad_response(access_token=access_token)) - # Mock the transport that would be created by HttpClientTransport mock_transport_instance = MagicMock(send=send) env_vars = { @@ -249,7 +247,7 @@ async def send(request, **kwargs): open_mock.assert_called_once_with(token_file_path, encoding="utf-8") def test_use_token_proxy_false_does_not_create_transport(self): - """Test that use_token_proxy=False (default) does not create HttpClientTransport.""" + """Test that use_token_proxy=False (default) does not create a custom transport.""" tenant_id = "tenant-id" client_id = "client-id" token_file_path = "foo-path" @@ -679,26 +677,6 @@ async def test_sni_with_custom_hostname(self): data = response.json() assert data["status"] == "healthy" - @pytest.mark.asyncio - async def test_connection_reuse(self): - """Test that connections are reused for multiple requests.""" - with TokenProxyTestServer(use_ssl=True) as server: - transport = _get_transport(sni=None, token_proxy_endpoint=None, ca_file=server.ca_file, ca_data=None) - assert transport is not None - - # Make multiple requests - requests_data = [] - for i in range(3): - request = HttpRequest("GET", f"{server.base_url}/health") - response = await transport.send(request) - assert response.status_code == 200 - requests_data.append(response.json()) - - # All requests should succeed - assert len(requests_data) == 3 - for data in requests_data: - assert data["status"] == "healthy" - @pytest.mark.asyncio async def test_ca_file_change_detection(self): """Test CA file change detection with real certificates.""" From 23824fa35f84edb6c1bd0d42cfa98d6d0faa96ab Mon Sep 17 00:00:00 2001 From: Paul Van Eck Date: Thu, 30 Oct 2025 03:12:43 +0000 Subject: [PATCH 06/12] Refactor Signed-off-by: Paul Van Eck --- .../pipeline/transport/_requests_basic.py | 1 - .../_credentials/workload_identity.py | 51 +-------------- .../_internal/token_binding_transport.py | 62 +++++++++++++++++++ .../token_binding_transport_mixin.py | 4 +- .../aio/_credentials/workload_identity.py | 42 +------------ .../aio/_internal/token_binding_transport.py | 38 ++++++++++++ .../azure-identity/tests/proxy_server.py | 2 +- 7 files changed, 109 insertions(+), 91 deletions(-) create mode 100644 sdk/identity/azure-identity/azure/identity/_internal/token_binding_transport.py create mode 100644 sdk/identity/azure-identity/azure/identity/aio/_internal/token_binding_transport.py diff --git a/sdk/core/azure-core/azure/core/pipeline/transport/_requests_basic.py b/sdk/core/azure-core/azure/core/pipeline/transport/_requests_basic.py index 44e150eddf7a..9f102f4b0b20 100644 --- a/sdk/core/azure-core/azure/core/pipeline/transport/_requests_basic.py +++ b/sdk/core/azure-core/azure/core/pipeline/transport/_requests_basic.py @@ -270,7 +270,6 @@ def __init__(self, **kwargs) -> None: self._use_env_settings = kwargs.pop("use_env_settings", True) # See https://github.com/Azure/azure-sdk-for-python/issues/25640 to understand why we track this self._has_been_opened = False - print('---done init---') def __enter__(self) -> "RequestsTransport": self.open() diff --git a/sdk/identity/azure-identity/azure/identity/_credentials/workload_identity.py b/sdk/identity/azure-identity/azure/identity/_credentials/workload_identity.py index 6216f021457c..309a2bbcc2b6 100644 --- a/sdk/identity/azure-identity/azure/identity/_credentials/workload_identity.py +++ b/sdk/identity/azure-identity/azure/identity/_credentials/workload_identity.py @@ -4,13 +4,11 @@ # ------------------------------------ import os import time -import ssl from typing import Any from typing import Optional from .client_assertion import ClientAssertionCredential from .._constants import EnvironmentVariables -from .._internal.token_binding_transport_mixin import TokenBindingTransportMixin WORKLOAD_CONFIG_ERROR = ( @@ -148,51 +146,9 @@ def __init__( def _get_transport(sni, token_proxy_endpoint, ca_file, ca_data): try: - from azure.core.pipeline.transport import ( # pylint: disable=non-abstract-transport-import, no-name-in-module - RequestsTransport, - ) - from requests.adapters import HTTPAdapter - from requests import Session - - class SNIAdapter(HTTPAdapter): - """A custom HTTPAdapter that allows setting a custom SNI hostname.""" - - def __init__(self, server_hostname, ca_data, **kwargs): - self.server_hostname = server_hostname - self.ca_data = ca_data - super().__init__(**kwargs) - - def init_poolmanager(self, connections, maxsize, block=False, **pool_kwargs): - if self.server_hostname: - pool_kwargs["server_hostname"] = self.server_hostname - pool_kwargs["ssl_context"] = ssl.create_default_context(cadata=self.ca_data) - super().init_poolmanager(connections, maxsize, block, **pool_kwargs) - - class CustomRequestsTransport(TokenBindingTransportMixin, RequestsTransport): - def __init__(self, *args, **kwargs): - super().__init__(*args, **kwargs) - self._create_session() - - def _create_session(self): - if self.session: # pylint: disable=access-member-before-definition - self.session.close() # pylint: disable=access-member-before-definition - - self.session = Session() - adapter = SNIAdapter(self._sni, self._ca_data) - self.session.mount("https://", adapter) - - def send(self, request, **kwargs): - self._update_request_url(request) - - # Check if CA file has changed and reload ca_data if needed - if self._ca_file and self._has_ca_file_changed(): - self._load_ca_file_to_data() - # If ca_data was updated, recreate SSL context with the new data - if self._ca_data: - self._create_session() - return super().send(request, **kwargs) - - transport = CustomRequestsTransport( + from .._internal.token_binding_transport import CustomRequestsTransport + + return CustomRequestsTransport( sni=sni, proxy_endpoint=token_proxy_endpoint, ca_file=ca_file, @@ -201,4 +157,3 @@ def send(self, request, **kwargs): except ImportError: return None - return transport diff --git a/sdk/identity/azure-identity/azure/identity/_internal/token_binding_transport.py b/sdk/identity/azure-identity/azure/identity/_internal/token_binding_transport.py new file mode 100644 index 000000000000..4eb546c84cc1 --- /dev/null +++ b/sdk/identity/azure-identity/azure/identity/_internal/token_binding_transport.py @@ -0,0 +1,62 @@ +# ------------------------------------ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. +# ------------------------------------ +""" +Transport class for WorkloadIdentityCredential with token proxy support. +""" +import ssl +from typing import Any, Optional + +from requests.adapters import HTTPAdapter +from requests import Session +from azure.core.pipeline.transport import ( # pylint: disable=non-abstract-transport-import, no-name-in-module + RequestsTransport, +) +from azure.core.rest import HttpRequest + +from .token_binding_transport_mixin import TokenBindingTransportMixin + + +class SNIAdapter(HTTPAdapter): + """A custom HTTPAdapter that allows setting a custom SNI hostname.""" + + def __init__(self, server_hostname: Optional[str], ca_data: Optional[str], **kwargs: Any) -> None: + self.server_hostname = server_hostname + self.ca_data = ca_data + super().__init__(**kwargs) + + def init_poolmanager(self, connections: int, maxsize: int, block: bool = False, **pool_kwargs: Any) -> None: + if self.server_hostname: + pool_kwargs["server_hostname"] = self.server_hostname + pool_kwargs["ssl_context"] = ssl.create_default_context(cadata=self.ca_data) + super().init_poolmanager(connections, maxsize, block, **pool_kwargs) + + +class CustomRequestsTransport(TokenBindingTransportMixin, RequestsTransport): + """Custom RequestsTransport with SNI and CA certificate support for WorkloadIdentityCredential.""" + + def __init__(self, *args: Any, **kwargs: Any) -> None: + self.session: Optional[Session] = None + super().__init__(*args, **kwargs) + self._create_session() + + def _create_session(self) -> None: + """Create a new requests session with custom SSL configuration.""" + if self.session: + self.session.close() + + self.session = Session() + adapter = SNIAdapter(self._sni, self._ca_data) + self.session.mount("https://", adapter) + + def send(self, request: HttpRequest, **kwargs: Any) -> Any: + self._update_request_url(request) + + # Check if CA file has changed and reload ca_data if needed + if self._ca_file and self._has_ca_file_changed(): + self._load_ca_file_to_data() + # If ca_data was updated, recreate SSL context with the new data + if self._ca_data: + self._create_session() + return super().send(request, **kwargs) diff --git a/sdk/identity/azure-identity/azure/identity/_internal/token_binding_transport_mixin.py b/sdk/identity/azure-identity/azure/identity/_internal/token_binding_transport_mixin.py index 44dbf426e98f..d790c25d6d24 100644 --- a/sdk/identity/azure-identity/azure/identity/_internal/token_binding_transport_mixin.py +++ b/sdk/identity/azure-identity/azure/identity/_internal/token_binding_transport_mixin.py @@ -5,7 +5,7 @@ # cspell:ignore cafile import os import urllib.parse -from typing import Optional +from typing import Optional, Any from azure.core.rest import HttpRequest @@ -13,7 +13,7 @@ class TokenBindingTransportMixin: """Mixin class providing URL validation, CA file tracking, and proxy URL functionality for transport classes.""" - def __init__(self, **kwargs): + def __init__(self, **kwargs: Any) -> None: """Initialize CA file tracking and proxy attributes.""" self._ca_file = kwargs.pop("ca_file", None) self._ca_data = kwargs.pop("ca_data", None) diff --git a/sdk/identity/azure-identity/azure/identity/aio/_credentials/workload_identity.py b/sdk/identity/azure-identity/azure/identity/aio/_credentials/workload_identity.py index 1801d4cf8d2b..d6d9a5732243 100644 --- a/sdk/identity/azure-identity/azure/identity/aio/_credentials/workload_identity.py +++ b/sdk/identity/azure-identity/azure/identity/aio/_credentials/workload_identity.py @@ -4,17 +4,12 @@ # ------------------------------------ # cspell:ignore cafile import os -import ssl import logging from typing import Any, Optional from .client_assertion import ClientAssertionCredential from ..._credentials.workload_identity import TokenFileMixin, WORKLOAD_CONFIG_ERROR from ..._constants import EnvironmentVariables -from ..._internal.token_binding_transport_mixin import TokenBindingTransportMixin - - -_LOGGER = logging.getLogger(__name__) class WorkloadIdentityCredential(ClientAssertionCredential, TokenFileMixin): @@ -127,44 +122,13 @@ def __init__( def _get_transport(sni, token_proxy_endpoint, ca_file, ca_data): try: - from azure.core.pipeline.transport import ( # pylint: disable=non-abstract-transport-import, no-name-in-module - AioHttpTransport, - ) - - class WorkloadIdentityAioHttpTransport(TokenBindingTransportMixin, AioHttpTransport): - - def __init__(self, *args, **kwargs): - super().__init__(*args, **kwargs) - self._ssl_context = ssl.create_default_context(cadata=self._ca_data) - - async def send(self, request, **kwargs): - self._update_request_url(request) - kwargs.setdefault("server_hostname", self._sni) - - # Check if CA file has changed and reload ca_data if needed - if self._ca_file and self._has_ca_file_changed(): - self._load_ca_file_to_data() - # If ca_data was updated, recreate SSL context with the new data - if self._ca_data: - self._ssl_context = ssl.create_default_context(cadata=self._ca_data) - - if self._ssl_context: - kwargs.setdefault("ssl", self._ssl_context) - return await super().send(request, **kwargs) - - async def __aenter__(self): - await super().__aenter__() - return self - - async def __aexit__(self, *args): - await super().__aexit__(*args) + from .._internal.token_binding_transport import CustomAioHttpTransport - transport = WorkloadIdentityAioHttpTransport( + return CustomAioHttpTransport( sni=sni, proxy_endpoint=token_proxy_endpoint, ca_file=ca_file, ca_data=ca_data, ) except ImportError: - transport = None - return transport + return None diff --git a/sdk/identity/azure-identity/azure/identity/aio/_internal/token_binding_transport.py b/sdk/identity/azure-identity/azure/identity/aio/_internal/token_binding_transport.py new file mode 100644 index 000000000000..bfc3cc53e494 --- /dev/null +++ b/sdk/identity/azure-identity/azure/identity/aio/_internal/token_binding_transport.py @@ -0,0 +1,38 @@ +# ------------------------------------ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. +# ------------------------------------ +""" +Transport class for the asynchronous WorkloadIdentityCredential with token proxy support. +""" +import ssl +from typing import Any + +from azure.core.pipeline.transport import ( # pylint: disable=non-abstract-transport-import, no-name-in-module + AioHttpTransport, +) +from azure.core.rest import HttpRequest + +from ..._internal.token_binding_transport_mixin import TokenBindingTransportMixin + + +class CustomAioHttpTransport(TokenBindingTransportMixin, AioHttpTransport): + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self._ssl_context = ssl.create_default_context(cadata=self._ca_data) + + async def send(self, request: HttpRequest, **kwargs: Any) -> Any: + self._update_request_url(request) + kwargs.setdefault("server_hostname", self._sni) + + # Check if CA file has changed and reload ca_data if needed + if self._ca_file and self._has_ca_file_changed(): + self._load_ca_file_to_data() + # If ca_data was updated, recreate SSL context with the new data + if self._ca_data: + self._ssl_context = ssl.create_default_context(cadata=self._ca_data) + + if self._ssl_context: + kwargs.setdefault("ssl", self._ssl_context) + return await super().send(request, **kwargs) diff --git a/sdk/identity/azure-identity/tests/proxy_server.py b/sdk/identity/azure-identity/tests/proxy_server.py index fb21e04c1bab..22a33973edff 100644 --- a/sdk/identity/azure-identity/tests/proxy_server.py +++ b/sdk/identity/azure-identity/tests/proxy_server.py @@ -4,7 +4,7 @@ # ------------------------------------ # cspell:ignore ests """ -Local test server for HttpClientTransport testing. +Local test server for token binding proxy testing. This server simulates a token proxy that can: 1. Accept HTTPS requests with custom SNI and CA certificates From 2071d34e838117f8c1612a69988f374ba3435f25 Mon Sep 17 00:00:00 2001 From: Paul Van Eck Date: Fri, 31 Oct 2025 01:25:06 +0000 Subject: [PATCH 07/12] Update changelog Signed-off-by: Paul Van Eck --- sdk/identity/azure-identity/CHANGELOG.md | 2 +- .../azure/identity/aio/_credentials/workload_identity.py | 1 - 2 files changed, 1 insertion(+), 2 deletions(-) diff --git a/sdk/identity/azure-identity/CHANGELOG.md b/sdk/identity/azure-identity/CHANGELOG.md index c0b31ed4f9b0..cd474e5b547d 100644 --- a/sdk/identity/azure-identity/CHANGELOG.md +++ b/sdk/identity/azure-identity/CHANGELOG.md @@ -4,7 +4,7 @@ ### Features Added -- Added support for WorkloadIdentityCredential identity binding mode in AKS. ([#43287](https://github.com/Azure/azure-sdk-for-python/pull/43287)) +- Added support for `WorkloadIdentityCredential` identity binding mode in AKS environments. This feature addresses Entra's limitation on the number of federated identity credentials (FICs) per managed identity by utilizing an AKS proxy that handles FIC exchanges on behalf of pods. When the environment is properly configured and `use_token_proxy=True` is specified in the `WorkloadIdentityCredential` constructor, all token requests are routed through this proxy service. ([#43287](https://github.com/Azure/azure-sdk-for-python/pull/43287)) ### Breaking Changes diff --git a/sdk/identity/azure-identity/azure/identity/aio/_credentials/workload_identity.py b/sdk/identity/azure-identity/azure/identity/aio/_credentials/workload_identity.py index d6d9a5732243..b62520e0f87e 100644 --- a/sdk/identity/azure-identity/azure/identity/aio/_credentials/workload_identity.py +++ b/sdk/identity/azure-identity/azure/identity/aio/_credentials/workload_identity.py @@ -4,7 +4,6 @@ # ------------------------------------ # cspell:ignore cafile import os -import logging from typing import Any, Optional from .client_assertion import ClientAssertionCredential From aec1604abed707922a2973100c4eca00b0c1e157 Mon Sep 17 00:00:00 2001 From: Paul Van Eck Date: Fri, 31 Oct 2025 05:05:01 +0000 Subject: [PATCH 08/12] test updates Signed-off-by: Paul Van Eck --- .../tests/test_workload_identity_credential.py | 4 +--- .../tests/test_workload_identity_credential_async.py | 9 +++------ 2 files changed, 4 insertions(+), 9 deletions(-) diff --git a/sdk/identity/azure-identity/tests/test_workload_identity_credential.py b/sdk/identity/azure-identity/tests/test_workload_identity_credential.py index 06ad11f984f4..b35ba07bf24d 100644 --- a/sdk/identity/azure-identity/tests/test_workload_identity_credential.py +++ b/sdk/identity/azure-identity/tests/test_workload_identity_credential.py @@ -10,6 +10,7 @@ import pytest from azure.core.rest import HttpRequest +from azure.core.exceptions import ServiceRequestError, ServiceResponseError from azure.identity import WorkloadIdentityCredential from azure.identity._credentials.workload_identity import _get_transport @@ -650,8 +651,6 @@ def test_ca_file_change_detection(self): def test_ssl_error_handling(self): """Test SSL error handling.""" - from azure.core.exceptions import ServiceRequestError, ServiceResponseError - with TokenProxyTestServer(use_ssl=True) as server: # Create transport without proper CA file (will cause SSL error) transport = _get_transport(sni=None, token_proxy_endpoint=None, ca_file=None, ca_data=None) @@ -678,7 +677,6 @@ def test_server_error_response(self): def test_slow_server_response(self): """Test handling of slow server responses.""" - import time with TokenProxyTestServer(use_ssl=True) as server: transport = _get_transport(sni=None, token_proxy_endpoint=None, ca_file=server.ca_file, ca_data=None) diff --git a/sdk/identity/azure-identity/tests/test_workload_identity_credential_async.py b/sdk/identity/azure-identity/tests/test_workload_identity_credential_async.py index 7fdc7c5d973d..9ab0a04a5591 100644 --- a/sdk/identity/azure-identity/tests/test_workload_identity_credential_async.py +++ b/sdk/identity/azure-identity/tests/test_workload_identity_credential_async.py @@ -5,11 +5,14 @@ # cspell:ignore cafile aexit ests import os import tempfile +import time +import asyncio from time import sleep as real_sleep from unittest.mock import mock_open, patch, MagicMock import pytest from azure.core.rest import HttpRequest +from azure.core.exceptions import ServiceRequestError, ServiceResponseError from azure.identity.aio import WorkloadIdentityCredential from azure.identity.aio._credentials.workload_identity import _get_transport @@ -717,8 +720,6 @@ async def test_ca_file_change_detection(self): @pytest.mark.asyncio async def test_ssl_error_handling(self): """Test SSL error handling.""" - from azure.core.exceptions import ServiceRequestError, ServiceResponseError - with TokenProxyTestServer(use_ssl=True) as server: # Create transport without proper CA file (will cause SSL error) transport = _get_transport(sni=None, token_proxy_endpoint=None, ca_file=None, ca_data=None) @@ -747,8 +748,6 @@ async def test_server_error_response(self): @pytest.mark.asyncio async def test_slow_server_response(self): """Test handling of slow server responses.""" - import time - with TokenProxyTestServer(use_ssl=True) as server: transport = _get_transport(sni=None, token_proxy_endpoint=None, ca_file=server.ca_file, ca_data=None) assert transport is not None @@ -817,8 +816,6 @@ async def test_query_parameters_preserved(self): @pytest.mark.asyncio async def test_concurrent_requests(self): """Test handling multiple concurrent requests.""" - import asyncio - with TokenProxyTestServer(use_ssl=True) as server: transport = _get_transport(sni=None, token_proxy_endpoint=None, ca_file=server.ca_file, ca_data=None) assert transport is not None From 6d0ec71bc0c1ab46054271ffb0974d1561aae17e Mon Sep 17 00:00:00 2001 From: Paul Van Eck Date: Tue, 4 Nov 2025 01:11:18 +0000 Subject: [PATCH 09/12] Updates Signed-off-by: Paul Van Eck --- .../_credentials/workload_identity.py | 52 +++++++++--------- .../aio/_credentials/workload_identity.py | 53 +++++++++---------- .../azure-identity/tests/proxy_server.py | 3 -- .../test_workload_identity_credential.py | 7 +-- ...test_workload_identity_credential_async.py | 7 +-- 5 files changed, 56 insertions(+), 66 deletions(-) diff --git a/sdk/identity/azure-identity/azure/identity/_credentials/workload_identity.py b/sdk/identity/azure-identity/azure/identity/_credentials/workload_identity.py index 309a2bbcc2b6..cc32a172d2bd 100644 --- a/sdk/identity/azure-identity/azure/identity/_credentials/workload_identity.py +++ b/sdk/identity/azure-identity/azure/identity/_credentials/workload_identity.py @@ -58,7 +58,8 @@ class WorkloadIdentityCredential(ClientAssertionCredential, TokenFileMixin): :keyword str token_file_path: The path to a file containing a Kubernetes service account token that authenticates the identity. :keyword bool use_token_proxy: Whether or not to read token proxy configuration from environment variables and use - a token proxy to acquire tokens. Defaults to False. + a token proxy to acquire tokens. If this value is True and proxy configuration isn't present or this value is + False, the credential will request tokens directly from Entra ID. Defaults to False. .. admonition:: Example: @@ -105,35 +106,30 @@ def __init__( if use_token_proxy: token_proxy_endpoint = os.environ.get(EnvironmentVariables.AZURE_KUBERNETES_TOKEN_PROXY) - if not token_proxy_endpoint: - raise ValueError( - "use_token_proxy is True, but no token proxy endpoint was found. " - f"Ensure that the {EnvironmentVariables.AZURE_KUBERNETES_TOKEN_PROXY} environment variable is set." + if token_proxy_endpoint: + sni = os.environ.get(EnvironmentVariables.AZURE_KUBERNETES_SNI_NAME) + ca_file = os.environ.get(EnvironmentVariables.AZURE_KUBERNETES_CA_FILE) + ca_data = os.environ.get(EnvironmentVariables.AZURE_KUBERNETES_CA_DATA) + + if ca_file and ca_data: + raise ValueError( + "Both AZURE_KUBERNETES_CA_FILE and AZURE_KUBERNETES_CA_DATA are set. Only one should be set." + ) + + transport = _get_transport( + sni=sni, + token_proxy_endpoint=token_proxy_endpoint, + ca_file=ca_file, + ca_data=ca_data, ) - sni = os.environ.get(EnvironmentVariables.AZURE_KUBERNETES_SNI_NAME) - ca_file = os.environ.get(EnvironmentVariables.AZURE_KUBERNETES_CA_FILE) - ca_data = os.environ.get(EnvironmentVariables.AZURE_KUBERNETES_CA_DATA) - - if ca_file and ca_data: - raise ValueError( - "Both AZURE_KUBERNETES_CA_FILE and AZURE_KUBERNETES_CA_DATA are set. Only one should be set." - ) - - transport = _get_transport( - sni=sni, - token_proxy_endpoint=token_proxy_endpoint, - ca_file=ca_file, - ca_data=ca_data, - ) - - if transport: - kwargs["transport"] = transport - else: - raise ValueError( - "Transport creation failed. Ensure that the requests package is installed to enable token " - "proxy usage in this credential." - ) + if transport: + kwargs["transport"] = transport + else: + raise ValueError( + "Transport creation failed. Ensure that the requests package is installed to enable token " + "proxy usage in this credential." + ) super(WorkloadIdentityCredential, self).__init__( tenant_id=tenant_id, diff --git a/sdk/identity/azure-identity/azure/identity/aio/_credentials/workload_identity.py b/sdk/identity/azure-identity/azure/identity/aio/_credentials/workload_identity.py index b62520e0f87e..66851d16c4d7 100644 --- a/sdk/identity/azure-identity/azure/identity/aio/_credentials/workload_identity.py +++ b/sdk/identity/azure-identity/azure/identity/aio/_credentials/workload_identity.py @@ -33,7 +33,8 @@ class WorkloadIdentityCredential(ClientAssertionCredential, TokenFileMixin): :keyword str token_file_path: The path to a file containing a Kubernetes service account token that authenticates the identity. :keyword bool use_token_proxy: Whether or not to read token proxy configuration from environment variables and use - a token proxy to acquire tokens. Defaults to False. + a token proxy to acquire tokens. If this value is True and proxy configuration isn't present or this value is + False, the credential will request tokens directly from Entra ID. Defaults to False. .. admonition:: Example: @@ -79,36 +80,30 @@ def __init__( self._token_file_path = token_file_path if use_token_proxy: - self._token_proxy_endpoint = os.environ.get(EnvironmentVariables.AZURE_KUBERNETES_TOKEN_PROXY) - if not self._token_proxy_endpoint: - raise ValueError( - "use_token_proxy is True, but no token proxy endpoint was found. " - f"Ensure that the {EnvironmentVariables.AZURE_KUBERNETES_TOKEN_PROXY} environment variable is set." + token_proxy_endpoint = os.environ.get(EnvironmentVariables.AZURE_KUBERNETES_TOKEN_PROXY) + if token_proxy_endpoint: + sni = os.environ.get(EnvironmentVariables.AZURE_KUBERNETES_SNI_NAME) + ca_file = os.environ.get(EnvironmentVariables.AZURE_KUBERNETES_CA_FILE) + ca_data = os.environ.get(EnvironmentVariables.AZURE_KUBERNETES_CA_DATA) + if ca_file and ca_data: + raise ValueError( + "Both AZURE_KUBERNETES_CA_FILE and AZURE_KUBERNETES_CA_DATA are set. Only one should be set." + ) + + transport = _get_transport( + sni=sni, + token_proxy_endpoint=token_proxy_endpoint, + ca_file=ca_file, + ca_data=ca_data, ) - self._sni = os.environ.get(EnvironmentVariables.AZURE_KUBERNETES_SNI_NAME) - ca_file = os.environ.get(EnvironmentVariables.AZURE_KUBERNETES_CA_FILE) - ca_data = os.environ.get(EnvironmentVariables.AZURE_KUBERNETES_CA_DATA) - - if ca_file and ca_data: - raise ValueError( - "Both AZURE_KUBERNETES_CA_FILE and AZURE_KUBERNETES_CA_DATA are set. Only one should be set." - ) - - transport = _get_transport( - sni=self._sni, - token_proxy_endpoint=self._token_proxy_endpoint, - ca_file=ca_file, - ca_data=ca_data, - ) - - if transport: - kwargs["transport"] = transport - else: - raise ValueError( - "Async transport creation failed. Ensure that the aiohttp package is installed to enable token " - "proxy usage in this credential." - ) + if transport: + kwargs["transport"] = transport + else: + raise ValueError( + "Async transport creation failed. Ensure that the aiohttp package is installed to enable " + "token proxy usage in this credential." + ) super().__init__( tenant_id=tenant_id, diff --git a/sdk/identity/azure-identity/tests/proxy_server.py b/sdk/identity/azure-identity/tests/proxy_server.py index 22a33973edff..0ef7e90c96bd 100644 --- a/sdk/identity/azure-identity/tests/proxy_server.py +++ b/sdk/identity/azure-identity/tests/proxy_server.py @@ -248,9 +248,6 @@ def start(self): context = ssl.create_default_context(ssl.Purpose.CLIENT_AUTH) if self.cert_file and self.key_file: context.load_cert_chain(self.cert_file, self.key_file) - # Disable certificate verification for testing - context.check_hostname = False - context.verify_mode = ssl.CERT_NONE self.server.socket = context.wrap_socket(self.server.socket, server_side=True) diff --git a/sdk/identity/azure-identity/tests/test_workload_identity_credential.py b/sdk/identity/azure-identity/tests/test_workload_identity_credential.py index b35ba07bf24d..d62cfa1508f6 100644 --- a/sdk/identity/azure-identity/tests/test_workload_identity_credential.py +++ b/sdk/identity/azure-identity/tests/test_workload_identity_credential.py @@ -163,8 +163,8 @@ def test_use_token_proxy_minimal_config(self): ca_data=None, ) - def test_use_token_proxy_missing_proxy_endpoint_raises_error(self): - """Test that use_token_proxy=True without proxy endpoint raises ValueError.""" + def test_use_token_proxy_missing_proxy_endpoint(self): + """Test that use_token_proxy=True without proxy endpoint uses the normal transport.""" tenant_id = "tenant-id" client_id = "client-id" token_file_path = "foo-path" @@ -174,13 +174,14 @@ def test_use_token_proxy_missing_proxy_endpoint_raises_error(self): if "AZURE_KUBERNETES_TOKEN_PROXY" in os.environ: del os.environ["AZURE_KUBERNETES_TOKEN_PROXY"] - with pytest.raises(ValueError, match="use_token_proxy is True, but no token proxy endpoint was found"): + with patch("azure.identity._credentials.workload_identity._get_transport") as mock_get_transport: WorkloadIdentityCredential( tenant_id=tenant_id, client_id=client_id, token_file_path=token_file_path, use_token_proxy=True, ) + mock_get_transport.assert_not_called() def test_use_token_proxy_both_ca_file_and_data_raises_error(self): """Test that setting both CA file and CA data raises ValueError.""" diff --git a/sdk/identity/azure-identity/tests/test_workload_identity_credential_async.py b/sdk/identity/azure-identity/tests/test_workload_identity_credential_async.py index 9ab0a04a5591..ae4371dbfbbb 100644 --- a/sdk/identity/azure-identity/tests/test_workload_identity_credential_async.py +++ b/sdk/identity/azure-identity/tests/test_workload_identity_credential_async.py @@ -165,8 +165,8 @@ def test_use_token_proxy_minimal_config(self): ca_data=None, ) - def test_use_token_proxy_missing_proxy_endpoint_raises_error(self): - """Test that use_token_proxy=True without proxy endpoint raises ValueError.""" + def test_use_token_proxy_missing_proxy_endpoint(self): + """Test that use_token_proxy=True without proxy endpoint uses the normal transport.""" tenant_id = "tenant-id" client_id = "client-id" token_file_path = "foo-path" @@ -176,13 +176,14 @@ def test_use_token_proxy_missing_proxy_endpoint_raises_error(self): if "AZURE_KUBERNETES_TOKEN_PROXY" in os.environ: del os.environ["AZURE_KUBERNETES_TOKEN_PROXY"] - with pytest.raises(ValueError, match="use_token_proxy is True, but no token proxy endpoint was found"): + with patch("azure.identity.aio._credentials.workload_identity._get_transport") as mock_get_transport: WorkloadIdentityCredential( tenant_id=tenant_id, client_id=client_id, token_file_path=token_file_path, use_token_proxy=True, ) + mock_get_transport.assert_not_called() def test_use_token_proxy_both_ca_file_and_data_raises_error(self): """Test that setting both CA file and CA data raises ValueError.""" From 93ae93486195defb9c8a72bfd305db95e6185b59 Mon Sep 17 00:00:00 2001 From: Paul Van Eck Date: Tue, 4 Nov 2025 20:14:56 +0000 Subject: [PATCH 10/12] kwarg update Signed-off-by: Paul Van Eck --- sdk/identity/azure-identity/CHANGELOG.md | 2 +- .../azure/identity/_credentials/workload_identity.py | 6 +----- .../azure/identity/aio/_credentials/workload_identity.py | 6 +----- 3 files changed, 3 insertions(+), 11 deletions(-) diff --git a/sdk/identity/azure-identity/CHANGELOG.md b/sdk/identity/azure-identity/CHANGELOG.md index cd474e5b547d..d434d7017122 100644 --- a/sdk/identity/azure-identity/CHANGELOG.md +++ b/sdk/identity/azure-identity/CHANGELOG.md @@ -4,7 +4,7 @@ ### Features Added -- Added support for `WorkloadIdentityCredential` identity binding mode in AKS environments. This feature addresses Entra's limitation on the number of federated identity credentials (FICs) per managed identity by utilizing an AKS proxy that handles FIC exchanges on behalf of pods. When the environment is properly configured and `use_token_proxy=True` is specified in the `WorkloadIdentityCredential` constructor, all token requests are routed through this proxy service. ([#43287](https://github.com/Azure/azure-sdk-for-python/pull/43287)) +- Added support for `WorkloadIdentityCredential` identity binding mode in AKS environments. This feature addresses Entra's limitation on the number of federated identity credentials (FICs) per managed identity by utilizing an AKS proxy that handles FIC exchanges on behalf of pods. ([#43287](https://github.com/Azure/azure-sdk-for-python/pull/43287)) ### Breaking Changes diff --git a/sdk/identity/azure-identity/azure/identity/_credentials/workload_identity.py b/sdk/identity/azure-identity/azure/identity/_credentials/workload_identity.py index cc32a172d2bd..cabeaed986ef 100644 --- a/sdk/identity/azure-identity/azure/identity/_credentials/workload_identity.py +++ b/sdk/identity/azure-identity/azure/identity/_credentials/workload_identity.py @@ -57,9 +57,6 @@ class WorkloadIdentityCredential(ClientAssertionCredential, TokenFileMixin): :keyword str client_id: The client ID of a Microsoft Entra app registration. :keyword str token_file_path: The path to a file containing a Kubernetes service account token that authenticates the identity. - :keyword bool use_token_proxy: Whether or not to read token proxy configuration from environment variables and use - a token proxy to acquire tokens. If this value is True and proxy configuration isn't present or this value is - False, the credential will request tokens directly from Entra ID. Defaults to False. .. admonition:: Example: @@ -77,7 +74,6 @@ def __init__( tenant_id: Optional[str] = None, client_id: Optional[str] = None, token_file_path: Optional[str] = None, - use_token_proxy: bool = False, **kwargs: Any, ) -> None: tenant_id = tenant_id or os.environ.get(EnvironmentVariables.AZURE_TENANT_ID) @@ -104,7 +100,7 @@ def __init__( self._token_file_path = token_file_path - if use_token_proxy: + if kwargs.pop("use_token_proxy", False): token_proxy_endpoint = os.environ.get(EnvironmentVariables.AZURE_KUBERNETES_TOKEN_PROXY) if token_proxy_endpoint: sni = os.environ.get(EnvironmentVariables.AZURE_KUBERNETES_SNI_NAME) diff --git a/sdk/identity/azure-identity/azure/identity/aio/_credentials/workload_identity.py b/sdk/identity/azure-identity/azure/identity/aio/_credentials/workload_identity.py index 66851d16c4d7..5b70d03bc0f8 100644 --- a/sdk/identity/azure-identity/azure/identity/aio/_credentials/workload_identity.py +++ b/sdk/identity/azure-identity/azure/identity/aio/_credentials/workload_identity.py @@ -32,9 +32,6 @@ class WorkloadIdentityCredential(ClientAssertionCredential, TokenFileMixin): :keyword str client_id: The client ID of a Microsoft Entra app registration. :keyword str token_file_path: The path to a file containing a Kubernetes service account token that authenticates the identity. - :keyword bool use_token_proxy: Whether or not to read token proxy configuration from environment variables and use - a token proxy to acquire tokens. If this value is True and proxy configuration isn't present or this value is - False, the credential will request tokens directly from Entra ID. Defaults to False. .. admonition:: Example: @@ -52,7 +49,6 @@ def __init__( tenant_id: Optional[str] = None, client_id: Optional[str] = None, token_file_path: Optional[str] = None, - use_token_proxy: bool = False, **kwargs: Any, ) -> None: tenant_id = tenant_id or os.environ.get(EnvironmentVariables.AZURE_TENANT_ID) @@ -79,7 +75,7 @@ def __init__( self._token_file_path = token_file_path - if use_token_proxy: + if kwargs.pop("use_token_proxy", False): token_proxy_endpoint = os.environ.get(EnvironmentVariables.AZURE_KUBERNETES_TOKEN_PROXY) if token_proxy_endpoint: sni = os.environ.get(EnvironmentVariables.AZURE_KUBERNETES_SNI_NAME) From 92182b4d6e538808ff4bbf4e22a92551d5246e14 Mon Sep 17 00:00:00 2001 From: Paul Van Eck Date: Wed, 5 Nov 2025 20:26:23 +0000 Subject: [PATCH 11/12] Updates - Add fallback AsyncioRequestsTransport - Add additonal error raise conditions - Ensure token binding kwarg isn't respected in chained credentials Signed-off-by: Paul Van Eck --- .../_credentials/workload_identity.py | 22 +- ...py => token_binding_transport_requests.py} | 15 +- .../aio/_credentials/workload_identity.py | 41 ++- ....py => token_binding_transport_aiohttp.py} | 2 +- .../token_binding_transport_asyncio.py | 45 +++ .../test_workload_identity_credential.py | 67 ++++- ...test_workload_identity_credential_async.py | 257 +++++++++++++++++- 7 files changed, 414 insertions(+), 35 deletions(-) rename sdk/identity/azure-identity/azure/identity/_internal/{token_binding_transport.py => token_binding_transport_requests.py} (85%) rename sdk/identity/azure-identity/azure/identity/aio/_internal/{token_binding_transport.py => token_binding_transport_aiohttp.py} (93%) create mode 100644 sdk/identity/azure-identity/azure/identity/aio/_internal/token_binding_transport_asyncio.py diff --git a/sdk/identity/azure-identity/azure/identity/_credentials/workload_identity.py b/sdk/identity/azure-identity/azure/identity/_credentials/workload_identity.py index cabeaed986ef..ef2453e2079d 100644 --- a/sdk/identity/azure-identity/azure/identity/_credentials/workload_identity.py +++ b/sdk/identity/azure-identity/azure/identity/_credentials/workload_identity.py @@ -9,6 +9,7 @@ from .client_assertion import ClientAssertionCredential from .._constants import EnvironmentVariables +from .._internal import within_credential_chain WORKLOAD_CONFIG_ERROR = ( @@ -16,6 +17,10 @@ "configured. See the troubleshooting guide for more information: " "https://aka.ms/azsdk/python/identity/workloadidentitycredential/troubleshoot" ) +CA_DATA_FILE_ERROR = "Both AZURE_KUBERNETES_CA_FILE and AZURE_KUBERNETES_CA_DATA are set. Only one should be set." +CUSTOM_PROXY_ENV_ERROR = ( + "AZURE_KUBERNETES_TOKEN_PROXY is not set but other custom endpoint-related environment variables are present." +) class TokenFileMixin: @@ -100,17 +105,14 @@ def __init__( self._token_file_path = token_file_path - if kwargs.pop("use_token_proxy", False): + if kwargs.pop("use_token_proxy", False) and not within_credential_chain.get(): token_proxy_endpoint = os.environ.get(EnvironmentVariables.AZURE_KUBERNETES_TOKEN_PROXY) + sni = os.environ.get(EnvironmentVariables.AZURE_KUBERNETES_SNI_NAME) + ca_file = os.environ.get(EnvironmentVariables.AZURE_KUBERNETES_CA_FILE) + ca_data = os.environ.get(EnvironmentVariables.AZURE_KUBERNETES_CA_DATA) if token_proxy_endpoint: - sni = os.environ.get(EnvironmentVariables.AZURE_KUBERNETES_SNI_NAME) - ca_file = os.environ.get(EnvironmentVariables.AZURE_KUBERNETES_CA_FILE) - ca_data = os.environ.get(EnvironmentVariables.AZURE_KUBERNETES_CA_DATA) - if ca_file and ca_data: - raise ValueError( - "Both AZURE_KUBERNETES_CA_FILE and AZURE_KUBERNETES_CA_DATA are set. Only one should be set." - ) + raise ValueError(CA_DATA_FILE_ERROR) transport = _get_transport( sni=sni, @@ -126,6 +128,8 @@ def __init__( "Transport creation failed. Ensure that the requests package is installed to enable token " "proxy usage in this credential." ) + elif sni or ca_file or ca_data: + raise ValueError(CUSTOM_PROXY_ENV_ERROR) super(WorkloadIdentityCredential, self).__init__( tenant_id=tenant_id, @@ -138,7 +142,7 @@ def __init__( def _get_transport(sni, token_proxy_endpoint, ca_file, ca_data): try: - from .._internal.token_binding_transport import CustomRequestsTransport + from .._internal.token_binding_transport_requests import CustomRequestsTransport return CustomRequestsTransport( sni=sni, diff --git a/sdk/identity/azure-identity/azure/identity/_internal/token_binding_transport.py b/sdk/identity/azure-identity/azure/identity/_internal/token_binding_transport_requests.py similarity index 85% rename from sdk/identity/azure-identity/azure/identity/_internal/token_binding_transport.py rename to sdk/identity/azure-identity/azure/identity/_internal/token_binding_transport_requests.py index 4eb546c84cc1..19d15abae1d9 100644 --- a/sdk/identity/azure-identity/azure/identity/_internal/token_binding_transport.py +++ b/sdk/identity/azure-identity/azure/identity/_internal/token_binding_transport_requests.py @@ -3,7 +3,7 @@ # Licensed under the MIT License. # ------------------------------------ """ -Transport class for WorkloadIdentityCredential with token proxy support. +Requests transport class for WorkloadIdentityCredential with token proxy support. """ import ssl from typing import Any, Optional @@ -39,14 +39,13 @@ class CustomRequestsTransport(TokenBindingTransportMixin, RequestsTransport): def __init__(self, *args: Any, **kwargs: Any) -> None: self.session: Optional[Session] = None super().__init__(*args, **kwargs) - self._create_session() + self._update_adaptor() - def _create_session(self) -> None: - """Create a new requests session with custom SSL configuration.""" - if self.session: - self.session.close() + def _update_adaptor(self) -> None: + """Update the session's adapter with the current SNI and CA data.""" + if not self.session: + self.session = Session() - self.session = Session() adapter = SNIAdapter(self._sni, self._ca_data) self.session.mount("https://", adapter) @@ -58,5 +57,5 @@ def send(self, request: HttpRequest, **kwargs: Any) -> Any: self._load_ca_file_to_data() # If ca_data was updated, recreate SSL context with the new data if self._ca_data: - self._create_session() + self._update_adaptor() return super().send(request, **kwargs) diff --git a/sdk/identity/azure-identity/azure/identity/aio/_credentials/workload_identity.py b/sdk/identity/azure-identity/azure/identity/aio/_credentials/workload_identity.py index 5b70d03bc0f8..47fe90a7d8fc 100644 --- a/sdk/identity/azure-identity/azure/identity/aio/_credentials/workload_identity.py +++ b/sdk/identity/azure-identity/azure/identity/aio/_credentials/workload_identity.py @@ -7,8 +7,14 @@ from typing import Any, Optional from .client_assertion import ClientAssertionCredential -from ..._credentials.workload_identity import TokenFileMixin, WORKLOAD_CONFIG_ERROR +from ..._credentials.workload_identity import ( + TokenFileMixin, + WORKLOAD_CONFIG_ERROR, + CA_DATA_FILE_ERROR, + CUSTOM_PROXY_ENV_ERROR, +) from ..._constants import EnvironmentVariables +from ..._internal import within_credential_chain class WorkloadIdentityCredential(ClientAssertionCredential, TokenFileMixin): @@ -75,16 +81,14 @@ def __init__( self._token_file_path = token_file_path - if kwargs.pop("use_token_proxy", False): + if kwargs.pop("use_token_proxy", False) and not within_credential_chain.get(): token_proxy_endpoint = os.environ.get(EnvironmentVariables.AZURE_KUBERNETES_TOKEN_PROXY) + sni = os.environ.get(EnvironmentVariables.AZURE_KUBERNETES_SNI_NAME) + ca_file = os.environ.get(EnvironmentVariables.AZURE_KUBERNETES_CA_FILE) + ca_data = os.environ.get(EnvironmentVariables.AZURE_KUBERNETES_CA_DATA) if token_proxy_endpoint: - sni = os.environ.get(EnvironmentVariables.AZURE_KUBERNETES_SNI_NAME) - ca_file = os.environ.get(EnvironmentVariables.AZURE_KUBERNETES_CA_FILE) - ca_data = os.environ.get(EnvironmentVariables.AZURE_KUBERNETES_CA_DATA) if ca_file and ca_data: - raise ValueError( - "Both AZURE_KUBERNETES_CA_FILE and AZURE_KUBERNETES_CA_DATA are set. Only one should be set." - ) + raise ValueError(CA_DATA_FILE_ERROR) transport = _get_transport( sni=sni, @@ -97,9 +101,11 @@ def __init__( kwargs["transport"] = transport else: raise ValueError( - "Async transport creation failed. Ensure that the aiohttp package is installed to enable " - "token proxy usage in this credential." + "Async transport creation failed. Ensure that the aiohttp or requests package is installed to " + "enable token proxy usage in this credential." ) + elif sni or ca_file or ca_data: + raise ValueError(CUSTOM_PROXY_ENV_ERROR) super().__init__( tenant_id=tenant_id, @@ -112,7 +118,7 @@ def __init__( def _get_transport(sni, token_proxy_endpoint, ca_file, ca_data): try: - from .._internal.token_binding_transport import CustomAioHttpTransport + from .._internal.token_binding_transport_aiohttp import CustomAioHttpTransport return CustomAioHttpTransport( sni=sni, @@ -121,4 +127,15 @@ def _get_transport(sni, token_proxy_endpoint, ca_file, ca_data): ca_data=ca_data, ) except ImportError: - return None + # Fallback to async-wrapped requests transport + try: + from .._internal.token_binding_transport_asyncio import CustomAsyncioRequestsTransport + + return CustomAsyncioRequestsTransport( + sni=sni, + proxy_endpoint=token_proxy_endpoint, + ca_file=ca_file, + ca_data=ca_data, + ) + except ImportError: + return None diff --git a/sdk/identity/azure-identity/azure/identity/aio/_internal/token_binding_transport.py b/sdk/identity/azure-identity/azure/identity/aio/_internal/token_binding_transport_aiohttp.py similarity index 93% rename from sdk/identity/azure-identity/azure/identity/aio/_internal/token_binding_transport.py rename to sdk/identity/azure-identity/azure/identity/aio/_internal/token_binding_transport_aiohttp.py index bfc3cc53e494..85a1031f522e 100644 --- a/sdk/identity/azure-identity/azure/identity/aio/_internal/token_binding_transport.py +++ b/sdk/identity/azure-identity/azure/identity/aio/_internal/token_binding_transport_aiohttp.py @@ -3,7 +3,7 @@ # Licensed under the MIT License. # ------------------------------------ """ -Transport class for the asynchronous WorkloadIdentityCredential with token proxy support. +Aiohttp transport class for the asynchronous WorkloadIdentityCredential with token proxy support. """ import ssl from typing import Any diff --git a/sdk/identity/azure-identity/azure/identity/aio/_internal/token_binding_transport_asyncio.py b/sdk/identity/azure-identity/azure/identity/aio/_internal/token_binding_transport_asyncio.py new file mode 100644 index 000000000000..fa236c2e4d4f --- /dev/null +++ b/sdk/identity/azure-identity/azure/identity/aio/_internal/token_binding_transport_asyncio.py @@ -0,0 +1,45 @@ +# ------------------------------------ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. +# ------------------------------------ +""" +Asyncio Requests transport class for the asynchronous WorkloadIdentityCredential with token proxy support. +""" +from typing import Any, Optional + +from requests import Session +from azure.core.pipeline.transport import ( # pylint: disable=non-abstract-transport-import, no-name-in-module + AsyncioRequestsTransport, +) +from azure.core.rest import HttpRequest + +from ..._internal.token_binding_transport_mixin import TokenBindingTransportMixin +from ..._internal.token_binding_transport_requests import SNIAdapter + + +class CustomAsyncioRequestsTransport(TokenBindingTransportMixin, AsyncioRequestsTransport): + + def __init__(self, *args, **kwargs): + self.session: Optional[Session] = None + super().__init__(*args, **kwargs) + self._update_adaptor() + + def _update_adaptor(self) -> None: + """Update the session's adapter with the current SNI and CA data.""" + if not self.session: + self.session = Session() + + adapter = SNIAdapter(self._sni, self._ca_data) + self.session.mount("https://", adapter) + + async def send(self, request: HttpRequest, **kwargs: Any) -> Any: + self._update_request_url(request) + + # Check if CA file has changed and reload ca_data if needed + if self._ca_file and self._has_ca_file_changed(): + self._load_ca_file_to_data() + # If ca_data was updated, recreate SSL context with the new data + if self._ca_data: + self._update_adaptor() + + return await super().send(request, **kwargs) diff --git a/sdk/identity/azure-identity/tests/test_workload_identity_credential.py b/sdk/identity/azure-identity/tests/test_workload_identity_credential.py index d62cfa1508f6..4ef048319eb7 100644 --- a/sdk/identity/azure-identity/tests/test_workload_identity_credential.py +++ b/sdk/identity/azure-identity/tests/test_workload_identity_credential.py @@ -207,6 +207,58 @@ def test_use_token_proxy_both_ca_file_and_data_raises_error(self): use_token_proxy=True, ) + def test_use_token_proxy_missing_endpoint_with_custom_env_vars_raises_error(self): + """Test that use_token_proxy=True without proxy endpoint but with other custom env vars raises ValueError.""" + tenant_id = "tenant-id" + client_id = "client-id" + token_file_path = "foo-path" + sni_hostname = "sni.example.com" + ca_file_path = "/path/to/ca.pem" + ca_data = "-----BEGIN CERTIFICATE-----\nTest CA data\n-----END CERTIFICATE-----" + + # Ensure proxy endpoint is not set + if "AZURE_KUBERNETES_TOKEN_PROXY" in os.environ: + del os.environ["AZURE_KUBERNETES_TOKEN_PROXY"] + + # Test with SNI set but no proxy endpoint + env_vars_sni = { + "AZURE_KUBERNETES_SNI_NAME": sni_hostname, + } + with patch.dict(os.environ, env_vars_sni, clear=False): + with pytest.raises(ValueError): + WorkloadIdentityCredential( + tenant_id=tenant_id, + client_id=client_id, + token_file_path=token_file_path, + use_token_proxy=True, + ) + + # Test with CA file set but no proxy endpoint + env_vars_ca_file = { + "AZURE_KUBERNETES_CA_FILE": ca_file_path, + } + with patch.dict(os.environ, env_vars_ca_file, clear=False): + with pytest.raises(ValueError): + WorkloadIdentityCredential( + tenant_id=tenant_id, + client_id=client_id, + token_file_path=token_file_path, + use_token_proxy=True, + ) + + # Test with CA data set but no proxy endpoint + env_vars_ca_data = { + "AZURE_KUBERNETES_CA_DATA": ca_data, + } + with patch.dict(os.environ, env_vars_ca_data, clear=False): + with pytest.raises(ValueError): + WorkloadIdentityCredential( + tenant_id=tenant_id, + client_id=client_id, + token_file_path=token_file_path, + use_token_proxy=True, + ) + @pytest.mark.parametrize("get_token_method", GET_TOKEN_METHODS) def test_use_token_proxy_get_token_success(self, get_token_method): """Test successful token acquisition when using token proxy.""" @@ -413,12 +465,11 @@ def test_custom_requests_transport_send_with_ca_file_reload(self, ca_data): # Mock the CA file change detection with patch.object(transport, "_has_ca_file_changed", return_value=True): with patch.object(transport, "_load_ca_file_to_data") as mock_load_ca: - with patch.object(transport, "_create_session") as mock_create_session: + with patch.object(transport, "_update_adaptor") as mock_update_adaptor: transport.send(mock_request) mock_load_ca.assert_called_once() - # _create_session should be called if _ca_data is truthy after loading - # This depends on the implementation details + mock_update_adaptor.assert_called_once() assert mock_parent_send.call_count == 2 @@ -607,13 +658,21 @@ def test_sni_with_custom_hostname(self): assert transport is not None request = HttpRequest("GET", f"{server.base_url}/health") - response = transport.send(request) assert response.status_code == 200 data = response.json() assert data["status"] == "healthy" + # Check an invalid SNI hostname + transport = _get_transport( + sni="unmatched.sni.hostname", token_proxy_endpoint=None, ca_file=server.ca_file, ca_data=None + ) + assert transport is not None + request = HttpRequest("GET", f"{server.base_url}/health") + with pytest.raises(ServiceRequestError): + transport.send(request) + def test_ca_file_change_detection(self): """Test CA file change detection with real certificates.""" with TokenProxyTestServer(use_ssl=True) as server: diff --git a/sdk/identity/azure-identity/tests/test_workload_identity_credential_async.py b/sdk/identity/azure-identity/tests/test_workload_identity_credential_async.py index ae4371dbfbbb..50237c166b05 100644 --- a/sdk/identity/azure-identity/tests/test_workload_identity_credential_async.py +++ b/sdk/identity/azure-identity/tests/test_workload_identity_credential_async.py @@ -15,6 +15,8 @@ from azure.core.exceptions import ServiceRequestError, ServiceResponseError from azure.identity.aio import WorkloadIdentityCredential from azure.identity.aio._credentials.workload_identity import _get_transport +from azure.identity.aio._internal.token_binding_transport_aiohttp import CustomAioHttpTransport +from azure.identity.aio._internal.token_binding_transport_asyncio import CustomAsyncioRequestsTransport from helpers import mock_response, build_aad_response, GET_TOKEN_METHODS from proxy_server import TokenProxyTestServer @@ -209,6 +211,58 @@ def test_use_token_proxy_both_ca_file_and_data_raises_error(self): use_token_proxy=True, ) + def test_use_token_proxy_missing_endpoint_with_custom_env_vars_raises_error(self): + """Test that use_token_proxy=True without proxy endpoint but with other custom env vars raises ValueError.""" + tenant_id = "tenant-id" + client_id = "client-id" + token_file_path = "foo-path" + sni_hostname = "sni.example.com" + ca_file_path = "/path/to/ca.pem" + ca_data = "-----BEGIN CERTIFICATE-----\nTest CA data\n-----END CERTIFICATE-----" + + # Ensure proxy endpoint is not set + if "AZURE_KUBERNETES_TOKEN_PROXY" in os.environ: + del os.environ["AZURE_KUBERNETES_TOKEN_PROXY"] + + # Test with SNI set but no proxy endpoint + env_vars_sni = { + "AZURE_KUBERNETES_SNI_NAME": sni_hostname, + } + with patch.dict(os.environ, env_vars_sni, clear=False): + with pytest.raises(ValueError): + WorkloadIdentityCredential( + tenant_id=tenant_id, + client_id=client_id, + token_file_path=token_file_path, + use_token_proxy=True, + ) + + # Test with CA file set but no proxy endpoint + env_vars_ca_file = { + "AZURE_KUBERNETES_CA_FILE": ca_file_path, + } + with patch.dict(os.environ, env_vars_ca_file, clear=False): + with pytest.raises(ValueError): + WorkloadIdentityCredential( + tenant_id=tenant_id, + client_id=client_id, + token_file_path=token_file_path, + use_token_proxy=True, + ) + + # Test with CA data set but no proxy endpoint + env_vars_ca_data = { + "AZURE_KUBERNETES_CA_DATA": ca_data, + } + with patch.dict(os.environ, env_vars_ca_data, clear=False): + with pytest.raises(ValueError): + WorkloadIdentityCredential( + tenant_id=tenant_id, + client_id=client_id, + token_file_path=token_file_path, + use_token_proxy=True, + ) + @pytest.mark.asyncio @pytest.mark.parametrize("get_token_method", GET_TOKEN_METHODS) async def test_use_token_proxy_get_token_success(self, get_token_method): @@ -282,7 +336,7 @@ def test_get_transport_creates_workload_identity_aiohttp_transport(self, ca_data ca_data=None, ) - assert transport is not None + assert type(transport) is CustomAioHttpTransport assert hasattr(transport, "_sni") assert hasattr(transport, "_proxy_endpoint") assert hasattr(transport, "_ca_file") @@ -681,6 +735,14 @@ async def test_sni_with_custom_hostname(self): data = response.json() assert data["status"] == "healthy" + transport = _get_transport( + sni="unmatched.sni.hostname", token_proxy_endpoint=None, ca_file=server.ca_file, ca_data=None + ) + assert transport is not None + request = HttpRequest("GET", f"{server.base_url}/health") + with pytest.raises(ServiceRequestError): + await transport.send(request) + @pytest.mark.asyncio async def test_ca_file_change_detection(self): """Test CA file change detection with real certificates.""" @@ -835,3 +897,196 @@ async def make_request(request_id): for request_id, status_code, data in results: assert status_code == 200 assert data["status"] == "healthy" + + +class TestCustomAsyncioRequestsTransportFallback: + """Test cases for the custom AsyncioRequestsTransport used by WorkloadIdentityCredential.""" + + def test_get_transport_creates_workload_identity_asyncio_requests_transport(self, ca_data): + """Test that _get_transport creates WorkloadIdentityAsyncioRequestsTransport with correct parameters.""" + sni = "test.sni.com" + proxy_endpoint = "https://proxy.example.com:8080" + ca_file = PEM_CERT_PATH + + with patch.dict("sys.modules", {"azure.identity.aio._internal.token_binding_transport_aiohttp": None}): + transport = _get_transport( + sni=sni, + token_proxy_endpoint=proxy_endpoint, + ca_file=ca_file, + ca_data=None, + ) + + assert type(transport) is CustomAsyncioRequestsTransport + assert hasattr(transport, "_sni") + assert hasattr(transport, "_proxy_endpoint") + assert hasattr(transport, "_ca_file") + assert hasattr(transport, "_ca_data") + assert transport._sni == sni + assert transport._proxy_endpoint == proxy_endpoint + assert transport._ca_file == ca_file + assert transport._ca_data == ca_data + + @pytest.mark.asyncio + async def test_basic_https_request(self): + """Test basic HTTPS request to test server.""" + with TokenProxyTestServer(use_ssl=True) as server: + # Create transport with server's CA certificate + with patch.dict("sys.modules", {"azure.identity.aio._internal.token_binding_transport_aiohttp": None}): + transport = _get_transport(sni=None, token_proxy_endpoint=None, ca_file=server.ca_file, ca_data=None) + assert type(transport) is CustomAsyncioRequestsTransport + request = HttpRequest("GET", f"{server.base_url}/health") + + response = await transport.send(request) + + assert response.status_code == 200 + data = response.json() + assert data["status"] == "healthy" + assert "timestamp" in data + + @pytest.mark.asyncio + async def test_proxy_endpoint_comprehensive(self): + """Test comprehensive proxy endpoint functionality with various HTTP methods and scenarios.""" + with TokenProxyTestServer(use_ssl=True) as server: + with patch.dict("sys.modules", {"azure.identity.aio._internal.token_binding_transport_aiohttp": None}): + transport = _get_transport( + sni=None, token_proxy_endpoint=server.base_url, ca_file=server.ca_file, ca_data=None + ) + assert type(transport) is CustomAsyncioRequestsTransport + + # Test 1: POST request with JSON body through proxy + post_data = {"grant_type": "client_credentials", "scope": "https://graph.microsoft.com/.default"} + post_request = HttpRequest( + "POST", + "https://login.microsoftonline.com/tenant/oauth2/v2.0/token2", + headers={"Content-Type": "application/json"}, + json=post_data, + ) + + post_response = await transport.send(post_request) + assert post_response.status_code == 200 + post_data_response = post_response.json() + assert post_data_response["method"] == "POST" + assert post_data_response["proxied_path"] == "/tenant/oauth2/v2.0/token2" + + # Test 2: PUT request through proxy + put_request = HttpRequest( + "PUT", + "https://graph.microsoft.com/v1.0/me/profile", + headers={"Content-Type": "application/json"}, + json={"displayName": "Test User"}, + ) + + put_response = await transport.send(put_request) + assert put_response.status_code == 200 + put_data_response = put_response.json() + assert put_data_response["method"] == "PUT" + assert put_data_response["proxied_path"] == "/v1.0/me/profile" + + # Test 3: DELETE request through proxy + delete_request = HttpRequest("DELETE", "https://graph.microsoft.com/v1.0/applications/app-id") + + delete_response = await transport.send(delete_request) + assert delete_response.status_code == 200 + delete_data_response = delete_response.json() + assert delete_data_response["method"] == "DELETE" + assert delete_data_response["proxied_path"] == "/v1.0/applications/app-id" + + # Test 4: Complex URL with multiple path segments and query parameters + complex_url = "https://management.azure.com/subscriptions/sub-id/resourceGroups/rg/providers/Microsoft.Storage/storageAccounts/account?api-version=2021-04-01&expand=properties" + complex_request = HttpRequest("GET", complex_url) + + complex_response = await transport.send(complex_request) + assert complex_response.status_code == 200 + complex_data_response = complex_response.json() + assert complex_data_response["method"] == "GET" + expected_path = "/subscriptions/sub-id/resourceGroups/rg/providers/Microsoft.Storage/storageAccounts/account?api-version=2021-04-01&expand=properties" + assert complex_data_response["proxied_path"] == expected_path + + # Test 5: Request with custom headers through proxy + headers_request = HttpRequest( + "GET", + "https://vault.azure.net/secrets/test-secret?api-version=7.3", + headers={ + "Authorization": "Bearer test-token", + "X-Custom-Header": "proxy-test-value", + "User-Agent": "Azure-SDK-For-Python", + }, + ) + + headers_response = await transport.send(headers_request) + assert headers_response.status_code == 200 + headers_data_response = headers_response.json() + assert headers_data_response["method"] == "GET" + assert headers_data_response["proxied_path"] == "/secrets/test-secret?api-version=7.3" + + # Verify headers were forwarded through proxy + received_headers = headers_data_response["headers_received"] + assert "Authorization" in received_headers + assert "X-Custom-Header" in received_headers + assert received_headers["Authorization"] == "Bearer test-token" + assert received_headers["X-Custom-Header"] == "proxy-test-value" + + @pytest.mark.asyncio + async def test_sni_with_custom_hostname(self): + """Test SNI (Server Name Indication) with custom hostname.""" + with TokenProxyTestServer(use_ssl=True) as server: + with patch.dict("sys.modules", {"azure.identity.aio._internal.token_binding_transport_aiohttp": None}): + # Use SNI with a different hostname than the server + transport = _get_transport( + sni="1234.ests.aks", token_proxy_endpoint=None, ca_file=server.ca_file, ca_data=None + ) + assert type(transport) is CustomAsyncioRequestsTransport + + request = HttpRequest("GET", f"{server.base_url}/health") + response = await transport.send(request) + + assert response.status_code == 200 + data = response.json() + assert data["status"] == "healthy" + + # Check an invalid SNI hostname + transport = _get_transport( + sni="unmatched.sni.hostname", token_proxy_endpoint=None, ca_file=server.ca_file, ca_data=None + ) + assert transport is not None + request = HttpRequest("GET", f"{server.base_url}/health") + with pytest.raises(ServiceRequestError): + await transport.send(request) + + @pytest.mark.asyncio + async def test_ca_file_change_detection(self): + """Test CA file change detection with real certificates.""" + with TokenProxyTestServer(use_ssl=True) as server: + with patch.dict("sys.modules", {"azure.identity.aio._internal.token_binding_transport_aiohttp": None}): + # Create a copy of the CA file that we can modify + ca_file = server.ca_file + if ca_file is None: + pytest.skip("CA file not available") + + assert ca_file is not None # Type hint for mypy + with tempfile.NamedTemporaryFile(mode="w", delete=False, suffix=".pem") as temp_ca: + with open(ca_file, "r") as original: + original_content = original.read() + temp_ca.write(original_content) + temp_ca_path = temp_ca.name + + try: + transport = _get_transport(sni=None, token_proxy_endpoint=None, ca_file=temp_ca_path, ca_data=None) + assert type(transport) is CustomAsyncioRequestsTransport + + # First request should work + request = HttpRequest("GET", f"{server.base_url}/health") + response1 = await transport.send(request) + assert response1.status_code == 200 + + # Modify the CA file (add some content) + real_sleep(0.1) + with open(temp_ca_path, "a") as f: + f.write("\n# Modified for testing\n") + + # Second request should still work (using the same cert content) + response2 = await transport.send(request) + assert response2.status_code == 200 + + finally: + os.unlink(temp_ca_path) From 088f46e2093141fb6ac39d797cae319a3e5a6012 Mon Sep 17 00:00:00 2001 From: Paul Van Eck Date: Thu, 6 Nov 2025 22:56:31 +0000 Subject: [PATCH 12/12] Clean up tests Signed-off-by: Paul Van Eck --- .../token_binding_transport_aiohttp.py | 3 +- .../azure-identity/tests/proxy_server.py | 7 - .../test_workload_identity_credential.py | 205 +--------------- ...test_workload_identity_credential_async.py | 226 +----------------- 4 files changed, 8 insertions(+), 433 deletions(-) diff --git a/sdk/identity/azure-identity/azure/identity/aio/_internal/token_binding_transport_aiohttp.py b/sdk/identity/azure-identity/azure/identity/aio/_internal/token_binding_transport_aiohttp.py index 85a1031f522e..1a21c5b13d54 100644 --- a/sdk/identity/azure-identity/azure/identity/aio/_internal/token_binding_transport_aiohttp.py +++ b/sdk/identity/azure-identity/azure/identity/aio/_internal/token_binding_transport_aiohttp.py @@ -33,6 +33,5 @@ async def send(self, request: HttpRequest, **kwargs: Any) -> Any: if self._ca_data: self._ssl_context = ssl.create_default_context(cadata=self._ca_data) - if self._ssl_context: - kwargs.setdefault("ssl", self._ssl_context) + kwargs.setdefault("ssl", self._ssl_context) return await super().send(request, **kwargs) diff --git a/sdk/identity/azure-identity/tests/proxy_server.py b/sdk/identity/azure-identity/tests/proxy_server.py index 0ef7e90c96bd..a6e6e1e3ae61 100644 --- a/sdk/identity/azure-identity/tests/proxy_server.py +++ b/sdk/identity/azure-identity/tests/proxy_server.py @@ -83,12 +83,6 @@ def _handle_request(self): # Simulate SSL error by closing connection self.wfile.close() return - elif path == "/slow": - # Simulate slow response using threading.Event instead of time.sleep - # to avoid being mocked by conftest.py - event = threading.Event() - event.wait(timeout=2) - self._send_json_response({"message": "slow response"}) else: self._send_proxy_response(path, body, headers) @@ -320,7 +314,6 @@ def main(): print(f" {server.base_url}/oauth2/v2.0/token - Mock OAuth token endpoint") print(f" {server.base_url}/error/500 - Simulate server error") print(f" {server.base_url}/error/ssl - Simulate SSL error") - print(f" {server.base_url}/slow - Simulate slow response") print(f" {server.base_url}/ - Generic proxy response") print("\nPress Ctrl+C to stop") diff --git a/sdk/identity/azure-identity/tests/test_workload_identity_credential.py b/sdk/identity/azure-identity/tests/test_workload_identity_credential.py index 4ef048319eb7..1bcdcefbe06d 100644 --- a/sdk/identity/azure-identity/tests/test_workload_identity_credential.py +++ b/sdk/identity/azure-identity/tests/test_workload_identity_credential.py @@ -259,46 +259,6 @@ def test_use_token_proxy_missing_endpoint_with_custom_env_vars_raises_error(self use_token_proxy=True, ) - @pytest.mark.parametrize("get_token_method", GET_TOKEN_METHODS) - def test_use_token_proxy_get_token_success(self, get_token_method): - """Test successful token acquisition when using token proxy.""" - tenant_id = "tenant-id" - client_id = "client-id" - access_token = "foo-access-token" - token_file_path = "foo-path" - assertion = "foo-assertion" - proxy_endpoint = "https://proxy.example.com:8080" - - def send(request, **kwargs): - assert "claims" not in kwargs - assert "tenant_id" not in kwargs - assert request.data.get("client_assertion") == assertion - return mock_response(json_payload=build_aad_response(access_token=access_token)) - - mock_transport_instance = MagicMock(send=send) - - env_vars = { - "AZURE_KUBERNETES_TOKEN_PROXY": proxy_endpoint, - } - - with patch.dict(os.environ, env_vars, clear=False): - with patch("azure.identity._credentials.workload_identity._get_transport") as mock_get_transport: - mock_get_transport.return_value = mock_transport_instance - - credential = WorkloadIdentityCredential( - tenant_id=tenant_id, - client_id=client_id, - token_file_path=token_file_path, - use_token_proxy=True, - ) - - open_mock = mock_open(read_data=assertion) - with patch("builtins.open", open_mock): - token = getattr(credential, get_token_method)("scope") - assert token.token == access_token - - open_mock.assert_called_once_with(token_file_path, encoding="utf-8") - def test_use_token_proxy_false_does_not_create_transport(self): """Test that use_token_proxy=False (default) does not create a custom transport.""" tenant_id = "tenant-id" @@ -358,150 +318,6 @@ def test_get_transport_with_minimal_config(self): assert transport._ca_file is None assert transport._ca_data is None - def test_custom_requests_transport_send_with_sni(self): - """Test that CustomRequestsTransport.send works with SNI configuration.""" - sni = "1234.ests.aks" - proxy_endpoint = "https://proxy.example.com:8080" - - transport = _get_transport( - sni=sni, - token_proxy_endpoint=proxy_endpoint, - ca_file=None, - ca_data=None, - ) - assert transport is not None - - # Mock request - mock_request = HttpRequest("POST", "https://login.microsoftonline.com/tenant/oauth2/v2.0/token") - mock_request.data = {"grant_type": "client_credentials"} - - with patch.object(transport.__class__.__bases__[1], "send") as mock_parent_send: - mock_response_obj = MagicMock() - mock_parent_send.return_value = mock_response_obj - - response = transport.send(mock_request) - - assert response == mock_response_obj - mock_parent_send.assert_called_once_with(mock_request) - - def test_custom_requests_transport_send_updates_url(self): - """Test that CustomRequestsTransport.send updates request URL through proxy endpoint.""" - proxy_endpoint = "https://proxy.example.com:8080" - - transport = _get_transport( - sni=None, - token_proxy_endpoint=proxy_endpoint, - ca_file=None, - ca_data=None, - ) - assert transport is not None - - # Mock request with original URL - original_url = "https://login.microsoftonline.com/tenant/oauth2/v2.0/token" - mock_request = HttpRequest("POST", original_url) - - with patch.object(transport.__class__.__bases__[1], "send") as mock_parent_send: - mock_response_obj = MagicMock() - mock_parent_send.return_value = mock_response_obj - - transport.send(mock_request) - - # Verify that _update_request_url was called (URL should be modified) - # The exact URL modification is tested in the mixin tests - mock_parent_send.assert_called_once_with(mock_request) - - def test_custom_requests_transport_send_with_ca_data(self, ca_data): - """Test that CustomRequestsTransport.send works with CA data.""" - transport = _get_transport( - sni=None, - token_proxy_endpoint="https://proxy.example.com:8080", - ca_file=None, - ca_data=ca_data, - ) - assert transport is not None - - mock_request = HttpRequest("GET", "https://login.microsoftonline.com/tenant/oauth2/v2.0/token") - - with patch.object(transport.__class__.__bases__[1], "send") as mock_parent_send: - mock_response_obj = MagicMock() - mock_parent_send.return_value = mock_response_obj - - response = transport.send(mock_request) - - assert response == mock_response_obj - mock_parent_send.assert_called_once_with(mock_request) - - def test_custom_requests_transport_send_with_ca_file_reload(self, ca_data): - """Test that CustomRequestsTransport.send reloads CA file when changed.""" - # Create a temporary CA file - with tempfile.NamedTemporaryFile(mode="w", delete=False, suffix=".pem") as ca_file: - ca_file.write(ca_data) - ca_file_path = ca_file.name - - try: - transport = _get_transport( - sni=None, - token_proxy_endpoint="https://proxy.example.com:8080", - ca_file=ca_file_path, - ca_data=None, - ) - assert transport is not None - - mock_request = HttpRequest("GET", "https://login.microsoftonline.com/tenant/oauth2/v2.0/token") - - # Mock the parent send method and file modification time tracking - with patch.object(transport.__class__.__bases__[1], "send") as mock_parent_send: - mock_response_obj = MagicMock() - mock_parent_send.return_value = mock_response_obj - - # First request - transport.send(mock_request) - - # Simulate file change by modifying the file - time.sleep(0.1) # Ensure different modification time - with open(ca_file_path, "a") as f: - f.write("\n# Modified for testing\n") - - # Mock the CA file change detection - with patch.object(transport, "_has_ca_file_changed", return_value=True): - with patch.object(transport, "_load_ca_file_to_data") as mock_load_ca: - with patch.object(transport, "_update_adaptor") as mock_update_adaptor: - transport.send(mock_request) - - mock_load_ca.assert_called_once() - mock_update_adaptor.assert_called_once() - - assert mock_parent_send.call_count == 2 - - finally: - os.unlink(ca_file_path) - - def test_custom_requests_transport_initialization_with_ca_data(self, ca_data): - """Test CustomRequestsTransport initialization with CA data creates SSL context.""" - transport = _get_transport( - sni=None, - token_proxy_endpoint="https://proxy.example.com:8080", - ca_file=None, - ca_data=ca_data, - ) - assert transport is not None - - # Verify transport was created properly - assert transport._ca_data == ca_data - - def test_custom_requests_transport_initialization_without_ca_data(self): - """Test CustomRequestsTransport initialization without CA data.""" - transport = _get_transport( - sni=None, - token_proxy_endpoint="https://proxy.example.com:8080", - ca_file=None, - ca_data=None, - ) - assert transport is not None - - # Verify transport was created properly - assert transport._ca_data is None - def test_custom_requests_transport_inherits_from_token_binding_mixin(self): """Test that CustomRequestsTransport inherits from TokenBindingTransportMixin.""" transport = _get_transport( @@ -704,6 +520,8 @@ def test_ca_file_change_detection(self): # Second request should still work (using the same cert content) response2 = transport.send(request) + assert transport._ca_data != original_content + assert "Modified for testing" in transport._ca_data assert response2.status_code == 200 finally: @@ -735,25 +553,6 @@ def test_server_error_response(self): assert response.status_code == 500 # Should not raise exception, just return error response - def test_slow_server_response(self): - """Test handling of slow server responses.""" - - with TokenProxyTestServer(use_ssl=True) as server: - transport = _get_transport(sni=None, token_proxy_endpoint=None, ca_file=server.ca_file, ca_data=None) - assert transport is not None - - request = HttpRequest("GET", f"{server.base_url}/slow") - - start_time = time.time() - response = transport.send(request) - elapsed_time = time.time() - start_time - - assert response.status_code == 200 - # Should take at least 2 seconds (server waits for 2s) - assert elapsed_time >= 2.0 - data = response.json() - assert data["message"] == "slow response" - def test_custom_headers_preserved(self): """Test that custom headers are preserved and sent to server.""" with TokenProxyTestServer(use_ssl=True) as server: diff --git a/sdk/identity/azure-identity/tests/test_workload_identity_credential_async.py b/sdk/identity/azure-identity/tests/test_workload_identity_credential_async.py index 50237c166b05..6580af09648b 100644 --- a/sdk/identity/azure-identity/tests/test_workload_identity_credential_async.py +++ b/sdk/identity/azure-identity/tests/test_workload_identity_credential_async.py @@ -346,6 +346,9 @@ def test_get_transport_creates_workload_identity_aiohttp_transport(self, ca_data assert transport._ca_file == ca_file assert transport._ca_data == ca_data + assert hasattr(transport, "_ssl_context") + assert transport._ssl_context is not None + def test_get_transport_with_minimal_config(self): """Test _get_transport with minimal configuration.""" proxy_endpoint = "https://proxy.example.com:8080" @@ -363,211 +366,9 @@ def test_get_transport_with_minimal_config(self): assert transport._ca_file is None assert transport._ca_data is None - @pytest.mark.asyncio - async def test_workload_identity_aiohttp_transport_send_with_sni(self): - """Test that WorkloadIdentityAioHttpTransport.send sets server_hostname correctly.""" - sni = "test.sni.com" - proxy_endpoint = "https://proxy.example.com:8080" - - transport = _get_transport( - sni=sni, - token_proxy_endpoint=proxy_endpoint, - ca_file=None, - ca_data=None, - ) - assert transport is not None - - # Mock the parent send method - mock_request = MagicMock() - mock_request.url = "https://login.microsoftonline.com/tenant/oauth2/v2.0/token" - - with patch.object(transport.__class__.__bases__[1], "send") as mock_parent_send: - mock_parent_send.return_value = MagicMock() - - await transport.send(mock_request) - - # Verify parent send was called with server_hostname set - mock_parent_send.assert_called_once() - call_args = mock_parent_send.call_args - assert call_args[1]["server_hostname"] == sni - - @pytest.mark.asyncio - async def test_workload_identity_aiohttp_transport_send_updates_url(self): - """Test that WorkloadIdentityAioHttpTransport.send updates request URL with proxy endpoint.""" - proxy_endpoint = "https://proxy.example.com:8080/path" - original_url = "https://login.microsoftonline.com/tenant/oauth2/v2.0/token" - expected_url = "https://proxy.example.com:8080/path/tenant/oauth2/v2.0/token" - - transport = _get_transport( - sni=None, - token_proxy_endpoint=proxy_endpoint, - ca_file=None, - ca_data=None, - ) - assert transport is not None - - mock_request = MagicMock() - mock_request.url = original_url - - with patch.object(transport.__class__.__bases__[1], "send") as mock_parent_send: - mock_parent_send.return_value = MagicMock() - - await transport.send(mock_request) - - # Verify URL was updated to use proxy endpoint - assert mock_request.url == expected_url - - @pytest.mark.asyncio - async def test_workload_identity_aiohttp_transport_send_with_ca_data(self, ca_data): - """Test that WorkloadIdentityAioHttpTransport.send creates SSL context from CA data.""" - proxy_endpoint = "https://proxy.example.com:8080" - - transport = _get_transport( - sni=None, - token_proxy_endpoint=proxy_endpoint, - ca_file=None, - ca_data=ca_data, - ) - assert transport is not None - - mock_request = MagicMock() - mock_request.url = "https://login.microsoftonline.com/tenant/oauth2/v2.0/token" - - with patch.object(transport.__class__.__bases__[1], "send") as mock_parent_send: - mock_parent_send.return_value = MagicMock() - - await transport.send(mock_request) - - # Verify SSL context was set - mock_parent_send.assert_called_once() - call_args = mock_parent_send.call_args - assert "ssl" in call_args[1] - assert call_args[1]["ssl"] is not None - - @pytest.mark.asyncio - async def test_workload_identity_aiohttp_transport_send_with_ca_file_reload(self, ca_data): - """Test that WorkloadIdentityAioHttpTransport.send reloads CA file when changed.""" - - # Create a temporary CA file - with tempfile.NamedTemporaryFile(mode="w", delete=False, suffix=".pem") as ca_file: - ca_file.write(ca_data) - ca_file_path = ca_file.name - - try: - proxy_endpoint = "https://proxy.example.com:8080" - - transport = _get_transport( - sni=None, - token_proxy_endpoint=proxy_endpoint, - ca_file=ca_file_path, - ca_data=None, - ) - - assert transport is not None - # Store original CA data - original_ca_data = transport._ca_data - - # Simulate file change by modifying mtime tracking - real_sleep(0.1) # Ensure different mtime - with open(ca_file_path, "a") as f: - f.write("\n") - - mock_request = MagicMock() - mock_request.url = "https://login.microsoftonline.com/tenant/oauth2/v2.0/token" - - with patch.object(transport.__class__.__bases__[1], "send") as mock_parent_send: - mock_parent_send.return_value = MagicMock() - - await transport.send(mock_request) - - # Verify CA data was reloaded - assert transport._ca_data != original_ca_data - - finally: - # Clean up temporary file - os.unlink(ca_file_path) - - @pytest.mark.asyncio - async def test_workload_identity_aiohttp_transport_context_manager(self): - """Test that WorkloadIdentityAioHttpTransport works as async context manager.""" - transport = _get_transport( - sni=None, - token_proxy_endpoint="https://proxy.example.com:8080", - ca_file=None, - ca_data=None, - ) - assert transport is not None - - # Mock the parent context manager methods - with patch.object(transport.__class__.__bases__[1], "__aenter__") as mock_aenter, patch.object( - transport.__class__.__bases__[1], "__aexit__" - ) as mock_aexit: - - mock_aenter.return_value = transport - mock_aexit.return_value = None - - async with transport as ctx_transport: - assert ctx_transport == transport - - mock_aenter.assert_called_once() - mock_aexit.assert_called_once() - - def test_workload_identity_aiohttp_transport_initialization_with_ca_data(self, ca_data): - """Test WorkloadIdentityAioHttpTransport initialization with CA data creates SSL context.""" - transport = _get_transport( - sni=None, - token_proxy_endpoint="https://proxy.example.com:8080", - ca_file=None, - ca_data=ca_data, - ) - assert transport is not None - - # Verify SSL context was created during initialization assert hasattr(transport, "_ssl_context") assert transport._ssl_context is not None - def test_workload_identity_aiohttp_transport_initialization_without_ca_data(self): - """Test WorkloadIdentityAioHttpTransport initialization without CA data.""" - transport = _get_transport( - sni=None, - token_proxy_endpoint="https://proxy.example.com:8080", - ca_file=None, - ca_data=None, - ) - assert transport is not None - - # Verify SSL context is created with None ca_data (creates default context) - assert hasattr(transport, "_ssl_context") - # SSL context should still be created even with None ca_data - - @pytest.mark.asyncio - async def test_workload_identity_aiohttp_transport_send_no_ssl_context_when_no_ca_data(self): - """Test that no SSL context is passed when ca_data is None and SSL context creation fails.""" - transport = _get_transport( - sni=None, - token_proxy_endpoint="https://proxy.example.com:8080", - ca_file=None, - ca_data=None, - ) - assert transport is not None - - # Mock SSL context creation to return None - with patch("ssl.create_default_context", return_value=None): - # Manually set ssl_context to None to test the conditional logic - with patch.object(transport, "_ssl_context", None): - mock_request = MagicMock() - mock_request.url = "https://login.microsoftonline.com/tenant/oauth2/v2.0/token" - - with patch.object(transport.__class__.__bases__[1], "send") as mock_parent_send: - mock_parent_send.return_value = MagicMock() - - await transport.send(mock_request) - - # Verify SSL context was not set when None - mock_parent_send.assert_called_once() - call_args = mock_parent_send.call_args - assert "ssl" not in call_args[1] or call_args[1].get("ssl") is None - def test_workload_identity_aiohttp_transport_inherits_from_token_binding_mixin(self): """Test that WorkloadIdentityAioHttpTransport inherits from TokenBindingTransportMixin.""" transport = _get_transport( @@ -775,6 +576,7 @@ async def test_ca_file_change_detection(self): # Second request should still work (using the same cert content) response2 = await transport.send(request) + assert transport._ca_data != original_content assert response2.status_code == 200 finally: @@ -808,25 +610,6 @@ async def test_server_error_response(self): assert response.status_code == 500 # Should not raise exception, just return error response - @pytest.mark.asyncio - async def test_slow_server_response(self): - """Test handling of slow server responses.""" - with TokenProxyTestServer(use_ssl=True) as server: - transport = _get_transport(sni=None, token_proxy_endpoint=None, ca_file=server.ca_file, ca_data=None) - assert transport is not None - - request = HttpRequest("GET", f"{server.base_url}/slow") - - start_time = time.time() - response = await transport.send(request) - elapsed_time = time.time() - start_time - - assert response.status_code == 200 - # Should take at least 2 seconds (server waits for 2s) - assert elapsed_time >= 2.0 - data = response.json() - assert data["message"] == "slow response" - @pytest.mark.asyncio async def test_custom_headers_preserved(self): """Test that custom headers are preserved and sent to server.""" @@ -1086,6 +869,7 @@ async def test_ca_file_change_detection(self): # Second request should still work (using the same cert content) response2 = await transport.send(request) + assert transport._ca_data != original_content assert response2.status_code == 200 finally: