fixing http api caching cases

Author: joeyzhao2018

datadog_lambda/constants.py

@@ -46,3 +46,4 @@ class XrayDaemon(object):
 
 class OtherConsts(object):
     parentSpanFinishTimeHeader = "x-datadog-parent-span-finish-time"
+    originalAuthorizerRequestEpoch = "x-datadog-original-epoch"

datadog_lambda/tracing.py

@@ -353,13 +353,17 @@ def get_injected_authorizer_data(event, event_source: _EventSource) -> dict:
 
         if event_source.equals(EventTypes.API_GATEWAY, subtype=EventSubtypes.HTTP_API):
             dd_data_raw = authorizer_headers.get("lambda", {}).get("_datadog")
+            dd_data = json.loads(base64.b64decode(dd_data_raw))
+            # [API_GATEWAY V2]if original authorizer request epoch is different, then it's cached
+            if event.get("requestContext", {}).get("timeEpoch") == dd_data.get(
+                OtherConsts.originalAuthorizerRequestEpoch
+            ):
+                return dd_data
+            else:
+                return None
         else:
             dd_data_raw = authorizer_headers.get("_datadog")
-
-        if dd_data_raw:
-            dd_data = json.loads(dd_data_raw)
-            return dd_data
-        return None
+            return json.loads(dd_data_raw)
 
     except Exception as e:
         logger.debug("Failed to check if invocated by an authorizer. error %s", e)

datadog_lambda/wrapper.py

@@ -3,6 +3,7 @@
 # This product includes software developed at Datadog (https://www.datadoghq.com/).
 # Copyright 2019 Datadog, Inc.
 
+import base64
 import os
 import logging
 import traceback
@@ -37,7 +38,12 @@
     create_inferred_span,
     InferredSpanInfo,
 )
-from datadog_lambda.trigger import extract_trigger_tags, extract_http_status_code_tag
+from datadog_lambda.trigger import (
+    extract_trigger_tags,
+    extract_http_status_code_tag,
+    EventTypes,
+    EventSubtypes,
+)
 from datadog_lambda.tag_object import tag_object
 
 profiling_env_var = os.environ.get("DD_PROFILING_ENABLED", "false").lower() == "true"
@@ -155,21 +161,7 @@ def __call__(self, event, context, **kwargs):
         self._before(event, context)
         try:
             self.response = self.func(event, context, **kwargs)
-            try:
-                if (
-                    self.make_authorizer_span
-                    and self.response
-                    and self.response.get("principalId")
-                    and self.response.get("policyDocument")
-                ):
-                    self._inject_authorizer_span_headers()
-            except Exception as e:
-                traceback.print_exc()
-                logger.debug(
-                    "Unable to inject the authorizer headers. \
-                    Continue to return the original response. Reason: %s",
-                    e,
-                )
+            self._ending(event)
             return self.response
         except Exception:
             submit_errors_metric(context)
@@ -179,19 +171,31 @@ def __call__(self, event, context, **kwargs):
         finally:
             self._after(event, context)
 
-    def _inject_authorizer_span_headers(self):
+    def _inject_authorizer_span_headers(self, event):
         finish_time_ns = (
             self.span.start_ns
             if InferredSpanInfo.is_async(self.inferred_span) and self.span
             else time_ns()
         )
-        self.response.setdefault("context", {})
-        self.response["context"].setdefault("_datadog", {})
         injected_headers = {}
         source_span = self.inferred_span if self.inferred_span else self.span
         HTTPPropagator.inject(source_span.context, injected_headers)
         injected_headers[OtherConsts.parentSpanFinishTimeHeader] = finish_time_ns / 1e6
-        self.response["context"]["_datadog"] = json.dumps(injected_headers)
+        encode_datadog_data = False
+        if self.event_source and self.event_source.equals(
+            EventTypes.API_GATEWAY, EventSubtypes.HTTP_API
+        ):
+            injected_headers[OtherConsts.originalAuthorizerRequestEpoch] = event.get(
+                "requestContext", {}
+            ).get("timeEpoch")
+            encode_datadog_data = True
+        datadog_data = json.dumps(injected_headers).encode()
+        if encode_datadog_data:
+            datadog_data = base64.b64encode(
+                datadog_data
+            )  # special characters can corrupt the response
+        self.response.setdefault("context", {})
+        self.response["context"]["_datadog"] = datadog_data
 
     def _before(self, event, context):
         try:
@@ -203,6 +207,7 @@ def _before(self, event, context):
             dd_context, trace_context_source, event_source = extract_dd_trace_context(
                 event, context, extractor=self.trace_extractor
             )
+            self.event_source = event_source
             # Create a Datadog X-Ray subsegment with the trace context
             if dd_context and trace_context_source == TraceContextSource.EVENT:
                 create_dd_dummy_metadata_subsegment(
@@ -232,6 +237,19 @@ def _before(self, event, context):
         except Exception:
             traceback.print_exc()
 
+    def _ending(self, event):
+        try:
+            if (
+                self.make_authorizer_span
+                and self.response
+                and self.response.get("principalId")
+                and self.response.get("policyDocument")
+                and self.event_source
+            ):
+                self._inject_authorizer_span_headers(event)
+        except Exception:
+            traceback.print_exc()
+
     def _after(self, event, context):
         try:
             status_code = extract_http_status_code_tag(self.trigger_tags, self.response)

tests/event_samples/authorizer-request-api-gateway-v2.json

@@ -11,9 +11,10 @@
     "cache-control": "no-cache",
     "content-length": "0",
     "host": "amddr1rix9.execute-api.sa-east-1.amazonaws.com",
-    "postman-token": "a9231288-1486-4753-a589-7ad3ac853c78",
+    "postman-token": "5e0efdb2-d93f-4d44-8792-45404d40dae2",
+    "user": "12345",
     "user-agent": "curl/7.64.1",
-    "x-amzn-trace-id": "Root=1-632a6082-6d32606a1e652ce274e8f055",
+    "x-amzn-trace-id": "Root=1-63321d1e-6b80ef0d2c2cf49600b9c28e",
     "x-forwarded-for": "38.122.226.210",
     "x-forwarded-port": "443",
     "x-forwarded-proto": "https"
@@ -23,7 +24,7 @@
     "apiId": "amddr1rix9",
     "authorizer": {
       "lambda": {
-        "_datadog": "{\"x-datadog-trace-id\": \"12056652649662348062\", \"x-datadog-parent-id\": \"12647339963998804799\", \"x-datadog-sampling-priority\": \"1\", \"x-datadog-tags\": \"_dd.p.dm=-0\", \"x-datadog-parent-span-finish-time\": 1663721602440.8286}",
+        "_datadog": "eyJ4LWRhdGFkb2ctdHJhY2UtaWQiOiAiMTQzNTY5ODM2MTk4NTI5MzMzNTQiLCAieC1kYXRhZG9nLXBhcmVudC1pZCI6ICIxMjY1ODYyMTA4MzUwNTQxMzgwOSIsICJ4LWRhdGFkb2ctc2FtcGxpbmctcHJpb3JpdHkiOiAiMSIsICJ4LWRhdGFkb2ctdGFncyI6ICJfZGQucC5kbT0tMCIsICJ4LWRhdGFkb2ctcGFyZW50LXNwYW4tZmluaXNoLXRpbWUiOiAxNjY0MjI4NjM5NTMzLjc3NTQsICJ4LWRhdGFkb2ctb3JpZ2luYWwtZXBvY2giOiAxNjY0MjI4NjM4MDg0fQ==",
         "scope": "this is just a string"
       }
     },
@@ -39,8 +40,8 @@
     "requestId": "1234567",
     "routeKey": "GET /hello",
     "stage": "$default",
-    "time": "21/Sep/2022:00:53:22 +0000",
-    "timeEpoch": 1663295021832
+    "time": "26/Sep/2022:21:43:58 +0000",
+    "timeEpoch": 1664228638084
   },
   "isBase64Encoded": false
 }

tests/test_tracing.py

@@ -617,7 +617,7 @@ def test_create_inferred_span_from_authorizer_token_api_gateway_v1_event(self):
     def test_create_inferred_span_from_authorizer_request_api_gateway_v2_event(self):
         event_sample_source = "authorizer-request-api-gateway-v2"
         span = self._http_common_testing_items(
-            event_sample_source, "aws.httpapi", 1663721602.44
+            event_sample_source, "aws.httpapi", 1664228639.533
         )
 
         # http-api specific checks:
@@ -628,7 +628,7 @@ def test_create_inferred_span_from_authorizer_request_api_gateway_v2_event(self)
         self.assertEqual(span.get_tag("http.source_ip"), "38.122.226.210")
         self.assertEqual(span.get_tag("http.user_agent"), "curl/7.64.1")
         self.assertEqual(
-            span.start, 1663721602.44
+            span.start, 1664228639.533
         )  # use the injected parent span finish time as an approximation
 
     def _http_common_testing_items(