enhance logic and tests for CloudWatch proxy (#121)

Co-authored-by: Claude Opus 4.5 <noreply@anthropic.com>

Author: whummer

aws-proxy/AGENTS.md

@@ -40,10 +40,13 @@ Some services have operations that are functionally read-only (don't modify stat
 
 If you find such operations, add them to the service-specific rules in `aws_proxy/server/aws_request_forwarder.py` in the `_is_read_request` method. This ensures that read-only proxy configurations correctly forward these operations rather than blocking them.
 
+**IMPORTANT**: This step is mandatory when adding a new service. Failure to identify non-standard read-only operations will cause `read_only: true` configurations to incorrectly block legitimate read requests.
+
 Example services with non-standard read-only operations:
 - **AppSync**: `EvaluateCode`, `EvaluateMappingTemplate`
 - **IAM**: `SimulateCustomPolicy`, `SimulatePrincipalPolicy`
 - **Cognito**: `InitiateAuth`
+- **DynamoDB**: `Scan`, `BatchGetItem`, `PartiQLSelect`
 
 When adding new integration tests, consider the following:
 * Include a mix of positive and negative assertions (i.e., presence and absence of resources).
@@ -52,3 +55,10 @@ When adding new integration tests, consider the following:
 * Make sure to either use fixtures (preferred), or reliable cleanups for removing the resources; several fixtures for creating AWS resources are available in the `localstack.testing.pytest.fixtures` module
 * If a test uses multiple resources with interdependencies (e.g., an SQS queue connected to an SNS topic), then the test needs to ensure that both resource types are proxied (i.e., created in real AWS), to avoid a situation where a resource in AWS is attempting to reference a local resource in LocalStack (using account ID `000000000000` in their ARN).
 * When waiting for the creation status of a resource, use the `localstack.utils.sync.retry(..)` utility function, rather than a manual `for` loop.
+* Avoid using `time.sleep()` in tests. Instead, use `localstack.utils.sync.retry(..)` to poll for the expected state. This makes tests more robust and avoids unnecessary delays when resources become available faster than expected.
+
+## Fixing or Enhancing Logic in the Proxy
+
+Notes:
+* The AWS proxy is running as a LocalStack Extension, and the tests are currently set up in a way that they assume the container to be running with the Extension in dev mode. Hence, in order to make actual changes to the proxy logic, we'll need to restart the LocalStack main container. You can either ask me (the user) to restart the container whenever you're making changes in the core logic, or alternatively remove the `localstack-main` container, and then run `EXTENSION_DEV_MODE=1 DEBUG=1 localstack start -d` again to restart the container, which may reveal some error logs, stack traces, etc.
+* If the proxy raises errors or something seems off, you can grab and parse the output of the LocalStack container via `localstack logs`.

aws-proxy/aws_proxy/client/auth_proxy.py

@@ -1,3 +1,4 @@
+# Note/disclosure: This file has been partially modified by an AI agent.
 import json
 import logging
 import os
@@ -11,9 +12,11 @@
 import boto3
 import requests
 from botocore.awsrequest import AWSPreparedRequest
-from botocore.model import OperationModel
-from localstack import config as localstack_config
+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 (
     AWS_REGION_US_EAST_1,
@@ -43,14 +46,15 @@
 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)
 
+
 # TODO make configurable
 CLI_PIP_PACKAGE = "localstack-extension-aws-proxy"
 # note: enable the line below temporarily for testing:
@@ -86,6 +90,8 @@ def proxy_request(self, request: Request, data: bytes) -> Response:
         if not parsed:
             return requests_response("", status_code=400)
         region_name, service_name = parsed
+        # Map AWS signing names to boto3 client names
+        service_name = SERVICE_NAME_MAPPING.get(service_name, service_name)
         query_string = to_str(request.query_string or "")
 
         LOG.debug(
@@ -97,10 +103,12 @@ def proxy_request(self, request: Request, data: bytes) -> Response:
             query_string,
         )
 
+        # Convert Quart headers to a dict for the LocalStack Request
+        headers_dict = dict(request.headers)
         request = Request(
             body=data,
             method=request.method,
-            headers=request.headers,
+            headers=headers_dict,
             path=request.path,
             query_string=query_string,
         )
@@ -172,13 +180,46 @@ def register_in_instance(self):
             )
             raise
 
+    def deregister_from_instance(self):
+        """Deregister this proxy from the LocalStack instance."""
+        port = getattr(self, "port", None)
+        if not port:
+            return
+        url = f"{external_service_url()}{HANDLER_PATH_PROXIES}/{port}"
+        LOG.debug("Deregistering proxy from main container via: %s", url)
+        try:
+            response = requests.delete(url)
+            return response
+        except Exception as e:
+            LOG.debug("Unable to deregister auth proxy: %s", e)
+
     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
+        # Get botocore's service model for making the actual AWS request
+        botocore_service_model = self._get_botocore_service_model(service_name)
+
+        # Check if request uses JSON protocol (X-Amz-Target header) while service model
+        # uses RPC v2 CBOR. In this case, we need to parse the request manually since
+        # create_parser would reject the X-Amz-Target header for RPC v2 services.
+        x_amz_target = request.headers.get("X-Amz-Target") or request.headers.get(
+            "X-Amzn-Target"
+        )
+        if x_amz_target and botocore_service_model.protocol == "smithy-rpc-v2-cbor":
+            # Extract operation name from X-Amz-Target (format: "ServiceName.OperationName")
+            operation_name = x_amz_target.split(".")[-1]
+            operation_model = botocore_service_model.operation_model(operation_name)
+            # Parse JSON body
+            parsed_request = json.loads(to_str(request.data)) if request.data else {}
+        else:
+            # Use LocalStack's parser for other protocols
+            localstack_service_model = load_service(service_name)
+            parser = create_parser(localstack_service_model)
+            ls_operation_model, parsed_request = parser.parse(request)
+            operation_model = botocore_service_model.operation_model(
+                ls_operation_model.name
+            )
 
-        parser = create_parser(load_service(service_name))
-        operation_model, parsed_request = parser.parse(request)
         request_context = {
             "client_region": region_name,
             "has_streaming_input": operation_model.has_streaming_input,
@@ -315,6 +356,19 @@ def _query_account_id_from_aws(self) -> str:
         result = sts_client.get_caller_identity()
         return result["Account"]
 
+    @staticmethod
+    @cache
+    def _get_botocore_service_model(service_name: str):
+        """
+        Get the botocore service model for a service. This is used instead of LocalStack's
+        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).
+        """
+        session = get_botocore_session()
+        loader = session.get_component("data_loader")
+        api_data = loader.load_service_model(service_name, "service-2")
+        return ServiceModel(api_data)
+
 
 def start_aws_auth_proxy(config: ProxyConfig, port: int = None) -> AuthProxyAWS:
     setup_logging()

aws-proxy/aws_proxy/server/aws_request_forwarder.py

@@ -1,9 +1,12 @@
+# Note/disclosure: This file has been partially modified by an AI agent.
 import json
 import logging
 import re
 from typing import Dict, Optional
+from urllib.parse import urlencode
 
 import requests
+from botocore.serialize import create_serializer
 from localstack.aws.api import RequestContext
 from localstack.aws.chain import Handler, HandlerChain
 from localstack.constants import APPLICATION_JSON, LOCALHOST, LOCALHOST_HOSTNAME
@@ -22,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__)
@@ -134,7 +137,41 @@ def _request_matches_resource(
                     secret_id, account_id=context.account_id, region_name=context.region
                 )
                 return bool(re.match(resource_name_pattern, secret_arn))
-            # TODO: add more resource patterns
+            if service_name == "cloudwatch":
+                # CloudWatch alarm ARN format: arn:aws:cloudwatch:{region}:{account}:alarm:{alarm_name}
+                alarm_name = context.service_request.get("AlarmName") or ""
+                alarm_names = context.service_request.get("AlarmNames") or []
+                if alarm_name:
+                    alarm_names = [alarm_name]
+                if alarm_names:
+                    for name in alarm_names:
+                        alarm_arn = f"arn:aws:cloudwatch:{context.region}:{context.account_id}:alarm:{name}"
+                        if re.match(resource_name_pattern, alarm_arn):
+                            return True
+                    return False
+                # For metric operations without alarm names, check if pattern is generic
+                return bool(re.match(resource_name_pattern, ".*"))
+            if service_name == "logs":
+                # CloudWatch Logs ARN format: arn:aws:logs:{region}:{account}:log-group:{name}:*
+                log_group_name = context.service_request.get("logGroupName") or ""
+                log_group_prefix = (
+                    context.service_request.get("logGroupNamePrefix") or ""
+                )
+                name = log_group_name or log_group_prefix
+                if name:
+                    log_group_arn = f"arn:aws:logs:{context.region}:{context.account_id}:log-group:{name}:*"
+                    return bool(re.match(resource_name_pattern, log_group_arn))
+                # Operations that don't have a log group name but should still be proxied
+                # (e.g., GetQueryResults uses queryId, not logGroupName)
+                operation_name = context.operation.name if context.operation else ""
+                if operation_name in {
+                    "GetQueryResults",
+                    "StopQuery",
+                    "DescribeQueries",
+                }:
+                    return True
+                # No log group name specified - check if pattern is generic
+                return bool(re.match(resource_name_pattern, ".*"))
         except re.error as e:
             raise Exception(
                 "Error evaluating regular expression - please verify proxy configuration"
@@ -168,6 +205,12 @@ def forward_request(
                 data = request.form
             elif request.data:
                 data = request.data
+
+            # Fallback: if data is empty and we have parsed service_request,
+            # reconstruct the request body (handles cases where form data was consumed)
+            if not data and context.service_request:
+                data = self._reconstruct_request_body(context, ctype)
+
             LOG.debug(
                 "Forward request: %s %s - %s - %s",
                 request.method,
@@ -228,6 +271,18 @@ def _is_read_request(self, context: RequestContext) -> bool:
             "EvaluateMappingTemplate",
         }:
             return True
+        if context.service.service_name == "logs" and operation_name in {
+            "FilterLogEvents",
+            "StartQuery",
+            "GetQueryResults",
+            "TestMetricFilter",
+        }:
+            return True
+        if context.service.service_name == "monitoring" and operation_name in {
+            "BatchGetServiceLevelObjectiveBudgetReport",
+            "BatchGetServiceLevelIndicatorReport",
+        }:
+            return True
         # TODO: add more rules
         return False
 
@@ -261,6 +316,30 @@ def _get_resource_names(cls, service_config: ProxyServiceConfig) -> list[str]:
 
     @classmethod
     def _get_canonical_service_name(cls, service_name: str) -> str:
-        if service_name == "sqs-query":
-            return "sqs"
-        return service_name
+        return SERVICE_NAME_MAPPING.get(service_name, service_name)
+
+    def _reconstruct_request_body(
+        self, context: RequestContext, content_type: str
+    ) -> bytes:
+        """
+        Reconstruct the request body from the parsed service_request.
+        This is used when the original request body was consumed during parsing.
+        """
+        try:
+            protocol = context.service.protocol
+            if protocol == "query" or "x-www-form-urlencoded" in (content_type or ""):
+                # For Query protocol, serialize using botocore serializer
+                serializer = create_serializer(protocol)
+                operation_model = context.operation
+                serialized = serializer.serialize_to_request(
+                    context.service_request, operation_model
+                )
+                body = serialized.get("body", {})
+                if isinstance(body, dict):
+                    return urlencode(body, doseq=True)
+                return body
+            elif protocol == "json" or protocol == "rest-json":
+                return json.dumps(context.service_request)
+        except Exception as e:
+            LOG.debug("Failed to reconstruct request body: %s", e)
+        return b""

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/server/request_handler.py

@@ -1,3 +1,4 @@
+# Note/disclosure: This file has been partially modified by an AI agent.
 import json
 import logging
 import os.path
@@ -43,6 +44,11 @@ def add_proxy(self, request: Request, **kwargs):
         result = handle_proxies_request(req)
         return result or {}
 
+    @route(f"{HANDLER_PATH_PROXIES}/<int:port>", methods=["DELETE"])
+    def delete_proxy(self, request: Request, port: int, **kwargs):
+        removed = AwsProxyHandler.PROXY_INSTANCES.pop(port, None)
+        return {"removed": removed is not None}
+
     @route(f"{HANDLER_PATH_PROXIES}/status", methods=["GET"])
     def get_status(self, request: Request, **kwargs):
         containers = get_proxy_containers()

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",
+}

aws-proxy/tests/conftest.py

@@ -1,3 +1,4 @@
+# Note/disclosure: This file has been partially modified by an AI agent.
 import os
 
 import pytest
@@ -51,4 +52,6 @@ def _start(config: dict = None):
     yield _start
 
     for proxy in proxies:
+        # Deregister from LocalStack instance before shutting down
+        proxy.deregister_from_instance()
         proxy.shutdown()