feat: wrap fetch and add tracing headers (#1186)

Co-authored-by: Tom Owers <owerstom@gmail.com>

rollup.config.js

@@ -79,6 +79,18 @@ export default [
         ],
         plugins: [...plugins],
     },
+    {
+        input: 'src/loader-tracing-headers.ts',
+        output: [
+            {
+                file: 'dist/tracing-headers.js',
+                sourcemap: true,
+                format: 'iife',
+                name: 'posthog',
+            },
+        ],
+        plugins: [...plugins],
+    },
     {
         input: 'src/loader-globals.ts',
         output: [

src/extensions/replay/rrweb-plugins/patch.ts

@@ -0,0 +1,46 @@
+// import { patch } from 'rrweb/typings/utils'
+// copied from https://github.com/rrweb-io/rrweb/blob/8aea5b00a4dfe5a6f59bd2ae72bb624f45e51e81/packages/rrweb/src/utils.ts#L129
+// which was copied from https://github.com/getsentry/sentry-javascript/blob/b2109071975af8bf0316d3b5b38f519bdaf5dc15/packages/utils/src/object.ts
+import { isFunction } from '../../../utils/type-utils'
+
+export function patch(
+    source: { [key: string]: any },
+    name: string,
+    replacement: (...args: unknown[]) => unknown
+): () => void {
+    try {
+        if (!(name in source)) {
+            return () => {
+                //
+            }
+        }
+
+        const original = source[name] as () => unknown
+        const wrapped = replacement(original)
+
+        // Make sure it's a function first, as we need to attach an empty prototype for `defineProperties` to work
+        // otherwise it'll throw "TypeError: Object.defineProperties called on non-object"
+        if (isFunction(wrapped)) {
+            // eslint-disable-next-line @typescript-eslint/no-unsafe-assignment
+            wrapped.prototype = wrapped.prototype || {}
+            Object.defineProperties(wrapped, {
+                __posthog_wrapped__: {
+                    enumerable: false,
+                    value: true,
+                },
+            })
+        }
+
+        source[name] = wrapped
+
+        return () => {
+            source[name] = original
+        }
+    } catch {
+        return () => {
+            //
+        }
+        // This can throw if multiple fill happens on a global object like XMLHttpRequest
+        // Fixes https://github.com/getsentry/sentry-javascript/issues/2043
+    }
+}

src/extensions/tracing-headers.ts

@@ -0,0 +1,53 @@
+import { PostHog } from '../posthog-core'
+import { assignableWindow } from '../utils/globals'
+import { logger } from '../utils/logger'
+import { loadScript } from '../utils'
+import Config from '../config'
+import { isUndefined } from '../utils/type-utils'
+
+const LOGGER_PREFIX = '[TRACING-HEADERS]'
+
+export class TracingHeaders {
+    private _restoreXHRPatch: (() => void) | undefined = undefined
+    private _restoreFetchPatch: (() => void) | undefined = undefined
+
+    constructor(private readonly instance: PostHog) {}
+
+    private _loadScript(cb: () => void): void {
+        if (assignableWindow.postHogTracingHeadersPatchFns) {
+            // already loaded
+            cb()
+        }
+
+        loadScript(
+            this.instance.requestRouter.endpointFor('assets', `/static/tracing-headers.js?v=${Config.LIB_VERSION}`),
+            (err) => {
+                if (err) {
+                    logger.error(LOGGER_PREFIX + ' failed to load script', err)
+                }
+                cb()
+            }
+        )
+    }
+    public startIfEnabledOrStop() {
+        if (this.instance.config.__add_tracing_headers) {
+            this._loadScript(this._startCapturing)
+        } else {
+            this._restoreXHRPatch?.()
+            this._restoreFetchPatch?.()
+            // we don't want to call these twice so we reset them
+            this._restoreXHRPatch = undefined
+            this._restoreFetchPatch = undefined
+        }
+    }
+
+    private _startCapturing = () => {
+        // NB: we can assert sessionManager is present only because we've checked previously
+        if (isUndefined(this._restoreXHRPatch)) {
+            assignableWindow.postHogTracingHeadersPatchFns._patchXHR(this.instance.sessionManager!)
+        }
+        if (isUndefined(this._restoreFetchPatch)) {
+            assignableWindow.postHogTracingHeadersPatchFns._patchFetch(this.instance.sessionManager!)
+        }
+    }
+}

src/loader-recorder.ts

@@ -23,7 +23,6 @@ import {
     isBoolean,
     isDocument,
     isFormData,
-    isFunction,
     isNull,
     isNullish,
     isObject,
@@ -34,6 +33,7 @@ import { logger } from './utils/logger'
 import { window } from './utils/globals'
 import { defaultNetworkOptions } from './extensions/replay/config'
 import { formDataToQuery } from './utils/request-utils'
+import { patch } from './extensions/replay/rrweb-plugins/patch'
 
 export type NetworkData = {
     requests: CapturedNetworkRequest[]
@@ -50,55 +50,6 @@ type ObservedPerformanceEntry = (PerformanceNavigationTiming | PerformanceResour
     responseStatus?: number
 }
 
-// import { patch } from 'rrweb/typings/utils'
-// copied from https://github.com/rrweb-io/rrweb/blob/8aea5b00a4dfe5a6f59bd2ae72bb624f45e51e81/packages/rrweb/src/utils.ts#L129
-// which was copied from https://github.com/getsentry/sentry-javascript/blob/b2109071975af8bf0316d3b5b38f519bdaf5dc15/packages/utils/src/object.ts
-export function patch(
-    source: { [key: string]: any },
-    name: string,
-    replacement: (...args: unknown[]) => unknown
-): () => void {
-    try {
-        if (!(name in source)) {
-            return () => {
-                //
-            }
-        }
-
-        const original = source[name] as () => unknown
-        const wrapped = replacement(original)
-
-        // Make sure it's a function first, as we need to attach an empty prototype for `defineProperties` to work
-        // otherwise it'll throw "TypeError: Object.defineProperties called on non-object"
-        if (isFunction(wrapped)) {
-            // eslint-disable-next-line @typescript-eslint/no-unsafe-assignment
-            wrapped.prototype = wrapped.prototype || {}
-            Object.defineProperties(wrapped, {
-                __rrweb_original__: {
-                    enumerable: false,
-                    value: original,
-                },
-                __posthog_wrapped__: {
-                    enumerable: false,
-                    value: true,
-                },
-            })
-        }
-
-        source[name] = wrapped
-
-        return () => {
-            source[name] = original
-        }
-    } catch {
-        return () => {
-            //
-        }
-        // This can throw if multiple fill happens on a global object like XMLHttpRequest
-        // Fixes https://github.com/getsentry/sentry-javascript/issues/2043
-    }
-}
-
 export function findLast<T>(array: Array<T>, predicate: (value: T) => boolean): T | undefined {
     const length = array.length
     for (let i = length - 1; i >= 0; i -= 1) {

src/loader-tracing-headers.ts

@@ -0,0 +1,64 @@
+import { SessionIdManager } from './sessionid'
+import { patch } from './extensions/replay/rrweb-plugins/patch'
+import { assignableWindow, window } from './utils/globals'
+
+const addTracingHeaders = (sessionManager: SessionIdManager, req: Request) => {
+    const { sessionId, windowId } = sessionManager.checkAndGetSessionAndWindowId(true)
+    req.headers.set('X-POSTHOG-SESSION-ID', sessionId)
+    req.headers.set('X-POSTHOG-WINDOW-ID', windowId)
+}
+
+const patchFetch = (sessionManager: SessionIdManager): (() => void) => {
+    // eslint-disable-next-line @typescript-eslint/ban-ts-comment
+    // @ts-ignore
+    return patch(window, 'fetch', (originalFetch: typeof fetch) => {
+        return async function (url: URL | RequestInfo, init?: RequestInit | undefined) {
+            // check IE earlier than this, we only initialize if Request is present
+            // eslint-disable-next-line compat/compat
+            const req = new Request(url, init)
+
+            addTracingHeaders(sessionManager, req)
+
+            return originalFetch(req)
+        }
+    })
+}
+
+const patchXHR = (sessionManager: SessionIdManager): (() => void) => {
+    return patch(
+        // we can assert this is present because we've checked previously
+        window!.XMLHttpRequest.prototype,
+        'open',
+        // eslint-disable-next-line @typescript-eslint/ban-ts-comment
+        // @ts-ignore
+        (originalOpen: typeof XMLHttpRequest.prototype.open) => {
+            return function (
+                method: string,
+                url: string | URL,
+                async = true,
+                username?: string | null,
+                password?: string | null
+            ) {
+                // because this function is returned in its actual context `this` _is_ an XMLHttpRequest
+                // eslint-disable-next-line @typescript-eslint/ban-ts-comment
+                // @ts-ignore
+                const xhr = this as XMLHttpRequest
+
+                // check IE earlier than this, we only initialize if Request is present
+                // eslint-disable-next-line compat/compat
+                const req = new Request(url)
+
+                addTracingHeaders(sessionManager, req)
+
+                return originalOpen.call(xhr, method, req.url, async, username, password)
+            }
+        }
+    )
+}
+
+if (assignableWindow) {
+    assignableWindow.postHogTracingHeadersPatchFns = {
+        _patchFetch: patchFetch,
+        _patchXHR: patchXHR,
+    }
+}

src/posthog-core.ts

@@ -72,6 +72,7 @@ import { Heatmaps } from './heatmaps'
 import { ScrollManager } from './scroll-manager'
 import { SimpleEventEmitter } from './utils/simple-event-emitter'
 import { Autocapture } from './autocapture'
+import { TracingHeaders } from './extensions/tracing-headers'
 import { ConsentManager } from './consent'
 
 /*
@@ -173,6 +174,7 @@ export const defaultConfig = (): PostHogConfig => ({
     disable_compression: false,
     session_idle_timeout_seconds: 30 * 60, // 30 minutes
     person_profiles: 'always',
+    __add_tracing_headers: false,
 })
 
 export const configRenames = (origConfig: Partial<PostHogConfig>): Partial<PostHogConfig> => {
@@ -271,6 +273,7 @@ export class PostHog {
 
     constructor() {
         this.config = defaultConfig()
+
         this.decideEndpointWasHit = false
         this.SentryIntegration = SentryIntegration
         this.__request_queue = []
@@ -392,6 +395,8 @@ export class PostHog {
         this.sessionManager = new SessionIdManager(this.config, this.persistence)
         this.sessionPropsManager = new SessionPropsManager(this.sessionManager, this.persistence)
 
+        new TracingHeaders(this).startIfEnabledOrStop()
+
         this.sessionRecording = new SessionRecording(this)
         this.sessionRecording.startIfEnabledOrStop()
 

src/types.ts

@@ -169,6 +169,12 @@ export interface PostHogConfig {
         /** How many events can be captured in a burst. This defaults to 10 times the events_per_second count  */
         events_burst_limit?: number
     }
+
+    /**
+     * PREVIEW - MAY CHANGE WITHOUT WARNING - DO NOT USE IN PRODUCTION
+     * whether to wrap fetch and add tracing headers to the request
+     * */
+    __add_tracing_headers?: boolean
 }
 
 export interface OptInOutCapturingOptions {