Skip to content

Commit 82d49f8

Browse files
authored
Add support for ContainerCredentialResolver (#516)
1 parent 0476832 commit 82d49f8

File tree

2 files changed

+571
-0
lines changed

2 files changed

+571
-0
lines changed
Lines changed: 206 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,206 @@
1+
# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
2+
# SPDX-License-Identifier: Apache-2.0
3+
import asyncio
4+
import ipaddress
5+
import json
6+
import os
7+
from dataclasses import dataclass
8+
from datetime import UTC, datetime
9+
from urllib.parse import urlparse
10+
11+
from smithy_core import URI
12+
from smithy_core.aio.interfaces.identity import IdentityResolver
13+
from smithy_core.exceptions import SmithyIdentityException
14+
from smithy_http import Field, Fields
15+
from smithy_http.aio import HTTPRequest
16+
from smithy_http.aio.interfaces import HTTPClient, HTTPResponse
17+
18+
from smithy_aws_core.identity import AWSCredentialsIdentity, IdentityProperties
19+
20+
_CONTAINER_METADATA_IP = "169.254.170.2"
21+
_CONTAINER_METADATA_ALLOWED_HOSTS = {
22+
_CONTAINER_METADATA_IP,
23+
"169.254.170.23",
24+
"fd00:ec2::23",
25+
"localhost",
26+
}
27+
_DEFAULT_TIMEOUT = 2
28+
_DEFAULT_RETRIES = 3
29+
_SLEEP_SECONDS = 1
30+
31+
32+
@dataclass
33+
class ContainerCredentialConfig:
34+
"""Configuration for container credential retrieval operations."""
35+
36+
timeout: int = _DEFAULT_TIMEOUT
37+
retries: int = _DEFAULT_RETRIES
38+
39+
40+
class ContainerMetadataClient:
41+
"""Client for remote credential retrieval in Container environments like ECS/EKS."""
42+
43+
def __init__(self, http_client: HTTPClient, config: ContainerCredentialConfig):
44+
self._http_client = http_client
45+
self._config = config
46+
47+
def _validate_allowed_url(self, uri: URI) -> None:
48+
if self._is_loopback(uri.host):
49+
return
50+
51+
if not self._is_allowed_container_metadata_host(uri.host):
52+
raise SmithyIdentityException(
53+
f"Unsupported host '{uri.host}'. "
54+
f"Can only retrieve metadata from a loopback address or "
55+
f"one of: {', '.join(_CONTAINER_METADATA_ALLOWED_HOSTS)}"
56+
)
57+
58+
async def get_credentials(self, uri: URI, fields: Fields) -> dict[str, str]:
59+
self._validate_allowed_url(uri)
60+
fields.set_field(Field(name="Accept", values=["application/json"]))
61+
62+
attempts = 0
63+
last_exc = None
64+
while attempts < self._config.retries:
65+
try:
66+
request = HTTPRequest(
67+
method="GET",
68+
destination=uri,
69+
fields=fields,
70+
)
71+
response: HTTPResponse = await self._http_client.send(request)
72+
body = await response.consume_body_async()
73+
if response.status != 200:
74+
raise SmithyIdentityException(
75+
f"Container metadata service returned {response.status}: "
76+
f"{body.decode('utf-8')}"
77+
)
78+
try:
79+
return json.loads(body.decode("utf-8"))
80+
except Exception as e:
81+
raise SmithyIdentityException(
82+
f"Unable to parse JSON from container metadata: {body.decode('utf-8')}"
83+
) from e
84+
except Exception as e:
85+
last_exc = e
86+
await asyncio.sleep(_SLEEP_SECONDS)
87+
attempts += 1
88+
89+
raise SmithyIdentityException(
90+
f"Failed to retrieve container metadata after {self._config.retries} attempt(s)"
91+
) from last_exc
92+
93+
def _is_loopback(self, hostname: str) -> bool:
94+
try:
95+
return ipaddress.ip_address(hostname).is_loopback
96+
except ValueError:
97+
return False
98+
99+
def _is_allowed_container_metadata_host(self, hostname: str) -> bool:
100+
return hostname in _CONTAINER_METADATA_ALLOWED_HOSTS
101+
102+
103+
class ContainerCredentialResolver(
104+
IdentityResolver[AWSCredentialsIdentity, IdentityProperties]
105+
):
106+
"""Resolves AWS Credentials from container credential sources."""
107+
108+
ENV_VAR = "AWS_CONTAINER_CREDENTIALS_RELATIVE_URI"
109+
ENV_VAR_FULL = "AWS_CONTAINER_CREDENTIALS_FULL_URI"
110+
ENV_VAR_AUTH_TOKEN = "AWS_CONTAINER_AUTHORIZATION_TOKEN" # noqa: S105
111+
ENV_VAR_AUTH_TOKEN_FILE = "AWS_CONTAINER_AUTHORIZATION_TOKEN_FILE" # noqa: S105
112+
113+
def __init__(
114+
self,
115+
http_client: HTTPClient,
116+
config: ContainerCredentialConfig | None = None,
117+
):
118+
self._http_client = http_client
119+
self._config = config or ContainerCredentialConfig()
120+
self._client = ContainerMetadataClient(http_client, self._config)
121+
self._credentials = None
122+
123+
async def _resolve_uri_from_env(self) -> URI:
124+
if self.ENV_VAR in os.environ:
125+
return URI(
126+
scheme="http",
127+
host=_CONTAINER_METADATA_IP,
128+
path=os.environ[self.ENV_VAR],
129+
)
130+
elif self.ENV_VAR_FULL in os.environ:
131+
parsed = urlparse(os.environ[self.ENV_VAR_FULL])
132+
return URI(
133+
scheme=parsed.scheme,
134+
host=parsed.hostname or "",
135+
port=parsed.port,
136+
path=parsed.path,
137+
)
138+
else:
139+
raise SmithyIdentityException(
140+
f"Neither {self.ENV_VAR} or {self.ENV_VAR_FULL} environment "
141+
"variables are set. Unable to resolve credentials."
142+
)
143+
144+
async def _resolve_fields_from_env(self) -> Fields:
145+
fields = Fields()
146+
if self.ENV_VAR_AUTH_TOKEN_FILE in os.environ:
147+
try:
148+
filename = os.environ[self.ENV_VAR_AUTH_TOKEN_FILE]
149+
auth_token = await asyncio.to_thread(self._read_file, filename)
150+
except (FileNotFoundError, PermissionError) as e:
151+
raise SmithyIdentityException(
152+
f"Unable to open {os.environ[self.ENV_VAR_AUTH_TOKEN_FILE]}."
153+
) from e
154+
155+
fields.set_field(Field(name="Authorization", values=[auth_token]))
156+
elif self.ENV_VAR_AUTH_TOKEN in os.environ:
157+
auth_token = os.environ[self.ENV_VAR_AUTH_TOKEN]
158+
fields.set_field(Field(name="Authorization", values=[auth_token]))
159+
160+
return fields
161+
162+
def _read_file(self, filename: str) -> str:
163+
with open(filename) as f:
164+
try:
165+
return f.read().strip()
166+
except UnicodeDecodeError as e:
167+
raise SmithyIdentityException(
168+
f"Unable to read valid utf-8 bytes from {filename}."
169+
) from e
170+
171+
async def get_identity(
172+
self, *, identity_properties: IdentityProperties
173+
) -> AWSCredentialsIdentity:
174+
if (
175+
self._credentials is not None
176+
and self._credentials.expiration
177+
and datetime.now(UTC) < self._credentials.expiration
178+
):
179+
return self._credentials
180+
181+
uri = await self._resolve_uri_from_env()
182+
fields = await self._resolve_fields_from_env()
183+
creds = await self._client.get_credentials(uri, fields)
184+
185+
access_key_id = creds.get("AccessKeyId")
186+
secret_access_key = creds.get("SecretAccessKey")
187+
session_token = creds.get("Token")
188+
expiration = creds.get("Expiration")
189+
account_id = creds.get("AccountId")
190+
191+
if isinstance(expiration, str):
192+
expiration = datetime.fromisoformat(expiration).replace(tzinfo=UTC)
193+
194+
if access_key_id is None or secret_access_key is None:
195+
raise SmithyIdentityException(
196+
"AccessKeyId and SecretAccessKey are required for container credentials"
197+
)
198+
199+
self._credentials = AWSCredentialsIdentity(
200+
access_key_id=access_key_id,
201+
secret_access_key=secret_access_key,
202+
session_token=session_token,
203+
expiration=expiration,
204+
account_id=account_id,
205+
)
206+
return self._credentials

0 commit comments

Comments
 (0)