Instrument Mysql

Support FunctionAssignment with sub properties

Author: timokoessler

instrumentation-wasm/src/js_transformer/helpers/get_name_str_for_member_expr.rs

@@ -1,4 +1,4 @@
-use oxc_allocator::Allocator;
+use oxc_allocator::{Allocator, CloneIn};
 use oxc_ast::ast::{Expression, MemberExpression};
 
 pub fn get_name_str_for_member_expr<'a>(
@@ -26,6 +26,10 @@ pub fn get_name_str_for_member_expr<'a>(
 
     let obj_name_str = match member_expr.object() {
         Expression::Identifier(identifier_ref) => identifier_ref.name.as_str(),
+        Expression::StaticMemberExpression(static_member_expr) => get_name_str_for_member_expr(
+            allocator,
+            &MemberExpression::StaticMemberExpression(static_member_expr.clone_in(allocator)),
+        )?,
         _ => {
             // Unsupported AST type
             return None;

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

@@ -1253,3 +1253,50 @@ t.test("Modify function declaration (CJS)", async (t) => {
     };`
   );
 });
+
+t.test(
+  "add inspectArgs to multi level function assignment expression (CJS)",
+  async (t) => {
+    const result = transformCode(
+      "express",
+      "1.0.0",
+      "application.js",
+      `
+        const app = require("example");
+        app.prototype.use = function (fn) {
+            console.log("test");
+        };
+        `,
+      "commonjs",
+      {
+        path: "application.js",
+        versionRange: "^1.0.0",
+        functions: [
+          {
+            nodeType: "FunctionAssignment",
+            name: "app.prototype.use",
+            identifier:
+              "express.application.js.app.prototype.use.MethodDefinition.v1.0.0",
+            inspectArgs: true,
+            modifyArgs: false,
+            modifyReturnValue: false,
+            modifyArgumentsObject: false,
+          },
+        ],
+      }
+    );
+
+    t.same(
+      compareCodeStrings(
+        result,
+        `const { __instrumentInspectArgs } = require("@aikidosec/firewall/instrument/internals");
+        const app = require("example");
+        app.prototype.use = function (fn) {
+            __instrumentInspectArgs("express.application.js.app.prototype.use.MethodDefinition.v1.0.0", arguments, "1.0.0", this);
+            console.log("test");
+        };`
+      ),
+      true
+    );
+  }
+);

library/sinks/MySQL.ts

@@ -57,6 +57,17 @@ export class MySQL implements Wrapper {
           kind: "sql_op",
           inspectArgs: (args) => this.inspectQuery(args),
         });
+      })
+      .addFileInstrumentation({
+        path: "lib/Connection.js",
+        functions: [
+          {
+            name: "Connection.prototype.query",
+            nodeType: "FunctionAssignment",
+            operationKind: "sql_op",
+            inspectArgs: (args) => this.inspectQuery(args),
+          },
+        ],
       });
   }
 }