Add basic IMDS credentials resolver

Author: jonathan343

packages/smithy-aws-core/src/smithy_aws_core/credentials_resolvers/__init__.py

@@ -2,5 +2,10 @@
 #  SPDX-License-Identifier: Apache-2.0
 from .environment import EnvironmentCredentialsResolver
 from .static import StaticCredentialsResolver
+from .imds import IMDSCredentialsResolver
 
-__all__ = ("EnvironmentCredentialsResolver", "StaticCredentialsResolver")
+__all__ = (
+    "EnvironmentCredentialsResolver",
+    "StaticCredentialsResolver",
+    "IMDSCredentialsResolver",
+)

packages/smithy-aws-core/src/smithy_aws_core/credentials_resolvers/imds.py

@@ -0,0 +1,230 @@
+#  Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#  SPDX-License-Identifier: Apache-2.0
+import json
+import threading
+from dataclasses import dataclass
+from datetime import datetime, timedelta
+from typing import Literal
+
+from smithy_core import URI
+from smithy_core.aio.interfaces.identity import IdentityResolver
+from smithy_core.exceptions import SmithyIdentityException
+from smithy_core.interfaces.identity import IdentityProperties
+from smithy_core.interfaces.retries import RetryStrategy
+from smithy_core.retries import SimpleRetryStrategy
+from smithy_http import Field, Fields
+from smithy_http.aio import HTTPRequest
+from smithy_http.aio.interfaces import HTTPClient
+
+from smithy_aws_core.identity import AWSCredentialsIdentity
+
+
+class Token:
+    """Represents an IMDSv2 session token with a value and method for checking
+    expiration."""
+
+    def __init__(self, value: bytes, ttl: int):
+        self._value = value
+        self._ttl = ttl
+        self._created_time = datetime.now()
+
+    def is_expired(self) -> bool:
+        """Check if the token has expired."""
+        return datetime.now() - self._created_time >= timedelta(seconds=self._ttl)
+
+    @property
+    def value(self) -> bytes:
+        return self._value
+
+
+class TokenCache:
+    """Holds the token needed to fetch instance metadata. In addition, it knows how to
+    refresh itself.
+
+    :param HTTPClient http_client: The client used for making http requests.
+    :param int token_ttl: The time in seconds before a token expires.
+    """
+
+    _MIN_TTL = 5
+    _MAX_TTL = 21600
+    _TOKEN_PATH = "/latest/api/token"
+
+    def __init__(
+        self, http_client: HTTPClient, base_uri: URI, token_ttl: int = _MAX_TTL
+    ):
+        self._http_client = http_client
+        self._base_uri = base_uri
+        self._token_ttl = self._validate_token_ttl(token_ttl)
+        self._refresh_lock = threading.Lock()
+        self._token = None
+
+    def _validate_token_ttl(self, ttl: int) -> int:
+        """Validates the token TTL value."""
+        if not self._MIN_TTL <= ttl <= self._MAX_TTL:
+            raise ValueError(
+                f"Token TTL must be between {self._MIN_TTL} and {self._MAX_TTL} seconds."
+            )
+        return ttl
+
+    def _should_refresh(self) -> bool:
+        """Determines if the token should be refreshed."""
+        return self._token is None or self._token.is_expired()
+
+    async def _refresh(self) -> None:
+        """Refreshes the token if needed, with thread safety."""
+        with self._refresh_lock:
+            if not self._should_refresh():
+                return
+            headers = Fields(
+                [
+                    # TODO: Add user-agent
+                    Field(
+                        name="x-aws-ec2-metadata-token-ttl-seconds",
+                        values=[str(self._token_ttl)],
+                    ),
+                ]
+            )
+            request = HTTPRequest(
+                method="PUT",
+                destination=URI(
+                    scheme=self._base_uri.scheme,
+                    host=self._base_uri.host,
+                    path=self._TOKEN_PATH,
+                ),
+                fields=headers,
+            )
+            response = await self._http_client.send(request)
+            token_value = await response.consume_body_async()
+            self._token = Token(token_value, self._token_ttl)
+
+    async def get_token(self) -> Token:
+        """Get the current token, refreshing it if expired."""
+        if self._should_refresh():
+            await self._refresh()
+        assert self._token is not None
+        return self._token
+
+
+@dataclass(init=False)
+class Config:
+    """Configuration for EC2Metadata."""
+
+    retry_strategy: RetryStrategy
+    endpoint_uri: URI
+    endpoint_mode: Literal["IPv4", "IPv6"]
+    port: int
+    token_ttl: int
+
+    def __init__(
+        self,
+        *,
+        retry_strategy: RetryStrategy | None = None,
+        endpoint_uri: URI | None = None,
+        endpoint_mode: Literal["IPv4", "IPv6"] = "IPv4",
+        port: int = 80,
+        token_ttl: int = 21600,
+        ec2_instance_profile_name: str | None = None,
+    ):
+        self.retry_strategy = retry_strategy or SimpleRetryStrategy(max_attempts=3)
+        self.endpoint_mode = endpoint_mode
+        self.endpoint_uri = self._resolve_endpoint(endpoint_uri, endpoint_mode)
+        self.port = port
+        self.token_ttl = token_ttl
+        self.ec2_instance_profile_name = ec2_instance_profile_name
+
+    def _resolve_endpoint(
+        self, endpoint_uri: URI | None, endpoint_mode: Literal["IPv4", "IPv6"]
+    ) -> URI:
+        if endpoint_uri is not None:
+            return endpoint_uri
+
+        host_mapping = {"IPv4": "169.254.169.254", "IPv6": "[fd00:ec2::254]"}
+
+        return URI(
+            scheme="http", host=host_mapping.get(endpoint_mode, host_mapping["IPv4"])
+        )
+
+
+class EC2Metadata:
+    def __init__(self, http_client: HTTPClient, config: Config | None = None):
+        self._http_client = http_client
+        self._config = config or Config()
+        self._token_cache = TokenCache(
+            http_client=self._http_client,
+            base_uri=self._config.endpoint_uri,
+            token_ttl=self._config.token_ttl,
+        )
+
+    async def get(self, *, path: str) -> str:
+        token = await self._token_cache.get_token()
+        headers = Fields(
+            [
+                # TODO: Add user-agent
+                Field(
+                    name="x-aws-ec2-metadata-token",
+                    values=[token.value.decode("utf-8")],
+                )
+            ]
+        )
+        request = HTTPRequest(
+            method="GET",
+            destination=URI(
+                scheme=self._config.endpoint_uri.scheme,
+                host=self._config.endpoint_uri.host,
+                port=self._config.port,
+                path=path,
+            ),
+            fields=headers,
+        )
+        response = await self._http_client.send(request=request)
+        body = await response.consume_body_async()
+        return body.decode("utf-8")
+
+
+class IMDSCredentialsResolver(
+    IdentityResolver[AWSCredentialsIdentity, IdentityProperties]
+):
+    """Resolves AWS Credentials from an EC2 Instance Metadata Service (IMDS) client."""
+
+    # TODO: Handle fallback to legacy path when a 404 is received.
+    _METADATA_PATH_BASE = "/latest/meta-data/iam/security-credentials-extended/"
+
+    def __init__(self, http_client: HTTPClient, config: Config | None = None):
+        self._http_client = http_client
+        self._ec2_metadata_client = EC2Metadata(http_client=http_client, config=config)
+        self._config = config or Config()
+        self._credentials = None
+        self._profile_name = self._config.ec2_instance_profile_name
+
+    async def get_identity(
+        self, *, identity_properties: IdentityProperties
+    ) -> AWSCredentialsIdentity:
+        if self._credentials is not None:
+            return self._credentials
+
+        profile = self._profile_name
+        if profile is None:
+            profile = await self._ec2_metadata_client.get(path=self._METADATA_PATH_BASE)
+
+        creds_str = await self._ec2_metadata_client.get(
+            path=f"{self._METADATA_PATH_BASE}/{profile}"
+        )
+        creds = json.loads(creds_str)
+
+        access_key_id = creds.get("AccessKeyId")
+        secret_access_key = creds.get("SecretAccessKey")
+        session_token = creds.get("Token")
+        account_id = creds.get("AccountId")
+
+        if access_key_id is None or secret_access_key is None:
+            raise SmithyIdentityException(
+                "AccessKeyId and SecretAccessKey are required"
+            )
+
+        self._credentials = AWSCredentialsIdentity(
+            access_key_id=access_key_id,
+            secret_access_key=secret_access_key,
+            session_token=session_token,
+            account_id=account_id,
+        )
+        return self._credentials