Refactor: use generic addHook/removeHook functions

Author: hansott

docs/outbound-requests.md

@@ -1,21 +1,39 @@
 # Monitoring Outbound Requests
 
-To monitor outbound HTTP/HTTPS requests made by your application, you can use the `onOutboundRequest` function. This is useful when you want to track external API calls, log outbound traffic, or analyze what domains your application connects to.
+To monitor outbound HTTP/HTTPS requests made by your application, you can use the `addHook` function with the `beforeOutboundRequest` hook. This is useful when you want to track external API calls, log outbound traffic, or analyze what domains your application connects to.
 
 ## Basic Usage
 
 ```js
-const { onOutboundRequest } = require("@aikidosec/firewall");
+const { addHook } = require("@aikidosec/firewall");
 
-onOutboundRequest(({ url, port, method }) => {
+addHook("beforeOutboundRequest", ({ url, port, method }) => {
   // url is a URL object: https://nodejs.org/api/url.html#class-url
   console.log(`${new Date().toISOString()} - ${method} ${url.href}`);
 });
 ```
 
+## Removing Hooks
+
+You can remove a previously registered hook using the `removeHook` function:
+
+```js
+const { addHook, removeHook } = require("@aikidosec/firewall");
+
+function myHook({ url, port, method }) {
+  console.log(`${method} ${url.href}`);
+}
+
+addHook("beforeOutboundRequest", myHook);
+
+// Later, when you want to remove it:
+removeHook("beforeOutboundRequest", myHook);
+```
+
 ## Important Notes
 
-- You can register multiple callbacks by calling `onOutboundRequest` multiple times.
-- Callbacks are triggered for all HTTP/HTTPS requests made through Node.js built-in modules (`http`, `https`), builtin fetch function, undici and anything that uses that.
-- Callbacks are called when the connection is initiated, before knowing if Zen will block the request.
-- Errors thrown in callbacks (both sync and async) are silently caught and not logged to prevent breaking your application.
+- You can register multiple hooks by calling `addHook` multiple times.
+- The same hook function can only be registered once (duplicates are automatically prevented).
+- Hooks are triggered for all HTTP/HTTPS requests made through Node.js built-in modules (`http`, `https`), builtin fetch function, undici and anything that uses that.
+- Hooks are called when the connection is initiated, before knowing if Zen will block the request.
+- Errors thrown in hooks (both sync and async) are silently caught and not logged to prevent breaking your application.

library/agent/Agent.ts

@@ -27,7 +27,7 @@ import { wrapInstalledPackages } from "./wrapInstalledPackages";
 import { Wrapper } from "./Wrapper";
 import { isAikidoCI } from "../helpers/isAikidoCI";
 import { AttackLogger } from "./AttackLogger";
-import { triggerOutboundRequestHooks } from "./hooks/outboundRequest";
+import { executeHooks } from "./hooks";
 import { Packages } from "./Packages";
 import { AIStatistics } from "./AIStatistics";
 import { isNewInstrumentationUnitTest } from "../helpers/isNewInstrumentationUnitTest";
@@ -570,7 +570,7 @@ export class Agent {
   }
 
   onConnectHTTP(url: URL, port: number, method: string) {
-    triggerOutboundRequestHooks({ url, port, method });
+    executeHooks("beforeOutboundRequest", { url, port, method });
   }
 
   onRouteExecute(context: Context) {

library/agent/hooks.test.ts

@@ -0,0 +1,145 @@
+import * as t from "tap";
+import { addHook, removeHook, executeHooks, OutboundRequestInfo } from "./hooks";
+
+t.test("it works", async (t) => {
+  let hookOneCalls = 0;
+  let hookTwoCalls = 0;
+
+  const testRequest: OutboundRequestInfo = {
+    url: new URL("https://example.com"),
+    port: 443,
+    method: "GET",
+  };
+
+  function hook1(request: OutboundRequestInfo) {
+    t.equal(request.url.href, "https://example.com/");
+    t.equal(request.port, 443);
+    t.equal(request.method, "GET");
+    hookOneCalls++;
+  }
+
+  function hook2(request: OutboundRequestInfo) {
+    t.equal(request.url.href, "https://example.com/");
+    t.equal(request.port, 443);
+    t.equal(request.method, "GET");
+    hookTwoCalls++;
+  }
+
+  function hook3() {
+    throw new Error("hook3 should not be called");
+  }
+
+  t.same(hookOneCalls, 0, "hookOneCalls starts at 0");
+  t.same(hookTwoCalls, 0, "hookTwoCalls starts at 0");
+
+  executeHooks("beforeOutboundRequest", testRequest);
+
+  t.same(hookOneCalls, 0, "hookOneCalls still at 0");
+  t.same(hookTwoCalls, 0, "hookTwoCalls still at 0");
+
+  addHook("beforeOutboundRequest", hook1);
+  // @ts-expect-error some other hook is not defined in the types
+  addHook("someOtherHook", hook3);
+  executeHooks("beforeOutboundRequest", testRequest);
+
+  t.equal(hookOneCalls, 1, "hook1 called once");
+  t.equal(hookTwoCalls, 0, "hook2 not called");
+
+  addHook("beforeOutboundRequest", hook2);
+  executeHooks("beforeOutboundRequest", testRequest);
+
+  t.equal(hookOneCalls, 2, "hook1 called twice");
+  t.equal(hookTwoCalls, 1, "hook2 called once");
+
+  removeHook("beforeOutboundRequest", hook1);
+  executeHooks("beforeOutboundRequest", testRequest);
+
+  t.equal(hookOneCalls, 2, "hook1 still called twice");
+  t.equal(hookTwoCalls, 2, "hook2 called twice");
+
+  removeHook("beforeOutboundRequest", hook2);
+  executeHooks("beforeOutboundRequest", testRequest);
+
+  t.equal(hookOneCalls, 2, "hook1 still called twice");
+  t.equal(hookTwoCalls, 2, "hook2 still called twice");
+});
+
+t.test("it handles errors gracefully", async (t) => {
+  let successCalls = 0;
+
+  function throwingHook() {
+    throw new Error("This should be caught");
+  }
+
+  function successHook() {
+    successCalls++;
+  }
+
+  const testRequest: OutboundRequestInfo = {
+    url: new URL("https://example.com"),
+    port: 443,
+    method: "POST",
+  };
+
+  addHook("beforeOutboundRequest", throwingHook);
+  addHook("beforeOutboundRequest", successHook);
+
+  // Should not throw even though one hook throws
+  t.doesNotThrow(() => {
+    executeHooks("beforeOutboundRequest", testRequest);
+  });
+
+  t.equal(successCalls, 1, "success hook still called despite error in other hook");
+
+  removeHook("beforeOutboundRequest", throwingHook);
+  removeHook("beforeOutboundRequest", successHook);
+});
+
+t.test("it handles async hooks with rejected promises", async (t) => {
+  let asyncCalls = 0;
+
+  async function asyncHook() {
+    asyncCalls++;
+    throw new Error("Async error");
+  }
+
+  const testRequest: OutboundRequestInfo = {
+    url: new URL("https://example.com"),
+    port: 443,
+    method: "DELETE",
+  };
+
+  addHook("beforeOutboundRequest", asyncHook);
+
+  // Should not throw even though async hook rejects
+  t.doesNotThrow(() => {
+    executeHooks("beforeOutboundRequest", testRequest);
+  });
+
+  t.equal(asyncCalls, 1, "async hook was called");
+
+  removeHook("beforeOutboundRequest", asyncHook);
+});
+
+t.test("it prevents duplicate hooks using Set", async (t) => {
+  let hookCalls = 0;
+
+  function hook() {
+    hookCalls++;
+  }
+
+  const testRequest: OutboundRequestInfo = {
+    url: new URL("https://example.com"),
+    port: 443,
+    method: "GET",
+  };
+
+  addHook("beforeOutboundRequest", hook);
+  addHook("beforeOutboundRequest", hook); // Try to add the same hook again
+
+  executeHooks("beforeOutboundRequest", testRequest);
+
+  t.equal(hookCalls, 1, "hook only called once despite being added twice");
+
+  removeHook("beforeOutboundRequest", hook);
+});

library/agent/hooks.ts

@@ -0,0 +1,60 @@
+export type OutboundRequestInfo = {
+  url: URL;
+  port: number;
+  method: string;
+};
+
+type HookName = "beforeOutboundRequest";
+
+// Map hook names to argument types
+interface HookTypes {
+  beforeOutboundRequest: {
+    args: [data: OutboundRequestInfo];
+  };
+}
+
+const hooks = new Map<
+  HookName,
+  Set<(...args: HookTypes[HookName]["args"]) => void | Promise<void>>
+>();
+
+export function addHook<N extends HookName>(
+  name: N,
+  fn: (...args: HookTypes[N]["args"]) => void | Promise<void>
+) {
+  if (!hooks.has(name)) {
+    hooks.set(name, new Set([fn]));
+  } else {
+    hooks.get(name)!.add(fn);
+  }
+}
+
+export function removeHook<N extends HookName>(
+  name: N,
+  fn: (...args: HookTypes[N]["args"]) => void | Promise<void>
+) {
+  if (hooks.has(name)) {
+    hooks.get(name)!.delete(fn);
+  }
+}
+
+export function executeHooks<N extends HookName>(
+  name: N,
+  ...args: [...HookTypes[N]["args"]]
+): void {
+  const hookSet = hooks.get(name);
+
+  for (const fn of hookSet ?? []) {
+    try {
+      const result = (fn as (...args: HookTypes[N]["args"]) => void | Promise<void>)(...args);
+      // If it returns a promise, catch any errors but don't wait
+      if (result instanceof Promise) {
+        result.catch(() => {
+          // Silently ignore errors from user hooks
+        });
+      }
+    } catch {
+      // Silently ignore errors from user hooks
+    }
+  }
+}

library/agent/hooks/outboundRequest.ts

@@ -15,7 +15,7 @@ import { isESM } from "./helpers/isESM";
 import { checkIndexImportGuard } from "./helpers/indexImportGuard";
 import { setRateLimitGroup } from "./ratelimiting/group";
 import { isLibBundled } from "./helpers/isLibBundled";
-import { onOutboundRequest } from "./agent/hooks/outboundRequest";
+import { addHook, removeHook } from "./agent/hooks";
 
 // Prevent logging twice / trying to start agent twice
 if (!isNewHookSystemUsed()) {
@@ -52,7 +52,8 @@ export {
   addKoaMiddleware,
   addRestifyMiddleware,
   setRateLimitGroup,
-  onOutboundRequest,
+  addHook,
+  removeHook,
 };
 
 // Required for ESM / TypeScript default export support
@@ -69,5 +70,6 @@ export default {
   addKoaMiddleware,
   addRestifyMiddleware,
   setRateLimitGroup,
-  onOutboundRequest,
+  addHook,
+  removeHook,
 };