Skip to content

Commit 84b2f5f

Browse files
bentskubentsku
authored
APIGW: add Binary Media support for response in AWS_PROXY (#12199)
1 parent d7ba30f commit 84b2f5f

File tree

4 files changed

+262
-2
lines changed

4 files changed

+262
-2
lines changed

‎localstack-core/localstack/services/apigateway/next_gen/execute_api/helpers.py

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -148,3 +148,29 @@ def parse_trace_id(trace_id: str) -> dict[str, str]:
148148
trace_values[key_value[0].capitalize()] = key_value[1]
149149

150150
return trace_values
151+
152+
153+
def mime_type_matches_binary_media_types(mime_type: str | None, binary_media_types: list[str]):
154+
if not mime_type or not binary_media_types:
155+
return False
156+
157+
mime_type_and_subtype = mime_type.split(",")[0].split(";")[0].split("/")
158+
if len(mime_type_and_subtype) != 2:
159+
return False
160+
mime_type, mime_subtype = mime_type_and_subtype
161+
162+
for bmt in binary_media_types:
163+
type_and_subtype = bmt.split(";")[0].split("/")
164+
if len(type_and_subtype) != 2:
165+
continue
166+
_type, subtype = type_and_subtype
167+
if _type == "*":
168+
continue
169+
170+
if subtype == "*" and mime_type == _type:
171+
return True
172+
173+
if mime_type == _type and mime_subtype == subtype:
174+
return True
175+
176+
return False

‎localstack-core/localstack/services/apigateway/next_gen/execute_api/integrations/aws.py

Lines changed: 26 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,7 @@
3434
from ..helpers import (
3535
get_lambda_function_arn_from_invocation_uri,
3636
get_source_arn,
37+
mime_type_matches_binary_media_types,
3738
render_uri_with_stage_variables,
3839
validate_sub_dict_of_typed_dict,
3940
)
@@ -392,9 +393,20 @@ def invoke(self, context: RestApiInvocationContext) -> EndpointResponse:
392393
response_headers = self._merge_lambda_response_headers(lambda_response)
393394
headers.update(response_headers)
394395

396+
# TODO: maybe centralize this flag inside the context, when we are also using it for other integration types
397+
# AWS_PROXY behaves a bit differently, but this could checked only once earlier
398+
binary_response_accepted = mime_type_matches_binary_media_types(
399+
context.invocation_request["headers"].get("Accept"),
400+
context.deployment.rest_api.rest_api.get("binaryMediaTypes", []),
401+
)
402+
body = self._parse_body(
403+
body=lambda_response.get("body"),
404+
is_base64_encoded=binary_response_accepted and lambda_response.get("isBase64Encoded"),
405+
)
406+
395407
return EndpointResponse(
396408
headers=headers,
397-
body=to_bytes(lambda_response.get("body") or ""),
409+
body=body,
398410
status_code=int(lambda_response.get("statusCode") or 200),
399411
)
400412

@@ -552,6 +564,19 @@ def _format_body(body: bytes) -> tuple[str, bool]:
552564
except UnicodeDecodeError:
553565
return to_str(base64.b64encode(body)), True
554566

567+
@staticmethod
568+
def _parse_body(body: str | None, is_base64_encoded: bool) -> bytes:
569+
if not body:
570+
return b""
571+
572+
if is_base64_encoded:
573+
try:
574+
return base64.b64decode(body)
575+
except Exception:
576+
raise InternalServerError("Internal server error", status_code=500)
577+
578+
return to_bytes(body)
579+
555580
@staticmethod
556581
def _merge_lambda_response_headers(lambda_response: LambdaProxyResponse) -> dict:
557582
headers = lambda_response.get("headers") or {}

‎tests/aws/services/apigateway/test_apigateway_lambda.py

Lines changed: 207 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,17 +1,19 @@
11
import base64
22
import json
33
import os
4+
import time
45

56
import pytest
67
import requests
78
from botocore.exceptions import ClientError
89

910
from localstack.aws.api.lambda_ import Runtime
11+
from localstack.testing.aws.util import is_aws_cloud
1012
from localstack.testing.pytest import markers
1113
from localstack.utils.aws import arns
1214
from localstack.utils.files import load_file
1315
from localstack.utils.strings import short_uid
14-
from localstack.utils.sync import retry
16+
from localstack.utils.sync import poll_condition, retry
1517
from tests.aws.services.apigateway.apigateway_fixtures import api_invoke_url, create_rest_resource
1618
from tests.aws.services.apigateway.conftest import (
1719
APIGATEWAY_ASSUME_ROLE_POLICY,
@@ -1312,3 +1314,207 @@ def invoke_api(url):
13121314
# retry is necessary against AWS, probably IAM permission delay
13131315
invoke_response = retry(invoke_api, sleep=2, retries=10, url=invocation_url)
13141316
snapshot.match("http-proxy-invocation-data-mapping", invoke_response)
1317+
1318+
1319+
@markers.aws.validated
1320+
def test_aws_proxy_binary_response(
1321+
create_rest_apigw,
1322+
create_lambda_function,
1323+
create_role_with_policy,
1324+
aws_client,
1325+
region_name,
1326+
):
1327+
_, role_arn = create_role_with_policy(
1328+
"Allow", "lambda:InvokeFunction", json.dumps(APIGATEWAY_ASSUME_ROLE_POLICY), "*"
1329+
)
1330+
timeout = 30 if is_aws_cloud() else 3
1331+
1332+
function_name = f"response-format-apigw-{short_uid()}"
1333+
create_function_response = create_lambda_function(
1334+
handler_file=LAMBDA_RESPONSE_FROM_BODY,
1335+
func_name=function_name,
1336+
runtime=Runtime.python3_12,
1337+
)
1338+
# create invocation role
1339+
lambda_arn = create_function_response["CreateFunctionResponse"]["FunctionArn"]
1340+
1341+
# create rest api
1342+
api_id, _, root = create_rest_apigw(
1343+
name=f"test-api-{short_uid()}",
1344+
description="Integration test API",
1345+
)
1346+
1347+
resource_id = aws_client.apigateway.create_resource(
1348+
restApiId=api_id, parentId=root, pathPart="{proxy+}"
1349+
)["id"]
1350+
1351+
aws_client.apigateway.put_method(
1352+
restApiId=api_id,
1353+
resourceId=resource_id,
1354+
httpMethod="ANY",
1355+
authorizationType="NONE",
1356+
)
1357+
1358+
# Lambda AWS_PROXY integration
1359+
aws_client.apigateway.put_integration(
1360+
restApiId=api_id,
1361+
resourceId=resource_id,
1362+
httpMethod="ANY",
1363+
type="AWS_PROXY",
1364+
integrationHttpMethod="POST",
1365+
uri=f"arn:aws:apigateway:{region_name}:lambda:path/2015-03-31/functions/{lambda_arn}/invocations",
1366+
credentials=role_arn,
1367+
)
1368+
1369+
# this deployment does not have any `binaryMediaTypes` configured, so it should not return any binary data
1370+
stage_1 = "test"
1371+
aws_client.apigateway.create_deployment(restApiId=api_id, stageName=stage_1)
1372+
endpoint = api_invoke_url(api_id=api_id, path="/test", stage=stage_1)
1373+
# Base64-encoded PNG image (example: 1x1 pixel transparent PNG)
1374+
image_base64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mP8/wIAAgkBAyMAlYwAAAAASUVORK5CYII="
1375+
binary_data = base64.b64decode(image_base64)
1376+
1377+
decoded_response = {
1378+
"statusCode": 200,
1379+
"body": image_base64,
1380+
"isBase64Encoded": True,
1381+
"headers": {
1382+
"Content-Type": "image/png",
1383+
"Cache-Control": "no-cache",
1384+
},
1385+
}
1386+
1387+
def _assert_invoke(accept: str | None, expect_binary: bool) -> bool:
1388+
headers = {"User-Agent": "python/test"}
1389+
if accept:
1390+
headers["Accept"] = accept
1391+
1392+
_response = requests.post(
1393+
url=endpoint,
1394+
data=json.dumps(decoded_response),
1395+
headers=headers,
1396+
)
1397+
if not _response.status_code == 200:
1398+
return False
1399+
1400+
if expect_binary:
1401+
return _response.content == binary_data
1402+
else:
1403+
return _response.text == image_base64
1404+
1405+
# we poll that the API is returning the right data after deployment
1406+
poll_condition(
1407+
lambda: _assert_invoke(accept="image/png", expect_binary=False), timeout=timeout, interval=1
1408+
)
1409+
if is_aws_cloud():
1410+
time.sleep(5)
1411+
1412+
# we did not configure binaryMedias so the API is not returning binary data even if all conditions are met
1413+
assert _assert_invoke(accept="image/png", expect_binary=False)
1414+
1415+
patch_operations = [
1416+
{"op": "add", "path": "/binaryMediaTypes/image~1png"},
1417+
# seems like wildcard with star on the left is not supported
1418+
{"op": "add", "path": "/binaryMediaTypes/*~1test"},
1419+
]
1420+
aws_client.apigateway.update_rest_api(restApiId=api_id, patchOperations=patch_operations)
1421+
# this deployment has `binaryMediaTypes` configured, so it should now return binary data if the client sends the
1422+
# right `Accept` header and the lambda returns the Content-Type
1423+
if is_aws_cloud():
1424+
time.sleep(10)
1425+
stage_2 = "test2"
1426+
endpoint = api_invoke_url(api_id=api_id, path="/test", stage=stage_2)
1427+
aws_client.apigateway.create_deployment(restApiId=api_id, stageName=stage_2)
1428+
1429+
# we poll that the API is returning the right data after deployment
1430+
poll_condition(
1431+
lambda: _assert_invoke(accept="image/png", expect_binary=True), timeout=timeout, interval=1
1432+
)
1433+
if is_aws_cloud():
1434+
time.sleep(10)
1435+
1436+
# all conditions are met
1437+
assert _assert_invoke(accept="image/png", expect_binary=True)
1438+
1439+
# client is sending the wrong accept, so the API returns the base64 data
1440+
assert _assert_invoke(accept="image/jpg", expect_binary=False)
1441+
1442+
# client is sending the wrong accept (wildcard), so the API returns the base64 data
1443+
assert _assert_invoke(accept="image/*", expect_binary=False)
1444+
1445+
# wildcard on the left is not supported
1446+
assert _assert_invoke(accept="*/test", expect_binary=False)
1447+
1448+
# client is sending an accept that matches the wildcard, but it does not work
1449+
assert _assert_invoke(accept="random/test", expect_binary=False)
1450+
1451+
# Accept has to exactly match what is configured
1452+
assert _assert_invoke(accept="*/*", expect_binary=False)
1453+
1454+
# client is sending a multiple accept, but AWS only checks the first one
1455+
assert _assert_invoke(accept="image/webp,image/png,*/*;q=0.8", expect_binary=False)
1456+
1457+
# client is sending a multiple accept, but AWS only checks the first one, which is right
1458+
assert _assert_invoke(accept="image/png,image/*,*/*;q=0.8", expect_binary=True)
1459+
1460+
# lambda is returning that the payload is not b64 encoded
1461+
decoded_response["isBase64Encoded"] = False
1462+
assert _assert_invoke(accept="image/png", expect_binary=False)
1463+
1464+
patch_operations = [
1465+
{"op": "add", "path": "/binaryMediaTypes/application~1*"},
1466+
{"op": "add", "path": "/binaryMediaTypes/image~1jpg"},
1467+
{"op": "remove", "path": "/binaryMediaTypes/*~1test"},
1468+
]
1469+
aws_client.apigateway.update_rest_api(restApiId=api_id, patchOperations=patch_operations)
1470+
if is_aws_cloud():
1471+
# AWS starts returning 200, but then fails again with 403. Wait a bit for it to be stable
1472+
time.sleep(10)
1473+
1474+
# this deployment has `binaryMediaTypes` configured, so it should now return binary data if the client sends the
1475+
# right `Accept` header
1476+
stage_3 = "test3"
1477+
endpoint = api_invoke_url(api_id=api_id, path="/test", stage=stage_3)
1478+
aws_client.apigateway.create_deployment(restApiId=api_id, stageName=stage_3)
1479+
decoded_response["isBase64Encoded"] = True
1480+
1481+
# we poll that the API is returning the right data after deployment
1482+
poll_condition(
1483+
lambda: _assert_invoke(accept="image/png", expect_binary=True), timeout=timeout, interval=1
1484+
)
1485+
if is_aws_cloud():
1486+
time.sleep(10)
1487+
1488+
# different scenario with right side wildcard, all working
1489+
decoded_response["headers"]["Content-Type"] = "application/test"
1490+
assert _assert_invoke(accept="application/whatever", expect_binary=True)
1491+
assert _assert_invoke(accept="application/test", expect_binary=True)
1492+
assert _assert_invoke(accept="application/*", expect_binary=True)
1493+
1494+
# lambda is returning a content-type that matches one binaryMediaType, but Accept matches another binaryMediaType
1495+
# it seems it does not matter, only Accept is checked
1496+
decoded_response["headers"]["Content-Type"] = "image/png"
1497+
assert _assert_invoke(accept="image/jpg", expect_binary=True)
1498+
1499+
# lambda is returning a content-type that matches the wildcard, but Accept matches another binaryMediaType
1500+
decoded_response["headers"]["Content-Type"] = "application/whatever"
1501+
assert _assert_invoke(accept="image/png", expect_binary=True)
1502+
1503+
# ContentType does not matter at all
1504+
decoded_response["headers"].pop("Content-Type")
1505+
assert _assert_invoke(accept="image/png", expect_binary=True)
1506+
1507+
# bad Accept
1508+
assert _assert_invoke(accept="application", expect_binary=False)
1509+
1510+
# no Accept
1511+
assert _assert_invoke(accept=None, expect_binary=False)
1512+
1513+
# bad base64
1514+
decoded_response["body"] = "èé+à)("
1515+
bad_b64_response = requests.post(
1516+
url=endpoint,
1517+
data=json.dumps(decoded_response),
1518+
headers={"User-Agent": "python/test", "Accept": "image/png"},
1519+
)
1520+
assert bad_b64_response.status_code == 500

‎tests/aws/services/apigateway/test_apigateway_lambda.validation.json

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,7 @@
11
{
2+
"tests/aws/services/apigateway/test_apigateway_lambda.py::test_aws_proxy_binary_response": {
3+
"last_validated_date": "2025-01-29T00:14:36+00:00"
4+
},
25
"tests/aws/services/apigateway/test_apigateway_lambda.py::test_aws_proxy_response_payload_format_validation": {
36
"last_validated_date": "2024-11-15T17:48:06+00:00"
47
},

0 commit comments

Comments
 (0)