Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 6 additions & 0 deletions aws-proxy/AGENTS.md
Original file line number Diff line number Diff line change
Expand Up @@ -52,3 +52,9 @@ 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.

## 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`.
68 changes: 61 additions & 7 deletions aws-proxy/aws_proxy/client/auth_proxy.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
# Note/disclosure: This file has been partially modified by an AI agent.
import json
import logging
import os
Expand All @@ -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,
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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(
Expand All @@ -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,
)
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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()
Expand Down
68 changes: 63 additions & 5 deletions aws-proxy/aws_proxy/server/aws_request_forwarder.py
Original file line number Diff line number Diff line change
@@ -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
Expand All @@ -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__)
Expand Down Expand Up @@ -134,7 +137,32 @@ 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))
# 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"
Expand Down Expand Up @@ -168,6 +196,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,
Expand Down Expand Up @@ -261,6 +295,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""
7 changes: 2 additions & 5 deletions aws-proxy/aws_proxy/server/extension.py
Original file line number Diff line number Diff line change
Expand Up @@ -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__)

Expand All @@ -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")
Expand All @@ -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())
6 changes: 6 additions & 0 deletions aws-proxy/aws_proxy/server/request_handler.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
# Note/disclosure: This file has been partially modified by an AI agent.
import json
import logging
import os.path
Expand Down Expand Up @@ -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()
Expand Down
6 changes: 6 additions & 0 deletions aws-proxy/aws_proxy/shared/constants.py
Original file line number Diff line number Diff line change
@@ -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",
}
3 changes: 3 additions & 0 deletions aws-proxy/tests/conftest.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
# Note/disclosure: This file has been partially modified by an AI agent.
import os

import pytest
Expand Down Expand Up @@ -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()
Loading