localstack · viren-nadkarni · Nov 24, 2022 · Jan 2, 2023 · Jan 4, 2023 · Jan 4, 2023
diff --git a/localstack/aws/connect.py b/localstack/aws/connect.py
@@ -0,0 +1,243 @@
+"""
+LocalStack client stack.
+
+This module provides the interface to perform cross-service communication between
+LocalStack providers.
+
+    from localstack.aws.connect import connect_to
+
+    key_pairs = connect_to('ec2').describe_key_pairs()
+    buckets = connect_to('s3', region='ap-south-1').list_buckets()
+"""
+import json
+from datetime import datetime, timezone
+from functools import cache
+from typing import Optional, TypedDict
+
+from boto3.session import Session
+from botocore.awsrequest import AWSPreparedRequest
+from botocore.client import BaseClient
+from botocore.config import Config
+
+from localstack import config
+from localstack.aws.api import RequestContext
+from localstack.constants import (
+    INTERNAL_AWS_ACCESS_KEY_ID,
+    INTERNAL_AWS_SECRET_ACCESS_KEY,
+    MAX_POOL_CONNECTIONS,
+)
+from localstack.utils.aws.arns import extract_region_from_arn
+from localstack.utils.aws.aws_stack import get_local_service_url
+from localstack.utils.aws.request_context import get_region_from_request_context
+
+#
+# Data transfer object
+#
+
+LOCALSTACK_DATA_HEADER = "x-localstack-data"
+"""Request header which contains the data transfer object."""
+
+
+class LocalStackData(TypedDict):
+    """
+    LocalStack Data Transfer Object.
+
+    This is sent with every internal request and contains any additional information
+    LocalStack might need for the purpose of policy enforcement. It is serialised
+    into text and sent in the request header.
+
+    The keys approximately correspond to:
+    https://docs.aws.amazon.com/IAM/latest/UserGuide/reference_policies_condition-keys.html
+    """
+
+    current_time: str
+    """Request datetime in ISO8601 format"""
+
+    source_arn: str
+    """ARN of resource which is triggering the call"""
+
+    source_service: str
+    """Service principal where the call originates, eg. `ec2`"""
+
+    target_arn: str
+    """ARN of the resource being targeted."""
+
+
+def dump_dto(data: LocalStackData) -> str:
+    # TODO@viren: Improve minification using custom JSONEncoder that use shortened keys
+
+    # To produce a compact JSON representation of DTO, remove spaces from separators
+    return json.dumps(data, separators=(",", ":"))
+
+
+def load_dto(data: str) -> LocalStackData:
+    return json.loads(data)
+
+
+#
+# Client
+#
+
+
+class ConnectFactory:
+    """
+    Factory to build the AWS client.
+    """
+
+    def __init__(
+        self,
+        use_ssl: bool = False,
+        verify: bool = False,
+        aws_access_key_id: Optional[str] = INTERNAL_AWS_ACCESS_KEY_ID,
+        aws_secret_access_key: Optional[str] = INTERNAL_AWS_SECRET_ACCESS_KEY,
+    ):
+        """
+        If either of the access keys are set to None, they are loaded from following
+        locations:
+        - AWS environment variables
+        - Credentials file `~/.aws/credentials`
+        - Config file `~/.aws/config`
+
+        :param use_ssl: Whether to use SSL
+        :param verify: Whether to verify SSL certificates
+        :param aws_access_key_id: Access key to use for the client.
+            If set to None, loads them from botocore session. See above.
+        :param aws_secret_access_key: Secret key to use for the client.
+            If set to None, loads them from botocore session. See above.
+        :param localstack_data: LocalStack data transfer object
+        """
+        self._use_ssl = use_ssl
        self._verify = verify
+        self._aws_access_key_id = aws_access_key_id
+        self._aws_secret_access_key = aws_secret_access_key
+        self._aws_session_token = None
+        self._session = Session()
+        self._config = Config(max_pool_connections=MAX_POOL_CONNECTIONS)
+
+    def get_partition_for_region(self, region_name: str) -> str:
+        """
+        Return the AWS partition name for a given region, eg. `aws`, `aws-cn`, etc.
+        """
+        return self._session.get_partition_for_region(region_name)
+
+    def get_session_region_name(self) -> str:
+        """
+        Return AWS region as set in the Boto session.
+        """
+        return self._session.region_name
+
+    def get_region_name(self) -> str:
+        """
+        Return the AWS region name from following sources, in order of availability.
+        - LocalStack request context
+        - LocalStack default region
+        - Boto session
+        """
+        return (
+            get_region_from_request_context()
+            or config.DEFAULT_REGION
+            or self.get_session_region_name()
+        )
+
+    @cache
+    def get_client(
+        self,
+        service_name: str,
+        region_name: str,
+        use_ssl: bool,
+        verify: bool,
+        endpoint_url: str,
+        aws_access_key_id: str,
+        aws_secret_access_key: str,
+        aws_session_token: str,
+        config: Config,
+    ) -> BaseClient:
+        return self._session.client(
+            service_name=service_name,
+            region_name=region_name,
+            use_ssl=use_ssl,
+            verify=verify,
+            endpoint_url=endpoint_url,
+            aws_access_key_id=aws_access_key_id,
+            aws_secret_access_key=aws_secret_access_key,
+            aws_session_token=aws_session_token,
+            config=config,
+        )
+
+    def __call__(
+        self,
+            target_service: str,
+        region_name: str = None,
+        endpoint_url: str = None,
+        config: Config = None,
+        source_arn: str = None,
+        source_service: str = None,
+        target_arn: str = None,
+    ) -> BaseClient:
+        """
+        Build and return the client.
+
+        Presence of any attribute apart from `source_*` or `target_*` argument
+        indicates that this is a client meant for internal calls.
+
+        :param target_service: Service to build the client for, eg. `s3`
+        :param region_name: Name of the AWS region to be associated with the client
+        :param endpoint_url: Full endpoint URL to be used by the client.
+            Defaults to appropriate LocalStack endpoint.
+        :param config: Boto config for advanced use.
+        :param source_arn: ARN of resource which triggers the call. Required for
+            internal calls.
+        :param source_service: Service name where call originates. Required for
+            internal calls.
+        :param target_arn: ARN of targeted resource. Overrides `region_name`.
+            Required for internal calls.
+        """
+        localstack_data = LocalStackData()
+
+        if source_arn:
+            localstack_data["source_arn"] = source_arn
+
+        if source_service:
+            localstack_data["source_service"] = source_service
+
+        if target_arn:
+            # Attention: region is overriden here
+            region_name = extract_region_from_arn(target_arn)
+            localstack_data["target_arn"] = target_arn
+
+        client = self.get_client(
+            service_name=target_service,
+            region_name=region_name or self.get_region_name(),
+            use_ssl=self._use_ssl,
+            verify=self._verify,
+            endpoint_url=endpoint_url or get_local_service_url(target_service),
+            aws_access_key_id=self._aws_access_key_id,
+            aws_secret_access_key=self._aws_secret_access_key,
+            aws_session_token=self._aws_session_token,
+            config=config or self._config,
+        )
+
+        def _handler(request: AWSPreparedRequest, **_):
+            data = localstack_data | LocalStackData(
+                current_time=datetime.now(timezone.utc).isoformat()
+            )
+            request.headers[LOCALSTACK_DATA_HEADER] = dump_dto(data)
+
+        if len(localstack_data):
+            client.meta.events.register("befor
6D40
e-send.*.*", handler=_handler)
+
+        return client
+
+
+connect_to = ConnectFactory()
+
+#
+# Utilities
+#
+
+
+def is_internal_call(context: RequestContext) -> bool:
+    """
+    Whether a given request is an internal LocalStack cross-service call.
+    """
+    return LOCALSTACK_DATA_HEADER in context.request.headers
diff --git a/localstack/aws/handlers/auth.py b/localstack/aws/handlers/auth.py
@@ -5,7 +5,11 @@
     set_aws_access_key_id,
     set_aws_account_id,
 )
-from localstack.constants import TEST_AWS_ACCESS_KEY_ID
+from localstack.constants import (
+    INTERNAL_AWS_ACCESS_KEY_ID,
+    INTERNAL_AWS_ACCOUNT_ID,
+    TEST_AWS_ACCESS_KEY_ID,
+)
 from localstack.http import Response
 from localstack.utils.aws.aws_stack import extract_access_key_id_from_auth_header
 
@@ -45,7 +49,17 @@ def __call__(self, chain: HandlerChain, context: RequestContext, response: Respo
         set_aws_access_key_id(access_key_id)
 
         # Obtain the account ID and save it in the request context
-        context.account_id = get_account_id_from_access_key_id(access_key_id)
+        if access_key_id == INTERNAL_AWS_ACCESS_KEY_ID:
+            # For internal calls, a special account ID is used for request context
+            # Cross account calls don't have the same auth flows as user-originating calls
+            # which means there is no true Account ID.
+            # The invocations of `get_aws_account_id()` used to resolve the stores must not break.
+            # We don't use the DEFAULT_AWS_ACCOUNT_ID either to help identify bugs.
+            # If correctly implemented with CrossAccountAttribute and ARNs, the provider
+            # will work with this internal AWS account ID.
+            context.account_id = INTERNAL_AWS_ACCOUNT_ID
+        else:
+            context.account_id = get_account_id_from_access_key_id(access_key_id)
 
         # Save the same account ID in the thread context
         set_aws_account_id(context.account_id)

diff --git a/localstack/constants.py b/localstack/constants.py
@@ -45,6 +45,9 @@
 # Fallback Account ID if not available in the client request
 DEFAULT_AWS_ACCOUNT_ID = "000000000000"
 
+# Fallback Account ID for internal calls
+INTERNAL_AWS_ACCOUNT_ID = "000012210000"
+
 # AWS user account ID used for tests - TODO move to config.py
 if "TEST_AWS_ACCOUNT_ID" not in os.environ:
     os.environ["TEST_AWS_ACCOUNT_ID"] = DEFAULT_AWS_ACCOUNT_ID

diff --git a/localstack/services/awslambda/event_source_listeners/sqs_event_source_listener.py b/localstack/services/awslambda/event_source_listeners/sqs_event_source_listener.py
@@ -63,6 +63,16 @@ def _listener_loop(self, *args):
                 for source in sources:
                     queue_arn = source["EventSourceArn"]
                     region_name = extract_region_from_arn(queue_arn)
+                    ## sqs_client = aws_stack.connect_to_service("sqs", region_name=region_name)
+                    # `sqs` is the factory
+                    # config args include stuff going into boto like retries, max conn pool,.
+                    # could be called `boto_config`
+                    # all options will map to boto_config args
+                    # aws.sqs.configure().credentials(aws_access_key_id=...)
+                    # client can be created from ARNs
+                    # can communicate with Queue ARN
+                    # all methods could be prefixed with `set_` eg. set_target_arn()
+                    # sqs_client = aws_client().target_arn(queue_arn).credentials(sts_call_credentials).sqs()
                     sqs_client = aws_stack.connect_to_service("sqs", region_name=region_name)
                     batch_size = max(min(source.get("BatchSize", 1), 10), 1)
 

diff --git a/localstack/services/s3/notifications.py b/localstack/services/s3/notifications.py
@@ -28,6 +28,7 @@
     TopicArn,
     TopicConfiguration,
 )
+from localstack.aws.connect import connect_to
 from localstack.config import DEFAULT_REGION
 from localstack.services.s3.models import get_moto_s3_backend
 from localstack.services.s3.utils import (
@@ -350,7 +351,7 @@ def _get_arn_value_and_name(topic_configuration: TopicConfiguration) -> [TopicAr
         return topic_configuration.get("TopicArn", ""), "TopicArn"
 
     def _verify_target_exists(self, arn: str, arn_data: ArnData) -> None:
-        client = aws_stack.connect_to_service(self.service_name, region_name=arn_data["region"])
+        client = connect_to(self.service_name, target_arn=arn, source_service="s3")
         try:
             client.get_topic_attributes(TopicArn=arn)
         except ClientError:

diff --git a/localstack/services/s3/s3_listener.py b/localstack/services/s3/s3_listener.py
@@ -21,6 +21,7 @@
 
 from localstack import config, constants
 from localstack.aws.api import CommonServiceException
+from localstack.aws.connect import connect_to
 from localstack.config import get_protocol as get_service_protocol
 from localstack.services.generic_proxy import ProxyListener
 from localstack.services.generic_proxy import append_cors_headers as _append_default_cors_headers
@@ -368,10 +369,11 @@ def send_notification_for_subscriber(
             )
     if notification.get("Topic"):
         region = arns.extract_region_from_arn(notification["Topic"])
-        sns_client = aws_stack.connect_to_service("sns", region_name=region)
+        topic_arn = notification["Topic"]
+        sns_client = connect_to("sns", target_arn=topic_arn, source_service="s3")
         try:
             sns_client.publish(
-                TopicArn=notification["Topic"],
+                TopicArn=topic_arn,
                 Message=message,
                 Subject="Amazon S3 Notification",
             )

diff --git a/localstack/services/ses/provider.py b/localstack/services/ses/provider.py
@@ -50,12 +50,11 @@
     VerificationAttributes,
     VerificationStatus,
 )
-from localstack.constants import TEST_AWS_SECRET_ACCESS_KEY
+from localstack.aws.connect import connect_to
 from localstack.services.internal import get_internal_apis
 from localstack.services.moto import call_moto
 from localstack.services.plugins import ServiceLifecycleHook
 from localstack.services.ses.models import SentEmail, SentEmailBody
-from localstack.utils.aws import arns, aws_stack
 from localstack.utils.files import mkdir
 from localstack.utils.strings import long_uid, to_str
 from localstack.utils.time import timestamp, timestamp_millis
@@ -577,13 +576,4 @@ def emit_delivery_event(self, payload: SNSPayload, sns_topic_arn: str):
 
     @staticmethod
     def _client_for_topic(topic_arn: str) -> "SNSClient":
-        arn_parameters = arns.parse_arn(topic_arn)
-        region = arn_parameters["region"]
-        access_key_id = arn_parameters["account"]
-
-        return aws_stack.connect_to_service(
-            "sn
4E25
s",
-            region_name=region,
-            aws_access_key_id=access_key_id,
-            aws_secret_access_key=TEST_AWS_SECRET_ACCESS_KEY,
-        )
+        return connect_to("sns", target_arn=topic_arn, source_service="ses")