Move event idenity to sign function

Author: JordonPhillips

packages/aws-sdk-signers/src/aws_sdk_signers/signers.py

@@ -803,27 +803,25 @@ class AsyncEventSigner:
     def __init__(
         self,
         *,
-        properties: SigV4SigningProperties,
-        identity: _AWSCredentialsIdentity,
         initial_signature: bytes,
         event_encoder_cls: type["EventHeaderEncoder"],
     ):
         self._prior_signature = initial_signature
         self._signing_lock = asyncio.Lock()
         self._event_encoder_cls = event_encoder_cls
-        self._properties = properties
-        self._identity = identity
 
     async def sign(
         self,
         *,
         event: "EventMessage",
+        identity: _AWSCredentialsIdentity,
+        properties: SigV4SigningProperties,
     ) -> "EventMessage":
         async with self._signing_lock:
             # Copy and prepopulate any missing values in the
             # signing properties.
             new_signing_properties = SigV4SigningProperties(  # type: ignore
-                **self._properties
+                **properties
             )
             # TODO: If date is in properties, parse a datetime from it.
             date_obj = datetime.datetime.now(datetime.UTC)
@@ -848,7 +846,7 @@ async def sign(
                 prior_signature=self._prior_signature,
             )
             event_signature = await self._sign_event(
-                identity=self._identity,
+                identity=identity,
                 timestamp=timestamp,
                 string_to_sign=string_to_sign,
                 properties=new_signing_properties,

packages/smithy-aws-core/src/smithy_aws_core/auth/sigv4.py

@@ -100,20 +100,15 @@ def signer(self) -> SigV4Signer:
         return self._signer
 
     def event_signer(
-        self,
-        request: HTTPRequest,
-        identity: AWSCredentialsIdentity,
-        properties: SigV4SigningProperties,
-    ) -> EventSigner | None:
+        self, request: HTTPRequest
+    ) -> EventSigner[AWSCredentialsIdentity, SigV4SigningProperties] | None:
         if not HAS_EVENT_STREAM:
             return None
 
         auth_value = request.fields["Authorization"].as_string()
         signature: str = re.split("Signature=", auth_value)[-1]
 
         return AsyncEventSigner(
-            identity=identity,
-            properties=properties,
             initial_signature=signature.encode("utf-8"),
             event_encoder_cls=EventHeaderEncoder,
         )

packages/smithy-aws-event-stream/src/smithy_aws_event_stream/aio/__init__.py

@@ -3,10 +3,13 @@
 import asyncio
 import logging
 from collections.abc import Callable
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
 
 from smithy_core.aio.interfaces import AsyncByteStream, AsyncWriter
 from smithy_core.aio.interfaces.auth import EventSigner
 from smithy_core.aio.interfaces.eventstream import EventPublisher, EventReceiver
+from smithy_core.aio.interfaces.identity import IdentityResolver
 from smithy_core.codecs import Codec
 from smithy_core.deserializers import DeserializeableShape, ShapeDeserializer
 from smithy_core.exceptions import ExpectationNotMetException
@@ -20,16 +23,29 @@
 logger = logging.getLogger(__name__)
 
 
+if TYPE_CHECKING:
+    from aws_sdk_signers import SigV4SigningProperties
+    from smithy_aws_core.identity import AWSCredentialsIdentity, AWSIdentityProperties
+
+
+@dataclass
+class SigningConfig:
+    signer: "EventSigner[AWSCredentialsIdentity, SigV4SigningProperties]"
+    signing_properties: "SigV4SigningProperties"
+    identity_resolver: "IdentityResolver[AWSCredentialsIdentity, AWSIdentityProperties]"
+    identity_properties: "AWSIdentityProperties"
+
+
 class AWSEventPublisher[E: SerializeableShape](EventPublisher[E]):
     def __init__(
         self,
         payload_codec: Codec,
         async_writer: AsyncWriter,
-        signer: EventSigner | None = None,
+        signing_config: SigningConfig | None = None,
         is_client_mode: bool = True,
     ):
         self._writer = async_writer
-        self._signer = signer
+        self._signing_config = signing_config
         self._serializer = _EventSerializer(
             payload_codec=payload_codec, is_client_mode=is_client_mode
         )
@@ -45,8 +61,16 @@ async def send(self, event: E) -> None:
             raise ExpectationNotMetException(
                 "Expected an event message to be serialized, but was None."
             )
-        if self._signer is not None:
-            result = await self._signer.sign(event=result)
+
+        if self._signing_config is not None:
+            identity = await self._signing_config.identity_resolver.get_identity(
+                properties=self._signing_config.identity_properties
+            )
+            result = await self._signing_config.signer.sign(
+                event=event,
+                identity=identity,
+                properties=self._signing_config.signing_properties,
+            )
 
         encoded_result = result.encode()
         try:

packages/smithy-core/src/smithy_core/aio/identity.py

@@ -1,16 +1,35 @@
 #  Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 #  SPDX-License-Identifier: Apache-2.0
 import logging
-from collections.abc import Sequence
-from typing import Final
+from collections.abc import Mapping, Sequence
+from typing import Any, Final
 
 from ..exceptions import SmithyIdentityException
+from ..interfaces.identity import Identity
 from .interfaces.identity import IdentityResolver
 
 logger: Final = logging.getLogger(__name__)
 
 
-class ChainedIdentityResolver[I, IP](IdentityResolver[I, IP]):
+# TODO: turn this into a decorator
+class CachingIdentityResolver[I: Identity, IP: Mapping[str, Any]](
+    IdentityResolver[I, IP]
+):
+    def __init__(self) -> None:
+        self._cached: I | None = None
+
+    async def get_identity(self, *, properties: IP) -> I:
+        if self._cached is None or self._cached.is_expired:
+            self._cached = await self._get_identity(properties=properties)
+        return self._cached
+
+    async def _get_identity(self, *, properties: IP) -> I:
+        raise NotImplementedError
+
+
+class ChainedIdentityResolver[I: Identity, IP: Mapping[str, Any]](
+    CachingIdentityResolver[I, IP]
+):
     """Attempts to resolve an identity by checking a sequence of sub-resolvers.
 
     If a nested resolver raises a :py:class:`SmithyIdentityException`, the next
@@ -22,9 +41,10 @@ def __init__(self, resolvers: Sequence[IdentityResolver[I, IP]]) -> None:
 
         :param resolvers: The sequence of resolvers to resolve identity from.
         """
+        super().__init__()
         self._resolvers = resolvers
 
-    async def get_identity(self, *, properties: IP) -> I:
+    async def _get_identity(self, *, properties: IP) -> I:
         logger.debug("Attempting to resolve identity from resolver chain.")
         for resolver in self._resolvers:
             try:

packages/smithy-core/src/smithy_core/aio/interfaces/auth.py

@@ -4,6 +4,7 @@
 from typing import Any, Protocol
 
 from ...interfaces import TypedProperties as _TypedProperties
+from ...interfaces.identity import Identity
 from ...shapes import ShapeID
 from . import Request
 from .identity import IdentityResolver
@@ -22,19 +23,21 @@ async def sign(self, *, request: R, identity: I, properties: SP) -> R:
         ...
 
 
-class EventSigner(Protocol):
+class EventSigner[I, SP: Mapping[str, Any]](Protocol):
     """A class that signs requests before they are sent."""
 
     # TODO: add a protocol type for events
-    async def sign(self, *, event: Any) -> Any:
+    async def sign(self, *, event: Any, identity: I, properties: SP) -> Any:
         """Get a signed version of the event.
 
         :param event: The event to be signed.
         """
         ...
 
 
-class AuthScheme[R: Request, I, IP: Mapping[str, Any], SP: Mapping[str, Any]](Protocol):
+class AuthScheme[R: Request, I: Identity, IP: Mapping[str, Any], SP: Mapping[str, Any]](
+    Protocol
+):
     """A class that coordinates identity and auth."""
 
     scheme_id: ShapeID
@@ -74,15 +77,11 @@ def signer(self) -> Signer[R, I, SP]:
         """Get a signer for the request."""
         ...
 
-    def event_signer(
-        self, request: R, identity: I, properties: SP
-    ) -> EventSigner | None:
+    def event_signer(self, request: R) -> EventSigner[I, SP] | None:
         """Construct a signer for event stream events.
 
         :param request: The request that will initiate the event stream. The request
             will not have been sent when this method is called.
-        :param identity: The identity to use to sign events.
-        :param properties: Additional properties used to sign events.
         :returns: An event signer if the scheme supports signing events, otherwise None.
         """
         return None

packages/smithy-core/src/smithy_core/aio/interfaces/identity.py

@@ -1,9 +1,12 @@
 #  Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 #  SPDX-License-Identifier: Apache-2.0
-from typing import Protocol
+from collections.abc import Mapping
+from typing import Any, Protocol
 
+from ...interfaces.identity import Identity
 
-class IdentityResolver[I, IP](Protocol):
+
+class IdentityResolver[I: Identity, IP: Mapping[str, Any]](Protocol):
     """Used to load a user's `Identity` from a given source.
 
     Each `Identity` may have one or more resolver implementations.