Hide require-in-the-middle and Shimmer libraries

This gives us a clean API to wrap things from a package: Allows for wrapping multiple packages at once and even based on versions

library/src/agent/Wrapper.ts

@@ -1,3 +1,174 @@
+import { Hook } from "require-in-the-middle";
+import { wrap } from "shimmer";
+import { getPackageVersion } from "../helpers/getPackageVersion";
+import { satisfiesVersion } from "../helpers/satisfiesVersion";
+
+type Interceptor = (args: unknown[], subject: unknown) => void | unknown[];
+
+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 function wrapPackages(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) => {
+          wrap(subject, method.getName(), function (original) {
+            return function (this: unknown, ...args: unknown[]) {
+              const updatedArgs = method.getInterceptor()(args, this);
+
+              return original.apply(
+                this,
+                Array.isArray(updatedArgs) ? updatedArgs : args
+              );
+            };
+          });
+        });
+      });
+
+      return exports;
+    });
+  });
+
+  return wrapped;
+}
+
 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,18 @@ 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 { Hooks, wrapPackages } 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 wrapPackages(hooks);
 }
 
 function getLogger(options: Options): Logger {

library/src/sinks/MongoDB.test.ts

@@ -4,6 +4,7 @@ import { setInstance } from "../agent/AgentSingleton";
 import { APIForTesting, Token } from "../agent/API";
 import { LoggerNoop } from "../agent/Logger";
 import { Context, runWithContext } from "../agent/Context";
+import { Hooks, wrapPackages } from "../agent/Wrapper";
 import { MongoDB } from "./MongoDB";
 
 const context: Context = {
@@ -21,7 +22,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);
+  wrapPackages(hooks);
 
   const agent = new Agent(
     true,