Move ttl validation to Config. Make config a parameter for TokenCache.

Author: jonathan343

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

@@ -19,6 +19,58 @@
 from smithy_aws_core.identity import AWSCredentialsIdentity
 
 
+@dataclass(init=False)
+class Config:
+    """Configuration for EC2Metadata."""
+
+    _HOST_MAPPING = {"IPv4": "169.254.169.254", "IPv6": "[fd00:ec2::254]"}
+    _MIN_TTL = 5
+    _MAX_TTL = 21600
+
+    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 = _MAX_TTL,
+        ec2_instance_profile_name: str | None = None,
+    ):
+        #  TODO: Implement retries.
+        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 = self._validate_token_ttl(token_ttl)
+        self.ec2_instance_profile_name = ec2_instance_profile_name
+
+    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 _resolve_endpoint(
+        self, endpoint_uri: URI | None, endpoint_mode: Literal["IPv4", "IPv6"]
+    ) -> URI:
+        if endpoint_uri is not None:
+            return endpoint_uri
+
+        return URI(
+            scheme="http",
+            host=self._HOST_MAPPING.get(endpoint_mode, self._HOST_MAPPING["IPv4"]),
+        )
+
+
 class Token:
     """Represents an IMDSv2 session token with a value and method for checking
     expiration."""
@@ -43,27 +95,15 @@ class TokenCache:
     In addition, it knows how to refresh itself.
     """
 
-    _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
-    ):
+    def __init__(self, http_client: HTTPClient, config: Config):
         self._http_client = http_client
-        self._base_uri = base_uri
-        self._token_ttl = self._validate_token_ttl(token_ttl)
+        self._config = config
+        self._base_uri = config.endpoint_uri
         self._refresh_lock = asyncio.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()
@@ -78,7 +118,7 @@ async def _refresh(self) -> None:
                     # TODO: Add user-agent
                     Field(
                         name="x-aws-ec2-metadata-token-ttl-seconds",
-                        values=[str(self._token_ttl)],
+                        values=[str(self._config.token_ttl)],
                     ),
                 ]
             )
@@ -93,7 +133,7 @@ async def _refresh(self) -> None:
             )
             response = await self._http_client.send(request)
             token_value = await response.consume_body_async()
-            self._token = Token(token_value, self._token_ttl)
+            self._token = Token(token_value, self._config.token_ttl)
 
     async def get_token(self) -> Token:
         """Get the current token, refreshing it if expired."""
@@ -103,55 +143,12 @@ async def get_token(self) -> Token:
         return self._token
 
 
-@dataclass(init=False)
-class Config:
-    """Configuration for EC2Metadata."""
-
-    _HOST_MAPPING = {"IPv4": "169.254.169.254", "IPv6": "[fd00:ec2::254]"}
-
-    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
-
-        return URI(
-            scheme="http",
-            host=self._HOST_MAPPING.get(endpoint_mode, self._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,
+            http_client=self._http_client, config=self._config
         )
 
     async def get(self, *, path: str) -> str: