Fix handling of Maps and other non plain object types

library/helpers/extractStringsFromUserInput.test.ts

@@ -345,3 +345,100 @@ t.test("it works with objects containing constructor key", async () => {
     fromArr(["test", "value", "constructor", "constructor value"])
   );
 });
+
+t.test("it works with objects containing prototype key", async () => {
+  t.same(
+    extractStringsFromUserInput({
+      test: "value",
+      prototype: "prototype value",
+    }),
+    fromArr(["test", "value", "prototype", "prototype value"])
+  );
+
+  t.same(
+    extractStringsFromUserInput({
+      test: "value",
+      __proto__: { protoKey: "protoValue" },
+    }),
+    fromArr(["test", "value", "protoKey", "protoValue"])
+  );
+});
+
+t.test("it works with Map objects", async () => {
+  const map = new Map<string | number, string | object>();
+  map.set("key1", "value1");
+  map.set("key2", { nestedKey: "nestedValue" });
+  map.set(5, "value3");
+
+  t.same(
+    extractStringsFromUserInput(map),
+    fromArr(["key1", "value1", "key2", "nestedKey", "nestedValue", "value3"])
+  );
+});
+
+t.test("it works with Sets", async () => {
+  const set = new Set<string | object>();
+  set.add("value1");
+  set.add({ nestedKey: "nestedValue" });
+
+  t.same(
+    extractStringsFromUserInput(set),
+    fromArr(["value1", "nestedKey", "nestedValue"])
+  );
+});
+
+t.test("it works with URLSearchParams", async () => {
+  const params = new URLSearchParams();
+  params.append("key1", "value1");
+  params.append("key2", "value2");
+
+  t.same(
+    extractStringsFromUserInput(params),
+    fromArr(["key1", "value1", "key2", "value2"])
+  );
+});
+
+t.test(
+  "it works with FormData",
+  {
+    skip:
+      typeof globalThis.FormData === "undefined"
+        ? "FormData is not supported in this environment"
+        : false,
+  },
+  async () => {
+    const formData = new globalThis.FormData();
+    formData.append("key1", "value1");
+    formData.append("key2", "value2");
+
+    t.same(
+      extractStringsFromUserInput(formData),
+      fromArr(["key1", "value1", "key2", "value2"])
+    );
+  }
+);
+
+t.test(
+  "it works with headers object",
+  {
+    skip:
+      typeof Headers === "undefined"
+        ? "Headers is not supported in this environment"
+        : false,
+  },
+  async () => {
+    const headers = new Headers();
+    headers.append("Content-Type", "application/json");
+    headers.append("Authorization", "Bearer token");
+
+    t.same(
+      extractStringsFromUserInput(headers),
+      fromArr([
+        "content-type",
+        "application/json",
+        "authorization",
+        "Bearer token",
+      ])
+    );
+  }
+);

library/helpers/extractStringsFromUserInput.ts

@@ -30,6 +30,48 @@ export function extractStringsFromUserInput(
     }
   }
 
+  if (obj instanceof Map) {
+    for (const [key, value] of obj.entries()) {
+      if (typeof key === "string") {
+        results.add(key);
+      }
+      extractStringsFromUserInput(value, depth + 1).forEach((v) =>
+        results.add(v)
+      );
+    }
+  }
+
+  if (obj instanceof Set) {
+    for (const value of obj.values()) {
+      extractStringsFromUserInput(value, depth + 1).forEach((v) =>
+        results.add(v)
+      );
+    }
+  }
+
+  if (obj instanceof URLSearchParams) {
+    for (const [key, value] of obj.entries()) {
+      results.add(key);
+      results.add(value);
+    }
+  }
+
+  if (globalThis.FormData && obj instanceof globalThis.FormData) {
+    obj.forEach((value, key) => {
+      results.add(key);
+      if (typeof value === "string") {
+        results.add(value);
+      }
+    });
+  }
+
+  if (globalThis.Headers && obj instanceof globalThis.Headers) {
+    obj.forEach((value, key) => {
+      results.add(key);
+      results.add(value);
+    });
+  }
+
   if (Array.isArray(obj)) {
     for (let i = 0; i < obj.length; i++) {
       extractStringsFromUserInput(obj[i], depth + 1).forEach((value) =>

library/sinks/MongoDB.tests.ts

@@ -282,6 +282,117 @@ export function createMongoDBTests(
           "Zen has blocked a NoSQL injection: MongoDB.Collection.find(...) originating from body.name"
         );
       }
+
+      await collection.insertMany([
+        { title: "abc" },
+        { title: "another title" },
+        { title: "yet another title" },
+      ]);
+
+      {
+        // Confirms that filtering with Maps is possible
+        const mapFilter = new Map<string, string>();
+        mapFilter.set("title", "abc");
+        const mapFilteredDocuments = await collection.find(mapFilter).toArray();
+        t.same(
+          mapFilteredDocuments.map((document) => document.title),
+          ["abc"]
+        );
+
+        // Confirms that filtering with nested Maps is possible
+        const mapSubFilter = new Map<string, any>();
+        mapSubFilter.set("$ne", "abc");
+        const mainMapFilter = new Map<string, any>();
+        mainMapFilter.set("title", mapSubFilter);
+        const mapSubFilteredDocuments = await collection
+          .find(mainMapFilter)
+          .toArray();
+        t.same(
+          mapSubFilteredDocuments.map((document) => document.title),
+          ["another title", "yet another title"]
+        );
+
+        // Confirms that map inside plain object is also working
+        const plainObjectWithMapFilter = {
+          title: new Map<string, any>([["$ne", "abc"]]),
+        };
+        const plainObjectWithMapFilteredDocuments = await collection
+          .find(plainObjectWithMapFilter)
+          .toArray();
+        t.same(
+          plainObjectWithMapFilteredDocuments.map((document) => document.title),
+          ["another title", "yet another title"]
+        );
+      }
+
+      const mapError = await t.rejects(async () => {
+        await runWithContext(unsafeContext, () => {
+          const filter = new Map<string, any>();
+          filter.set("title", { $ne: null });
+          return collection.find(filter).toArray();
+        });
+      });
+      t.ok(mapError instanceof Error);
+      if (mapError instanceof Error) {
+        t.same(
+          mapError.message,
+          "Zen has blocked a NoSQL injection: MongoDB.Collection.find(...) originating from body.myTitle"
+        );
+      }
+
+      const mapError2 = await t.rejects(async () => {
+        await runWithContext(
+          {
+            ...unsafeContext,
+            body: {
+              title: new Map<string, any>([["$ne", null]]),
+            },
+          },
+          () => {
+            const filter = new Map<string, any>();
+            filter.set("title", { $ne: null });
+            return collection.find(filter).toArray();
+          }
+        );
+      });
+      t.ok(mapError2 instanceof Error);
+      if (mapError2 instanceof Error) {
+        t.same(
+          mapError2.message,
+          "Zen has blocked a NoSQL injection: MongoDB.Collection.find(...) originating from body.title"
+        );
+      }
+
+      const mapError3 = await t.rejects(async () => {
+        await runWithContext(unsafeContext, () => {
+          const filter = {
+            title: new Map<string, any>([["$ne", null]]),
+          };
+          return collection.find(filter).toArray();
+        });
+      });
+      t.ok(mapError3 instanceof Error);
+      if (mapError3 instanceof Error) {
+        t.same(
+          mapError3.message,
+          "Zen has blocked a NoSQL injection: MongoDB.Collection.find(...) originating from body.myTitle"
+        );
+      }
+
+      const mapError4 = await t.rejects(async () => {
+        await runWithContext(unsafeContext, () => {
+          const filter = new Map<string, any>();
+          filter.set("title", new Map<string, any>([["$ne", null]]));
+          return collection.find(filter).toArray();
+        });
+      });
+      t.ok(mapError4 instanceof Error);
+      if (mapError4 instanceof Error) {
+        t.same(
+          mapError4.message,
+          "Zen has blocked a NoSQL injection: MongoDB.Collection.find(...) originating from body.myTitle"
+        );
+      }
     } catch (error: any) {
       t.fail(error.message);
     } finally {

library/sinks/MongoDB.ts

@@ -3,7 +3,6 @@ import { Hooks } from "../agent/hooks/Hooks";
 import { InterceptorResult } from "../agent/hooks/InterceptorResult";
 import type { WrapPackageInfo } from "../agent/hooks/WrapPackageInfo";
 import { detectNoSQLInjection } from "../vulnerabilities/nosql-injection/detectNoSQLInjection";
-import { isPlainObject } from "../helpers/isPlainObject";
 import { Context, getContext } from "../agent/Context";
 import { Wrapper } from "../agent/Wrapper";
 import { wrapExport } from "../agent/hooks/wrapExport";
@@ -147,7 +146,7 @@ export class MongoDB implements Wrapper {
       return undefined;
     }
 
-    if (args.length > 0 && isPlainObject(args[0])) {
+    if (args.length > 0) {
       const filter = args[0];
 
       return this.inspectFilter(
@@ -172,7 +171,7 @@ export class MongoDB implements Wrapper {
       return undefined;
     }
 
-    if (args.length > 1 && isPlainObject(args[1])) {
+    if (args.length > 1) {
       const filter = args[1];
 
       return this.inspectFilter(

library/vulnerabilities/nosql-injection/detectNoSQLInjection.test.ts

@@ -879,3 +879,32 @@ t.test("not a valid injection attempt", async (t) => {
     }
   );
 });
+
+t.test("it works with Maps", async (t) => {
+  t.same(
+    detectNoSQLInjection(
+      createContext({
+        body: new Map<string, any>([
+          ["username", "admin"],
+          [
+            "test",
+            new Map<string, any>([
+              ["$ne", ""],
+              ["hello", "world"],
+            ]),
+          ],
+        ]),
+      }),
+      {
+        username: "admin",
+        test: { $ne: "", hello: "world" },
+      }
+    ),
+    {
+      injection: true,
+      source: "body",
+      pathsToPayload: [".test"],
+      payload: { $ne: "" },
+    }
+  );
+});

library/vulnerabilities/nosql-injection/detectNoSQLInjection.ts

@@ -65,6 +65,14 @@ function matchFilterPartInUser(
     return matchFilterPartInUser(userInput.join(), filterPart, pathToPayload);
   }
 
+  if (userInput instanceof Map) {
+    return matchFilterPartInUser(
+      mapToPlainObject(userInput),
+      filterPart,
+      pathToPayload
+    );
+  }
+
   return {
     match: false,
   };
@@ -127,6 +135,13 @@ function findFilterPartWithOperators(
     }
   }
 
+  if (partOfFilter instanceof Map) {
+    return findFilterPartWithOperators(
+      userInput,
+      mapToPlainObject(partOfFilter)
+    );
+  }
+
   return { found: false };
 }
 
@@ -143,6 +158,10 @@ export function detectNoSQLInjection(
   request: Context,
   filter: unknown
 ): DetectionResult {
+  if (filter instanceof Map) {
+    return detectNoSQLInjection(request, mapToPlainObject(filter));
+  }
+
   if (!isPlainObject(filter) && !Array.isArray(filter)) {
     return { injection: false };
   }
@@ -164,3 +183,13 @@ export function detectNoSQLInjection(
 
   return { injection: false };
 }
+
+function mapToPlainObject(map: Map<unknown, unknown>): Record<string, unknown> {
+  const obj: Record<string, unknown> = {};
+  for (const [key, value] of map.entries()) {
+    if (typeof key === "string") {
+      obj[key] = value;
+    }
+  }
+  return obj;
+}