Avoid for loop over sources

It's a tiny improvement but makes it a bit more readable too.

Author: hansott

library/agent/Context.test.ts

@@ -107,27 +107,21 @@ t.test("it clears cache when context is mutated", async (t) => {
   const context = { ...sampleContext };
 
   runWithContext(context, () => {
-    t.same(extractStringsFromUserInputCached(getContext()!, "body"), undefined);
     t.same(
-      extractStringsFromUserInputCached(getContext()!, "query"),
-      new Set(["abc", "def"])
+      extractStringsFromUserInputCached(getContext()!),
+      new Set(["abc", "def", "http://localhost:4000"])
     );
 
     updateContext(getContext()!, "query", {});
-    t.same(extractStringsFromUserInputCached(getContext()!, "body"), undefined);
     t.same(
-      extractStringsFromUserInputCached(getContext()!, "query"),
-      new Set()
+      extractStringsFromUserInputCached(getContext()!),
+      new Set(["http://localhost:4000"])
     );
 
     runWithContext({ ...context, body: { a: "z" }, query: { b: "y" } }, () => {
       t.same(
-        extractStringsFromUserInputCached(getContext()!, "body"),
-        new Set(["a", "z"])
-      );
-      t.same(
-        extractStringsFromUserInputCached(getContext()!, "query"),
-        new Set(["b", "y"])
+        extractStringsFromUserInputCached(getContext()!),
+        new Set(["a", "z", "b", "y", "http://localhost:4000"])
       );
     });
   });

library/agent/Context.ts

@@ -2,7 +2,6 @@ import type { ParsedQs } from "qs";
 import { extractStringsFromUserInput } from "../helpers/extractStringsFromUserInput";
 import { ContextStorage } from "./context/ContextStorage";
 import { AsyncResource } from "async_hooks";
-import { Source, SOURCES } from "./Source";
 import type { Endpoint } from "./Config";
 
 export type User = { id: string; name?: string };
@@ -25,7 +24,7 @@ export type Context = {
   xml?: unknown[];
   subdomains?: string[]; // https://expressjs.com/en/5x/api.html#req.subdomains
   markUnsafe?: unknown[];
-  cache?: Map<Source, ReturnType<typeof extractStringsFromUserInput>>;
+  cache?: ReturnType<typeof extractStringsFromUserInput>;
   /**
    * Used to store redirects in outgoing http(s) requests that are started by a user-supplied input (hostname and port / url) to prevent SSRF redirect attacks.
    */
@@ -44,10 +43,6 @@ export function getContext(): Readonly<Context> | undefined {
   return ContextStorage.getStore();
 }
 
-function isSourceKey(key: string): key is Source {
-  return SOURCES.includes(key as Source);
-}
-
 // We need to use a function to mutate the context because we need to clear the cache when the user input changes
 export function updateContext<K extends keyof Context>(
   context: Context,
@@ -56,9 +51,8 @@ export function updateContext<K extends keyof Context>(
 ) {
   context[key] = value;
 
-  if (context.cache && isSourceKey(key)) {
-    context.cache.delete(key);
-  }
+  // Clear all the cached user input strings
+  delete context.cache;
 }
 
 /**

library/helpers/extractStringsFromUserInputCached.ts

@@ -1,25 +1,29 @@
 import { Context } from "../agent/Context";
-import { Source } from "../agent/Source";
+import { SOURCES } from "../agent/Source";
 import { extractStringsFromUserInput } from "./extractStringsFromUserInput";
 
+type ReturnValue = ReturnType<typeof extractStringsFromUserInput>;
+
 export function extractStringsFromUserInputCached(
-  context: Context,
-  source: Source
-): ReturnType<typeof extractStringsFromUserInput> | undefined {
-  if (!context[source]) {
-    return undefined;
+  context: Context
+): ReturnValue {
+  if (context.cache) {
+    return context.cache;
   }
 
-  if (!context.cache) {
-    context.cache = new Map();
-  }
+  const userStrings: ReturnValue = new Set();
 
-  let result = context.cache.get(source);
+  for (const source of SOURCES) {
+    if (!context[source]) {
+      continue;
+    }
 
-  if (!result) {
-    result = extractStringsFromUserInput(context[source]);
-    context.cache.set(source, result);
+    for (const item of extractStringsFromUserInput(context[source])) {
+      userStrings.add(item);
+    }
   }
 
-  return result;
+  context.cache = userStrings;
+
+  return userStrings;
 }

library/helpers/getSourceForUserString.test.ts

@@ -0,0 +1,40 @@
+import * as t from "tap";
+import { Context } from "../agent/Context";
+import { getSourceForUserString } from "./getSourceForUserString";
+
+function createContext(): Context {
+  return {
+    remoteAddress: "::1",
+    method: "POST",
+    url: "http://local.aikido.io",
+    query: {},
+    headers: {},
+    body: {
+      image: "http://localhost:4000/api/internal",
+    },
+    cookies: {},
+    routeParams: {},
+    source: "express",
+    route: "/posts/:id",
+  };
+}
+
+t.test(
+  "it returns undefined if the user string cannot be found in the context",
+  async () => {
+    t.same(getSourceForUserString(createContext(), "unknown"), undefined);
+  }
+);
+
+t.test(
+  "it returns source if the user string is found in the context",
+  async () => {
+    t.same(
+      getSourceForUserString(
+        createContext(),
+        "http://localhost:4000/api/internal"
+      ),
+      "body"
+    );
+  }
+);

library/helpers/getSourceForUserString.ts

@@ -0,0 +1,22 @@
+import { Context } from "../agent/Context";
+import { Source, SOURCES } from "../agent/Source";
+import { extractStringsFromUserInput } from "./extractStringsFromUserInput";
+
+export function getSourceForUserString(
+  context: Context,
+  str: string
+): Source | undefined {
+  for (const source of SOURCES) {
+    if (!context[source]) {
+      continue;
+    }
+
+    const userStrings = extractStringsFromUserInput(context[source]);
+
+    if (userStrings.has(str)) {
+      return source;
+    }
+  }
+
+  return undefined;
+}

library/sources/xml/isXmlInContext.ts

@@ -1,6 +1,6 @@
 import { Context } from "../../agent/Context";
 import { SOURCES } from "../../agent/Source";
-import { extractStringsFromUserInputCached } from "../../helpers/extractStringsFromUserInputCached";
+import { extractStringsFromUserInput } from "../../helpers/extractStringsFromUserInput";
 
 /**
  * Checks if the XML string can be found in the context.
@@ -11,15 +11,14 @@ export function isXmlInContext(xml: string, context: Context): boolean {
       // Skip parsed XML
       continue;
     }
-    const userInput = extractStringsFromUserInputCached(context, source);
-    if (!userInput) {
+
+    if (!context[source]) {
       continue;
     }
 
-    for (const str of userInput) {
-      if (str === xml) {
-        return true;
-      }
+    const userInput = extractStringsFromUserInput(context[source]);
+    if (userInput.has(xml)) {
+      return true;
     }
   }
 

library/vulnerabilities/attack-wave-detection/queryParamsContainDangerousPayload.ts

@@ -1,5 +1,5 @@
 import type { Context } from "../../agent/Context";
-import { extractStringsFromUserInputCached } from "../../helpers/extractStringsFromUserInputCached";
+import { extractStringsFromUserInput } from "../../helpers/extractStringsFromUserInput";
 
 const keywords = [
   "SELECT (CASE WHEN",
@@ -24,10 +24,15 @@ const keywords = [
  * Check the query for some common SQL or path traversal patterns.
  */
 export function queryParamsContainDangerousPayload(context: Context): boolean {
-  const queryStrings = extractStringsFromUserInputCached(context, "query");
+  if (!context.query) {
+    return false;
+  }
+
+  const queryStrings = extractStringsFromUserInput(context.query);
   if (!queryStrings) {
     return false;
   }
+
   for (const str of queryStrings) {
     // Performance optimization
     // Some keywords like ../ are shorter than this min length check

library/vulnerabilities/js-injection/checkContextForJsInjection.ts

@@ -1,8 +1,8 @@
 import { Context } from "../../agent/Context";
 import { InterceptorResult } from "../../agent/hooks/InterceptorResult";
-import { SOURCES } from "../../agent/Source";
 import { getPathsToPayload } from "../../helpers/attackPath";
 import { extractStringsFromUserInputCached } from "../../helpers/extractStringsFromUserInputCached";
+import { getSourceForUserString } from "../../helpers/getSourceForUserString";
 import { detectJsInjection } from "./detectJsInjection";
 
 /**
@@ -18,14 +18,10 @@ export function checkContextForJsInjection({
   operation: string;
   context: Context;
 }): InterceptorResult {
-  for (const source of SOURCES) {
-    const userInput = extractStringsFromUserInputCached(context, source);
-    if (!userInput) {
-      continue;
-    }
-
-    for (const str of userInput) {
-      if (detectJsInjection(js, str)) {
+  for (const str of extractStringsFromUserInputCached(context)) {
+    if (detectJsInjection(js, str)) {
+      const source = getSourceForUserString(context, str);
+      if (source) {
         return {
           operation: operation,
           kind: "code_injection",

library/vulnerabilities/path-traversal/checkContextForPathTraversal.ts

@@ -1,8 +1,8 @@
 import { Context } from "../../agent/Context";
 import { InterceptorResult } from "../../agent/hooks/InterceptorResult";
-import { SOURCES } from "../../agent/Source";
 import { getPathsToPayload } from "../../helpers/attackPath";
 import { extractStringsFromUserInputCached } from "../../helpers/extractStringsFromUserInputCached";
+import { getSourceForUserString } from "../../helpers/getSourceForUserString";
 import { detectPathTraversal } from "./detectPathTraversal";
 
 /**
@@ -26,14 +26,10 @@ export function checkContextForPathTraversal({
     return;
   }
 
-  for (const source of SOURCES) {
-    const userInput = extractStringsFromUserInputCached(context, source);
-    if (!userInput) {
-      continue;
-    }
-
-    for (const str of userInput) {
-      if (detectPathTraversal(pathString, str, checkPathStart, isUrl)) {
+  for (const str of extractStringsFromUserInputCached(context)) {
+    if (detectPathTraversal(pathString, str, checkPathStart, isUrl)) {
+      const source = getSourceForUserString(context, str);
+      if (source) {
         return {
           operation: operation,
           kind: "path_traversal",

library/vulnerabilities/shell-injection/checkContextForShellInjection.ts

@@ -1,8 +1,8 @@
 import { Context } from "../../agent/Context";
 import { InterceptorResult } from "../../agent/hooks/InterceptorResult";
-import { SOURCES } from "../../agent/Source";
 import { getPathsToPayload } from "../../helpers/attackPath";
 import { extractStringsFromUserInputCached } from "../../helpers/extractStringsFromUserInputCached";
+import { getSourceForUserString } from "../../helpers/getSourceForUserString";
 import { detectShellInjection } from "./detectShellInjection";
 
 /**
@@ -18,14 +18,10 @@ export function checkContextForShellInjection({
   operation: string;
   context: Context;
 }): InterceptorResult {
-  for (const source of SOURCES) {
-    const userInput = extractStringsFromUserInputCached(context, source);
-    if (!userInput) {
-      continue;
-    }
-
-    for (const str of userInput) {
-      if (detectShellInjection(command, str)) {
+  for (const str of extractStringsFromUserInputCached(context)) {
+    if (detectShellInjection(command, str)) {
+      const source = getSourceForUserString(context, str);
+      if (source) {
         return {
           operation: operation,
           kind: "shell_injection",