From 02e5035b42aa7a24a1d6b6219f784f25cc665ff2 Mon Sep 17 00:00:00 2001
From: Abhijeet Prasad <aprasad@sentry.io>
Date: Tue, 4 Jul 2023 09:43:26 -0400
Subject: [PATCH] ref(serverless): Extract propagation context (#8429)

---
 packages/serverless/src/awslambda.ts          | 23 +++++++-------
 packages/serverless/src/gcpfunction/http.ts   | 30 +++++++------------
 .../serverless/test/__mocks__/@sentry/node.ts |  1 +
 3 files changed, 24 insertions(+), 30 deletions(-)

diff --git a/packages/serverless/src/awslambda.ts b/packages/serverless/src/awslambda.ts
index 888b57852327..b5cf41af5e90 100644
--- a/packages/serverless/src/awslambda.ts
+++ b/packages/serverless/src/awslambda.ts
@@ -3,7 +3,7 @@ import type { Scope } from '@sentry/node';
 import * as Sentry from '@sentry/node';
 import { captureException, captureMessage, flush, getCurrentHub, withScope } from '@sentry/node';
 import type { Integration } from '@sentry/types';
-import { baggageHeaderToDynamicSamplingContext, extractTraceparentData, isString, logger } from '@sentry/utils';
+import { isString, logger, tracingContextFromHeaders } from '@sentry/utils';
 // NOTE: I have no idea how to fix this right now, and don't want to waste more time, as it builds just fine — Kamil
 // eslint-disable-next-line import/no-unresolved
 import type { Context, Handler } from 'aws-lambda';
@@ -274,17 +274,20 @@ export function wrapHandler<TEvent, TResult>(
       }, timeoutWarningDelay) as unknown as NodeJS.Timeout;
     }
 
-    // Applying `sentry-trace` to context
-    let traceparentData;
-    const eventWithHeaders = event as { headers?: { [key: string]: string } };
-    if (eventWithHeaders.headers && isString(eventWithHeaders.headers['sentry-trace'])) {
-      traceparentData = extractTraceparentData(eventWithHeaders.headers['sentry-trace']);
-    }
+    const hub = getCurrentHub();
 
-    const baggageHeader = eventWithHeaders.headers && eventWithHeaders.headers.baggage;
-    const dynamicSamplingContext = baggageHeaderToDynamicSamplingContext(baggageHeader);
+    const eventWithHeaders = event as { headers?: { [key: string]: string } };
 
-    const hub = getCurrentHub();
+    const sentryTrace =
+      eventWithHeaders.headers && isString(eventWithHeaders.headers['sentry-trace'])
+        ? eventWithHeaders.headers['sentry-trace']
+        : undefined;
+    const baggage = eventWithHeaders.headers?.baggage;
+    const { traceparentData, dynamicSamplingContext, propagationContext } = tracingContextFromHeaders(
+      sentryTrace,
+      baggage,
+    );
+    hub.getScope().setPropagationContext(propagationContext);
 
     const transaction = hub.startTransaction({
       name: context.functionName,
diff --git a/packages/serverless/src/gcpfunction/http.ts b/packages/serverless/src/gcpfunction/http.ts
index 9bc9052d179d..d3b24f5ba743 100644
--- a/packages/serverless/src/gcpfunction/http.ts
+++ b/packages/serverless/src/gcpfunction/http.ts
@@ -1,13 +1,6 @@
 import type { AddRequestDataToEventOptions } from '@sentry/node';
 import { captureException, flush, getCurrentHub } from '@sentry/node';
-import {
-  baggageHeaderToDynamicSamplingContext,
-  extractTraceparentData,
-  isString,
-  isThenable,
-  logger,
-  stripUrlQueryAndFragment,
-} from '@sentry/utils';
+import { isString, isThenable, logger, stripUrlQueryAndFragment, tracingContextFromHeaders } from '@sentry/utils';
 
 import { domainify, proxyFunction } from './../utils';
 import type { HttpFunction, WrapperOptions } from './general';
@@ -68,21 +61,18 @@ function _wrapHttpFunction(fn: HttpFunction, wrapOptions: Partial<HttpFunctionWr
     ...wrapOptions,
   };
   return (req, res) => {
+    const hub = getCurrentHub();
+
     const reqMethod = (req.method || '').toUpperCase();
     const reqUrl = stripUrlQueryAndFragment(req.originalUrl || req.url || '');
 
-    // Applying `sentry-trace` to context
-    let traceparentData;
-    const reqWithHeaders = req as { headers?: { [key: string]: string } };
-    if (reqWithHeaders.headers && isString(reqWithHeaders.headers['sentry-trace'])) {
-      traceparentData = extractTraceparentData(reqWithHeaders.headers['sentry-trace']);
-    }
-
-    const baggageHeader = reqWithHeaders.headers?.baggage;
-
-    const dynamicSamplingContext = baggageHeaderToDynamicSamplingContext(baggageHeader);
-
-    const hub = getCurrentHub();
+    const sentryTrace = req.headers && isString(req.headers['sentry-trace']) ? req.headers['sentry-trace'] : undefined;
+    const baggage = req.headers?.baggage;
+    const { traceparentData, dynamicSamplingContext, propagationContext } = tracingContextFromHeaders(
+      sentryTrace,
+      baggage,
+    );
+    hub.getScope().setPropagationContext(propagationContext);
 
     const transaction = hub.startTransaction({
       name: `${reqMethod} ${reqUrl}`,
diff --git a/packages/serverless/test/__mocks__/@sentry/node.ts b/packages/serverless/test/__mocks__/@sentry/node.ts
index 355065317ea7..6da20c091780 100644
--- a/packages/serverless/test/__mocks__/@sentry/node.ts
+++ b/packages/serverless/test/__mocks__/@sentry/node.ts
@@ -22,6 +22,7 @@ export const fakeScope = {
   setSpan: jest.fn(),
   getTransaction: jest.fn(() => fakeTransaction),
   setSDKProcessingMetadata: jest.fn(),
+  setPropagationContext: jest.fn(),
 };
 export const fakeSpan = {
   finish: jest.fn(),