Skip to content

Commit d1643ce

Browse files
timokoesslertimokoessler
committed
Support argument modification
1 parent ddcf9cd commit d1643ce

File tree

10 files changed

+336
-42
lines changed

10 files changed

+336
-42
lines changed

‎instrumentation-wasm/src/js_transformer/helpers/get_import_code_str.rs‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,10 @@
1+
use oxc_span::SourceType;
2+
3+
pub fn get_import_code_str(source_type: &SourceType, has_module_syntax: bool) -> &'static str {
4+
// Todo maybe import only needed functions, not sure if it is worth it
5+
if source_type.is_script() || source_type.is_unambiguous() && !has_module_syntax {
6+
return "const { __instrumentInspectArgs, __instrumentModifyArgs } = require('@aikidosec/firewall/instrument/internals');";
7+
}
8+
9+
"import { __instrumentInspectArgs, __instrumentModifyArgs } from '@aikidosec/firewall/instrument/internals';"
10+
}

‎instrumentation-wasm/src/js_transformer/helpers/get_method_arg_names.rs‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,18 @@
1+
use oxc_ast::ast::MethodDefinition;
2+
3+
pub fn get_method_arg_names(method_definition: &MethodDefinition) -> Vec<String> {
4+
let mut arg_names = Vec::new();
5+
6+
// Todo check whats appens if rest parameter is used
7+
8+
method_definition
9+
.value
10+
.params
11+
.items
12+
.iter()
13+
.for_each(|param| {
14+
arg_names.push(param.pattern.get_identifier_name().unwrap().to_string());
15+
});
16+
17+
arg_names
18+
}

‎instrumentation-wasm/src/js_transformer/helpers/mod.rs‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,2 +1,4 @@
1+
pub mod get_import_code_str;
2+
pub mod get_method_arg_names;
13
pub mod parse_js_code_to_statements;
24
pub mod select_sourcetype_based_on_enum;

‎instrumentation-wasm/src/js_transformer/transformer.rs‎

Lines changed: 46 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ use oxc_traverse::{traverse_mut, Traverse, TraverseCtx};
88

99
use super::{
1010
helpers::{
11+
get_import_code_str::get_import_code_str, get_method_arg_names::get_method_arg_names,
1112
parse_js_code_to_statements::parse_js_code_to_statements,
1213
select_sourcetype_based_on_enum::select_sourcetype_based_on_enum,
1314
},
@@ -54,33 +55,17 @@ pub fn transform_code_str(code: &str, instructions_json: &str, src_type: i32) ->
5455

5556
traverse_mut(t, &allocator, program, symbols, scopes);
5657

57-
if source_type.is_script()
58-
|| source_type.is_unambiguous() && parser_result.module_record.has_module_syntax == false
59-
{
60-
// Add require statement
61-
program.body.insert(
62-
0,
63-
parse_js_code_to_statements(
64-
&allocator,
65-
"const { __instrumentInspectArgs } = require('@aikidosec/firewall/instrument/internals');",
66-
SourceType::cjs(),
67-
)
68-
.pop()
69-
.unwrap(),
70-
);
71-
} else {
72-
// Add import statement
73-
program.body.insert(
74-
0,
75-
parse_js_code_to_statements(
76-
&allocator,
77-
"import { __instrumentInspectArgs } from '@aikidosec/firewall/instrument/internals';",
78-
SourceType::mjs(),
79-
)
80-
.pop()
81-
.unwrap(),
82-
);
83-
}
58+
// Add import / require statement
59+
program.body.insert(
60+
0,
61+
parse_js_code_to_statements(
62+
&allocator,
63+
get_import_code_str(&source_type, parser_result.module_record.has_module_syntax),
64+
SourceType::cjs(),
65+
)
66+
.pop()
67+
.unwrap(),
68+
);
8469

8570
// Todo: Update source map?
8671
let js = Codegen::new()
@@ -116,6 +101,11 @@ impl<'a> Traverse<'a> for Transformer<'a> {
116101
return;
117102
}
118103

104+
if !node.kind.is_method() {
105+
// Ignore constructor, getters and setters for now
106+
return;
107+
}
108+
119109
// Todo implement submethod counting for nested functions for supporting modifications of return value
120110

121111
let method_name = node.key.name().unwrap().to_string();
@@ -138,9 +128,38 @@ impl<'a> Traverse<'a> for Transformer<'a> {
138128
self.current_function_identifier = Some(function_identifier.clone());
139129
self.modify_return_value = instruction.modify_return_value;*/
140130

131+
// We need to collect the arg names before we make the body mutable
132+
let arg_names = if instruction.modify_args {
133+
// Todo check whats appens if rest parameter is used
134+
135+
get_method_arg_names(node)
136+
} else {
137+
Vec::new()
138+
};
139+
141140
let body = node.value.body.as_mut().unwrap();
142141

142+
if instruction.modify_args {
143+
// Modify the arguments by adding a statement to the beginning of the function
144+
// [arg1, arg2, ...] = __instrumentModifyArgs('function_identifier', [arg1, arg2, ...]);
145+
146+
let arg_names_str = arg_names.join(", ");
147+
let source_text: &'a str = self.allocator.alloc_str(&format!(
148+
"[{}] = __instrumentModifyArgs('{}', [{}]);",
149+
arg_names_str, instruction.identifier, arg_names_str
150+
));
151+
152+
body.statements.insert(
153+
0,
154+
parse_js_code_to_statements(self.allocator, &source_text, SourceType::mjs())
155+
.into_iter()
156+
.next()
157+
.unwrap(),
158+
);
159+
}
160+
143161
if instruction.inspect_args {
162+
// Add a statement to the beginning of the function: __instrumentInspectArgs('function_identifier', arguments);
144163
let source_text: &'a str = self.allocator.alloc_str(&format!(
145164
"__instrumentInspectArgs('{}', arguments);",
146165
instruction.identifier

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

Lines changed: 119 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -42,7 +42,7 @@ t.test("add inspectArgs to method definition (ESM)", async (t) => {
4242
t.same(
4343
compareCodeStrings(
4444
result,
45-
`import { __instrumentInspectArgs } from "@aikidosec/firewall/instrument/internals";
45+
`import { __instrumentInspectArgs, __instrumentModifyArgs } from "@aikidosec/firewall/instrument/internals";
4646
import { test } from "test";
4747
class Test {
4848
private testValue = 42;
@@ -97,7 +97,7 @@ t.test("add inspectArgs to method definition (CJS)", async (t) => {
9797
t.same(
9898
compareCodeStrings(
9999
result,
100-
`const { __instrumentInspectArgs } = require("@aikidosec/firewall/instrument/internals");
100+
`const { __instrumentInspectArgs, __instrumentModifyArgs } = require("@aikidosec/firewall/instrument/internals");
101101
import { test } from "test";
102102
class Test {
103103
private testValue = 42;
@@ -152,7 +152,7 @@ t.test("wrong function name", async (t) => {
152152
t.same(
153153
compareCodeStrings(
154154
result,
155-
`const { __instrumentInspectArgs } = require("@aikidosec/firewall/instrument/internals");
155+
`const { __instrumentInspectArgs, __instrumentModifyArgs } = require("@aikidosec/firewall/instrument/internals");
156156
import { test } from "test";
157157
class Test {
158158
private testValue = 42;
@@ -206,7 +206,7 @@ t.test("typescript code", async (t) => {
206206
t.same(
207207
compareCodeStrings(
208208
result,
209-
`import { __instrumentInspectArgs } from "@aikidosec/firewall/instrument/internals";
209+
`import { __instrumentInspectArgs, __instrumentModifyArgs } from "@aikidosec/firewall/instrument/internals";
210210
import { test } from "test";
211211
class Test {
212212
private testValue: number = 42;
@@ -289,8 +289,122 @@ t.test("empty code", async (t) => {
289289
t.same(
290290
compareCodeStrings(
291291
result,
292-
`import { __instrumentInspectArgs } from "@aikidosec/firewall/instrument/internals";`
292+
`import { __instrumentInspectArgs, __instrumentModifyArgs } from "@aikidosec/firewall/instrument/internals";`
293293
),
294294
true
295295
);
296296
});
297+
298+
t.test("add modifyArgs to method definition (ESM)", async (t) => {
299+
const result = transformCode(
300+
"test.js",
301+
`
302+
import { test } from "test";
303+
class Test {
304+
305+
private testValue = 42;
306+
307+
constructor() {
308+
this.testFunction(testValue);
309+
}
310+
testFunction(arg1) {
311+
console.log("test");
312+
}
313+
}
314+
`,
315+
true,
316+
{
317+
path: "test.js",
318+
versionRange: "^1.0.0",
319+
functions: [
320+
{
321+
nodeType: "MethodDefinition",
322+
name: "testFunction",
323+
identifier: "testmodule.test.js.testFunction.v1.0.0",
324+
inspectArgs: false,
325+
modifyArgs: true,
326+
modifyReturnValue: false,
327+
},
328+
],
329+
}
330+
);
331+
332+
t.same(
333+
compareCodeStrings(
334+
result,
335+
`import { __instrumentInspectArgs, __instrumentModifyArgs } from "@aikidosec/firewall/instrument/internals";
336+
import { test } from "test";
337+
class Test {
338+
private testValue = 42;
339+
340+
constructor() {
341+
this.testFunction(testValue);
342+
}
343+
testFunction(arg1) {
344+
[arg1] = __instrumentModifyArgs("testmodule.test.js.testFunction.v1.0.0", [arg1]);
345+
console.log("test");
346+
}
347+
}`
348+
),
349+
true
350+
);
351+
});
352+
353+
t.test(
354+
"add modifyArgs and inspectArgs to method definition (ESM)",
355+
async (t) => {
356+
const result = transformCode(
357+
"test.js",
358+
`
359+
import { test } from "test";
360+
class Test {
361+
362+
private testValue = 42;
363+
364+
constructor() {
365+
this.testFunction(testValue);
366+
}
367+
testFunction(arg1) {
368+
console.log("test");
369+
}
370+
}
371+
`,
372+
true,
373+
{
374+
path: "test.js",
375+
versionRange: "^1.0.0",
376+
functions: [
377+
{
378+
nodeType: "MethodDefinition",
379+
name: "testFunction",
380+
identifier: "testmodule.test.js.testFunction.v1.0.0",
381+
inspectArgs: true,
382+
modifyArgs: true,
383+
modifyReturnValue: false,
384+
},
385+
],
386+
}
387+
);
388+
389+
t.same(
390+
compareCodeStrings(
391+
result,
392+
`import { __instrumentInspectArgs, __instrumentModifyArgs } from "@aikidosec/firewall/instrument/internals";
393+
import { test } from "test";
394+
class Test {
395+
private testValue = 42;
396+
397+
constructor() {
398+
this.testFunction(testValue);
399+
}
400+
testFunction(arg1) {
401+
__instrumentInspectArgs("testmodule.test.js.testFunction.v1.0.0", arguments);
402+
[arg1] = __instrumentModifyArgs("testmodule.test.js.testFunction.v1.0.0", [arg1]);
403+
console.log("test");
404+
}
405+
}`
406+
),
407+
true
408+
);
409+
}
410+
);

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

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3,8 +3,6 @@ import type { PackageFileInstrumentationInstructionJSON } from "./types";
33
import { wasm_transform_code_str } from "./wasm/node_code_instrumentation";
44
import { getSourceType } from "./getSourceType";
55

6-
// Todo check if caching is done by Node or if we need to cache the result
7-
86
export function transformCode(
97
path: string,
108
code: string,

‎library/agent/hooks/instrumentation/injectedFunctions.ts‎

Lines changed: 33 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -7,12 +7,16 @@ export function __instrumentInspectArgs(id: string, args: unknown[]): void {
77
if (!agent) {
88
return;
99
}
10+
try {
11+
const cbFuncs = getPackageCallbacks(id);
1012

11-
const cbFuncs = getPackageCallbacks(id);
12-
13-
if (typeof cbFuncs.inspectArgs === "function") {
14-
// Todo support subject?
15-
cbFuncs.inspectArgs(args, agent, undefined);
13+
if (typeof cbFuncs.inspectArgs === "function") {
14+
// Todo support subject?
15+
cbFuncs.inspectArgs(args, agent, undefined);
16+
}
17+
} catch (error) {
18+
// Do not crash the application if an error occurs
19+
console.error(error); // We don't have a logger yet :(
1620
}
1721
}
1822

@@ -50,3 +54,27 @@ export function __wrapBuiltinExports(id: string, exports: unknown): unknown {
5054

5155
return exports;
5256
}
57+
58+
export function __instrumentModifyArgs(id: string, args: unknown[]): unknown[] {
59+
const agent = getInstance();
60+
if (!agent) {
61+
return args;
62+
}
63+
64+
try {
65+
const cbFuncs = getPackageCallbacks(id);
66+
67+
if (typeof cbFuncs.modifyArgs === "function") {
68+
const newArgs = cbFuncs.modifyArgs(args, agent);
69+
// Only return the new arguments if they are an array
70+
if (Array.isArray(newArgs)) {
71+
return newArgs;
72+
}
73+
}
74+
} catch (error) {
75+
// Do not crash the application if an error occurs
76+
console.error(error); // We don't have a logger yet :(
77+
}
78+
79+
return args;
80+
}

0 commit comments

Comments
 (0)