Catch gql header in XHR proxy

DataDog · Sep 21, 2023 · b45680a · b45680a
1 parent 76c715d
commit b45680a
Show file tree

Hide file tree

Showing 5 changed files with 100 additions and 8 deletions.
diff --git a/packages/core/src/index.tsx b/packages/core/src/index.tsx
@@ -19,6 +19,11 @@ import { SdkVerbosity } from './SdkVerbosity';
 import { TrackingConsent } from './TrackingConsent';
 import { DdLogs } from './logs/DdLogs';
 import { DdRum } from './rum/DdRum';
+import {
+    DATADOG_GRAPH_QL_OPERATION_TYPE_HEADER,
+    DATADOG_GRAPH_QL_OPERATION_NAME_HEADER,
+    DATADOG_GRAPH_QL_VARIABLES_HEADER
+} from './rum/instrumentation/resourceTracking/graphql/graphqlHeaders';
 import { RumActionType, ErrorSource, PropagatorType } from './rum/types';
 import { DatadogProvider } from './sdk/DatadogProvider/DatadogProvider';
 import { DdTrace } from './trace/DdTrace';
@@ -42,5 +47,8 @@ export {
     VitalsUpdateFrequency,
     PropagatorType,
     UploadFrequency,
-    BatchSize
+    BatchSize,
+    DATADOG_GRAPH_QL_OPERATION_TYPE_HEADER,
+    DATADOG_GRAPH_QL_OPERATION_NAME_HEADER,
+    DATADOG_GRAPH_QL_VARIABLES_HEADER
 };
diff --git a/...es/core/src/rum/instrumentation/resourceTracking/graphql/__tests__/graphqlHeaders.test.ts b/...es/core/src/rum/instrumentation/resourceTracking/graphql/__tests__/graphqlHeaders.test.ts
@@ -0,0 +1,22 @@
+import {
+    DATADOG_GRAPH_QL_OPERATION_NAME_HEADER,
+    DATADOG_GRAPH_QL_OPERATION_TYPE_HEADER,
+    DATADOG_GRAPH_QL_VARIABLES_HEADER,
+    isDatadogCustomHeader
+} from '../graphqlHeaders';
+
+describe('GraphQL custom headers', () => {
+    it.each([
+        DATADOG_GRAPH_QL_OPERATION_NAME_HEADER,
+        DATADOG_GRAPH_QL_OPERATION_TYPE_HEADER,
+        DATADOG_GRAPH_QL_VARIABLES_HEADER
+    ])('%s matches the custom header pattern', header => {
+        expect(isDatadogCustomHeader(header)).toBeTruthy();
+    });
+
+    describe('isDatadogCustomHeader', () => {
+        it('returns false for non-custom headers', () => {
+            expect(isDatadogCustomHeader('non-custom-header')).toBeFalsy();
+        });
+    });
+});
diff --git a/packages/core/src/rum/instrumentation/resourceTracking/graphql/graphqlHeaders.ts b/packages/core/src/rum/instrumentation/resourceTracking/graphql/graphqlHeaders.ts
@@ -0,0 +1,9 @@
+const DATADOG_CUSTOM_HEADER_PREFIX = '_dd-custom-header';
+
+export const DATADOG_GRAPH_QL_OPERATION_NAME_HEADER = `${DATADOG_CUSTOM_HEADER_PREFIX}-graph-ql-operation-name`;
+export const DATADOG_GRAPH_QL_VARIABLES_HEADER = `${DATADOG_CUSTOM_HEADER_PREFIX}-graph-ql-variables`;
+export const DATADOG_GRAPH_QL_OPERATION_TYPE_HEADER = `${DATADOG_CUSTOM_HEADER_PREFIX}-graph-ql-operation-type`;
+
+export const isDatadogCustomHeader = (header: string) => {
+    return header.match(new RegExp(`^${DATADOG_CUSTOM_HEADER_PREFIX}`));
+};
diff --git a/packages/core/src/rum/instrumentation/resourceTracking/requestProxy/XHRProxy/XHRProxy.ts b/packages/core/src/rum/instrumentation/resourceTracking/requestProxy/XHRProxy/XHRProxy.ts
@@ -8,6 +8,12 @@ import Timer from '../../../../../Timer';
 import { getTracingHeaders } from '../../distributedTracing/distributedTracingHeaders';
 import type { DdRumResourceTracingAttributes } from '../../distributedTracing/distributedTracing';
 import { getTracingAttributes } from '../../distributedTracing/distributedTracing';
+import {
+    DATADOG_GRAPH_QL_OPERATION_NAME_HEADER,
+    DATADOG_GRAPH_QL_OPERATION_TYPE_HEADER,
+    DATADOG_GRAPH_QL_VARIABLES_HEADER,
+    isDatadogCustomHeader
+} from '../../graphql/graphqlHeaders';
 import type { RequestProxyOptions } from '../interfaces/RequestProxy';
 import { RequestProxy } from '../interfaces/RequestProxy';
 
@@ -41,6 +47,7 @@ export class XHRProxy extends RequestProxy {
     private providers: XHRProxyProviders;
     private static originalXhrOpen: typeof XMLHttpRequest.prototype.open;
     private static originalXhrSend: typeof XMLHttpRequest.prototype.send;
+    private static originalXhrSetRequestHeader: typeof XMLHttpRequest.prototype.setRequestHeader;
 
     constructor(providers: XHRProxyProviders) {
         super();
@@ -50,12 +57,15 @@ export class XHRProxy extends RequestProxy {
     onTrackingStart = (context: RequestProxyOptions) => {
         XHRProxy.originalXhrOpen = this.providers.xhrType.prototype.open;
         XHRProxy.originalXhrSend = this.providers.xhrType.prototype.send;
+        XHRProxy.originalXhrSetRequestHeader = this.providers.xhrType.prototype.setRequestHeader;
         proxyRequests(this.providers, context);
     };
 
     onTrackingStop = () => {
         this.providers.xhrType.prototype.open = XHRProxy.originalXhrOpen;
         this.providers.xhrType.prototype.send = XHRProxy.originalXhrSend;
+        this.providers.xhrType.prototype.setRequestHeader =
+            XHRProxy.originalXhrSetRequestHeader;
     };
 }
 
@@ -65,6 +75,7 @@ const proxyRequests = (
 ): void => {
     proxyOpen(providers, context);
     proxySend(providers);
+    proxySetRequestHeader(providers);
 };
 
 const proxyOpen = (
@@ -181,3 +192,32 @@ const reportXhr = async (
         resourceContext: xhrProxy
     });
 };
+
+const proxySetRequestHeader = (providers: XHRProxyProviders): void => {
+    const xhrType = providers.xhrType;
+    const originalXhrSetRequestHeader = xhrType.prototype.setRequestHeader;
+
+    xhrType.prototype.setRequestHeader = function (
+        this: DdRumXhr,
+        header: string,
+        value: string
+    ) {
+        if (isDatadogCustomHeader(header)) {
+            if (header === DATADOG_GRAPH_QL_OPERATION_NAME_HEADER) {
+                // TODO: add information to request
+                return;
+            }
+            if (header === DATADOG_GRAPH_QL_OPERATION_TYPE_HEADER) {
+                // TODO: add information to request
+                return;
+            }
+            if (header === DATADOG_GRAPH_QL_VARIABLES_HEADER) {
+                // TODO: add information to request
+                return;
+            }
+        }
+
+        // eslint-disable-next-line prefer-rest-params
+        return originalXhrSetRequestHeader.apply(this, arguments as any);
+    };
+};
diff --git a/packages/react-native-apollo-client/src/DatadogLink.ts b/packages/react-native-apollo-client/src/DatadogLink.ts
@@ -5,6 +5,11 @@
  */
 
 import { ApolloLink } from '@apollo/client';
+import {
+    DATADOG_GRAPH_QL_OPERATION_TYPE_HEADER,
+    DATADOG_GRAPH_QL_OPERATION_NAME_HEADER,
+    DATADOG_GRAPH_QL_VARIABLES_HEADER
+} from '@datadog/mobile-react-native';
 
 import { getOperationName, getVariables, getOperationType } from './helpers';
 
@@ -16,14 +21,22 @@ export class DatadogLink extends ApolloLink {
             const operationType = getOperationType(operation);
 
             operation.setContext(({ headers = {} }) => {
+                const newHeaders: Record<string, string | null> = {
+                    ...headers
+                };
+
+                newHeaders[
+                    DATADOG_GRAPH_QL_OPERATION_TYPE_HEADER
+                ] = operationType;
+                newHeaders[
+                    DATADOG_GRAPH_QL_OPERATION_NAME_HEADER
+                ] = operationName;
+                newHeaders[
+                    DATADOG_GRAPH_QL_VARIABLES_HEADER
+                ] = formattedVariables;
+
                 return {
-                    headers: {
-                        ...headers,
-                        // TODO: import headers from core
-                        '_dd-graph-ql-operation-name': operationName,
-                        '_dd-graph-ql-variables': formattedVariables,
-                        '_dd-graph-ql-operation-type': operationType
-                    }
+                    headers: newHeaders
                 };
             });