Modify instrumentation instruction system (WIP)

Author: timokoessler

instrumentation-wasm/src/js_transformer/transformer.rs

@@ -132,7 +132,7 @@ impl<'a> Traverse<'a> for Transformer<'a> {
 
         if instruction.inspect_args {
             let source_text: &'a str = self.allocator.alloc_str(&format!(
-                "__instrumentInspectArgs('{}', arguments);",
+                "__instrumentInspectArgs('{}', false, arguments);",
                 self.current_function_identifier.as_ref().unwrap()
             ));
 

library/agent/applyHooks.ts

@@ -9,7 +9,7 @@ import { registerNodeHooks } from "./hooks/instrumentation";
 import {
   setBuiltinsToInstrument,
   setPackagesToInstrument,
-} from "./hooks/instrumentation/loadHook";
+} from "./hooks/instrumentation/instructions";
 
 /**
  * Hooks allows you to register packages and then wrap specific methods on
@@ -26,7 +26,7 @@ export function applyHooks(hooks: Hooks, newInstrumentation: boolean = false) {
   } else {
     setPackagesToInstrument(hooks.getPackages());
     setBuiltinsToInstrument(hooks.getBuiltInModules());
-    registerNodeHooks(hooks);
+    registerNodeHooks();
   }
 
   hooks.getGlobals().forEach((g) => {

library/agent/hooks/BuiltinModule.ts

@@ -3,7 +3,9 @@ import { RequireInterceptor } from "./RequireInterceptor";
 
 export class BuiltinModule {
   private requireInterceptors: RequireInterceptor[] = [];
-  private instrumentationInstructions: BuiltinInstrumentationInstruction[] = [];
+  private instrumentationInstructions:
+    | BuiltinInstrumentationInstruction
+    | undefined;
 
   constructor(private readonly name: string) {
     if (!this.name) {
@@ -23,11 +25,13 @@ export class BuiltinModule {
     return this.requireInterceptors;
   }
 
-  addInstrumentation(instruction: BuiltinInstrumentationInstruction) {
-    this.instrumentationInstructions.push(instruction);
+  setInstrumentationInstruction(
+    instruction: BuiltinInstrumentationInstruction
+  ) {
+    this.instrumentationInstructions = instruction;
   }
 
-  getInstrumentationInstructions() {
+  getInstrumentationInstruction() {
     return this.instrumentationInstructions;
   }
 }

library/agent/hooks/instrumentation/builtinShim.ts

@@ -1,11 +1,15 @@
 import { getExportsForBuiltin } from "./getExportsForBuiltin";
-import { BuiltinInstrumentationInstruction } from "./types";
+import {
+  BuiltinInstrumentationInstruction,
+  BuiltinInstrumentationInstructionJSON,
+} from "./types";
 
 export function generateBuildinShim(
-  moduleName: string,
-  instructions: BuiltinInstrumentationInstruction["functions"]
+  builtinName: string,
+  builtinNameWithoutPrefix: string,
+  instructions: BuiltinInstrumentationInstructionJSON["functions"]
 ): string | undefined {
-  const exports = getExportsForBuiltin(moduleName);
+  const exports = getExportsForBuiltin(builtinName);
 
   // Filter out non-existing exports
   const methods = instructions.filter((m) => exports.has(m.name));
@@ -20,14 +24,14 @@ export function generateBuildinShim(
   );
 
   return `
-        const orig = process.getBuiltinModule(${JSON.stringify(moduleName)});
+        const orig = process.getBuiltinModule(${JSON.stringify(builtinName)});
         const { __instrumentInspectArgs } = require('@aikidosec/firewall/instrument/internals');
 
         ${methods
           .map(
             (method) => `
                 exports.${method.name} = function() {
-                    ${method.inspectArgs ? `__instrumentInspectArgs("${moduleName}.${method.name}", arguments);` : ""}
+                    ${method.inspectArgs ? `__instrumentInspectArgs("${builtinNameWithoutPrefix}.${method.name}", true, arguments);` : ""}
                     return orig.${method.name}(...arguments);
                 };`
           )

library/agent/hooks/instrumentation/codeTransformation.ts

@@ -1,14 +1,16 @@
-import type { PackageFileInstrumentationInstruction } from "./types";
+import type { PackageFileInstrumentationInstructionJSON } from "./types";
 // eslint-disable-next-line camelcase
 import { wasm_transform_code_str } from "./wasm/node_code_instrumentation";
 import { getSourceType } from "./getSourceType";
 
+// Todo check if caching is done by Node or if we need to cache the result
+
 export function transformCode(
   path: string,
   code: string,
   moduleName: string,
   isESM: boolean,
-  fileInstructions: PackageFileInstrumentationInstruction
+  fileInstructions: PackageFileInstrumentationInstructionJSON
 ): string {
   const result = wasm_transform_code_str(
     code,

library/agent/hooks/instrumentation/index.ts

@@ -1,19 +1,16 @@
 import { onModuleLoad } from "./loadHook";
-import * as module from "node:module";
+import * as mod from "node:module";
 import type { RegisterHookFunction } from "./types";
 import { Hooks } from "../Hooks";
 
-export function registerNodeHooks(hooks: Hooks) {
-  if (
-    !("registerHooks" in module) ||
-    typeof module.registerHooks !== "function"
-  ) {
+export function registerNodeHooks() {
+  if (!("registerHooks" in mod) || typeof mod.registerHooks !== "function") {
     throw new Error("This Node.js version is not supported");
   }
 
   // Hook into the ESM & CJS module loading process
   // Types are required because official Node.js typings are not up-to-date
-  (module.registerHooks as RegisterHookFunction)({
+  (mod.registerHooks as RegisterHookFunction)({
     load(url, context, nextLoad) {
       const result = nextLoad(url, context);
       return onModuleLoad(url, context, result);

library/agent/hooks/instrumentation/injectedFunctions.ts

@@ -1,3 +1,7 @@
-export function __instrumentInspectArgs(name: string, args: unknown[]): void {
+export function __instrumentInspectArgs(
+  id: string,
+  isBuiltin: boolean,
+  args: unknown[]
+): void {
   // Todo do something
 }

library/agent/hooks/instrumentation/instructions.ts

@@ -0,0 +1,126 @@
+import { Package } from "../Package";
+import { BuiltinModule } from "../BuiltinModule";
+import {
+  BuiltinInstrumentationInstructionJSON,
+  IntereptorFunctionsObj,
+  PackageFileInstrumentationInstructionJSON,
+} from "./types";
+import { satisfiesVersion } from "../../../helpers/satisfiesVersion";
+
+// Keys are package / builtin names
+let packages = new Map<string, PackageFileInstrumentationInstructionJSON[]>();
+let builtins = new Map<string, BuiltinInstrumentationInstructionJSON>();
+
+// Stores the callbacks for the instrumented functions of builtin modules
+// Identifier for builtin is: moduleName.functionName
+let builtinCallbacks = new Map<string, IntereptorFunctionsObj>();
+// Stores the callbacks for the instrumented functions of package files
+// Identifier for package file is: packageName.relativePath.functionName.matchingVersion
+let packageFileCallbacks = new Map<string, IntereptorFunctionsObj>();
+
+export function setPackagesToInstrument(_packages: Package[]) {
+  // Clear the previous packages
+  packages = new Map();
+  packageFileCallbacks = new Map();
+
+  for (const pkg of _packages) {
+    const packageInstructions = pkg
+      .getVersions()
+      .map((versionedPackage) => {
+        return versionedPackage
+          .getFileInstrumentationInstructions()
+          .map((file) => {
+            return {
+              path: file.path,
+              versionRange: versionedPackage.getRange(),
+              functions: file.functions.map((func) => {
+                const identifier = `${pkg.getName()}.${file.path}.${func.name}.${versionedPackage.getRange()}`;
+
+                return {
+                  nodeType: func.nodeType,
+                  name: func.name,
+                  identifier,
+                  inspectArgs: !!func.inspectArgs,
+                  modifyArgs: !!func.modifyArgs,
+                  modifyReturnValue: !!func.modifyReturnValue,
+                };
+              }),
+            };
+          })
+          .flat();
+      })
+      .flat();
+
+    packages.set(pkg.getName(), packageInstructions);
+  }
+}
+
+export function setBuiltinsToInstrument(builtinModules: BuiltinModule[]) {
+  // Clear the previous builtins
+  builtins = new Map();
+  builtinCallbacks = new Map();
+
+  for (const builtin of builtinModules) {
+    const instructions = builtin.getInstrumentationInstruction();
+
+    if (
+      !instructions ||
+      !instructions.functions ||
+      instructions.functions.length === 0
+    ) {
+      continue;
+    }
+
+    // Check if function is included twice
+    const functionNames = new Set<string>();
+    for (const f of instructions.functions) {
+      if (functionNames.has(f.name)) {
+        throw new Error(
+          `Function ${f.name} is included twice in the instrumentation instructions for ${builtin.getName()}`
+        );
+      }
+      functionNames.add(f.name);
+    }
+
+    const functions = instructions.functions.map((f) => {
+      builtinCallbacks.set(`${builtin.getName()}.${f.name}`, {
+        inspectArgs: f.inspectArgs,
+        modifyArgs: f.modifyArgs,
+        modifyReturnValue: f.modifyReturnValue,
+      });
+
+      return {
+        name: f.name,
+        inspectArgs: !!f.inspectArgs,
+        modifyArgs: !!f.modifyArgs,
+        modifyReturnValue: !!f.modifyReturnValue,
+      };
+    });
+
+    builtins.set(builtin.getName(), {
+      functions,
+    });
+  }
+}
+
+export function getBuiltinInstrumentationInstructions(
+  name: string
+): BuiltinInstrumentationInstructionJSON | undefined {
+  return builtins.get(name);
+}
+
+export function shouldPatchPackage(name: string): boolean {
+  return packages.has(name);
+}
+
+export function getPackageFileInstrumentationInstructions(
+  packageName: string,
+  version: string
+): PackageFileInstrumentationInstructionJSON | undefined {
+  const instructions = packages.get(packageName);
+  if (!instructions) {
+    return;
+  }
+
+  return instructions.find((f) => satisfiesVersion(f.versionRange, version));
+}