Merge pull request #458 from AikidoSec/custom-dispatcher

Immediately patch global dispatcher undici

library/sinks/Undici.custom.test.ts

@@ -0,0 +1,100 @@
+/* eslint-disable prefer-rest-params */
+import * as dns from "dns";
+import * as t from "tap";
+import { Token } from "../agent/api/Token";
+import { Context, runWithContext } from "../agent/Context";
+import { wrap } from "../helpers/wrap";
+import { getMajorNodeVersion } from "../helpers/getNodeVersion";
+import { Undici } from "./Undici";
+import { createTestAgent } from "../helpers/createTestAgent";
+
+function createContext(): Context {
+  return {
+    remoteAddress: "::1",
+    method: "POST",
+    url: "http://localhost:4003",
+    query: {},
+    headers: {},
+    body: {
+      image: "http://thisdomainpointstointernalip.com/path",
+    },
+    cookies: {},
+    routeParams: {},
+    source: "express",
+    route: "/posts/:id",
+  };
+}
+
+wrap(dns, "lookup", function lookup(original) {
+  return function lookup() {
+    original.apply(
+      // @ts-expect-error We don't know the type of `this`
+      this,
+      ["localhost", ...Array.from(arguments).slice(1)]
+    );
+  };
+});
+
+t.test(
+  "it works",
+  {
+    skip:
+      getMajorNodeVersion() <= 16 ? "ReadableStream is not available" : false,
+  },
+  async (t) => {
+    const agent = createTestAgent({
+      token: new Token("123"),
+    });
+    agent.start([new Undici()]);
+
+    const { request, Dispatcher, setGlobalDispatcher, getGlobalDispatcher } =
+      require("undici") as typeof import("undici");
+
+    // See https://www.npmjs.com/package/@n8n_io/license-sdk
+    // They set a custom dispatcher to proxy certain requests
+    const originalDispatcher = getGlobalDispatcher();
+
+    const kOptions = Object.getOwnPropertySymbols(originalDispatcher).find(
+      (symbol) => {
+        return symbol.description === "options";
+      }
+    );
+
+    if (!kOptions) {
+      throw new Error("Could not find the options symbol on the dispatcher");
+    }
+
+    // @ts-expect-error kOptions is a symbol
+    const originalOptions = originalDispatcher[kOptions];
+
+    t.ok(
+      "connect" in originalOptions &&
+        originalOptions.connect &&
+        "lookup" in originalOptions.connect
+    );
+
+    setGlobalDispatcher(
+      new (class CustomDispatcher extends Dispatcher {
+        // @ts-expect-error The types of options and handler are unknown
+        dispatch(options, handler) {
+          // Custom logic comes here
+
+          // Fallback to the original dispatcher
+          return originalDispatcher.dispatch(options, handler);
+        }
+      })()
+    );
+
+    await runWithContext(createContext(), async () => {
+      const error = await t.rejects(() =>
+        request("http://thisdomainpointstointernalip.com")
+      );
+      if (error instanceof Error) {
+        t.same(
+          error.message,
+          "Zen has blocked a server-side request forgery: undici.[method](...) originating from body.image"
+        );
+      }
+    });
+  }
+);

library/sinks/Undici.test.ts

@@ -345,7 +345,7 @@ t.test(
     logger.clear();
     setGlobalDispatcher(new UndiciAgent({}));
     t.same(logger.getMessages(), [
-      "undici.setGlobalDispatcher was called, we can't provide protection!",
+      "undici.setGlobalDispatcher(..) was called, we can't guarantee protection!",
     ]);
   }
 );

library/sinks/Undici.ts

@@ -1,5 +1,6 @@
 import { lookup } from "dns";
 import { Agent } from "../agent/Agent";
+import { getInstance } from "../agent/AgentSingleton";
 import { getContext } from "../agent/Context";
 import { Hooks } from "../agent/hooks/Hooks";
 import { InterceptorResult } from "../agent/hooks/InterceptorResult";
@@ -24,8 +25,6 @@ const methods = [
 ];
 
 export class Undici implements Wrapper {
-  private patchedGlobalDispatcher = false;
-
   private inspectHostname(
     agent: Agent,
     hostname: string,
@@ -109,26 +108,34 @@ export class Undici implements Wrapper {
       .addPackage("undici")
       .withVersion("^4.0.0 || ^5.0.0 || ^6.0.0")
       .onRequire((exports, pkgInfo) => {
+        const agent = getInstance();
+
+        if (!agent) {
+          // No agent, we can't do anything
+          return;
+        }
+
+        // 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
+        // (In case the user has a custom dispatcher that conditionally calls the original dispatcher)
+        this.patchGlobalDispatcher(agent, exports);
+
         // Print a warning that we can't provide protection if setGlobalDispatcher is called
         wrapExport(exports, "setGlobalDispatcher", pkgInfo, {
           inspectArgs: (args, agent) => {
-            if (this.patchedGlobalDispatcher) {
-              agent.log(
-                `undici.setGlobalDispatcher was called, we can't provide protection!`
-              );
-            }
+            agent.log(
+              `undici.setGlobalDispatcher(..) was called, we can't guarantee protection!`
+            );
           },
         });
+
         // Wrap all methods that can make requests
         for (const method of methods) {
           wrapExport(exports, method, pkgInfo, {
             // Whenever a request is made, we'll check the hostname whether it's a private IP
             // If global dispatcher is not patched, we'll patch it
             inspectArgs: (args, agent) => {
-              if (!this.patchedGlobalDispatcher) {
-                this.patchGlobalDispatcher(agent, exports);
-                this.patchedGlobalDispatcher = true;
-              }
               return this.inspect(args, agent, method);
             },
           });