From c49f12e7b9b6f909734cb4161043cf2a91f59df8 Mon Sep 17 00:00:00 2001
From: Mahad Janjua <majanjua@amazon.com>
Date: Thu, 6 Jun 2024 17:19:07 -0700
Subject: [PATCH] Allow lambda passthrough trace header

---
 packages/core/lib/patchers/aws3_p.ts | 4 +++-
 packages/core/lib/patchers/aws_p.js  | 3 +++
 packages/core/lib/patchers/http_p.js | 6 ++++--
 packages/core/lib/utils.js           | 4 ++++
 4 files changed, 14 insertions(+), 3 deletions(-)

diff --git a/packages/core/lib/patchers/aws3_p.ts b/packages/core/lib/patchers/aws3_p.ts
index 782317e4..c72d17a6 100644
--- a/packages/core/lib/patchers/aws3_p.ts
+++ b/packages/core/lib/patchers/aws3_p.ts
@@ -140,7 +140,9 @@ const getXRayMiddleware = (config: RegionResolvedConfig, manualSegment?: Segment
     }
   }
 
-  args.request.headers['X-Amzn-Trace-Id'] = traceHeader;
+  if (!segment.noOp) {
+    args.request.headers['X-Amzn-Trace-Id'] = traceHeader;
+  }
 
   let res;
   try {
diff --git a/packages/core/lib/patchers/aws_p.js b/packages/core/lib/patchers/aws_p.js
index 776d33ef..52c1f021 100644
--- a/packages/core/lib/patchers/aws_p.js
+++ b/packages/core/lib/patchers/aws_p.js
@@ -88,6 +88,9 @@ function captureAWSRequest(req) {
   const data = parent.segment ? parent.segment.additionalTraceData : parent.additionalTraceData;
 
   var buildListener = function(req) {
+    if (parent.noOp) {
+      return;
+    }
     let traceHeader = 'Root=' + traceId + ';Parent=' + subsegment.id +
       ';Sampled=' + (subsegment.notTraced ? '0' : '1');
     if (data != null) {
diff --git a/packages/core/lib/patchers/http_p.js b/packages/core/lib/patchers/http_p.js
index fcf2f16f..d84563e1 100644
--- a/packages/core/lib/patchers/http_p.js
+++ b/packages/core/lib/patchers/http_p.js
@@ -130,8 +130,10 @@ function enableCapture(module, downstreamXRayEnabled, subsegmentCallback) {
       options.headers = {};
     }
 
-    options.headers['X-Amzn-Trace-Id'] = 'Root=' + root.trace_id + ';Parent=' + subsegment.id +
-      ';Sampled=' + (subsegment.notTraced ? '0' : '1');
+    if (!parent.noOp) {
+      options.headers['X-Amzn-Trace-Id'] = 'Root=' + root.trace_id + ';Parent=' + subsegment.id +
+        ';Sampled=' + (subsegment.notTraced ? '0' : '1');
+    }
 
     const errorCapturer = function errorCapturer(e) {
       if (subsegmentCallback) {
diff --git a/packages/core/lib/utils.js b/packages/core/lib/utils.js
index 93efcd3b..e047b71a 100644
--- a/packages/core/lib/utils.js
+++ b/packages/core/lib/utils.js
@@ -168,6 +168,10 @@ var utils = {
       if (!traceData) {
         traceData = {};
         logger.getLogger().error('_X_AMZN_TRACE_ID is empty or has an invalid format');
+      } else if (traceData.root && !traceData.parent && !traceData.sampled) {
+        // Lambda PassThrough only has root, treat as valid in this case and mark the segment
+        segment.noOp = true;
+        valid = true;
       } else if (!traceData.root || !traceData.parent || !traceData.sampled) {
         logger.getLogger().error('_X_AMZN_TRACE_ID is missing required information');
       } else {