[Identity] Implement binding mode support in WorkloadIdentityCredential

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 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
 
 ### Bugs Fixed

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)

sdk/identity/azure-identity/azure/identity/_credentials/workload_identity.py

@@ -9,13 +9,18 @@
 
 from .client_assertion import ClientAssertionCredential
 from .._constants import EnvironmentVariables
+from .._internal import within_credential_chain
 
 
 WORKLOAD_CONFIG_ERROR = (
     "WorkloadIdentityCredential authentication unavailable. The workload options are not fully "
     "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:
@@ -99,10 +104,52 @@ def __init__(
         assert token_file_path is not None
 
         self._token_file_path = token_file_path
+
+        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:
+                if ca_file and ca_data:
+                    raise ValueError(CA_DATA_FILE_ERROR)
+
+                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."
+                    )
+            elif sni or ca_file or ca_data:
+                raise ValueError(CUSTOM_PROXY_ENV_ERROR)
+
         super(WorkloadIdentityCredential, self).__init__(
             tenant_id=tenant_id,
             client_id=client_id,
             func=self._get_service_account_token,
             token_file_path=token_file_path,
             **kwargs,
         )
+
+
+def _get_transport(sni, token_proxy_endpoint, ca_file, ca_data):
+    try:
+        from .._internal.token_binding_transport_requests import CustomRequestsTransport
+
+        return CustomRequestsTransport(
+            sni=sni,
+            proxy_endpoint=token_proxy_endpoint,
+            ca_file=ca_file,
+            ca_data=ca_data,
+        )
+
+    except ImportError:
+        return None

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, Any
+
+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: 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)
+        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
+        """
+        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(
+                (
+                    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

sdk/identity/azure-identity/azure/identity/_internal/token_binding_transport_requests.py

@@ -0,0 +1,61 @@
+# ------------------------------------
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+# ------------------------------------
+"""
+Requests 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._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)
+
+    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 super().send(request, **kwargs)

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"

sdk/identity/azure-identity/azure/identity/aio/_credentials/workload_identity.py

@@ -2,11 +2,19 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
 # ------------------------------------
+# cspell:ignore cafile
 import os
 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):
@@ -72,10 +80,62 @@ def __init__(
         assert token_file_path is not None
 
         self._token_file_path = token_file_path
+
+        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:
+                if ca_file and ca_data:
+                    raise ValueError(CA_DATA_FILE_ERROR)
+
+                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(
+                        "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,
             client_id=client_id,
             func=self._get_service_account_token,
             token_file_path=token_file_path,
             **kwargs,
         )
+
+
+def _get_transport(sni, token_proxy_endpoint, ca_file, ca_data):
+    try:
+        from .._internal.token_binding_transport_aiohttp import CustomAioHttpTransport
+
+        return CustomAioHttpTransport(
+            sni=sni,
+            proxy_endpoint=token_proxy_endpoint,
+            ca_file=ca_file,
+            ca_data=ca_data,
+        )
+    except ImportError:
+        # 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