Add support for API Gateway v2 Lambda payload (#288)

* Detect API Gateway v2 trigger events on AWS Lambda

* Extract trace context from API Gateway v2 events

* Extract tags from API Gateway v2 trigger events for the entry span

* Safely read HTTP headers from AWL Lambda proxy event payloads

Author: Andrey Slotin

instana/instrumentation/aws/triggers.py

@@ -13,8 +13,12 @@
 
 def get_context(tracer, event):
     # TODO: Search for more types of trigger context
-    if is_api_gateway_proxy_trigger(event) or is_application_load_balancer_trigger(event):
-        return tracer.extract('http_headers', event['headers'])
+    is_proxy_event = is_api_gateway_proxy_trigger(event) or \
+        is_api_gateway_v2_proxy_trigger(event) or \
+        is_application_load_balancer_trigger(event)
+
+    if is_proxy_event:
+        return tracer.extract('http_headers', event.get('headers', {}))
 
     return tracer.extract('http_headers', event)
 
@@ -26,6 +30,20 @@ def is_api_gateway_proxy_trigger(event):
     return True
 
 
+def is_api_gateway_v2_proxy_trigger(event):
+    for key in ["version", "requestContext"]:
+        if key not in event:
+            return False
+
+    if event["version"] != "2.0":
+        return False
+
+    for key in ["apiId", "stage", "http"]:
+        if key not in event["requestContext"]:
+            return False
+
+    return True
+
 def is_application_load_balancer_trigger(event):
     if 'requestContext' in event and 'elb' in event['requestContext']:
         return True
@@ -143,6 +161,23 @@ def enrich_lambda_span(agent, span, event, context):
             if agent.options.extra_http_headers is not None:
                 capture_extra_headers(event, span, agent.options.extra_http_headers)
 
+        elif is_api_gateway_v2_proxy_trigger(event):
+            logger.debug("Detected as API Gateway v2.0 Proxy Trigger")
+
+            reqCtx = event["requestContext"]
+
+            # trim optional HTTP method prefix
+            route_path = event["routeKey"].split(" ", 2)[-1]
+
+            span.set_tag(STR_LAMBDA_TRIGGER, 'aws:api.gateway')
+            span.set_tag('http.method', reqCtx["http"]["method"])
+            span.set_tag('http.url', reqCtx["http"]["path"])
+            span.set_tag('http.path_tpl', route_path)
+            span.set_tag('http.params', read_http_query_params(event))
+
+            if agent.options.extra_http_headers is not None:
+                capture_extra_headers(event, span, agent.options.extra_http_headers)
+
         elif is_application_load_balancer_trigger(event):
             logger.debug("Detected as Application Load Balancer Trigger")
             span.set_tag(STR_LAMBDA_TRIGGER, 'aws:application.load.balancer')

tests/data/lambda/api_gateway_v2_event.json

@@ -0,0 +1,75 @@
+{
+  "version": "2.0",
+  "routeKey": "ANY /my/{resource}",
+  "rawPath": "/my/path",
+  "rawQueryString": "parameter1=value1&parameter1=value2&parameter2=value",
+  "cookies": [
+    "cookie1",
+    "cookie2"
+  ],
+  "headers": {
+    "Header1": "value1",
+    "Header2": "value1,value2",
+    "X-Instana-T": "0000000000001234",
+    "X-Instana-S": "0000000000004567",
+    "X-Instana-L": "1",
+    "X-Instana-Synthetic": "1",
+    "X-Custom-Header-1": "value1",
+    "x-custom-header-2": "value2"
+  },
+  "queryStringParameters": {
+    "secret": "key",
+    "q": "term"
+  },
+  "requestContext": {
+    "accountId": "123456789012",
+    "apiId": "api-id",
+    "authentication": {
+      "clientCert": {
+        "clientCertPem": "CERT_CONTENT",
+        "subjectDN": "www.example.com",
+        "issuerDN": "Example issuer",
+        "serialNumber": "a1:a1:a1:a1:a1:a1:a1:a1:a1:a1:a1:a1:a1:a1:a1:a1",
+        "validity": {
+          "notBefore": "May 28 12:30:02 2019 GMT",
+          "notAfter": "Aug  5 09:36:04 2021 GMT"
+        }
+      }
+    },
+    "authorizer": {
+      "jwt": {
+        "claims": {
+          "claim1": "value1",
+          "claim2": "value2"
+        },
+        "scopes": [
+          "scope1",
+          "scope2"
+        ]
+      }
+    },
+    "domainName": "id.execute-api.us-east-1.amazonaws.com",
+    "domainPrefix": "id",
+    "http": {
+      "method": "POST",
+      "path": "/my/path",
+      "protocol": "HTTP/1.1",
+      "sourceIp": "IP",
+      "userAgent": "agent"
+    },
+    "requestId": "id",
+    "routeKey": "$default",
+    "stage": "$default",
+    "time": "12/Mar/2020:19:03:58 +0000",
+    "timeEpoch": 1583348638390
+  },
+  "body": "Hello from Lambda",
+  "pathParameters": {
+    "parameter1": "value1"
+  },
+  "isBase64Encoded": false,
+  "stageVariables": {
+    "stageVariable1": "value1",
+    "stageVariable2": "value2"
+  }
+}

tests/platforms/test_lambda.py

@@ -300,6 +300,73 @@ def test_api_gateway_trigger_tracing(self):
         else:
             self.assertEqual("foo=['bar']", span.data['http']['params'])
 
+    def test_api_gateway_v2_trigger_tracing(self):
+        with open(self.pwd + '/../data/lambda/api_gateway_v2_event.json', 'r') as json_file:
+            event = json.load(json_file)
+
+        self.create_agent_and_setup_tracer()
+
+        # Call the Instana Lambda Handler as we do in the real world.  It will initiate tracing and then
+        # figure out the original (the users') Lambda Handler and execute it.
+        # The original Lambda handler is set in os.environ["LAMBDA_HANDLER"]
+        result = lambda_handler(event, self.context)
+
+        assert isinstance(result, dict)
+        assert 'headers' in result
+        assert 'Server-Timing' in result['headers']
+
+        time.sleep(1)
+        payload = self.agent.collector.prepare_payload()
+
+        self.assertTrue("metrics" in payload)
+        self.assertTrue("spans" in payload)
+        self.assertEqual(2, len(payload.keys()))
+
+        self.assertTrue(isinstance(payload['metrics']['plugins'], list))
+        self.assertTrue(len(payload['metrics']['plugins']) == 1)
+        plugin_data = payload['metrics']['plugins'][0]
+
+        self.assertEqual('com.instana.plugin.aws.lambda', plugin_data['name'])
+        self.assertEqual('arn:aws:lambda:us-east-2:12345:function:TestPython:1', plugin_data['entityId'])
+
+        self.assertEqual(1, len(payload['spans']))
+
+        span = payload['spans'][0]
+        self.assertEqual('aws.lambda.entry', span.n)
+        self.assertEqual('0000000000001234', span.t)
+        self.assertIsNotNone(span.s)
+        self.assertEqual('0000000000004567', span.p)
+        self.assertIsNotNone(span.ts)
+        self.assertIsNotNone(span.d)
+
+        server_timing_value = "intid;desc=%s" % span.t
+        assert result['headers']['Server-Timing'] == server_timing_value
+
+        self.assertEqual({'hl': True, 'cp': 'aws', 'e': 'arn:aws:lambda:us-east-2:12345:function:TestPython:1'},
+                         span.f)
+
+        self.assertTrue(span.sy)
+
+        self.assertIsNone(span.ec)
+        self.assertIsNone(span.data['lambda']['error'])
+
+        self.assertEqual('arn:aws:lambda:us-east-2:12345:function:TestPython:1', span.data['lambda']['arn'])
+        self.assertEqual(None, span.data['lambda']['alias'])
+        self.assertEqual('python', span.data['lambda']['runtime'])
+        self.assertEqual('TestPython', span.data['lambda']['functionName'])
+        self.assertEqual('1', span.data['lambda']['functionVersion'])
+        self.assertIsNone(span.data['service'])
+
+        self.assertEqual('aws:api.gateway', span.data['lambda']['trigger'])
+        self.assertEqual('POST', span.data['http']['method'])
+        self.assertEqual('/my/path', span.data['http']['url'])
+        self.assertEqual('/my/{resource}', span.data['http']['path_tpl'])
+        if sys.version[:3] == '2.7':
+            self.assertEqual(u"q=term&secret=key", span.data['http']['params'])
+        else:
+            self.assertEqual("secret=key&q=term", span.data['http']['params'])
+
+
     def test_application_lb_trigger_tracing(self):
         with open(self.pwd + '/../data/lambda/api_gateway_event.json', 'r') as json_file:
             event = json.load(json_file)