Improve SSRF performance: re-use extracted hostname from URL

hansott · hansott · commit 4879c2945705 · 2025-02-21T11:29:32.000+01:00
diff --git a/benchmarks/operations/benchmark.js b/benchmarks/operations/benchmark.js
@@ -30,6 +30,10 @@ const modules = [
     module: "http-request",
     name: "Outgoing HTTP request (`http.request`)",
   },
+  {
+    module: "undici",
+    name: "Outgoing HTTP request (`undici.request`)",
+  },
 ];
 
 (async () => {
diff --git a/benchmarks/operations/http-request.js b/benchmarks/operations/http-request.js
@@ -3,13 +3,8 @@ const http = require("http");
 module.exports = {
   step: async function step() {
     return new Promise((resolve, reject) => {
-      const options = {
-        hostname: "localhost",
-        port: 10411,
-      };
-
-      const req = http.request(options, (res) => {
-        res.on("data", () => {});
+      const req = http.request("http://localhost:10411", (res) => {
+        res.resume();
         res.on("end", resolve);
       });
 
diff --git a/benchmarks/operations/package-lock.json b/benchmarks/operations/package-lock.json
diff --git a/benchmarks/operations/package.json b/benchmarks/operations/package.json
@@ -7,6 +7,7 @@
   "dependencies": {
     "@aikidosec/firewall": "file:../../build",
     "better-sqlite3": "^11.7.2",
-    "mongodb": "^6.9.0"
+    "mongodb": "^6.9.0",
+    "undici": "^7.3.0"
   }
 }
diff --git a/benchmarks/operations/undici.js b/benchmarks/operations/undici.js
@@ -0,0 +1,7 @@
+const undici = require("undici");
+
+module.exports = {
+  step: async function step() {
+    return await undici.request("http://localhost:10411");
+  },
+};
diff --git a/library/sinks/Fetch.ts b/library/sinks/Fetch.ts
@@ -8,6 +8,7 @@ import { Wrapper } from "../agent/Wrapper";
 import { getPortFromURL } from "../helpers/getPortFromURL";
 import { tryParseURL } from "../helpers/tryParseURL";
 import { checkContextForSSRF } from "../vulnerabilities/ssrf/checkContextForSSRF";
+import { Hostname } from "../vulnerabilities/ssrf/Hostname";
 import { inspectDNSLookupCalls } from "../vulnerabilities/ssrf/inspectDNSLookupCalls";
 import { wrapDispatch } from "./undici/wrapDispatch";
 
@@ -16,13 +17,13 @@ export class Fetch implements Wrapper {
 
   private inspectHostname(
     agent: Agent,
-    hostname: string,
+    hostname: URL,
     port: number | undefined
   ): InterceptorResult {
     // Let the agent know that we are connecting to this hostname
     // This is to build a list of all hostnames that the application is connecting to
     if (typeof port === "number" && port > 0) {
-      agent.onConnectHostname(hostname, port);
+      agent.onConnectHostname(hostname.hostname, port);
     }
     const context = getContext();
 
@@ -31,7 +32,7 @@ export class Fetch implements Wrapper {
     }
 
     return checkContextForSSRF({
-      hostname: hostname,
+      hostname: Hostname.fromURL(hostname),
       operation: "fetch",
       context: context,
       port: port,
@@ -44,11 +45,7 @@ export class Fetch implements Wrapper {
       if (typeof args[0] === "string" && args[0].length > 0) {
         const url = tryParseURL(args[0]);
         if (url) {
-          const attack = this.inspectHostname(
-            agent,
-            url.hostname,
-            getPortFromURL(url)
-          );
+          const attack = this.inspectHostname(agent, url, getPortFromURL(url));
           if (attack) {
             return attack;
           }
@@ -62,11 +59,7 @@ export class Fetch implements Wrapper {
       if (Array.isArray(args[0])) {
         const url = tryParseURL(args[0].toString());
         if (url) {
-          const attack = this.inspectHostname(
-            agent,
-            url.hostname,
-            getPortFromURL(url)
-          );
+          const attack = this.inspectHostname(agent, url, getPortFromURL(url));
           if (attack) {
             return attack;
           }
@@ -77,7 +70,7 @@ export class Fetch implements Wrapper {
       if (args[0] instanceof URL && args[0].hostname.length > 0) {
         const attack = this.inspectHostname(
           agent,
-          args[0].hostname,
+          args[0],
           getPortFromURL(args[0])
         );
         if (attack) {
@@ -89,11 +82,7 @@ export class Fetch implements Wrapper {
       if (args[0] instanceof Request) {
         const url = tryParseURL(args[0].url);
         if (url) {
-          const attack = this.inspectHostname(
-            agent,
-            url.hostname,
-            getPortFromURL(url)
-          );
+          const attack = this.inspectHostname(agent, url, getPortFromURL(url));
           if (attack) {
             return attack;
           }
diff --git a/library/sinks/HTTPRequest.ts b/library/sinks/HTTPRequest.ts
@@ -7,6 +7,7 @@ import { InterceptorResult } from "../agent/hooks/InterceptorResult";
 import { Wrapper } from "../agent/Wrapper";
 import { getPortFromURL } from "../helpers/getPortFromURL";
 import { checkContextForSSRF } from "../vulnerabilities/ssrf/checkContextForSSRF";
+import { Hostname } from "../vulnerabilities/ssrf/Hostname";
 import { inspectDNSLookupCalls } from "../vulnerabilities/ssrf/inspectDNSLookupCalls";
 import { isRedirectToPrivateIP } from "../vulnerabilities/ssrf/isRedirectToPrivateIP";
 import { getUrlFromHTTPRequestArgs } from "./http-request/getUrlFromHTTPRequestArgs";
@@ -34,7 +35,7 @@ export class HTTPRequest implements Wrapper {
 
     // Check if the hostname is inside the context
     const foundDirectSSRF = checkContextForSSRF({
-      hostname: url.hostname,
+      hostname: Hostname.fromURL(url),
       operation: `${module}.request`,
       context: context,
       port: port,
diff --git a/library/sinks/Undici.ts b/library/sinks/Undici.ts
@@ -7,7 +7,9 @@ import { InterceptorResult } from "../agent/hooks/InterceptorResult";
 import { Wrapper } from "../agent/Wrapper";
 import { getSemverNodeVersion } from "../helpers/getNodeVersion";
 import { isVersionGreaterOrEqual } from "../helpers/isVersionGreaterOrEqual";
+import { tryParseURL } from "../helpers/tryParseURL";
 import { checkContextForSSRF } from "../vulnerabilities/ssrf/checkContextForSSRF";
+import { Hostname } from "../vulnerabilities/ssrf/Hostname";
 import { inspectDNSLookupCalls } from "../vulnerabilities/ssrf/inspectDNSLookupCalls";
 import { wrapDispatch } from "./undici/wrapDispatch";
 import { wrapExport } from "../agent/hooks/wrapExport";
@@ -40,8 +42,13 @@ export class Undici implements Wrapper {
       return undefined;
     }
 
+    const hostnameURL = tryParseURL(`http://${hostname}`);
+    if (!hostnameURL) {
+      return undefined;
+    }
+
     return checkContextForSSRF({
-      hostname: hostname,
+      hostname: Hostname.fromURL(hostnameURL),
       operation: `undici.${method}`,
       context,
       port,
diff --git a/library/sinks/http-request/wrapResponseHandler.ts b/library/sinks/http-request/wrapResponseHandler.ts
@@ -6,6 +6,7 @@ import { isRedirectStatusCode } from "../../helpers/isRedirectStatusCode";
 import { tryParseURL } from "../../helpers/tryParseURL";
 import { findHostnameInContext } from "../../vulnerabilities/ssrf/findHostnameInContext";
 import { getRedirectOrigin } from "../../vulnerabilities/ssrf/getRedirectOrigin";
+import { Hostname } from "../../vulnerabilities/ssrf/Hostname";
 import { getUrlFromHTTPRequestArgs } from "./getUrlFromHTTPRequestArgs";
 
 /**
@@ -78,7 +79,11 @@ function addRedirectToContext(source: URL, destination: URL, context: Context) {
   const sourcePort = getPortFromURL(source);
 
   // Check if the source hostname is in the context - is true if it's the first redirect in the chain and the user input is the source
-  const found = findHostnameInContext(source.hostname, context, sourcePort);
+  const found = findHostnameInContext(
+    Hostname.fromURL(source),
+    context,
+    sourcePort
+  );
 
   // If the source hostname is not in the context, check if it's a redirect in a already existing chain
   if (!found && context.outgoingRequestRedirects) {
diff --git a/library/sinks/undici/onRedirect.ts b/library/sinks/undici/onRedirect.ts
@@ -1,6 +1,7 @@
 import { Context, updateContext } from "../../agent/Context";
 import { findHostnameInContext } from "../../vulnerabilities/ssrf/findHostnameInContext";
 import { getRedirectOrigin } from "../../vulnerabilities/ssrf/getRedirectOrigin";
+import { Hostname } from "../../vulnerabilities/ssrf/Hostname";
 import { RequestContextStorage } from "./RequestContextStorage";
 
 /**
@@ -20,7 +21,7 @@ export function onRedirect(
 
   // Check if the source hostname is in the context - is true if it's the first redirect in the chain and the user input is the source
   const found = findHostnameInContext(
-    requestContext.url.hostname,
+    Hostname.fromURL(requestContext.url),
     context,
     requestContext.port
   );
diff --git a/library/vulnerabilities/ssrf/Hostname.ts b/library/vulnerabilities/ssrf/Hostname.ts
@@ -0,0 +1,27 @@
+import { tryParseURL } from "../../helpers/tryParseURL";
+
+export class Hostname {
+  private constructor(private readonly url: URL) {}
+
+  static fromURL(url: URL) {
+    return new Hostname(url);
+  }
+
+  static fromString(str: string) {
+    const url = tryParseURL(`http://${str}`);
+
+    if (!url) {
+      return undefined;
+    }
+
+    return new Hostname(url);
+  }
+
+  asString() {
+    return this.url.hostname;
+  }
+
+  toString() {
+    throw new Error("Use asString() instead");
+  }
+}
diff --git a/library/vulnerabilities/ssrf/checkContextForSSRF.ts b/library/vulnerabilities/ssrf/checkContextForSSRF.ts
@@ -7,6 +7,7 @@ import { extractStringsFromUserInputCached } from "../../helpers/extractStringsF
 import { containsPrivateIPAddress } from "./containsPrivateIPAddress";
 import { findHostnameInUserInput } from "./findHostnameInUserInput";
 import { getMetadataForSSRFAttack } from "./getMetadataForSSRFAttack";
+import { Hostname } from "./Hostname";
 import { isRequestToItself } from "./isRequestToItself";
 
 /**
@@ -19,7 +20,7 @@ export function checkContextForSSRF({
   operation,
   context,
 }: {
-  hostname: string;
+  hostname: Hostname;
   port: number | undefined;
   operation: string;
   context: Context;
@@ -28,7 +29,7 @@ export function checkContextForSSRF({
   // DNS lookup calls will be inspected somewhere else
   // This is just to inspect direct invocations of `http.request` and similar
   // Where the hostname might be a private IP address (or localhost)
-  if (!containsPrivateIPAddress(hostname)) {
+  if (!containsPrivateIPAddress(hostname.asString())) {
     return;
   }
 
@@ -61,7 +62,10 @@ export function checkContextForSSRF({
           kind: "ssrf",
           source: source,
           pathsToPayload: paths,
-          metadata: getMetadataForSSRFAttack({ hostname, port }),
+          metadata: getMetadataForSSRFAttack({
+            hostname: hostname.asString(),
+            port,
+          }),
           payload: str,
         };
       }
diff --git a/library/vulnerabilities/ssrf/findHostnameInContext.ts b/library/vulnerabilities/ssrf/findHostnameInContext.ts
@@ -3,6 +3,7 @@ import { Source, SOURCES } from "../../agent/Source";
 import { getPathsToPayload } from "../../helpers/attackPath";
 import { extractStringsFromUserInputCached } from "../../helpers/extractStringsFromUserInputCached";
 import { findHostnameInUserInput } from "./findHostnameInUserInput";
+import { Hostname } from "./Hostname";
 import { isRequestToItself } from "./isRequestToItself";
 
 type HostnameLocation = {
@@ -14,7 +15,7 @@ type HostnameLocation = {
 };
 
 export function findHostnameInContext(
-  hostname: string,
+  hostname: Hostname,
   context: Context,
   port: number | undefined
 ): HostnameLocation | undefined {
@@ -48,7 +49,7 @@ export function findHostnameInContext(
           pathsToPayload: paths,
           payload: str,
           port: port,
-          hostname: hostname,
+          hostname: hostname.asString(),
         };
       }
     }
diff --git a/library/vulnerabilities/ssrf/findHostnameInUserInput.test.ts b/library/vulnerabilities/ssrf/findHostnameInUserInput.test.ts
diff --git a/library/vulnerabilities/ssrf/findHostnameInUserInput.ts b/library/vulnerabilities/ssrf/findHostnameInUserInput.ts
diff --git a/library/vulnerabilities/ssrf/inspectDNSLookupCalls.ts b/library/vulnerabilities/ssrf/inspectDNSLookupCalls.ts
diff --git a/library/vulnerabilities/ssrf/isRedirectToPrivateIP.ts b/library/vulnerabilities/ssrf/isRedirectToPrivateIP.ts

Original file line number	Diff line number	Diff line change
`@@ -7,6 +7,7 @@`
`7`	`7`	`"dependencies": {`
`8`	`8`	`"@aikidosec/firewall": "file:../../build",`
`9`	`9`	`"better-sqlite3": "^11.7.2",`
`10`		`- "mongodb": "^6.9.0"`
	`10`	`+ "mongodb": "^6.9.0",`
	`11`	`+ "undici": "^7.3.0"`
`11`	`12`	`}`
`12`	`13`	`}`