Skip to content

Commit e67710d

Browse files
cloutierMatcloutierMat
authored
Apigw ng implement header remapping (#11237)
1 parent a0a1ba0 commit e67710d

32 files changed

+1826
-416
lines changed

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

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@
55
from rolo.gateway import RequestContext
66
from werkzeug.datastructures import Headers
77

8-
from localstack.aws.api.apigateway import Method, Resource
8+
from localstack.aws.api.apigateway import Integration, Method, Resource
99
from localstack.services.apigateway.models import RestApiDeployment
1010

1111
from .variables import ContextVariables, LoggingContextVariables
@@ -72,10 +72,10 @@ class RestApiInvocationContext(RequestContext):
7272
This context is going to be used to pass relevant information across an API Gateway invocation.
7373
"""
7474

75-
invocation_request: Optional[InvocationRequest]
76-
"""Contains the data relative to the invocation request"""
7775
deployment: Optional[RestApiDeployment]
7876
"""Contains the invoked REST API Resources"""
77+
integration: Optional[Integration]
78+
"""The Method Integration for the invoked request"""
7979
api_id: Optional[str]
8080
"""The REST API identifier of the invoked API"""
8181
stage: Optional[str]
@@ -87,7 +87,7 @@ class RestApiInvocationContext(RequestContext):
8787
account_id: Optional[str]
8888
"""The account the REST API is living in."""
8989
resource: Optional[Resource]
90-
"""The resource the invocation matched""" # TODO: verify if needed through the invocation
90+
"""The resource the invocation matched"""
9191
resource_method: Optional[Method]
9292
"""The method of the resource the invocation matched"""
9393
stage_variables: Optional[dict[str, str]]
@@ -96,6 +96,8 @@ class RestApiInvocationContext(RequestContext):
9696
"""The $context used in data models, authorizers, mapping templates, and CloudWatch access logging"""
9797
logging_context_variables: Optional[LoggingContextVariables]
9898
"""Additional $context variables available only for access logging, not yet implemented"""
99+
invocation_request: Optional[InvocationRequest]
100+
"""Contains the data relative to the invocation request"""
99101
integration_request: Optional[IntegrationRequest]
100102
"""Contains the data needed to construct an HTTP request to an Integration"""
101103
endpoint_response: Optional[EndpointResponse]

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@ def __init__(self):
3333
)
3434
self.response_handlers.extend(
3535
[
36+
handlers.response_enricher
3637
# add composite response handlers?
3738
]
3839
)

‎localstack-core/localstack/services/apigateway/next_gen/execute_api/handlers/__init__.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010
from .method_response import MethodResponseHandler
1111
from .parse import InvocationRequestParser
1212
from .resource_router import InvocationRequestRouter
13+
from .response_enricher import InvocationResponseEnricher
1314

1415
legacy_handler = LegacyHandler()
1516
parse_request = InvocationRequestParser()
@@ -22,3 +23,4 @@
2223
method_response_handler = MethodResponseHandler()
2324
gateway_exception_handler = GatewayExceptionHandler()
2425
api_key_validation_handler = ApiKeyValidationHandler()
26+
response_enricher = InvocationResponseEnricher()

‎localstack-core/localstack/services/apigateway/next_gen/execute_api/handlers/gateway_exception.py

Lines changed: 23 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22
import logging
33

44
from rolo import Response
5+
from werkzeug.datastructures import Headers
56

67
from localstack.constants import APPLICATION_JSON
78
from localstack.services.apigateway.next_gen.execute_api.api import (
@@ -13,6 +14,9 @@
1314
BaseGatewayException,
1415
get_gateway_response_or_default,
1516
)
17+
from localstack.services.apigateway.next_gen.execute_api.variables import (
18+
GatewayResponseContextVarsError,
19+
)
1620

1721
LOG = logging.getLogger(__name__)
1822

@@ -40,11 +44,21 @@ def __call__(
4044
)
4145
return
4246

43-
LOG.debug("Error raised during invocation: %s", exception.type)
47+
LOG.info("Error raised during invocation: %s", exception.type)
48+
self.set_error_context(exception, context)
4449
error = self.create_exception_response(exception, context)
4550
if error:
4651
response.update_from(error)
4752

53+
@staticmethod
54+
def set_error_context(exception: BaseGatewayException, context: RestApiInvocationContext):
55+
context.context_variables["error"] = GatewayResponseContextVarsError(
56+
message=exception.message,
57+
messageString=exception.message,
58+
responseType=exception.type,
59+
validationErrorString="", # TODO
60+
)
61+
4862
def create_exception_response(
4963
self, exception: BaseGatewayException, context: RestApiInvocationContext
5064
):
@@ -60,13 +74,17 @@ def create_exception_response(
6074
if not status_code:
6175
status_code = exception.status_code or 500
6276

63-
return Response(response=content, headers=headers, status=status_code)
77+
response = Response(response=content, headers=headers, status=status_code)
78+
return response
6479

65-
def _build_response_content(self, exception: BaseGatewayException) -> str:
80+
@staticmethod
81+
def _build_response_content(exception: BaseGatewayException) -> str:
6682
# TODO apply responseTemplates to the content. We should also handle the default simply by managing the default
6783
# template body `{"message":$context.error.messageString}`
6884
return json.dumps({"message": exception.message})
6985

70-
def _build_response_headers(self, exception: BaseGatewayException) -> dict:
86+
@staticmethod
87+
def _build_response_headers(exception: BaseGatewayException) -> dict:
7188
# TODO apply responseParameters to the headers and get content-type from the gateway_response
72-
return {"content-type": APPLICATION_JSON, "x-amzn-ErrorType": exception.code}
89+
headers = Headers({"Content-Type": APPLICATION_JSON, "x-amzn-ErrorType": exception.code})
90+
return headers

‎localstack-core/localstack/services/apigateway/next_gen/execute_api/handlers/integration.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,7 @@ def __call__(
1818
context: RestApiInvocationContext,
1919
response: Response,
2020
):
21-
integration_type = context.resource_method["methodIntegration"]["type"]
21+
integration_type = context.integration["type"]
2222
is_proxy = "PROXY" in integration_type
2323

2424
integration = REST_API_INTEGRATIONS.get(integration_type)

‎localstack-core/localstack/services/apigateway/next_gen/execute_api/handlers/integration_request.py

Lines changed: 95 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -1,16 +1,18 @@
11
import logging
2+
from http import HTTPMethod
23

34
from werkzeug.datastructures import Headers
45

56
from localstack.aws.api.apigateway import Integration, IntegrationType
67
from localstack.constants import APPLICATION_JSON
7-
from localstack.http import Response
8+
from localstack.http import Request, Response
89
from localstack.utils.collections import merge_recursive
9-
from localstack.utils.strings import to_bytes, to_str
10+
from localstack.utils.strings import short_uid, to_bytes, to_str
1011

1112
from ..api import RestApiGatewayHandler, RestApiGatewayHandlerChain
1213
from ..context import IntegrationRequest, InvocationRequest, RestApiInvocationContext
13-
from ..gateway_response import UnsupportedMediaTypeError
14+
from ..gateway_response import InternalServerError, UnsupportedMediaTypeError
15+
from ..header_utils import drop_headers, set_default_headers
1416
from ..helpers import render_integration_uri
1517
from ..parameters_mapping import ParametersMapper, RequestDataMapping
1618
from ..template_mapping import (
@@ -23,6 +25,30 @@
2325

2426
LOG = logging.getLogger(__name__)
2527

28+
# Illegal headers to include in transformation
29+
ILLEGAL_INTEGRATION_REQUESTS_COMMON = [
30+
"content-length",
31+
"transfer-encoding",
32+
"x-amzn-trace-id",
33+
"X-Amzn-Apigateway-Api-Id",
34+
]
35+
ILLEGAL_INTEGRATION_REQUESTS_AWS = [
36+
*ILLEGAL_INTEGRATION_REQUESTS_COMMON,
37+
"authorization",
38+
"connection",
39+
"expect",
40+
"proxy-authenticate",
41+
"te",
42+
]
43+
44+
# These are dropped after the templates override were applied. they will never make it to the requests.
45+
DROPPED_FROM_INTEGRATION_REQUESTS_COMMON = ["Expect", "Proxy-Authenticate", "TE"]
46+
DROPPED_FROM_INTEGRATION_REQUESTS_AWS = [*DROPPED_FROM_INTEGRATION_REQUESTS_COMMON, "Referer"]
47+
DROPPED_FROM_INTEGRATION_REQUESTS_HTTP = [*DROPPED_FROM_INTEGRATION_REQUESTS_COMMON, "Via"]
48+
49+
# Default headers
50+
DEFAULT_REQUEST_HEADERS = {"Accept": APPLICATION_JSON, "Connection": "keep-alive"}
51+
2652

2753
class PassthroughBehavior(str):
2854
# TODO maybe this class should be moved where it can also be used for validation in
@@ -48,7 +74,7 @@ def __call__(
4874
context: RestApiInvocationContext,
4975
response: Response,
5076
):
51-
integration: Integration = context.resource_method["methodIntegration"]
77+
integration: Integration = context.integration
5278
integration_type = integration["type"]
5379

5480
integration_request_parameters = integration["requestParameters"] or {}
@@ -59,7 +85,8 @@ def __call__(
5985

6086
if integration_type in (IntegrationType.AWS_PROXY, IntegrationType.HTTP_PROXY):
6187
# `PROXY` types cannot use integration mapping templates, they pass most of the data straight
62-
headers = context.invocation_request["headers"]
88+
# We make a copy to avoid modifying the invocation headers and keep a cleaner history
89+
headers = context.invocation_request["headers"].copy()
6390
query_string_parameters: dict[str, list[str]] = context.invocation_request[
6491
"multi_value_query_string_parameters"
6592
]
@@ -68,31 +95,22 @@ def __call__(
6895
# HTTP_PROXY still make uses of the request data mappings, and merges it with the invocation request
6996
# this is undocumented but validated behavior
7097
if integration_type == IntegrationType.HTTP_PROXY:
71-
headers = headers.copy()
72-
to_merge = {
73-
k: v
74-
for k, v in request_data_mapping["header"].items()
75-
if k not in ("Content-Type", "Accept")
76-
}
77-
headers.update(to_merge)
98+
# These headers won't be passed through by default from the invocation.
99+
# They can however be added through request mappings.
100+
drop_headers(headers, ["Host", "Content-Encoding"])
101+
headers.update(request_data_mapping["header"])
78102

79103
query_string_parameters = self._merge_http_proxy_query_string(
80104
query_string_parameters, request_data_mapping["querystring"]
81105
)
82106

83107
else:
108+
self._set_proxy_headers(headers, context.request)
84109
# AWS_PROXY does not allow URI path rendering
85110
# TODO: verify this
86111
path_parameters = {}
87112

88113
else:
89-
# default values, can be overridden with right casing
90-
default_headers = {
91-
"Content-Type": APPLICATION_JSON,
92-
"Accept": APPLICATION_JSON,
93-
}
94-
request_data_mapping["header"] = default_headers | request_data_mapping["header"]
95-
96114
# find request template to raise UnsupportedMediaTypeError early
97115
request_template = self.get_request_template(
98116
integration=integration, request=context.invocation_request
@@ -107,6 +125,13 @@ def __call__(
107125
headers = Headers(request_data_mapping["header"])
108126
query_string_parameters = request_data_mapping["querystring"]
109127

128+
# Some headers can't be modified by parameter mappings or mapping templates.
129+
# Aws will raise in those were present. Even for AWS_PROXY, where it is not applying them.
130+
if header_mappings := request_data_mapping["header"]:
131+
self._validate_headers_mapping(header_mappings, integration_type)
132+
133+
self._apply_header_transforms(headers, integration_type, context)
134+
110135
# looks like the stageVariables rendering part is done in the Integration part in AWS
111136
# but we can avoid duplication by doing it here for now
112137
# TODO: if the integration if of AWS Lambda type and the Lambda is in another account, we cannot render
@@ -223,3 +248,54 @@ def _merge_http_proxy_query_string(
223248
new_query_string_parameters[param] = value
224249

225250
return new_query_string_parameters
251+
252+
@staticmethod
253+
def _set_proxy_headers(headers: Headers, request: Request):
254+
headers.set("X-Forwarded-For", request.remote_addr)
255+
headers.set("X-Forwarded-Port", request.environ.get("SERVER_PORT"))
256+
headers.set(
257+
"X-Forwarded-Proto",
258+
request.environ.get("SERVER_PROTOCOL", "").split("/")[0],
259+
)
260+
261+
@staticmethod
262+
def _apply_header_transforms(
263+
headers: Headers, integration_type: IntegrationType, context: RestApiInvocationContext
264+
):
265+
# Dropping matching headers for the provided integration type
266+
match integration_type:
267+
case IntegrationType.AWS:
268+
drop_headers(headers, DROPPED_FROM_INTEGRATION_REQUESTS_AWS)
269+
case IntegrationType.HTTP | IntegrationType.HTTP_PROXY:
270+
drop_headers(headers, DROPPED_FROM_INTEGRATION_REQUESTS_HTTP)
271+
case _:
272+
drop_headers(headers, DROPPED_FROM_INTEGRATION_REQUESTS_COMMON)
273+
274+
# Adding default headers to the requests headers
275+
default_headers = {
276+
**DEFAULT_REQUEST_HEADERS,
277+
"User-Agent": f"AmazonAPIGateway_{context.api_id}",
278+
}
279+
if (
280+
content_type := context.request.headers.get("Content-Type")
281+
) and context.request.method not in {HTTPMethod.OPTIONS, HTTPMethod.GET, HTTPMethod.HEAD}:
282+
default_headers["Content-Type"] = content_type
283+
284+
set_default_headers(headers, default_headers)
285+
headers.set("X-Amzn-Trace-Id", short_uid()) # TODO
286+
if integration_type not in (IntegrationType.AWS_PROXY, IntegrationType.AWS):
287+
headers.set("X-Amzn-Apigateway-Api-Id", context.api_id)
288+
289+
@staticmethod
290+
def _validate_headers_mapping(headers: dict[str, str], integration_type: IntegrationType):
291+
"""Validates and raises an error when attempting to set an illegal header"""
292+
to_validate = ILLEGAL_INTEGRATION_REQUESTS_COMMON
293+
if integration_type in {IntegrationType.AWS, IntegrationType.AWS_PROXY}:
294+
to_validate = ILLEGAL_INTEGRATION_REQUESTS_AWS
295+
296+
for header in headers.keys():
297+
if header.lower() in to_validate:
298+
LOG.debug(
299+
"Execution failed due to configuration error: %s header already present", header
300+
)
301+
raise InternalServerError("Internal server error")

‎localstack-core/localstack/services/apigateway/next_gen/execute_api/handlers/integration_response.py

Lines changed: 1 addition & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -47,7 +47,7 @@ def __call__(
4747
):
4848
# TODO: we should log the response coming in from the Integration, either in Integration or here.
4949
# before modification / after?
50-
integration: Integration = context.resource_method["methodIntegration"]
50+
integration: Integration = context.integration
5151
integration_type = integration["type"]
5252

5353
if integration_type in (IntegrationType.AWS_PROXY, IntegrationType.HTTP_PROXY):
@@ -58,7 +58,6 @@ def __call__(
5858
endpoint_response: EndpointResponse = context.endpoint_response
5959
status_code = endpoint_response["status_code"]
6060
body = endpoint_response["body"]
61-
headers = endpoint_response["headers"]
6261

6362
# we first need to find the right IntegrationResponse based on their selection template, linked to the status
6463
# code of the Response
@@ -104,26 +103,6 @@ def __call__(
104103
LOG.debug("Response header overrides: %s", header_override)
105104
response_headers.update(header_override)
106105

107-
# setting up default content-type
108-
if not response_headers.get("content-type"):
109-
response_headers.set("Content-Type", APPLICATION_JSON)
110-
111-
# TODO : refactor the following into method response using table from
112-
# https://docs.aws.amazon.com/apigateway/latest/developerguide/api-gateway-known-issues.html#api-gateway-known-issues-rest-apis
113-
# When trying to override certain headers, they will instead be remapped
114-
# There may be other headers, but these have been confirmed on aws
115-
remapped_headers = ("connection", "content-length", "date", "x-amzn-requestid")
116-
for header in remapped_headers:
117-
if value := response_headers.get(header):
118-
response_headers[f"x-amzn-remapped-{header}"] = value
119-
response_headers.remove(header)
120-
121-
# Those headers are passed through from the response headers, there might be more?
122-
passthrough_headers = ["connection"]
123-
for header in passthrough_headers:
124-
if values := headers.getlist(header):
125-
response_headers.setlist(header, values)
126-
127106
LOG.debug("Method response body after transformations: %s", body)
128107
context.invocation_response = InvocationResponse(
129108
body=body,

0 commit comments

Comments
 (0)