Remove diagnostics_channel usage

library/sinks/Undici.ts

@@ -1,7 +1,7 @@
 import { lookup } from "dns";
 import { Agent } from "../agent/Agent";
 import { getInstance } from "../agent/AgentSingleton";
-import { getContext } from "../agent/Context";
+import { bindContext, getContext } from "../agent/Context";
 import { Hooks } from "../agent/hooks/Hooks";
 import { InterceptorResult } from "../agent/hooks/InterceptorResult";
 import { Wrapper } from "../agent/Wrapper";
@@ -12,9 +12,8 @@ import { inspectDNSLookupCalls } from "../vulnerabilities/ssrf/inspectDNSLookupC
 import { wrapDispatch } from "./undici/wrapDispatch";
 import { wrapExport } from "../agent/hooks/wrapExport";
 import { wrapNewInstance } from "../agent/hooks/wrapNewInstance";
-import { onHeaders } from "./undici/onHeaders";
 import { getHostnameAndPortFromArgs } from "./undici/getHostnameAndPortFromArgs";
-import { subscribe } from "diagnostics_channel";
+import { wrapOnHeaders } from "./undici/wrapOnHeaders";
 
 const methods = [
   "request",
@@ -103,7 +102,13 @@ export class Undici implements Wrapper {
     }
 
     // @ts-expect-error No types for this
-    instance.dispatch = wrapDispatch(instance.dispatch, agent, false, context);
+    instance.dispatch = bindContext(
+      // @ts-expect-error No types for this
+      wrapDispatch(instance.dispatch, agent, false, context)
+    );
+
+    // @ts-expect-error No types for this
+    instance.onHeaders = wrapOnHeaders(instance.onHeaders, context);
 
     return instance;
   }
@@ -125,11 +130,6 @@ export class Undici implements Wrapper {
           return;
         }
 
-        // Subscribe to the undici:request:headers diagnostic channel to check for redirects
-        if (isVersionGreaterOrEqual("16.17.0", getSemverNodeVersion())) {
-          subscribe("undici:request:headers", onHeaders);
-        }
-
         // Immediately patch the global dispatcher before returning the module
         // The global dispatcher might be overwritten by the user
         // But at least they have a reference to our dispatcher instead of the original one

library/sinks/undici/wrapDispatch.ts

@@ -1,6 +1,9 @@
 /* eslint-disable max-lines-per-function */
 import type { Dispatcher } from "undici-v6";
-import { runWithUndiciRequestContext } from "./RequestContextStorage";
+import {
+  getUndiciRequestContext,
+  runWithUndiciRequestContext,
+} from "./RequestContextStorage";
 import { getMetadataForSSRFAttack } from "../../vulnerabilities/ssrf/getMetadataForSSRFAttack";
 import { Context, getContext } from "../../agent/Context";
 import { getPortFromURL } from "../../helpers/getPortFromURL";
@@ -9,6 +12,7 @@ import { attackKindHumanName } from "../../agent/Attack";
 import { escapeHTML } from "../../helpers/escapeHTML";
 import { isRedirectToPrivateIP } from "../../vulnerabilities/ssrf/isRedirectToPrivateIP";
 import { getUrlFromOptions } from "./getUrlFromOptions";
+import { wrapOnHeaders } from "./wrapOnHeaders";
 
 type Dispatch = Dispatcher["dispatch"];
 
@@ -60,6 +64,8 @@ export function wrapDispatch(
 
     blockRedirectToPrivateIP(url, context, agent, isFetch);
 
+    handler.onHeaders = wrapOnHeaders(handler.onHeaders, context, url);
+
     // We also pass the incoming context as part of the outgoing request context to prevent context mismatch, if the request is a redirect (argContext is set)
     return runWithUndiciRequestContext(
       {

library/sinks/undici/wrapOnHeaders.ts

@@ -0,0 +1,80 @@
+import { parseHeaders } from "./parseHeaders";
+import { isRedirectStatusCode } from "../../helpers/isRedirectStatusCode";
+import { addRedirectToContext } from "../../vulnerabilities/ssrf/addRedirectToContext";
+import { Context, getContext } from "../../agent/Context";
+import type { Dispatcher } from "undici-v6";
+import { getUrlFromOptions } from "./getUrlFromOptions";
+
+type OnHeaders = Dispatcher.DispatchHandlers["onHeaders"];
+
+/**
+ * Check if the response is a redirect. If yes, determine the destination URL.
+ */
+export function wrapOnHeaders(
+  orig: OnHeaders,
+  context: Context,
+  requestUrl?: URL
+): OnHeaders {
+  // @ts-expect-error We return undefined if there is no original function, thats fine because the onHeaders function is optional
+  return function onHeaders() {
+    // eslint-disable-next-line prefer-rest-params
+    const args = Array.from(arguments);
+
+    try {
+      let sourceURL = requestUrl;
+      if (!sourceURL) {
+        // @ts-expect-error No types for this
+        sourceURL = getUrlFromOptions(this.opts);
+      }
+
+      if (sourceURL) {
+        const destinationUrl = getRedirectDestination(args, sourceURL);
+        if (destinationUrl) {
+          addRedirectToContext(sourceURL, destinationUrl, context);
+        }
+      }
+    } catch {
+      // Log if we have a logger with log levels
+    }
+
+    // It's a optional function, so we need to check if it was originally defined
+    if (typeof orig === "function") {
+      return orig.apply(
+        // @ts-expect-error We don't know the type of this
+        this,
+        // @ts-expect-error Arguments are not typed
+        // eslint-disable-next-line prefer-rest-params
+        arguments
+      );
+    }
+  };
+}
+
+function getRedirectDestination(
+  args: unknown[],
+  sourceURL: URL
+): URL | undefined {
+  const statusCode = args[0];
+
+  // Check if the response is a redirect
+  if (typeof statusCode !== "number" || !isRedirectStatusCode(statusCode)) {
+    return;
+  }
+
+  // Get redirect destination
+  const headers = parseHeaders(args[1] as any);
+  if (typeof headers.location !== "string") {
+    return;
+  }
+
+  // Get the destination URL
+  return parseLocationHeader(headers.location, sourceURL.origin);
+}
+
+// The location header can be an absolute or relative URL
+function parseLocationHeader(header: string, origin: string) {
+  if (header.startsWith("/")) {
+    return new URL(header, origin);
+  }
+  return new URL(header);
+}