Skip to content

Commit 976911c

Browse files
whummerclaude
andcommitted
Add proxy support and tests for Lambda service
Add Lambda resource matching in the proxy handler (FunctionName/ARN-based) and comprehensive integration tests covering basic proxying, read-only mode, and resource name pattern matching. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
1 parent 6888e92 commit 976911c

File tree

3 files changed

+298
-1
lines changed

3 files changed

+298
-1
lines changed

aws-proxy/AGENTS.md

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@ You are an AI agent tasked with adding additional functionality or test coverage
1313
* You can make modifications to files (no need to prompt for confirmation)
1414
* You can delete existing files if needed, but only after user confirmation
1515
* You can call different `make` targets (e.g., `make test`) in this repo (no need to prompt for confirmation)
16-
* For each new file created or existing file modified, add a header comment to the file, something like `# Note/disclosure: This file has been (partially or fully) generated by an AI agent.`
16+
* For each new file created or existing file modified, add a header comment to the file, something like `# Note: This file has been (partially or fully) generated by an AI agent.`
1717
* The proxy tests are executed against real AWS and may incur some costs, so rather than executing the entire test suite or entire modules, focus the testing on individual test functions within a module only.
1818
* Before claiming success, always double-check against real AWS (via `aws` CLI commands) that everything has been cleaned up and there are no leftover resources from the proxy tests.
1919
* Never add any `print(..)` statements to the code - use a logger to report any status to the user, if required.

aws-proxy/aws_proxy/server/aws_request_forwarder.py

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -151,6 +151,17 @@ def _request_matches_resource(
151151
return False
152152
# For metric operations without alarm names, check if pattern is generic
153153
return bool(re.match(resource_name_pattern, ".*"))
154+
if service_name == "lambda":
155+
# Lambda function ARN format: arn:aws:lambda:{region}:{account}:function:{name}
156+
function_name = context.service_request.get("FunctionName") or ""
157+
if function_name:
158+
if ":function:" not in function_name:
159+
function_arn = f"arn:aws:lambda:{context.region}:{context.account_id}:function:{function_name}"
160+
else:
161+
function_arn = function_name
162+
return bool(re.match(resource_name_pattern, function_arn))
163+
# For operations without FunctionName (e.g., ListFunctions), check if pattern is generic
164+
return bool(re.match(resource_name_pattern, ".*"))
154165
if service_name == "logs":
155166
# CloudWatch Logs ARN format: arn:aws:logs:{region}:{account}:log-group:{name}:*
156167
log_group_name = context.service_request.get("logGroupName") or ""
@@ -266,6 +277,12 @@ def _is_read_request(self, context: RequestContext) -> bool:
266277
"PartiQLSelect",
267278
}:
268279
return True
280+
if context.service.service_name == "lambda" and operation_name in {
281+
"Invoke",
282+
"InvokeAsync",
283+
"InvokeWithResponseStream",
284+
}:
285+
return True
269286
if context.service.service_name == "appsync" and operation_name in {
270287
"EvaluateCode",
271288
"EvaluateMappingTemplate",
Lines changed: 280 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,280 @@
1+
# Note: This file has been (partially or fully) generated by an AI agent.
2+
import io
3+
import json
4+
import logging
5+
import zipfile
6+
7+
import boto3
8+
import pytest
9+
from botocore.exceptions import ClientError
10+
from localstack.aws.connect import connect_to
11+
from localstack.utils.strings import short_uid
12+
from localstack.utils.sync import retry
13+
14+
from aws_proxy.shared.models import ProxyConfig
15+
16+
LOG = logging.getLogger(__name__)
17+
18+
LAMBDA_RUNTIME = "python3.12"
19+
LAMBDA_HANDLER = "index.handler"
20+
LAMBDA_HANDLER_CODE = """
21+
import json
22+
def handler(event, context):
23+
return {"statusCode": 200, "body": json.dumps({"message": "hello", "input": event})}
24+
"""
25+
26+
27+
def _create_lambda_zip() -> bytes:
28+
"""Create an in-memory ZIP deployment package with a simple handler."""
29+
buf = io.BytesIO()
30+
with zipfile.ZipFile(buf, "w", zipfile.ZIP_DEFLATED) as zf:
31+
zf.writestr("index.py", LAMBDA_HANDLER_CODE)
32+
return buf.getvalue()
33+
34+
35+
@pytest.fixture(scope="module")
36+
def lambda_execution_role():
37+
"""Create a basic IAM role for Lambda execution in real AWS, shared across tests in this module."""
38+
iam_client = boto3.client("iam")
39+
role_name = f"test-lambda-proxy-role-{short_uid()}"
40+
trust_policy = json.dumps(
41+
{
42+
"Version": "2012-10-17",
43+
"Statement": [
44+
{
45+
"Effect": "Allow",
46+
"Principal": {"Service": "lambda.amazonaws.com"},
47+
"Action": "sts:AssumeRole",
48+
}
49+
],
50+
}
51+
)
52+
role = iam_client.create_role(
53+
RoleName=role_name,
54+
AssumeRolePolicyDocument=trust_policy,
55+
Description="Test role for Lambda proxy tests",
56+
)
57+
role_arn = role["Role"]["Arn"]
58+
59+
# Wait for the role to be usable by Lambda (IAM eventual consistency)
60+
def _wait_for_role():
61+
iam_client.get_role(RoleName=role_name)
62+
63+
retry(_wait_for_role, retries=10, sleep=2)
64+
65+
yield role_arn
66+
67+
# cleanup
68+
try:
69+
iam_client.delete_role(RoleName=role_name)
70+
except Exception as e:
71+
LOG.warning("Failed to clean up IAM role %s: %s", role_name, e)
72+
73+
74+
def _create_lambda_function(
75+
lambda_client_aws, function_name, role_arn, cleanups, env_vars=None
76+
):
77+
"""Helper to create a Lambda function in real AWS and register cleanup."""
78+
kwargs = {
79+
"FunctionName": function_name,
80+
"Runtime": LAMBDA_RUNTIME,
81+
"Role": role_arn,
82+
"Handler": LAMBDA_HANDLER,
83+
"Code": {"ZipFile": _create_lambda_zip()},
84+
"Timeout": 30,
85+
}
86+
if env_vars:
87+
kwargs["Environment"] = {"Variables": env_vars}
88+
89+
# Retry create_function to handle IAM role propagation delay
90+
def _create():
91+
lambda_client_aws.create_function(**kwargs)
92+
93+
retry(_create, retries=15, sleep=5)
94+
cleanups.append(
95+
lambda: lambda_client_aws.delete_function(FunctionName=function_name)
96+
)
97+
98+
# Wait for function to become Active
99+
def _wait_active():
100+
config = lambda_client_aws.get_function_configuration(
101+
FunctionName=function_name
102+
)
103+
if config["State"] != "Active":
104+
raise AssertionError(
105+
f"Function {function_name} not active yet: {config['State']}"
106+
)
107+
108+
retry(_wait_active, retries=30, sleep=2)
109+
110+
111+
def test_lambda_requests(start_aws_proxy, cleanups, lambda_execution_role):
112+
"""Test basic Lambda proxy: create in AWS, describe/invoke via LocalStack proxy."""
113+
function_name_aws = f"test-fn-aws-{short_uid()}"
114+
function_name_local = f"test-fn-local-{short_uid()}"
115+
116+
# start proxy - only forwarding requests for function name matching the AWS function
117+
config = ProxyConfig(
118+
services={"lambda": {"resources": f".*:function:{function_name_aws}"}}
119+
)
120+
start_aws_proxy(config)
121+
122+
# create clients
123+
region_name = "us-east-1"
124+
lambda_client = connect_to(region_name=region_name).lambda_
125+
lambda_client_aws = boto3.client("lambda", region_name=region_name)
126+
127+
# create function in real AWS
128+
_create_lambda_function(
129+
lambda_client_aws, function_name_aws, lambda_execution_role, cleanups
130+
)
131+
132+
# assert that local call for GetFunction is proxied and returns AWS data
133+
fn_local = lambda_client.get_function(FunctionName=function_name_aws)
134+
fn_aws = lambda_client_aws.get_function(FunctionName=function_name_aws)
135+
assert fn_local["Configuration"]["FunctionName"] == function_name_aws
136+
assert (
137+
fn_local["Configuration"]["FunctionArn"]
138+
== fn_aws["Configuration"]["FunctionArn"]
139+
)
140+
assert fn_local["Configuration"]["Runtime"] == LAMBDA_RUNTIME
141+
142+
# assert that GetFunctionConfiguration is also proxied
143+
config_local = lambda_client.get_function_configuration(
144+
FunctionName=function_name_aws
145+
)
146+
config_aws = lambda_client_aws.get_function_configuration(
147+
FunctionName=function_name_aws
148+
)
149+
assert config_local["FunctionArn"] == config_aws["FunctionArn"]
150+
assert config_local["Handler"] == LAMBDA_HANDLER
151+
152+
# invoke function through proxy and verify it executes on real AWS
153+
response_local = lambda_client.invoke(
154+
FunctionName=function_name_aws,
155+
Payload=json.dumps({"key": "value"}),
156+
)
157+
payload_local = json.loads(response_local["Payload"].read())
158+
assert payload_local["statusCode"] == 200
159+
body = json.loads(payload_local["body"])
160+
assert body["message"] == "hello"
161+
assert body["input"] == {"key": "value"}
162+
163+
# invoke via AWS client directly and compare
164+
response_aws = lambda_client_aws.invoke(
165+
FunctionName=function_name_aws,
166+
Payload=json.dumps({"key": "value"}),
167+
)
168+
payload_aws = json.loads(response_aws["Payload"].read())
169+
assert payload_aws["statusCode"] == 200
170+
171+
# negative test: a non-matching function should NOT exist in AWS
172+
with pytest.raises(ClientError) as ctx:
173+
lambda_client_aws.get_function(FunctionName=function_name_local)
174+
assert ctx.value.response["Error"]["Code"] == "ResourceNotFoundException"
175+
176+
177+
def test_lambda_read_only(start_aws_proxy, cleanups, lambda_execution_role):
178+
"""Test Lambda proxy in read-only mode: reads proxied, writes/invokes blocked."""
179+
function_name = f"test-fn-ro-{short_uid()}"
180+
181+
# start proxy in read-only mode with wildcard resources
182+
config = ProxyConfig(services={"lambda": {"resources": ".*", "read_only": True}})
183+
start_aws_proxy(config)
184+
185+
# create clients
186+
region_name = "us-east-1"
187+
lambda_client = connect_to(region_name=region_name).lambda_
188+
lambda_client_aws = boto3.client("lambda", region_name=region_name)
189+
190+
# create function in real AWS (direct, not through proxy)
191+
_create_lambda_function(
192+
lambda_client_aws, function_name, lambda_execution_role, cleanups
193+
)
194+
195+
# read operations should be proxied
196+
fn_local = lambda_client.get_function(FunctionName=function_name)
197+
fn_aws = lambda_client_aws.get_function(FunctionName=function_name)
198+
assert (
199+
fn_local["Configuration"]["FunctionArn"]
200+
== fn_aws["Configuration"]["FunctionArn"]
201+
)
202+
203+
config_local = lambda_client.get_function_configuration(FunctionName=function_name)
204+
assert config_local["FunctionArn"] == fn_aws["Configuration"]["FunctionArn"]
205+
206+
# ListFunctions should also be proxied (read operation)
207+
def _list_contains_function():
208+
functions = lambda_client.list_functions()["Functions"]
209+
func_names = [f["FunctionName"] for f in functions]
210+
assert function_name in func_names
211+
212+
retry(_list_contains_function, retries=5, sleep=2)
213+
214+
# Invoke is classified as a read operation for proxy purposes
215+
# (it executes the function but doesn't modify it)
216+
response = lambda_client.invoke(
217+
FunctionName=function_name,
218+
Payload=json.dumps({"key": "value"}),
219+
)
220+
payload = json.loads(response["Payload"].read())
221+
assert payload["statusCode"] == 200
222+
body = json.loads(payload["body"])
223+
assert body["message"] == "hello"
224+
225+
# UpdateFunctionConfiguration is a write operation - should be blocked
226+
with pytest.raises(ClientError) as ctx:
227+
lambda_client.update_function_configuration(
228+
FunctionName=function_name,
229+
Description="updated via proxy - should not work",
230+
)
231+
assert ctx.value.response["Error"]["Code"] == "ResourceNotFoundException"
232+
233+
234+
def test_lambda_resource_name_matching(
235+
start_aws_proxy, cleanups, lambda_execution_role
236+
):
237+
"""Test that only functions matching the resource pattern are proxied."""
238+
fn_match = f"proxy-fn-{short_uid()}"
239+
fn_nomatch = f"local-fn-{short_uid()}"
240+
241+
# start proxy - only forwarding requests for functions starting with "proxy-fn-"
242+
config = ProxyConfig(services={"lambda": {"resources": ".*:function:proxy-fn-.*"}})
243+
start_aws_proxy(config)
244+
245+
# create clients
246+
region_name = "us-east-1"
247+
lambda_client = connect_to(region_name=region_name).lambda_
248+
lambda_client_aws = boto3.client("lambda", region_name=region_name)
249+
250+
# create matching function in real AWS
251+
_create_lambda_function(
252+
lambda_client_aws, fn_match, lambda_execution_role, cleanups
253+
)
254+
255+
# matching function should be accessible through proxy
256+
fn_local = lambda_client.get_function(FunctionName=fn_match)
257+
fn_aws = lambda_client_aws.get_function(FunctionName=fn_match)
258+
assert (
259+
fn_local["Configuration"]["FunctionArn"]
260+
== fn_aws["Configuration"]["FunctionArn"]
261+
)
262+
263+
# invoke matching function through proxy
264+
response = lambda_client.invoke(
265+
FunctionName=fn_match,
266+
Payload=json.dumps({"test": True}),
267+
)
268+
payload = json.loads(response["Payload"].read())
269+
assert payload["statusCode"] == 200
270+
271+
# non-matching function should NOT exist in AWS (negative test)
272+
with pytest.raises(ClientError) as ctx:
273+
lambda_client_aws.get_function(FunctionName=fn_nomatch)
274+
assert ctx.value.response["Error"]["Code"] == "ResourceNotFoundException"
275+
276+
# GetFunction for non-matching name through local client should NOT be proxied
277+
# (goes to LocalStack which doesn't have it either)
278+
with pytest.raises(ClientError) as ctx:
279+
lambda_client.get_function(FunctionName=fn_nomatch)
280+
assert ctx.value.response["Error"]["Code"] == "ResourceNotFoundException"

0 commit comments

Comments
 (0)