add tracestate header to outgoing requests

lobsterkatie · lobsterkatie · commit ea6b9ef1085b · 2021-08-19T22:18:47.000-07:00
diff --git a/packages/node/src/integrations/http.ts b/packages/node/src/integrations/http.ts
@@ -1,5 +1,5 @@
 import { getCurrentHub } from '@sentry/core';
-import { Integration, Span } from '@sentry/types';
+import { Integration, Span, TraceHeaders } from '@sentry/types';
 import { fill, logger, parseSemver } from '@sentry/utils';
 import * as http from 'http';
 import * as https from 'https';
@@ -115,11 +115,9 @@ function _createWrappedRequestMethodFactory(
             op: 'request',
           });
 
-          const sentryTraceHeader = span.toTraceparent();
-          logger.log(
-            `[Tracing] Adding sentry-trace header ${sentryTraceHeader} to outgoing request to ${requestUrl}: `,
-          );
-          requestOptions.headers = { ...requestOptions.headers, 'sentry-trace': sentryTraceHeader };
+          const traceHeaders = span.getTraceHeaders();
+          logger.log(`[Tracing] Adding sentry-trace and tracestate headers to outgoing request to ${requestUrl}.`);
+          requestOptions.headers = { ...requestOptions.headers, ...(traceHeaders as TraceHeaders) };
         }
       }
 
diff --git a/packages/node/test/integrations/http.test.ts b/packages/node/test/integrations/http.test.ts
@@ -9,7 +9,7 @@ import { NodeClient } from '../../src/client';
 import { Http as HttpIntegration } from '../../src/integrations/http';
 
 describe('tracing', () => {
-  function createTransactionOnScope() {
+  function createTransactionOnScope(): Transaction {
     const hub = new Hub(
       new NodeClient({
         dsn: 'https://dogsarebadatkeepingsecrets@squirrelchasers.ingest.sentry.io/12312012',
@@ -24,7 +24,10 @@ describe('tracing', () => {
     jest.spyOn(sentryCore, 'getCurrentHub').mockReturnValue(hub);
     jest.spyOn(hubModule, 'getCurrentHub').mockReturnValue(hub);
 
-    const transaction = hub.startTransaction({ name: 'dogpark' });
+    // we have to cast this to a Transaction (the class) because hub.startTransaction only returns a Transaction (the
+    // interface, which doesn't have things like spanRecorder) (and because @sentry/hub can't depend on @sentry/tracing,
+    // we can't fix that)
+    const transaction = hub.startTransaction({ name: 'dogpark' }) as Transaction;
     hub.getScope()?.setSpan(transaction);
 
     return transaction;
@@ -36,7 +39,7 @@ describe('tracing', () => {
       .reply(200);
 
     const transaction = createTransactionOnScope();
-    const spans = (transaction as Span).spanRecorder?.spans as Span[];
+    const spans = transaction.spanRecorder?.spans as Span[];
 
     http.get('http://dogs.are.great/');
 
@@ -55,7 +58,7 @@ describe('tracing', () => {
       .reply(200);
 
     const transaction = createTransactionOnScope();
-    const spans = (transaction as Span).spanRecorder?.spans as Span[];
+    const spans = transaction.spanRecorder?.spans as Span[];
 
     http.get('http://squirrelchasers.ingest.sentry.io/api/12312012/store/');
 
@@ -64,21 +67,23 @@ describe('tracing', () => {
     expect((spans[0] as Transaction).name).toEqual('dogpark');
   });
 
-  it('attaches the sentry-trace header to outgoing non-sentry requests', async () => {
+  it('attaches tracing headers to outgoing non-sentry requests', async () => {
     nock('http://dogs.are.great')
       .get('/')
       .reply(200);
 
     createTransactionOnScope();
 
     const request = http.get('http://dogs.are.great/');
-    const sentryTraceHeader = request.getHeader('sentry-trace') as string;
+    const sentryTraceHeader = request.getHeader('sentry-trace');
+    const tracestateHeader = request.getHeader('tracestate');
 
     expect(sentryTraceHeader).toBeDefined();
-    expect(SENTRY_TRACE_REGEX.test(sentryTraceHeader)).toBe(true);
+    expect(tracestateHeader).toBeDefined();
+    expect(SENTRY_TRACE_REGEX.test(sentryTraceHeader as string)).toBe(true);
   });
 
-  it("doesn't attach the sentry-trace header to outgoing sentry requests", () => {
+  it("doesn't attach tracing headers to outgoing sentry requests", () => {
     nock('http://squirrelchasers.ingest.sentry.io')
       .get('/api/12312012/store/')
       .reply(200);
@@ -87,7 +92,9 @@ describe('tracing', () => {
 
     const request = http.get('http://squirrelchasers.ingest.sentry.io/api/12312012/store/');
     const sentryTraceHeader = request.getHeader('sentry-trace');
+    const tracestateHeader = request.getHeader('tracestate');
 
     expect(sentryTraceHeader).not.toBeDefined();
+    expect(tracestateHeader).not.toBeDefined();
   });
 });
diff --git a/packages/tracing/src/browser/request.ts b/packages/tracing/src/browser/request.ts
@@ -1,6 +1,6 @@
+import { Span } from '@sentry/types';
 import { addInstrumentationHandler, isInstanceOf, isMatchingPattern } from '@sentry/utils';
 
-import { Span } from '../span';
 import { SpanStatus } from '../spanstatus';
 import { getActiveTransaction, hasTracingEnabled } from '../utils';
 
@@ -59,6 +59,16 @@ export interface FetchData {
   endTimestamp?: number;
 }
 
+type PolymorphicRequestHeaders =
+  | Record<string, string>
+  | Array<[string, string]>
+  // the below is not preicsely the Header type used in Request, but it'll pass duck-typing
+  | {
+      // eslint-disable-next-line @typescript-eslint/no-explicit-any
+      [key: string]: any;
+      append: (key: string, value: string) => void;
+    };
+
 /** Data returned from XHR request */
 export interface XHRData {
   xhr?: {
@@ -183,22 +193,26 @@ export function fetchCallback(
     const request = (handlerData.args[0] = handlerData.args[0] as string | Request);
     // eslint-disable-next-line @typescript-eslint/no-explicit-any
     const options = (handlerData.args[1] = (handlerData.args[1] as { [key: string]: any }) || {});
-    let headers = options.headers;
+    let headers: PolymorphicRequestHeaders = options.headers;
     if (isInstanceOf(request, Request)) {
       headers = (request as Request).headers;
     }
+    const traceHeaders = span.getTraceHeaders() as Record<string, string>;
     if (headers) {
-      // eslint-disable-next-line @typescript-eslint/no-unsafe-member-access
-      if (typeof headers.append === 'function') {
-        // eslint-disable-next-line @typescript-eslint/no-unsafe-member-access
-        headers.append('sentry-trace', span.toTraceparent());
+      if ('append' in headers && typeof headers.append === 'function') {
+        headers.append('sentry-trace', traceHeaders['sentry-trace']);
+        if (traceHeaders.tracestate) {
+          headers.append('tracestate', traceHeaders.tracestate);
+        }
       } else if (Array.isArray(headers)) {
-        headers = [...headers, ['sentry-trace', span.toTraceparent()]];
+        // TODO use the nicer version below once we stop supporting Node 6
+        // headers = [...headers, ...Object.entries(traceHeaders)];
+        headers = [...headers, ['sentry-trace', traceHeaders['sentry-trace']], ['tracestate', traceHeaders.tracestate]];
       } else {
-        headers = { ...headers, 'sentry-trace': span.toTraceparent() };
+        headers = { ...headers, ...traceHeaders };
       }
     } else {
-      headers = { 'sentry-trace': span.toTraceparent() };
+      headers = traceHeaders;
     }
     options.headers = headers;
   }
@@ -254,7 +268,11 @@ export function xhrCallback(
 
     if (handlerData.xhr.setRequestHeader) {
       try {
-        handlerData.xhr.setRequestHeader('sentry-trace', span.toTraceparent());
+        const sentryHeaders = span.getTraceHeaders();
+        handlerData.xhr.setRequestHeader('sentry-trace', sentryHeaders['sentry-trace']);
+        if (sentryHeaders.tracestate) {
+          handlerData.xhr.setRequestHeader('tracestate', sentryHeaders.tracestate);
+        }
       } catch (_) {
         // Error: InvalidStateError: Failed to execute 'setRequestHeader' on 'XMLHttpRequest': The object's state must be OPENED.
       }
diff --git a/packages/tracing/test/browser/request.test.ts b/packages/tracing/test/browser/request.test.ts
@@ -6,12 +6,41 @@ import { Span, SpanStatus, Transaction } from '../../src';
 import { fetchCallback, FetchData, instrumentOutgoingRequests, xhrCallback, XHRData } from '../../src/browser/request';
 import { addExtensionMethods } from '../../src/hubextensions';
 import * as tracingUtils from '../../src/utils';
+import { objectFromEntries } from '../testutils';
+
+// This is a normal base64 regex, modified to reflect that fact that we strip the trailing = or == off
+const stripped_base64 = '([a-zA-Z0-9+/]{4})*([a-zA-Z0-9+/]{2,3})?';
+
+const TRACESTATE_HEADER_REGEX = new RegExp(
+  `sentry=(${stripped_base64})` +  // our part of the header - should be the only part or at least the first part
+    `(,\\w+=\\w+)*`, // any number of copies of a comma followed by `name=value`
+);
 
 beforeAll(() => {
   addExtensionMethods();
-  // @ts-ignore need to override global Request because it's not in the jest environment (even with an
-  // `@jest-environment jsdom` directive, for some reason)
-  global.Request = {};
+
+  // Add Request to the global scope (necessary because for some reason Request isn't in the jest environment, even with
+  // an `@jest-environment jsdom` directive)
+
+  type MockHeaders = {
+    [key: string]: any;
+    append: (key: string, value: string) => void;
+  };
+
+  class Request {
+    public headers: MockHeaders;
+    constructor() {
+      // We need our headers to act like an object for key-lookup purposes, but also have an append method that adds
+      // items as its siblings. This hack precludes a key named `append`, of course, but for our purposes it's enough.
+      const headers = {} as MockHeaders;
+      headers.append = (key: string, value: any): void => {
+        headers[key] = value;
+      };
+      this.headers = headers;
+    }
+  }
+
+  (global as any).Request = Request;
 });
 
 const hasTracingEnabled = jest.spyOn(tracingUtils, 'hasTracingEnabled');
@@ -49,15 +78,15 @@ describe('instrumentOutgoingRequests', () => {
   });
 });
 
-describe('callbacks', () => {
+describe('fetch and xhr callbacks', () => {
   let hub: Hub;
   let transaction: Transaction;
   const alwaysCreateSpan = () => true;
   const neverCreateSpan = () => false;
   const fetchHandlerData: FetchData = {
     args: ['http://dogs.are.great/', {}],
     fetchData: { url: 'http://dogs.are.great/', method: 'GET' },
-    startTimestamp: 1356996072000,
+    startTimestamp: 2012112120121231,
   };
   const xhrHandlerData: XHRData = {
     xhr: {
@@ -72,18 +101,27 @@ describe('callbacks', () => {
       // setRequestHeader: XMLHttpRequest.prototype.setRequestHeader,
       setRequestHeader,
     },
-    startTimestamp: 1353501072000,
+    startTimestamp: 2012112120121231,
   };
-  const endTimestamp = 1356996072000;
+  const endTimestamp = 2013041520130908;
 
   beforeAll(() => {
-    hub = new Hub(new BrowserClient({ tracesSampleRate: 1 }));
+    hub = new Hub(
+      new BrowserClient({
+        dsn: 'https://dogsarebadatkeepingsecrets@squirrelchasers.ingest.sentry.io/12312012',
+        environment: 'dogpark',
+        release: 'off.leash.park',
+
+        tracesSampleRate: 1,
+      }),
+    );
     makeMain(hub);
   });
 
   beforeEach(() => {
-    transaction = hub.startTransaction({ name: 'organizations/users/:userid', op: 'pageload' }) as Transaction;
+    transaction = hub.startTransaction({ name: 'meetNewDogFriend', op: 'wag.tail' }) as Transaction;
     hub.configureScope(scope => scope.setSpan(transaction));
+    jest.clearAllMocks();
   });
 
   describe('fetchCallback()', () => {
@@ -111,20 +149,17 @@ describe('callbacks', () => {
       expect(spans).toEqual({});
     });
 
-    it('does not add fetch request headers if tracing is disabled', () => {
+    it('does not add tracing headers if tracing is disabled', () => {
       hasTracingEnabled.mockReturnValueOnce(false);
 
       // make a local copy so the global one doesn't get mutated
-      const handlerData: FetchData = {
-        args: ['http://dogs.are.great/', {}],
-        fetchData: { url: 'http://dogs.are.great/', method: 'GET' },
-        startTimestamp: 1353501072000,
-      };
+      const handlerData = { ...fetchHandlerData };
 
       fetchCallback(handlerData, alwaysCreateSpan, {});
 
       const headers = (handlerData.args[1].headers as Record<string, string>) || {};
       expect(headers['sentry-trace']).not.toBeDefined();
+      expect(headers['tracestate']).not.toBeDefined();
     });
 
     it('creates and finishes fetch span on active transaction', () => {
@@ -179,8 +214,67 @@ describe('callbacks', () => {
       expect(newSpan!.status).toBe(SpanStatus.fromHttpCode(404));
     });
 
-    it('adds sentry-trace header to fetch requests', () => {
-      // TODO
+    describe('adding tracing headers to fetch requests', () => {
+      it('can handle headers added with an `append` method', () => {
+        const handlerData: FetchData = { ...fetchHandlerData, args: [new Request('http://dogs.are.great'), {}] };
+
+        fetchCallback(handlerData, alwaysCreateSpan, {});
+
+        const headers = handlerData.args[1].headers;
+        expect(headers['sentry-trace']).toBeDefined();
+        expect(headers['tracestate']).toBeDefined();
+      });
+
+      it('can handle existing headers in array form', () => {
+        const handlerData = {
+          ...fetchHandlerData,
+          args: [
+            'http://dogs.are.great/',
+            {
+              headers: [
+                ['GREETING_PROTOCOL', 'mutual butt sniffing'],
+                ['TAIL_ACTION', 'wagging'],
+              ],
+            },
+          ],
+        };
+
+        fetchCallback(handlerData, alwaysCreateSpan, {});
+
+        const headers = objectFromEntries((handlerData.args[1] as any).headers);
+        expect(headers['sentry-trace']).toBeDefined();
+        expect(headers['tracestate']).toBeDefined();
+      });
+
+      it('can handle existing headers in object form', () => {
+        const handlerData = {
+          ...fetchHandlerData,
+          args: [
+            'http://dogs.are.great/',
+            {
+              headers: { GREETING_PROTOCOL: 'mutual butt sniffing', TAIL_ACTION: 'wagging' },
+            },
+          ],
+        };
+
+        fetchCallback(handlerData, alwaysCreateSpan, {});
+
+        const headers = (handlerData.args[1] as any).headers;
+        expect(headers['sentry-trace']).toBeDefined();
+        expect(headers['tracestate']).toBeDefined();
+      });
+
+      it('can handle there being no existing headers', () => {
+        // override the value of `args`, even though we're overriding it with the same data, as a means of deep copying
+        // the one part which gets mutated
+        const handlerData = { ...fetchHandlerData, args: ['http://dogs.are.great/', {}] };
+
+        fetchCallback(handlerData, alwaysCreateSpan, {});
+
+        const headers = (handlerData.args[1] as any).headers;
+        expect(headers['sentry-trace']).toBeDefined();
+        expect(headers['tracestate']).toBeDefined();
+      });
     });
   });
 
@@ -201,21 +295,22 @@ describe('callbacks', () => {
       expect(spans).toEqual({});
     });
 
-    it('does not add xhr request headers if tracing is disabled', () => {
+    it('does not add tracing headers if tracing is disabled', () => {
       hasTracingEnabled.mockReturnValueOnce(false);
 
       xhrCallback(xhrHandlerData, alwaysCreateSpan, {});
 
       expect(setRequestHeader).not.toHaveBeenCalled();
     });
 
-    it('adds sentry-trace header to XHR requests', () => {
+    it('adds tracing headers to XHR requests', () => {
       xhrCallback(xhrHandlerData, alwaysCreateSpan, {});
 
       expect(setRequestHeader).toHaveBeenCalledWith(
         'sentry-trace',
         expect.stringMatching(tracingUtils.SENTRY_TRACE_REGEX),
       );
+      expect(setRequestHeader).toHaveBeenCalledWith('tracestate', expect.stringMatching(TRACESTATE_HEADER_REGEX));
     });
 
     it('creates and finishes XHR span on active transaction', () => {
