Skip to content

Commit 4250940

Browse files
pvaneckpvaneck
authored
[Identity] Implement binding mode support in WorkloadIdentityCredential (Azure#43287)
Signed-off-by: Paul Van Eck <[email protected]>
1 parent 9ceb7da commit 4250940

13 files changed

+2097
-5
lines changed

‎sdk/identity/azure-identity/CHANGELOG.md‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,11 @@
11
# Release History
22

3-
## 1.25.2 (Unreleased)
3+
## 1.26.0b1 (Unreleased)
44

55
### Features Added
66

7+
- 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))
8+
79
### Breaking Changes
810

911
### Bugs Fixed

‎sdk/identity/azure-identity/azure/identity/_constants.py‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -68,5 +68,10 @@ class EnvironmentVariables:
6868
AZURE_REGIONAL_AUTHORITY_NAME = "AZURE_REGIONAL_AUTHORITY_NAME"
6969

7070
AZURE_FEDERATED_TOKEN_FILE = "AZURE_FEDERATED_TOKEN_FILE"
71+
AZURE_KUBERNETES_SNI_NAME = "AZURE_KUBERNETES_SNI_NAME"
72+
AZURE_KUBERNETES_TOKEN_PROXY = "AZURE_KUBERNETES_TOKEN_PROXY"
73+
AZURE_KUBERNETES_CA_FILE = "AZURE_KUBERNETES_CA_FILE"
74+
AZURE_KUBERNETES_CA_DATA = "AZURE_KUBERNETES_CA_DATA"
75+
7176
AZURE_TOKEN_CREDENTIALS = "AZURE_TOKEN_CREDENTIALS"
7277
WORKLOAD_IDENTITY_VARS = (AZURE_AUTHORITY_HOST, AZURE_TENANT_ID, AZURE_FEDERATED_TOKEN_FILE)

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

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,13 +9,18 @@
99

1010
from .client_assertion import ClientAssertionCredential
1111
from .._constants import EnvironmentVariables
12+
from .._internal import within_credential_chain
1213

1314

1415
WORKLOAD_CONFIG_ERROR = (
1516
"WorkloadIdentityCredential authentication unavailable. The workload options are not fully "
1617
"configured. See the troubleshooting guide for more information: "
1718
"https://aka.ms/azsdk/python/identity/workloadidentitycredential/troubleshoot"
1819
)
20+
CA_DATA_FILE_ERROR = "Both AZURE_KUBERNETES_CA_FILE and AZURE_KUBERNETES_CA_DATA are set. Only one should be set."
21+
CUSTOM_PROXY_ENV_ERROR = (
22+
"AZURE_KUBERNETES_TOKEN_PROXY is not set but other custom endpoint-related environment variables are present."
23+
)
1924

2025

2126
class TokenFileMixin:
@@ -99,10 +104,52 @@ def __init__(
99104
assert token_file_path is not None
100105

101106
self._token_file_path = token_file_path
107+
108+
if kwargs.pop("use_token_proxy", False) and not within_credential_chain.get():
109+
token_proxy_endpoint = os.environ.get(EnvironmentVariables.AZURE_KUBERNETES_TOKEN_PROXY)
110+
sni = os.environ.get(EnvironmentVariables.AZURE_KUBERNETES_SNI_NAME)
111+
ca_file = os.environ.get(EnvironmentVariables.AZURE_KUBERNETES_CA_FILE)
112+
ca_data = os.environ.get(EnvironmentVariables.AZURE_KUBERNETES_CA_DATA)
113+
if token_proxy_endpoint:
114+
if ca_file and ca_data:
115+
raise ValueError(CA_DATA_FILE_ERROR)
116+
117+
transport = _get_transport(
118+
sni=sni,
119+
token_proxy_endpoint=token_proxy_endpoint,
120+
ca_file=ca_file,
121+
ca_data=ca_data,
122+
)
123+
124+
if transport:
125+
kwargs["transport"] = transport
126+
else:
127+
raise ValueError(
128+
"Transport creation failed. Ensure that the requests package is installed to enable token "
129+
"proxy usage in this credential."
130+
)
131+
elif sni or ca_file or ca_data:
132+
raise ValueError(CUSTOM_PROXY_ENV_ERROR)
133+
102134
super(WorkloadIdentityCredential, self).__init__(
103135
tenant_id=tenant_id,
104136
client_id=client_id,
105137
func=self._get_service_account_token,
106138
token_file_path=token_file_path,
107139
**kwargs,
108140
)
141+
142+
143+
def _get_transport(sni, token_proxy_endpoint, ca_file, ca_data):
144+
try:
145+
from .._internal.token_binding_transport_requests import CustomRequestsTransport
146+
147+
return CustomRequestsTransport(
148+
sni=sni,
149+
proxy_endpoint=token_proxy_endpoint,
150+
ca_file=ca_file,
151+
ca_data=ca_data,
152+
)
153+
154+
except ImportError:
155+
return None

‎sdk/identity/azure-identity/azure/identity/_internal/token_binding_transport_mixin.py‎

Lines changed: 123 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,123 @@
1+
# ------------------------------------
2+
# Copyright (c) Microsoft Corporation.
3+
# Licensed under the MIT License.
4+
# ------------------------------------
5+
# cspell:ignore cafile
6+
import os
7+
import urllib.parse
8+
from typing import Optional, Any
9+
10+
from azure.core.rest import HttpRequest
11+
12+
13+
class TokenBindingTransportMixin:
14+
"""Mixin class providing URL validation, CA file tracking, and proxy URL functionality for transport classes."""
15+
16+
def __init__(self, **kwargs: Any) -> None:
17+
"""Initialize CA file tracking and proxy attributes."""
18+
self._ca_file = kwargs.pop("ca_file", None)
19+
self._ca_data = kwargs.pop("ca_data", None)
20+
self._proxy_endpoint = kwargs.pop("proxy_endpoint", None)
21+
self._sni = kwargs.pop("sni", None)
22+
23+
self._ca_file_mtime: Optional[float] = None
24+
25+
if self._ca_file and self._ca_data:
26+
raise ValueError("Both ca_file and ca_data are set. Only one should be set")
27+
28+
if self._proxy_endpoint:
29+
self._validate_url(self._proxy_endpoint)
30+
31+
# If we have a ca_file, read it once and store as ca_data
32+
if self._ca_file:
33+
self._load_ca_file_to_data()
34+
35+
super().__init__()
36+
37+
def _validate_url(self, url: str) -> None:
38+
"""Validate that a URL meets security requirements for HTTPS connections.
39+
40+
:param url: The URL to validate.
41+
:type url: str
42+
:raises ValueError: If the URL does not meet security requirements.
43+
"""
44+
parsed_url = urllib.parse.urlparse(url)
45+
if parsed_url.scheme != "https":
46+
raise ValueError(f"Endpoint URL ({url}) must use the 'https' scheme. Got '{parsed_url.scheme}' instead.")
47+
if parsed_url.username or parsed_url.password:
48+
raise ValueError(f"Endpoint URL ({url}) must not contain username or password.")
49+
if parsed_url.fragment:
50+
raise ValueError(f"Endpoint URL ({url}) must not contain a fragment.")
51+
if parsed_url.query:
52+
raise ValueError(f"Endpoint URL ({url}) must not contain query parameters.")
53+
54+
def _load_ca_file_to_data(self) -> None:
55+
"""Load CA file content into ca_data and track modification time.
56+
57+
:raises ValueError: If the CA file is empty on first read.
58+
"""
59+
try:
60+
with open(self._ca_file, "r", encoding="utf-8") as f:
61+
content = f.read()
62+
63+
# Check if the file is empty
64+
if not content:
65+
# If no prior ca_data exists (first read), fail
66+
if self._ca_data is None:
67+
raise ValueError(f"CA file ({self._ca_file}) is empty. Cannot establish secure connection.")
68+
# If we had prior ca_data, keep it (mid-rotation scenario)
69+
return
70+
71+
# File has content, update ca_data and tracking
72+
self._ca_data = content
73+
self._ca_file_mtime = os.path.getmtime(self._ca_file)
74+
except (OSError, IOError) as e:
75+
# If no prior ca_data exists (first read), fail
76+
if self._ca_data is None:
77+
raise ValueError(f"Failed to read CA file ({self._ca_file}): {e}") from e
78+
# If we can't read the file, keep existing ca_data but clear mtime
79+
# so we'll try to reload on the next change check
80+
self._ca_file_mtime = None
81+
82+
def _has_ca_file_changed(self) -> bool:
83+
"""Check if the CA file has changed since last load.
84+
85+
:return: True if the CA file has changed, False otherwise.
86+
:rtype: bool
87+
"""
88+
if not self._ca_file:
89+
return False
90+
91+
if not os.path.exists(self._ca_file):
92+
# File was deleted, consider this a change if we had data before
93+
return self._ca_data is not None or self._ca_file_mtime is not None
94+
95+
try:
96+
# Check modification time
97+
current_mtime = os.path.getmtime(self._ca_file)
98+
return self._ca_file_mtime != current_mtime
99+
except (OSError, IOError):
100+
# If we can't read the file stats, assume it changed
101+
return True
102+
103+
def _update_request_url(self, request: HttpRequest) -> None:
104+
"""Update the request URL to use proxy endpoint if configured.
105+
106+
:param request: The HTTP request object to update.
107+
:type request: ~azure.core.rest.HttpRequest
108+
"""
109+
if self._proxy_endpoint:
110+
parsed_request_url = urllib.parse.urlparse(request.url)
111+
parsed_proxy_url = urllib.parse.urlparse(self._proxy_endpoint)
112+
combined_path = parsed_proxy_url.path.rstrip("/") + "/" + parsed_request_url.path.lstrip("/")
113+
new_url = urllib.parse.urlunparse(
114+
(
115+
parsed_proxy_url.scheme,
116+
parsed_proxy_url.netloc,
117+
combined_path,
118+
parsed_request_url.params,
119+
parsed_request_url.query,
120+
parsed_request_url.fragment,
121+
)
122+
)
123+
request.url = new_url

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

Lines changed: 61 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,61 @@
1+
# ------------------------------------
2+
# Copyright (c) Microsoft Corporation.
3+
# Licensed under the MIT License.
4+
# ------------------------------------
5+
"""
6+
Requests transport class for WorkloadIdentityCredential with token proxy support.
7+
"""
8+
import ssl
9+
from typing import Any, Optional
10+
11+
from requests.adapters import HTTPAdapter
12+
from requests import Session
13+
from azure.core.pipeline.transport import ( # pylint: disable=non-abstract-transport-import, no-name-in-module
14+
RequestsTransport,
15+
)
16+
from azure.core.rest import HttpRequest
17+
18+
from .token_binding_transport_mixin import TokenBindingTransportMixin
19+
20+
21+
class SNIAdapter(HTTPAdapter):
22+
"""A custom HTTPAdapter that allows setting a custom SNI hostname."""
23+
24+
def __init__(self, server_hostname: Optional[str], ca_data: Optional[str], **kwargs: Any) -> None:
25+
self.server_hostname = server_hostname
26+
self.ca_data = ca_data
27+
super().__init__(**kwargs)
28+
29+
def init_poolmanager(self, connections: int, maxsize: int, block: bool = False, **pool_kwargs: Any) -> None:
30+
if self.server_hostname:
31+
pool_kwargs["server_hostname"] = self.server_hostname
32+
pool_kwargs["ssl_context"] = ssl.create_default_context(cadata=self.ca_data)
33+
super().init_poolmanager(connections, maxsize, block, **pool_kwargs)
34+
35+
36+
class CustomRequestsTransport(TokenBindingTransportMixin, RequestsTransport):
37+
"""Custom RequestsTransport with SNI and CA certificate support for WorkloadIdentityCredential."""
38+
39+
def __init__(self, *args: Any, **kwargs: Any) -> None:
40+
self.session: Optional[Session] = None
41+
super().__init__(*args, **kwargs)
42+
self._update_adaptor()
43+
44+
def _update_adaptor(self) -> None:
45+
"""Update the session's adapter with the current SNI and CA data."""
46+
if not self.session:
47+
self.session = Session()
48+
49+
adapter = SNIAdapter(self._sni, self._ca_data)
50+
self.session.mount("https://", adapter)
51+
52+
def send(self, request: HttpRequest, **kwargs: Any) -> Any:
53+
self._update_request_url(request)
54+
55+
# Check if CA file has changed and reload ca_data if needed
56+
if self._ca_file and self._has_ca_file_changed():
57+
self._load_ca_file_to_data()
58+
# If ca_data was updated, recreate SSL context with the new data
59+
if self._ca_data:
60+
self._update_adaptor()
61+
return super().send(request, **kwargs)

‎sdk/identity/azure-identity/azure/identity/_version.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,4 +2,4 @@
22
# Copyright (c) Microsoft Corporation.
33
# Licensed under the MIT License.
44
# ------------------------------------
5-
VERSION = "1.25.2"
5+
VERSION = "1.26.0b1"

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

Lines changed: 61 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,11 +2,19 @@
22
# Copyright (c) Microsoft Corporation.
33
# Licensed under the MIT License.
44
# ------------------------------------
5+
# cspell:ignore cafile
56
import os
67
from typing import Any, Optional
8+
79
from .client_assertion import ClientAssertionCredential
8-
from ..._credentials.workload_identity import TokenFileMixin, WORKLOAD_CONFIG_ERROR
10+
from ..._credentials.workload_identity import (
11+
TokenFileMixin,
12+
WORKLOAD_CONFIG_ERROR,
13+
CA_DATA_FILE_ERROR,
14+
CUSTOM_PROXY_ENV_ERROR,
15+
)
916
from ..._constants import EnvironmentVariables
17+
from ..._internal import within_credential_chain
1018

1119

1220
class WorkloadIdentityCredential(ClientAssertionCredential, TokenFileMixin):
@@ -72,10 +80,62 @@ def __init__(
7280
assert token_file_path is not None
7381

7482
self._token_file_path = token_file_path
83+
84+
if kwargs.pop("use_token_proxy", False) and not within_credential_chain.get():
85+
token_proxy_endpoint = os.environ.get(EnvironmentVariables.AZURE_KUBERNETES_TOKEN_PROXY)
86+
sni = os.environ.get(EnvironmentVariables.AZURE_KUBERNETES_SNI_NAME)
87+
ca_file = os.environ.get(EnvironmentVariables.AZURE_KUBERNETES_CA_FILE)
88+
ca_data = os.environ.get(EnvironmentVariables.AZURE_KUBERNETES_CA_DATA)
89+
if token_proxy_endpoint:
90+
if ca_file and ca_data:
91+
raise ValueError(CA_DATA_FILE_ERROR)
92+
93+
transport = _get_transport(
94+
sni=sni,
95+
token_proxy_endpoint=token_proxy_endpoint,
96+
ca_file=ca_file,
97+
ca_data=ca_data,
98+
)
99+
100+
if transport:
101+
kwargs["transport"] = transport
102+
else:
103+
raise ValueError(
104+
"Async transport creation failed. Ensure that the aiohttp or requests package is installed to "
105+
"enable token proxy usage in this credential."
106+
)
107+
elif sni or ca_file or ca_data:
108+
raise ValueError(CUSTOM_PROXY_ENV_ERROR)
109+
75110
super().__init__(
76111
tenant_id=tenant_id,
77112
client_id=client_id,
78113
func=self._get_service_account_token,
79114
token_file_path=token_file_path,
80115
**kwargs,
81116
)
117+
118+
119+
def _get_transport(sni, token_proxy_endpoint, ca_file, ca_data):
120+
try:
121+
from .._internal.token_binding_transport_aiohttp import CustomAioHttpTransport
122+
123+
return CustomAioHttpTransport(
124+
sni=sni,
125+
proxy_endpoint=token_proxy_endpoint,
126+
ca_file=ca_file,
127+
ca_data=ca_data,
128+
)
129+
except ImportError:
130+
# Fallback to async-wrapped requests transport
131+
try:
132+
from .._internal.token_binding_transport_asyncio import CustomAsyncioRequestsTransport
133+
134+
return CustomAsyncioRequestsTransport(
135+
sni=sni,
136+
proxy_endpoint=token_proxy_endpoint,
137+
ca_file=ca_file,
138+
ca_data=ca_data,
139+
)
140+
except ImportError:
141+
return None

0 commit comments

Comments
 (0)