🚧 Add event source detection

mangum/adapter.py

@@ -6,7 +6,7 @@
 from dataclasses import dataclass, InitVar
 from contextlib import ExitStack
 
-from mangum.types import ASGIApp, Scope
+from mangum.types import ASGIApp, Scope, EventSource
 from mangum.protocols.lifespan import LifespanCycle
 from mangum.protocols.http import HTTPCycle
 from mangum.exceptions import ConfigurationError
@@ -84,6 +84,8 @@ def __post_init__(self, text_mime_types: typing.Optional[typing.List[str]]) -> N
     def __call__(self, event: dict, context: "LambdaContext") -> dict:
         self.logger.debug("Event received.")
 
+        event_source = EventSource.get_event_source(event)
+
         with ExitStack() as stack:
             if self.lifespan != "off":
                 lifespan_cycle: typing.ContextManager = LifespanCycle(
@@ -98,13 +100,18 @@ def __call__(self, event: dict, context: "LambdaContext") -> dict:
             elif not isinstance(initial_body, bytes):
                 initial_body = initial_body.encode()
 
-            scope = self._create_scope(event, context)
+            scope = self._create_scope(event, event_source, context)
+
             http_cycle = HTTPCycle(scope, text_mime_types=self.text_mime_types)
+            # TODO: we should construct response according to request event type.
+            # Take event_source into account in the future.
             response = http_cycle(self.app, initial_body)
 
         return response
 
-    def _create_scope(self, event: dict, context: "LambdaContext") -> Scope:
+    def _create_scope(
+        self, event: dict, event_source: EventSource, context: "LambdaContext"
+    ) -> Scope:
         """Creates a scope object according to ASGI specification from a Lambda Event.
 
         https://asgi.readthedocs.io/en/latest/specs/www.html#http-connection-scope
@@ -127,7 +134,7 @@ def _create_scope(self, event: dict, context: "LambdaContext") -> Scope:
             headers = {}
 
         # API Gateway v2
-        if event.get("version") == "2.0":
+        if event_source == EventSource.API_GW_V2:
             source_ip = request_context["http"]["sourceIp"]
             path = request_context["http"]["path"]
             http_method = request_context["http"]["method"]
@@ -138,7 +145,7 @@ def _create_scope(self, event: dict, context: "LambdaContext") -> Scope:
 
         # API Gateway v1 / ELB
         else:
-            if "elb" in request_context:
+            if event_source in {EventSource.ALB, EventSource.ALB_MULTIVALUEHEADERS}:
                 # NOTE: trust only the most right side value
                 source_ip = headers.get("x-forwarded-for", "").split(", ")[-1]
             else:

mangum/types.py

@@ -1,3 +1,4 @@
+import enum
 import typing
 from typing_extensions import Protocol
 
@@ -10,3 +11,23 @@
 class ASGIApp(Protocol):
     async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
         ...  # pragma: no cover
+
+
+class EventSource(enum.Enum):
+    ALB = enum.auto()
+    ALB_MULTIVALUEHEADERS = enum.auto()
+    API_GW_V1 = enum.auto()
+    API_GW_V2 = enum.auto()
+
+    @classmethod
+    def get_event_source(cls, event: dict) -> "EventSource":
+        version_val = event.get("version", None)
+        multi_value_headers_val = event.get("multiValueHeaders", None)
+        if version_val == "1.0":
+            return cls.API_GW_V1
+        elif version_val == "2.0":
+            return cls.API_GW_V2
+        elif multi_value_headers_val is not None:
+            return cls.ALB_MULTIVALUEHEADERS
+        else:
+            return cls.ALB

tests/conftest.py

@@ -7,6 +7,7 @@ def mock_http_event(request):
     body = request.param[1]
     multi_value_query_parameters = request.param[2]
     event = {
+        "version": "1.0",
         "path": "/test/hello",
         "body": body,
         "headers": {

tests/test_http.py

@@ -89,6 +89,7 @@ async def app(scope, receive, send):
                 },
                 "resource": "/{proxy+}",
                 "stageVariables": {"stageVarName": "stageVarValue"},
+                "version": "1.0",
             },
             "client": ("192.168.100.1", 0),
             "headers": [
@@ -207,6 +208,7 @@ async def app(scope, receive, send):
                 },
                 "resource": "/{proxy+}",
                 "stageVariables": {"stageVarName": "stageVarValue"},
+                "version": "1.0",
             },
             "client": ("192.168.100.1", 0),
             "headers": [

tests/test_types.py

@@ -0,0 +1,187 @@
+from mangum.types import EventSource
+
+import pytest
+
+
+@pytest.mark.parametrize(
+    "event,event_source",
+    [
+        (
+            {
+                "body": '{"username":"xyz","password":"xyz"}',
+                "headers": {
+                    "accept": "*/*",
+                    "content-length": "35",
+                    "content-type": "application/json",
+                    "head": "abc",
+                    "host": "test-755069476.eu-central-1.elb.amazonaws.com",
+                    "user-agent": "curl/7.54.0",
+                    "x-amzn-trace-id": "Root=1-5fff2c3d-02201fad1b80d4fe331120e7",
+                    "x-forwarded-for": "77.79.170.37",
+                    "x-forwarded-port": "80",
+                    "x-forwarded-proto": "http",
+                },
+                "httpMethod": "POST",
+                "isBase64Encoded": False,
+                "path": "/user/id/123123",
+                "queryStringParameters": {"a": "bar", "q": "baz"},
+                "requestContext": {
+                    "elb": {
+                        "targetGroupArn": (
+                            "arn:aws:elasticloadbalancing:eu-central-1:123:"
+                            "targetgroup/lambda/6a653f0ffbc88fec"
+                        )
+                    },
+                },
+            },
+            EventSource.ALB,
+        ),
+        (
+            {
+                "body": '{"username":"xyz","password":"xyz"}',
+                "httpMethod": "POST",
+                "isBase64Encoded": False,
+                "multiValueHeaders": {
+                    "accept": ["*/*"],
+                    "content-length": ["35"],
+                    "content-type": ["application/json"],
+                    "head": ["123", "abc"],
+                    "host": ["test-755069476.eu-central-1.elb.amazonaws.com"],
+                    "user-agent": ["curl/7.54.0"],
+                    "x-amzn-trace-id": ["Root=1-5fff2c9e-554de4de50c1c373369d2ba9"],
+                    "x-forwarded-for": ["77.79.170.37"],
+                    "x-forwarded-port": ["80"],
+                    "x-forwarded-proto": ["http"],
+                },
+                "multiValueQueryStringParameters": {"a": ["bar"], "q": ["foo", "baz"]},
+                "path": "/user/id/123123",
+                "requestContext": {
+                    "elb": {
+                        "targetGroupArn": (
+                            "arn:aws:elasticloadbalancing:eu-central-1:123:"
+                            "targetgroup/lambda/6a653f0ffbc88fec"
+                        )
+                    }
+                },
+            },
+            EventSource.ALB_MULTIVALUEHEADERS,
+        ),
+        (
+            {
+                "body": '{"username":"xyz","password":"xyz"}',
+                "headers": {
+                    "Content-Length": "35",
+                    "Content-Type": "application/json",
+                    "Cookie": "foo=bar",
+                    "Host": "pr3k5m9ob8.execute-api.eu-central-1.amazonaws.com",
+                    "User-Agent": "curl/7.54.0",
+                    "X-Amzn-Trace-Id": "Root=1-5fffc75d-272ca7e75cb8ec6a02f15752",
+                    "X-Forwarded-For": "77.79.170.37",
+                    "X-Forwarded-Port": "443",
+                    "X-Forwarded-Proto": "https",
+                    "accept": "*/*",
+                    "head": "abc",
+                },
+                "httpMethod": "POST",
+                "isBase64Encoded": False,
+                "multiValueHeaders": {
+                    "Content-Length": ["35"],
+                    "Content-Type": ["application/json"],
+                    "Cookie": ["foo=bar"],
+                    "Host": ["pr3k5m9ob8.execute-api.eu-central-1.amazonaws.com"],
+                    "User-Agent": ["curl/7.54.0"],
+                    "X-Amzn-Trace-Id": ["Root=1-5fffc75d-272ca7e75cb8ec6a02f15752"],
+                    "X-Forwarded-For": ["77.79.170.37"],
+                    "X-Forwarded-Port": ["443"],
+                    "X-Forwarded-Proto": ["https"],
+                    "accept": ["*/*"],
+                    "head": ["123", "abc"],
+                },
+                "multiValueQueryStringParameters": {"a": ["bar"], "q": ["foo", "baz"]},
+                "path": "/logger",
+                "pathParameters": None,
+                "queryStringParameters": {"a": "bar", "q": "baz"},
+                "requestContext": {
+                    "accountId": "386635533411",
+                    "apiId": "pr3k5m9ob8",
+                    "domainName": "pr3k5m9ob8.execute-api.eu-central-1.amazonaws.com",
+                    "domainPrefix": "pr3k5m9ob8",
+                    "extendedRequestId": "ZHwWkiwKliAEJNg=",
+                    "httpMethod": "POST",
+                    "identity": {
+                        "accessKey": None,
+                        "accountId": None,
+                        "caller": None,
+                        "cognitoAmr": None,
+                        "cognitoAuthenticationProvider": None,
+                        "cognitoAuthenticationType": None,
+                        "cognitoIdentityId": None,
+                        "cognitoIdentityPoolId": None,
+                        "principalOrgId": None,
+                        "sourceIp": "77.79.170.37",
+                        "user": None,
+                        "userAgent": "curl/7.54.0",
+                        "userArn": None,
+                    },
+                    "path": "/logger",
+                    "protocol": "HTTP/1.1",
+                    "requestId": "ZHwWkiwKliAEJNg=",
+                    "requestTime": "14/Jan/2021:04:23:57 +0000",
+                    "requestTimeEpoch": 1610598237125,
+                    "resourceId": "ANY /logger",
+                    "resourcePath": "/logger",
+                    "stage": "$default",
+                },
+                "resource": "/logger",
+                "stageVariables": None,
+                "version": "1.0",
+            },
+            EventSource.API_GW_V1,
+        ),
+        (
+            {
+                "body": '{"username":"xyz","password":"xyz"}',
+                "cookies": ["foo=bar"],
+                "headers": {
+                    "accept": "*/*",
+                    "content-length": "35",
+                    "content-type": "application/json",
+                    "head": "123,abc",
+                    "host": "pr3k5m9ob8.execute-api.eu-central-1.amazonaws.com",
+                    "user-agent": "curl/7.54.0",
+                    "x-amzn-trace-id": "Root=1-5fffc60e-1679fb5a1d62219d67db18e1",
+                    "x-forwarded-for": "77.79.170.37",
+                    "x-forwarded-port": "443",
+                    "x-forwarded-proto": "https",
+                },
+                "isBase64Encoded": False,
+                "queryStringParameters": {"a": "bar", "q": "foo,baz"},
+                "rawPath": "/logger",
+                "rawQueryString": "q=foo&a=bar&q=baz",
+                "requestContext": {
+                    "accountId": "386635533411",
+                    "apiId": "pr3k5m9ob8",
+                    "domainName": "pr3k5m9ob8.execute-api.eu-central-1.amazonaws.com",
+                    "domainPrefix": "pr3k5m9ob8",
+                    "http": {
+                        "method": "POST",
+                        "path": "/logger",
+                        "protocol": "HTTP/1.1",
+                        "sourceIp": "77.79.170.37",
+                        "userAgent": "curl/7.54.0",
+                    },
+                    "requestId": "ZHviWj0LliAEMig=",
+                    "routeKey": "ANY /logger",
+                    "stage": "$default",
+                    "time": "14/Jan/2021:04:18:22 +0000",
+                    "timeEpoch": 1610597902927,
+                },
+                "routeKey": "ANY /logger",
+                "version": "2.0",
+            },
+            EventSource.API_GW_V2,
+        ),
+    ],
+)
+def test_event_source_detection(event, event_source):
+    assert EventSource.get_event_source(event) == event_source