Support FunctionExpression

Author: timokoessler

instrumentation-wasm/src/js_transformer/transformer_impl.rs

@@ -132,10 +132,14 @@ impl<'a> Traverse<'a, TraverseState> for Transformer<'a> {
         node: &mut oxc_ast::ast::Function<'a>,
         _ctx: &mut TraverseCtx<'a, TraverseState>,
     ) {
-        if node.r#type != FunctionType::FunctionDeclaration {
-            // Only instrument function declarations here
-            return;
-        }
+        let node_type = match node.r#type {
+            FunctionType::FunctionDeclaration => "FunctionDeclaration",
+            FunctionType::FunctionExpression => "FunctionExpression",
+            _ => {
+                // Ignore unknown function types
+                return;
+            }
+        };
 
         if node.id.is_none() || node.body.is_none() {
             // No identifier or body, nothing to instrument
@@ -150,7 +154,7 @@ impl<'a> Traverse<'a, TraverseState> for Transformer<'a> {
             .file_instructions
             .functions
             .iter()
-            .find(|f| f.node_type == "FunctionDeclaration" && f.name == function_name);
+            .find(|f| f.node_type == node_type && f.name == function_name);
 
         if matching_instruction.is_none() {
             return;

library/agent/hooks/instrumentation/codeTransformation.test.ts

@@ -1568,3 +1568,57 @@ t.test(
     );
   }
 );
+
+t.test("Modify function expression (ESM)", async (t) => {
+  const result = transformCode(
+    "testpkg",
+    "1.0.0",
+    "application.js",
+    `
+      const x = 1;
+
+      const y = function test123(arg1, arg2 = "default") {
+        console.log("test123");
+
+        return "abc";
+      }
+
+      function test123() {
+        console.log("ignore");
+      }
+    `,
+    "module",
+    {
+      path: "application.js",
+      versionRange: "^1.0.0",
+      functions: [
+        {
+          nodeType: "FunctionExpression",
+          name: "test123",
+          identifier:
+            "testpkg.application.js.test123.FunctionExpression.v1.0.0",
+          inspectArgs: true,
+          modifyArgs: false,
+          modifyReturnValue: true,
+          modifyArgumentsObject: false,
+        },
+      ],
+    }
+  );
+
+  isSameCode(
+    result,
+    `import { __instrumentInspectArgs, __instrumentModifyReturnValue } from "@aikidosec/firewall/instrument/internals";
+    const x = 1;
+    
+    const y = function test123(arg1, arg2 = "default") {
+        __instrumentInspectArgs("testpkg.application.js.test123.FunctionExpression.v1.0.0", arguments, "1.0.0", this);
+        console.log("test123");
+        return __instrumentModifyReturnValue("testpkg.application.js.test123.FunctionExpression.v1.0.0", arguments, "abc", this);
+    };
+    
+    function test123() {
+        console.log("ignore");
+    }`
+  );
+});

library/agent/hooks/instrumentation/types.ts

@@ -68,7 +68,11 @@ export type IntereptorCallbackInfoObj = {
 };
 
 export type PackageFunctionInstrumentationInstruction = {
-  nodeType: "MethodDefinition" | "FunctionAssignment" | "FunctionDeclaration";
+  nodeType:
+    | "MethodDefinition"
+    | "FunctionAssignment"
+    | "FunctionDeclaration"
+    | "FunctionExpression";
   name: string;
   /**
    * Used for stats reporting to core, e.g. sql_op