Fork W3CTraceContextPropagator to workaround Edge bug with RegExp (#98)

* Fork W3CTraceContextPropagator to workaround Edge bug with RegExp

* changeset

.changeset/shy-monkeys-hug.md

@@ -0,0 +1,5 @@
+---
+"@vercel/otel": patch
+---
+
+Edge-comaptible W3CTraceContextPropagator

packages/otel/src/propagators/w3c-tracecontext-propagator.test.ts

@@ -0,0 +1,201 @@
+import { beforeEach, describe, expect, it } from "vitest";
+import type { Context, TextMapGetter, TextMapSetter } from "@opentelemetry/api";
+import {
+  context as contextApi,
+  createTraceState,
+  trace as traceApi,
+} from "@opentelemetry/api";
+import { W3CTraceContextPropagator } from "./w3c-tracecontext-propagator";
+
+describe("W3CTraceContextPropagator", () => {
+  let baseContext: Context;
+  let propagator: W3CTraceContextPropagator;
+
+  beforeEach(() => {
+    propagator = new W3CTraceContextPropagator();
+    baseContext = contextApi.active();
+  });
+
+  it("should return fields", () => {
+    expect(propagator.fields()).toEqual(["traceparent", "tracestate"]);
+  });
+
+  describe("exract", () => {
+    it("should extract the trace parent", () => {
+      const context = propagator.extract(
+        baseContext,
+        new Map([
+          [
+            "traceparent",
+            "00-6d2eac29c9283ece795b4fbaa2d57225-bad4e819c34d2cdb-01",
+          ],
+        ]),
+        GETTER
+      );
+      expect(traceApi.getSpanContext(context)).toEqual({
+        isRemote: true,
+        traceId: "6d2eac29c9283ece795b4fbaa2d57225",
+        spanId: "bad4e819c34d2cdb",
+        traceFlags: 1,
+      });
+    });
+
+    it("should extract the trace parent as an array", () => {
+      const context = propagator.extract(
+        baseContext,
+        new Map([
+          [
+            "traceparent",
+            ["00-6d2eac29c9283ece795b4fbaa2d57225-bad4e819c34d2cdb-01"],
+          ],
+        ]),
+        GETTER
+      );
+      expect(traceApi.getSpanContext(context)).toEqual({
+        isRemote: true,
+        traceId: "6d2eac29c9283ece795b4fbaa2d57225",
+        spanId: "bad4e819c34d2cdb",
+        traceFlags: 1,
+      });
+    });
+
+    it("should extract the trace state", () => {
+      const context = propagator.extract(
+        baseContext,
+        new Map([
+          [
+            "traceparent",
+            "00-6d2eac29c9283ece795b4fbaa2d57225-bad4e819c34d2cdb-01",
+          ],
+          ["tracestate", "foo=11,bar=12"],
+        ]),
+        GETTER
+      );
+      expect(traceApi.getSpanContext(context)).toContain({
+        isRemote: true,
+        traceId: "6d2eac29c9283ece795b4fbaa2d57225",
+        spanId: "bad4e819c34d2cdb",
+        traceFlags: 1,
+      });
+      expect(traceApi.getSpanContext(context)?.traceState?.get("foo")).toBe(
+        "11"
+      );
+      expect(traceApi.getSpanContext(context)?.traceState?.get("bar")).toBe(
+        "12"
+      );
+    });
+
+    describe("failure cases", () => {
+      it("should NOT extract invalid trace parent", () => {
+        const context = propagator.extract(
+          baseContext,
+          new Map([["traceparent", "00-bad-bad4e819c34d2cdb-01"]]),
+          GETTER
+        );
+        expect(traceApi.getSpanContext(context)).toBeUndefined();
+      });
+
+      it("should NOT extract invalid span", () => {
+        const context = propagator.extract(
+          baseContext,
+          new Map([
+            ["traceparent", "00-6d2eac29c9283ece795b4fbaa2d57225-bad-01"],
+          ]),
+          GETTER
+        );
+        expect(traceApi.getSpanContext(context)).toBeUndefined();
+      });
+
+      it("should NOT extract invalid version", () => {
+        const context = propagator.extract(
+          baseContext,
+          new Map([
+            [
+              "traceparent",
+              "bad-6d2eac29c9283ece795b4fbaa2d57225-bad4e819c34d2cdb-01",
+            ],
+          ]),
+          GETTER
+        );
+        expect(traceApi.getSpanContext(context)).toBeUndefined();
+      });
+
+      it("should NOT extract invalid version", () => {
+        const context = propagator.extract(
+          baseContext,
+          new Map([
+            [
+              "traceparent",
+              "00-6d2eac29c9283ece795b4fbaa2d57225-bad4e819c34d2cdb-01-unknown",
+            ],
+          ]),
+          GETTER
+        );
+        expect(traceApi.getSpanContext(context)).toBeUndefined();
+      });
+    });
+  });
+
+  describe("inject", () => {
+    it("should inject the trace parent", () => {
+      const carrier = new Map<string, string | string[]>();
+      const context = traceApi.setSpanContext(baseContext, {
+        traceId: "6d2eac29c9283ece795b4fbaa2d57225",
+        spanId: "bad4e819c34d2cdb",
+        traceFlags: 1,
+      });
+      propagator.inject(context, carrier, SETTER);
+      expect(carrier.get("traceparent")).toBe(
+        "00-6d2eac29c9283ece795b4fbaa2d57225-bad4e819c34d2cdb-01"
+      );
+      expect(carrier.get("tracestate")).toBeUndefined();
+    });
+
+    it("should inject the trace state", () => {
+      const carrier = new Map<string, string | string[]>();
+      const context = traceApi.setSpanContext(baseContext, {
+        traceId: "6d2eac29c9283ece795b4fbaa2d57225",
+        spanId: "bad4e819c34d2cdb",
+        traceFlags: 1,
+        traceState: createTraceState("foo=11,bar=12"),
+      });
+      propagator.inject(context, carrier, SETTER);
+      expect(carrier.get("traceparent")).toBe(
+        "00-6d2eac29c9283ece795b4fbaa2d57225-bad4e819c34d2cdb-01"
+      );
+      expect(carrier.get("tracestate")).toEqual("foo=11,bar=12");
+    });
+
+    describe("failure cases", () => {
+      it("should NOT inject empty span", () => {
+        const carrier = new Map<string, string | string[]>();
+        propagator.inject(baseContext, carrier, SETTER);
+        expect(carrier.get("traceparent")).toBeUndefined();
+        expect(carrier.get("tracestate")).toBeUndefined();
+      });
+
+      it("should inject an invalid span parent", () => {
+        const carrier = new Map<string, string | string[]>();
+        const context = traceApi.setSpanContext(baseContext, {
+          traceId: "bad",
+          spanId: "worse",
+          traceFlags: -1,
+        });
+        propagator.inject(context, carrier, SETTER);
+        expect(carrier.get("traceparent")).toBeUndefined();
+        expect(carrier.get("tracestate")).toBeUndefined();
+      });
+    });
+  });
+});
+
+type Carrier = Map<string, string | string[]>;
+
+const GETTER: TextMapGetter<Carrier> = {
+  get: (carrier, key) => carrier.get(key),
+  keys: (carrier) => Array.from(carrier.keys()),
+};
+
+const SETTER: TextMapSetter<Carrier> = {
+  set: (carrier, key, value) => carrier.set(key, value),
+};

packages/otel/src/propagators/w3c-tracecontext-propagator.ts

@@ -0,0 +1,105 @@
+import {
+  createTraceState,
+  isSpanContextValid,
+  trace as traceApi,
+} from "@opentelemetry/api";
+import type {
+  Context,
+  SpanContext,
+  TextMapGetter,
+  TextMapPropagator,
+  TextMapSetter,
+} from "@opentelemetry/api";
+import { isTracingSuppressed } from "@opentelemetry/core";
+
+const VERSION = "00";
+
+const TRACE_PARENT_HEADER = "traceparent";
+const TRACE_STATE_HEADER = "tracestate";
+
+/**
+ * Same as the `W3CTraceContextPropagator` from `@opentelemetry/core`, but with
+ * a workaround for RegExp issue in Edge.
+ */
+export class W3CTraceContextPropagator implements TextMapPropagator {
+  fields(): string[] {
+    return [TRACE_PARENT_HEADER, TRACE_STATE_HEADER];
+  }
+
+  // eslint-disable-next-line @typescript-eslint/no-explicit-any
+  inject(context: Context, carrier: any, setter: TextMapSetter): void {
+    const spanContext = traceApi.getSpanContext(context);
+    if (
+      !spanContext ||
+      isTracingSuppressed(context) ||
+      !isSpanContextValid(spanContext)
+    )
+      return;
+
+    const traceParent = `${VERSION}-${spanContext.traceId}-${
+      spanContext.spanId
+    }-0${Number(spanContext.traceFlags || 0).toString(16)}`;
+
+    setter.set(carrier, TRACE_PARENT_HEADER, traceParent);
+    if (spanContext.traceState) {
+      setter.set(
+        carrier,
+        TRACE_STATE_HEADER,
+        spanContext.traceState.serialize()
+      );
+    }
+  }
+
+  // eslint-disable-next-line @typescript-eslint/no-explicit-any
+  extract(context: Context, carrier: any, getter: TextMapGetter): Context {
+    const traceParentHeader = getter.get(carrier, TRACE_PARENT_HEADER);
+    if (!traceParentHeader) return context;
+    const traceParent = Array.isArray(traceParentHeader)
+      ? traceParentHeader[0]
+      : traceParentHeader;
+    if (typeof traceParent !== "string") return context;
+    const spanContext = parseTraceParent(traceParent);
+    if (!spanContext) return context;
+
+    spanContext.isRemote = true;
+
+    const traceStateHeader = getter.get(carrier, TRACE_STATE_HEADER);
+    if (traceStateHeader) {
+      // If more than one `tracestate` header is found, we merge them into a
+      // single header.
+      const state = Array.isArray(traceStateHeader)
+        ? traceStateHeader.join(",")
+        : traceStateHeader;
+      spanContext.traceState = createTraceState(
+        typeof state === "string" ? state : undefined
+      );
+    }
+    return traceApi.setSpanContext(context, spanContext);
+  }
+}
+
+function parseTraceParent(traceParent: string): SpanContext | null {
+  const [version, traceId, spanId, traceFlags, other] = traceParent.split("-");
+  if (
+    !version ||
+    !traceId ||
+    !spanId ||
+    !traceFlags ||
+    version.length !== 2 ||
+    traceId.length !== 32 ||
+    spanId.length !== 16 ||
+    traceFlags.length !== 2
+  )
+    return null;
+
+  // According to the specification the implementation should be compatible
+  // with future versions. If there are more parts, we only reject it if it's using version 00
+  // See https://www.w3.org/TR/trace-context/#versioning-of-traceparent
+  if (version === "00" && other) return null;
+
+  return {
+    traceId,
+    spanId,
+    traceFlags: parseInt(traceFlags, 16),
+  };
+}

packages/otel/src/sdk.ts

@@ -35,7 +35,6 @@ import { AsyncLocalStorageContextManager } from "@opentelemetry/context-async-ho
 import {
   CompositePropagator,
   W3CBaggagePropagator,
-  W3CTraceContextPropagator,
   baggageUtils,
 } from "@opentelemetry/core";
 import { CompositeSpanProcessor } from "./processor/composite-span-processor";
@@ -51,6 +50,7 @@ import type {
   SpanProcessorOrName,
 } from "./types";
 import { FetchInstrumentation } from "./instrumentations/fetch";
+import { W3CTraceContextPropagator } from "./propagators/w3c-tracecontext-propagator";
 
 type Env = ReturnType<typeof parseEnvironment>;