minor refactoring

Author: whummer

aws-proxy/aws_proxy/client/auth_proxy.py

@@ -12,7 +12,10 @@
 import boto3
 import requests
 from botocore.awsrequest import AWSPreparedRequest
-from botocore.model import OperationModel
+from botocore.model import OperationModel, ServiceModel
+from botocore.session import get_session as get_botocore_session
+from localstack.aws.protocol.parser import create_parser
+from localstack.aws.spec import load_service
 from localstack import config as localstack_config
 from localstack.config import external_service_url
 from localstack.constants import (
@@ -43,18 +46,14 @@
 from aws_proxy import config as repl_config
 from aws_proxy.client.utils import truncate_content
 from aws_proxy.config import HANDLER_PATH_PROXIES
-from aws_proxy.shared.constants import HEADER_HOST_ORIGINAL
+from aws_proxy.shared.constants import HEADER_HOST_ORIGINAL, SERVICE_NAME_MAPPING
 from aws_proxy.shared.models import AddProxyRequest, ProxyConfig
 
 LOG = logging.getLogger(__name__)
 LOG.setLevel(logging.INFO)
 if localstack_config.DEBUG:
     LOG.setLevel(logging.DEBUG)
 
-# Mapping from AWS service signing names to boto3 client names
-SERVICE_NAME_MAPPING = {
-    "monitoring": "cloudwatch",
-}
 
 # TODO make configurable
 CLI_PIP_PACKAGE = "localstack-extension-aws-proxy"
@@ -197,13 +196,19 @@ def deregister_from_instance(self):
     def _parse_aws_request(
         self, request: Request, service_name: str, region_name: str, client
     ) -> Tuple[OperationModel, AWSPreparedRequest, Dict]:
-        from localstack.aws.protocol.parser import create_parser
+        # Use LocalStack's load_service for parsing, as it handles multi-protocol detection
+        # from request headers (e.g., detecting JSON vs RPC v2 CBOR from content-type).
+        localstack_service_model = load_service(service_name)
+        parser = create_parser(localstack_service_model)
+        ls_operation_model, parsed_request = parser.parse(request)
+
+        # Get botocore's operation model for making the actual AWS request,
+        # as botocore client methods expect botocore model objects.
+        botocore_service_model = self._get_botocore_service_model(service_name)
+        operation_model = botocore_service_model.operation_model(
+            ls_operation_model.name
+        )
 
-        # Use botocore's service model to ensure protocol compatibility
-        # (LocalStack's load_service may return newer protocol versions that don't match the client)
-        service_model = self._get_botocore_service_model(service_name)
-        parser = create_parser(service_model)
-        operation_model, parsed_request = parser.parse(request)
         request_context = {
             "client_region": region_name,
             "has_streaming_input": operation_model.has_streaming_input,
@@ -348,10 +353,7 @@ def _get_botocore_service_model(service_name: str):
         load_service() to ensure protocol compatibility, as LocalStack may use newer protocol
         versions (e.g., smithy-rpc-v2-cbor) while clients use older protocols (e.g., query).
         """
-        import botocore.session
-        from botocore.model import ServiceModel
-
-        session = botocore.session.get_session()
+        session = get_botocore_session()
         loader = session.get_component("data_loader")
         api_data = loader.load_service_model(service_name, "service-2")
         return ServiceModel(api_data)

aws-proxy/aws_proxy/server/aws_request_forwarder.py

@@ -25,7 +25,7 @@
 except ImportError:
     from localstack.constants import TEST_AWS_ACCESS_KEY_ID
 
-from aws_proxy.shared.constants import HEADER_HOST_ORIGINAL
+from aws_proxy.shared.constants import HEADER_HOST_ORIGINAL, SERVICE_NAME_MAPPING
 from aws_proxy.shared.models import ProxyInstance, ProxyServiceConfig
 
 LOG = logging.getLogger(__name__)
@@ -295,12 +295,7 @@ def _get_resource_names(cls, service_config: ProxyServiceConfig) -> list[str]:
 
     @classmethod
     def _get_canonical_service_name(cls, service_name: str) -> str:
-        # Map internal/signing service names to boto3 client names
-        mapping = {
-            "sqs-query": "sqs",
-            "monitoring": "cloudwatch",
-        }
-        return mapping.get(service_name, service_name)
+        return SERVICE_NAME_MAPPING.get(service_name, service_name)
 
     def _reconstruct_request_body(
         self, context: RequestContext, content_type: str

aws-proxy/aws_proxy/server/extension.py

@@ -8,7 +8,8 @@
 
 from localstack.services.internal import get_internal_apis
 
-from aws_proxy.server.request_handler import WebApp
+from aws_proxy.server.aws_request_forwarder import AwsProxyHandler
+from aws_proxy.server.request_handler import RequestHandler, WebApp
 
 LOG = logging.getLogger(__name__)
 
@@ -27,8 +28,6 @@ def on_extension_load(self):
             )
 
     def update_gateway_routes(self, router: http.Router[http.RouteHandler]):
-        from aws_proxy.server.request_handler import RequestHandler
-
         super().update_gateway_routes(router)
 
         LOG.info("AWS Proxy: adding routes to activate extension")
@@ -38,7 +37,5 @@ def collect_routes(self, routes: list[t.Any]):
         routes.append(WebApp())
 
     def update_request_handlers(self, handlers: CompositeHandler):
-        from aws_proxy.server.aws_request_forwarder import AwsProxyHandler
-
         LOG.debug("AWS Proxy: adding AWS proxy handler to the request chain")
         handlers.handlers.append(AwsProxyHandler())

aws-proxy/aws_proxy/shared/constants.py

@@ -1,2 +1,8 @@
 # header name for the original request host name forwarded in the request to the target proxy handler
 HEADER_HOST_ORIGINAL = "x-ls-host-original"
+
+# Mapping from AWS service signing names to boto3 client names
+SERVICE_NAME_MAPPING = {
+    "monitoring": "cloudwatch",
+    "sqs-query": "sqs",
+}