Add support for argument mutation in complex argument functions

src/instrumentation.rs

@@ -8,7 +8,7 @@ use swc_core::common::{Span, SyntaxContext};
 use swc_core::ecma::{
     ast::{
         ArrowExpr, AssignExpr, AssignTarget, BlockStmt, ClassDecl, ClassMethod, Constructor, Expr,
-        FnDecl, FnExpr, Ident, Lit, MemberProp, MethodProp, Module, ModuleItem, Pat, PropName,
+        FnDecl, FnExpr, Ident, Lit, MemberProp, MethodProp, Module, ModuleItem, Param, Pat, PropName,
         Script, SimpleAssignTarget, Stmt, Str, VarDecl,
     },
     atoms::Atom,
@@ -49,9 +49,9 @@ impl Instrumentation {
         self.is_correct_class = false;
     }
 
-    fn new_fn(&self, body: BlockStmt) -> ArrowExpr {
+    fn new_fn(&self, body: BlockStmt, params: Vec<Pat>) -> ArrowExpr {
         ArrowExpr {
-            params: vec![],
+            params,
             body: Box::new(body.into()),
             is_async: self.config.function_query.kind().is_async(),
             is_generator: false,
@@ -81,7 +81,7 @@ impl Instrumentation {
         define_channel
     }
 
-    fn insert_tracing(&mut self, body: &mut BlockStmt) {
+    fn insert_tracing(&mut self, body: &mut BlockStmt, params: Vec<Param>) {
         self.count += 1;
 
         let original_stmts = std::mem::take(&mut body.stmts);
@@ -93,7 +93,20 @@ impl Instrumentation {
             ..body.clone()
         };
 
-        let traced_fn = self.new_fn(original_body);
+        let original_params: Vec<Pat> = params.into_iter().map(|p| p.pat).collect();
+
+        let wrapped_fn = self.new_fn(original_body, original_params);
+
+        let traced_body = BlockStmt {
+            span: Span::default(),
+            ctxt: SyntaxContext::empty(),
+            stmts: vec![
+                quote!("const __apm$wrapped = $wrapped;" as Stmt, wrapped: Expr = wrapped_fn.into()),
+                quote!("return __apm$wrapped.apply(null, __apm$original_args);" as Stmt),
+            ],
+        };
+
+        let traced_fn = self.new_fn(traced_body, vec![]);
 
         let ch_ident = ident!(format!("tr_ch_apm${}", &self.config.channel_name));
         let trace_ident = ident!(format!(
@@ -103,6 +116,7 @@ impl Instrumentation {
         ));
 
         body.stmts = vec![
+            quote!("const __apm$original_args = arguments" as Stmt),
             quote!("const __apm$traced = $traced;" as Stmt, traced: Expr = traced_fn.into()),
             quote!(
                 "if (!$ch.hasSubscribers) return __apm$traced();" as Stmt,
@@ -168,7 +182,7 @@ impl Instrumentation {
             && func_expr.function.body.is_some()
         {
             if let Some(body) = func_expr.function.body.as_mut() {
-                self.insert_tracing(body);
+                self.insert_tracing(body, func_expr.function.params.clone());
             }
             true
         } else {
@@ -206,7 +220,7 @@ impl Instrumentation {
             && node.function.body.is_some()
         {
             if let Some(body) = node.function.body.as_mut() {
-                self.insert_tracing(body);
+                self.insert_tracing(body, node.function.params.clone());
             }
         }
         false
@@ -253,7 +267,7 @@ impl Instrumentation {
             && node.function.body.is_some()
         {
             if let Some(body) = node.function.body.as_mut() {
-                self.insert_tracing(body);
+                self.insert_tracing(body, node.function.params.clone());
             }
         }
         true
@@ -286,7 +300,7 @@ impl Instrumentation {
             && node.function.body.is_some()
         {
             if let Some(body) = node.function.body.as_mut() {
-                self.insert_tracing(body);
+                self.insert_tracing(body, node.function.params.clone());
             }
         }
         false

tests/arguments_mutation/mod.js

@@ -0,0 +1,30 @@
+/**
+ * Unless explicitly stated otherwise all files in this repository are licensed under the Apache-2.0 License.
+ * This product includes software developed at Datadog (https://www.datadoghq.com/). Copyright 2025 Datadog, Inc.
+ **/
+const assert = require('assert');
+
+function fetch_simple (url, cb) {
+  assert.strictEqual(this.this, 'this');
+  assert.strictEqual(url, 'https://example.com');
+  assert.strictEqual(cb.length, 2);
+  const result = cb.apply(this, ['arg1', 'arg2']);
+  assert.strictEqual(result, 'result');
+  return 'return';
+}
+
+function fetch_complex ({ url, tuple: [a = 'a', b = 'b'] }, cb, optional = 'default', ...rest) {
+  assert.strictEqual(this.this, 'this');
+  assert.strictEqual(url, 'https://example.com');
+  assert.strictEqual(a, 'a');
+  assert.strictEqual(b, 'b');
+  assert.strictEqual(cb.length, 2);
+  assert.strictEqual(optional, 'default');
+  assert.deepStrictEqual(rest, []);
+  const result = cb.apply(this, ['arg1', 'arg2']);
+  assert.strictEqual(result, 'result');
+  return 'return';
+}
+
+
+module.exports = { fetch_simple, fetch_complex };

tests/arguments_mutation/mod.rs

@@ -0,0 +1,25 @@
+use crate::common::*;
+use orchestrion_js::*;
+
+#[test]
+fn arguments_mutation() {
+    transpile_and_test(
+        file!(),
+        false,
+        Config::new(
+            vec![
+                InstrumentationConfig::new(
+                    "fetch_simple",
+                    test_module_matcher(),
+                    FunctionQuery::function_declaration("fetch_simple", FunctionKind::Sync),
+                ),
+                InstrumentationConfig::new(
+                    "fetch_complex",
+                    test_module_matcher(),
+                    FunctionQuery::function_declaration("fetch_complex", FunctionKind::Sync),
+                ),
+            ],
+            None,
+        ),
+    );
+}

tests/arguments_mutation/test.js

@@ -0,0 +1,38 @@
+/**
+ * Unless explicitly stated otherwise all files in this repository are licensed under the Apache-2.0 License.
+ * This product includes software developed at Datadog (https://www.datadoghq.com/). Copyright 2025 Datadog, Inc.
+ **/
+const { fetch_simple, fetch_complex } = require('./instrumented.js');
+const assert = require('assert');
+const { tracingChannel } = require('diagnostics_channel');
+
+const handler = {
+  start (message) {
+    const originalCb = message.arguments[1];
+    const wrappedCb = function (a, b) {
+      assert.strictEqual(this.this, 'this');
+      assert.strictEqual(a, 'arg1');
+      assert.strictEqual(b, 'arg2');
+      arguments[1] = 'arg2_mutated';
+      return originalCb.apply(this, arguments);
+    }
+
+    message.arguments[1] = wrappedCb;
+  }
+};
+
+tracingChannel('orchestrion:undici:fetch_simple').subscribe(handler);
+tracingChannel('orchestrion:undici:fetch_complex').subscribe(handler);
+
+assert.strictEqual(fetch_simple.length, 2);
+assert.strictEqual(fetch_complex.length, 2);
+
+const cb = function (a, b) {
+  assert.strictEqual(this.this, 'this');
+  assert.strictEqual(a, 'arg1');
+  assert.strictEqual(b, 'arg2_mutated');
+  return 'result';
+};
+
+assert.strictEqual(fetch_simple.apply({ this: 'this' }, ['https://example.com', cb]), 'return');
+assert.strictEqual(fetch_complex.apply({ this: 'this' }, [{ url: 'https://example.com', tuple: [] }, cb]), 'return');

tests/instrumentor_test.rs

@@ -4,6 +4,7 @@
  **/
 mod common;
 
+mod arguments_mutation;
 mod class_method_cjs;
 mod constructor_cjs;
 mod constructor_mjs;