Merge pull request #47 from AikidoSec/patch-wrapper-api

Hide require-in-the-middle and Shimmer libraries

library/src/agent/Wrapper.test.ts

@@ -0,0 +1,25 @@
+import * as t from "tap";
+import { Hooks } from "./Wrapper";
+
+t.test("package throws error if name is empty", async (t) => {
+  const hooks = new Hooks();
+
+  t.throws(() => hooks.package(""));
+});
+
+t.test("withVersion throws error if version is empty", async (t) => {
+  const hooks = new Hooks();
+  const subject = hooks.package("package");
+
+  t.throws(() => subject.withVersion(""));
+});
+
+t.test("method throws error if name is empty", async (t) => {
+  const hooks = new Hooks();
+  const subject = hooks
+    .package("package")
+    .withVersion("^1.0.0")
+    .subject((exports) => exports);
+
+  t.throws(() => subject.method("", () => {}));
+});

library/src/agent/Wrapper.ts

@@ -1,3 +1,110 @@
+type Interceptor = (args: unknown[], subject: unknown) => void | unknown[];
+
+export class Method {
+  constructor(
+    private readonly name: string,
+    private readonly interceptor: Interceptor
+  ) {
+    if (!this.name) {
+      throw new Error("Method name is required");
+    }
+  }
+
+  getName() {
+    return this.name;
+  }
+
+  getInterceptor() {
+    return this.interceptor;
+  }
+}
+
+class Selector {
+  private methods: Method[] = [];
+
+  constructor(private readonly selector: (exports: unknown) => unknown) {}
+
+  method(name: string, interceptor: Interceptor) {
+    const method = new Method(name, interceptor);
+    this.methods.push(method);
+
+    return method;
+  }
+
+  getSelector() {
+    return this.selector;
+  }
+
+  getMethods() {
+    return this.methods;
+  }
+}
+
+class VersionedPackage {
+  private selectors: Selector[] = [];
+
+  constructor(private readonly range: string) {
+    if (!this.range) {
+      throw new Error("Version range is required");
+    }
+  }
+
+  getRange() {
+    return this.range;
+  }
+
+  subject(selector: (exports: any) => unknown): Selector {
+    const fn = new Selector(selector);
+    this.selectors.push(fn);
+
+    return fn;
+  }
+
+  getSelectors() {
+    return this.selectors;
+  }
+}
+
+class Package {
+  private versions: VersionedPackage[] = [];
+
+  constructor(private readonly packageName: string) {}
+
+  getName() {
+    return this.packageName;
+  }
+
+  withVersion(range: string): VersionedPackage {
+    const pkg = new VersionedPackage(range);
+    this.versions.push(pkg);
+
+    return pkg;
+  }
+
+  getVersions() {
+    return this.versions;
+  }
+}
+
+export class Hooks {
+  private readonly packages: Package[] = [];
+
+  package(packageName: string): Package {
+    if (!packageName) {
+      throw new Error("Package name is required");
+    }
+
+    const pkg = new Package(packageName);
+    this.packages.push(pkg);
+
+    return pkg;
+  }
+
+  getPackages() {
+    return this.packages;
+  }
+}
+
 export interface Wrapper {
-  wrap(): void;
+  wrap(hooks: Hooks): void;
 }

library/src/agent/protect.ts

@@ -1,6 +1,4 @@
 import type { APIGatewayProxyHandler } from "aws-lambda";
-import { getPackageVersion } from "../helpers/getPackageVersion";
-import { satisfiesVersion } from "../helpers/satisfiesVersion";
 import { Agent } from "./Agent";
 import { getInstance, setInstance } from "./AgentSingleton";
 import { API, APIFetch, APIThrottled, Token } from "./API";
@@ -10,40 +8,19 @@ import { MongoDB } from "../sinks/MongoDB";
 import { Postgres } from "../sinks/Postgres";
 import * as shimmer from "shimmer";
 import { Logger, LoggerConsole, LoggerNoop } from "./Logger";
-import { Wrapper } from "./Wrapper";
+import { wrap } from "./wrap";
+import { Hooks } from "./Wrapper";
 import { Options, getOptions } from "../helpers/getOptions";
 
 function wrapInstalledPackages() {
-  const packages: Record<string, { range: string; wrapper: Wrapper }> = {
-    express: {
-      range: "^4.0.0",
-      wrapper: new Express(),
-    },
-    mongodb: {
-      range: "^4.0.0 || ^5.0.0 || ^6.0.0",
-      wrapper: new MongoDB(),
-    },
-    pg: {
-      range: "^8.0.0",
-      wrapper: new Postgres(),
-    },
-  };
-
-  const wrapped: Record<string, { version: string; supported: boolean }> = {};
-  for (const packageName in packages) {
-    const { range, wrapper } = packages[packageName];
-    const version = getPackageVersion(packageName);
-    wrapped[packageName] = {
-      version,
-      supported: version ? satisfiesVersion(range, version) : false,
-    };
-
-    if (wrapped[packageName].supported) {
-      wrapper.wrap();
-    }
-  }
+  const wrappers = [new Express(), new MongoDB(), new Postgres()];
+
+  const hooks = new Hooks();
+  wrappers.forEach((wrapper) => {
+    wrapper.wrap(hooks);
+  });
 
-  return wrapped;
+  return wrap(hooks);
 }
 
 function getLogger(options: Options): Logger {

library/src/agent/wrap.test.ts

@@ -0,0 +1,47 @@
+import * as t from "tap";
+import { wrap } from "./wrap";
+import { Hooks } from "./Wrapper";
+
+t.test("it ignores if package is not installed", async (t) => {
+  const hooks = new Hooks();
+  hooks.package("unknown").withVersion("^1.0.0");
+
+  t.same(wrap(hooks), {});
+});
+
+t.test("it ignores if packages have empty selectors", async (t) => {
+  const hooks = new Hooks();
+  hooks.package("shimmer").withVersion("^1.0.0");
+
+  t.same(wrap(hooks), {});
+});
+
+t.test("it ignores unknown selectors", async (t) => {
+  const hooks = new Hooks();
+  hooks
+    .package("shimmer")
+    .withVersion("^1.0.0")
+    .subject((exports) => exports.doesNotExist)
+    .method("method", () => {});
+
+  t.same(wrap(hooks), {
+    shimmer: {
+      version: "1.2.1",
+      supported: true,
+    },
+  });
+
+  // Force require to load shimmer
+  require("shimmer");
+});
+
+t.test("it ignores if version is not supported", async (t) => {
+  const hooks = new Hooks();
+  hooks
+    .package("shimmer")
+    .withVersion("^2.0.0")
+    .subject((exports) => exports)
+    .method("method", () => {});
+
+  t.same(wrap(hooks), {});
+});

library/src/agent/wrap.ts

@@ -0,0 +1,72 @@
+import { Hook } from "require-in-the-middle";
+import { wrap as shim } from "shimmer";
+import { getPackageVersion } from "../helpers/getPackageVersion";
+import { satisfiesVersion } from "../helpers/satisfiesVersion";
+import { Hooks, Method } from "./Wrapper";
+
+function wrapMethod(subject: unknown, method: Method) {
+  // @ts-expect-error We don't now the type of the subject
+  shim(subject, method.getName(), function wrap(original: Function) {
+    return function wrap() {
+      // eslint-disable-next-line prefer-rest-params
+      const args = Array.from(arguments);
+      // @ts-expect-error We don't now the type of this
+      const updatedArgs = method.getInterceptor()(args, this);
+
+      return original.apply(
+        // @ts-expect-error We don't now the type of this
+        this,
+        // eslint-disable-next-line prefer-rest-params
+        Array.isArray(updatedArgs) ? updatedArgs : arguments
+      );
+    };
+  });
+}
+
+export function wrap(hooks: Hooks) {
+  const wrapped: Record<string, { version: string; supported: boolean }> = {};
+
+  hooks.getPackages().forEach((pkg) => {
+    const version = getPackageVersion(pkg.getName());
+
+    if (!version) {
+      return;
+    }
+
+    const selectors = pkg
+      .getVersions()
+      .map((versioned) => {
+        if (!satisfiesVersion(versioned.getRange(), version)) {
+          return [];
+        }
+
+        return versioned.getSelectors();
+      })
+      .flat();
+
+    if (selectors.length === 0) {
+      return;
+    }
+
+    wrapped[pkg.getName()] = {
+      version,
+      supported: true,
+    };
+
+    new Hook([pkg.getName()], (exports) => {
+      selectors.forEach((selector) => {
+        const subject = selector.getSelector()(exports);
+
+        if (!subject) {
+          return;
+        }
+
+        selector.getMethods().forEach((method) => wrapMethod(subject, method));
+      });
+
+      return exports;
+    });
+  });
+
+  return wrapped;
+}

library/src/sinks/MongoDB.test.ts

@@ -4,6 +4,8 @@ import { setInstance } from "../agent/AgentSingleton";
 import { APIForTesting, Token } from "../agent/API";
 import { LoggerNoop } from "../agent/Logger";
 import { Context, runWithContext } from "../agent/Context";
+import { wrap } from "../agent/wrap";
+import { Hooks } from "../agent/Wrapper";
 import { MongoDB } from "./MongoDB";
 
 const context: Context = {
@@ -21,7 +23,9 @@ const context: Context = {
 };
 
 t.test("we can highjack the MongoDB library", async (t) => {
-  new MongoDB().wrap();
+  const hooks = new Hooks();
+  new MongoDB().wrap(hooks);
+  wrap(hooks);
 
   const agent = new Agent(
     true,