Skip to content

Commit 352e624

Browse files
timokoesslertimokoessler
committed
Constructor modification for PubSub
+ Fix MySQL2
1 parent 157ceb5 commit 352e624

File tree

7 files changed

+247
-38
lines changed

7 files changed

+247
-38
lines changed

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

Lines changed: 34 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@ use oxc_ast::{
33
AstBuilder, NONE,
44
ast::{
55
Argument, ArrayExpressionElement, AssignmentOperator, AssignmentTarget, Expression,
6-
FunctionBody,
6+
FunctionBody, Statement,
77
},
88
};
99
use oxc_span::SPAN;
@@ -15,6 +15,7 @@ pub fn insert_inspect_args<'a>(
1515
identifier: &str,
1616
pkg_version: &'a str,
1717
body: &mut Box<'a, FunctionBody<'a>>,
18+
is_constructor: bool,
1819
) {
1920
let mut inspect_args: OxcVec<'a, Argument<'a>> = builder.vec_with_capacity(4);
2021

@@ -49,7 +50,8 @@ pub fn insert_inspect_args<'a>(
4950

5051
let stmt_expression = builder.statement_expression(SPAN, call_expr);
5152

52-
body.statements.insert(0, stmt_expression);
53+
let insert_pos = get_insert_pos(body, is_constructor);
54+
body.statements.insert(insert_pos, stmt_expression);
5355
}
5456

5557
// Modify the arguments by adding a statement to the beginning of the function
@@ -61,6 +63,7 @@ pub fn insert_modify_args<'a>(
6163
arg_names: &Vec<String>,
6264
body: &mut Box<'a, FunctionBody<'a>>,
6365
modify_arguments_object: bool,
66+
is_constructor: bool,
6467
) {
6568
if modify_arguments_object {
6669
// If we are modifying the arguments object, we need to use the arguments object directly
@@ -129,8 +132,8 @@ pub fn insert_modify_args<'a>(
129132

130133
let stmt_expression = builder.statement_expression(SPAN, obj_assign_call_expr);
131134

132-
body.statements.insert(0, stmt_expression);
133-
135+
let insert_pos = get_insert_pos(body, is_constructor);
136+
body.statements.insert(insert_pos, stmt_expression);
134137
return;
135138
}
136139

@@ -199,5 +202,31 @@ pub fn insert_modify_args<'a>(
199202

200203
let stmt_expression = builder.statement_expression(SPAN, arr_assignment_expr);
201204

202-
body.statements.insert(0, stmt_expression);
205+
let insert_pos = get_insert_pos(body, is_constructor);
206+
body.statements.insert(insert_pos, stmt_expression);
207+
}
208+
209+
// Determine the position to insert the code in the function body.
210+
// If it's a constructor, we look for the first super call and insert after it.
211+
fn get_insert_pos(body: &FunctionBody, is_constructor: bool) -> usize {
212+
if !is_constructor || body.statements.is_empty() {
213+
0
214+
} else {
215+
let mut index = 0;
216+
for statement in &body.statements {
217+
match statement {
218+
Statement::ExpressionStatement(expr_stmt) => {
219+
if let Expression::CallExpression(call_expr) = &expr_stmt.expression {
220+
if let Expression::Super(_) = &call_expr.callee {
221+
// Found a super call, insert after this statement
222+
return index + 1; // Insert after the super call
223+
}
224+
}
225+
}
226+
_ => {}
227+
};
228+
index += 1;
229+
}
230+
0 // No super call found, insert at the beginning
231+
}
203232
}

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

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@ pub fn insert_instrument_method_calls<'a>(
1616
arg_names: &Vec<String>,
1717
pkg_version: &'a str,
1818
body: &mut Box<'a, FunctionBody<'a>>,
19+
is_constructor: bool,
1920
) {
2021
if instruction.modify_args && !arg_names.is_empty() {
2122
insert_modify_args(
@@ -25,6 +26,7 @@ pub fn insert_instrument_method_calls<'a>(
2526
arg_names,
2627
body,
2728
instruction.modify_arguments_object,
29+
is_constructor,
2830
);
2931
}
3032

@@ -35,10 +37,11 @@ pub fn insert_instrument_method_calls<'a>(
3537
&instruction.identifier,
3638
pkg_version,
3739
body,
40+
is_constructor,
3841
);
3942
}
4043

41-
if instruction.modify_return_value {
44+
if instruction.modify_return_value && !is_constructor {
4245
transform_return_statements(
4346
allocator,
4447
builder,

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

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -28,8 +28,8 @@ impl<'a> Traverse<'a, TraverseState> for Transformer<'a> {
2828
return;
2929
}
3030

31-
if !node.kind.is_method() {
32-
// Ignore constructor, getters and setters for now
31+
if !node.kind.is_method() && !node.kind.is_constructor() {
32+
// Ignore getters and setters for now
3333
return;
3434
}
3535

@@ -63,6 +63,7 @@ impl<'a> Traverse<'a, TraverseState> for Transformer<'a> {
6363
&arg_names,
6464
self.pkg_version,
6565
body,
66+
node.kind.is_constructor(),
6667
);
6768
}
6869

@@ -124,6 +125,7 @@ impl<'a> Traverse<'a, TraverseState> for Transformer<'a> {
124125
&arg_names,
125126
self.pkg_version,
126127
body,
128+
false,
127129
);
128130
}
129131

@@ -178,6 +180,7 @@ impl<'a> Traverse<'a, TraverseState> for Transformer<'a> {
178180
&arg_names,
179181
self.pkg_version,
180182
body,
183+
false,
181184
);
182185
}
183186
}

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

Lines changed: 141 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1622,3 +1622,144 @@ t.test("Modify function expression (ESM)", async (t) => {
16221622
}`
16231623
);
16241624
});
1625+
1626+
t.test("Modify constructor (ESM)", async (t) => {
1627+
const result = transformCode(
1628+
"testpkg",
1629+
"1.0.0",
1630+
"application.js",
1631+
`
1632+
class Test {
1633+
constructor() {
1634+
console.log("ignore");
1635+
const x = 1;
1636+
}
1637+
}
1638+
`,
1639+
"module",
1640+
{
1641+
path: "application.js",
1642+
versionRange: "^1.0.0",
1643+
functions: [
1644+
{
1645+
nodeType: "MethodDefinition",
1646+
name: "constructor",
1647+
identifier:
1648+
"testpkg.application.js.constructor.MethodDefinition.v1.0.0",
1649+
inspectArgs: true,
1650+
modifyArgs: false,
1651+
modifyReturnValue: false,
1652+
modifyArgumentsObject: false,
1653+
},
1654+
],
1655+
}
1656+
);
1657+
1658+
isSameCode(
1659+
result,
1660+
`import { __instrumentInspectArgs } from "@aikidosec/firewall/instrument/internals";
1661+
class Test {
1662+
constructor() {
1663+
__instrumentInspectArgs("testpkg.application.js.constructor.MethodDefinition.v1.0.0", arguments, "1.0.0", this);
1664+
console.log("ignore");
1665+
const x = 1;
1666+
}
1667+
}
1668+
`
1669+
);
1670+
});
1671+
1672+
t.test("Modify constructor with super (ESM)", async (t) => {
1673+
const result = transformCode(
1674+
"testpkg",
1675+
"1.0.0",
1676+
"application.js",
1677+
`
1678+
class Test {
1679+
constructor() {
1680+
super();
1681+
console.log("ignore");
1682+
}
1683+
}
1684+
`,
1685+
"module",
1686+
{
1687+
path: "application.js",
1688+
versionRange: "^1.0.0",
1689+
functions: [
1690+
{
1691+
nodeType: "MethodDefinition",
1692+
name: "constructor",
1693+
identifier:
1694+
"testpkg.application.js.constructor.MethodDefinition.v1.0.0",
1695+
inspectArgs: true,
1696+
modifyArgs: false,
1697+
modifyReturnValue: false,
1698+
modifyArgumentsObject: false,
1699+
},
1700+
],
1701+
}
1702+
);
1703+
1704+
isSameCode(
1705+
result,
1706+
`import { __instrumentInspectArgs } from "@aikidosec/firewall/instrument/internals";
1707+
class Test {
1708+
constructor() {
1709+
super();
1710+
__instrumentInspectArgs("testpkg.application.js.constructor.MethodDefinition.v1.0.0", arguments, "1.0.0", this);
1711+
console.log("ignore");
1712+
}
1713+
}
1714+
`
1715+
);
1716+
});
1717+
1718+
t.test("Modify constructor with super (ESM)", async (t) => {
1719+
const result = transformCode(
1720+
"testpkg",
1721+
"1.0.0",
1722+
"application.js",
1723+
`
1724+
class Test {
1725+
constructor(arg1) {
1726+
test();
1727+
super();
1728+
console.log("ignore");
1729+
}
1730+
}
1731+
`,
1732+
"module",
1733+
{
1734+
path: "application.js",
1735+
versionRange: "^1.0.0",
1736+
functions: [
1737+
{
1738+
nodeType: "MethodDefinition",
1739+
name: "constructor",
1740+
identifier:
1741+
"testpkg.application.js.constructor.MethodDefinition.v1.0.0",
1742+
inspectArgs: true,
1743+
modifyArgs: true,
1744+
modifyReturnValue: true,
1745+
modifyArgumentsObject: false,
1746+
},
1747+
],
1748+
}
1749+
);
1750+
1751+
isSameCode(
1752+
result,
1753+
`import { __instrumentInspectArgs, __instrumentModifyArgs, __instrumentModifyReturnValue } from "@aikidosec/firewall/instrument/internals";
1754+
class Test {
1755+
constructor(arg1) {
1756+
test();
1757+
super();
1758+
__instrumentInspectArgs("testpkg.application.js.constructor.MethodDefinition.v1.0.0", arguments, "1.0.0", this);
1759+
[arg1] = __instrumentModifyArgs("testpkg.application.js.constructor.MethodDefinition.v1.0.0", [arg1], this);
1760+
console.log("ignore");
1761+
}
1762+
}
1763+
`
1764+
);
1765+
});

‎library/sinks/MySQL2.ts‎

Lines changed: 24 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
import { getContext } from "../agent/Context";
22
import { Hooks } from "../agent/hooks/Hooks";
3+
import { PackageFunctionInstrumentationInstruction } from "../agent/hooks/instrumentation/types";
34
import { InterceptorResult } from "../agent/hooks/InterceptorResult";
45
import { wrapExport } from "../agent/hooks/wrapExport";
56
import { WrapPackageInfo } from "../agent/hooks/WrapPackageInfo";
@@ -72,6 +73,23 @@ export class MySQL2 implements Wrapper {
7273
return connectionPrototype;
7374
}
7475

76+
private getFunctionInstructions(): PackageFunctionInstrumentationInstruction[] {
77+
return [
78+
{
79+
nodeType: "MethodDefinition",
80+
name: "query",
81+
inspectArgs: (args) => this.inspectQuery("mysql2.query", args),
82+
operationKind: "sql_op",
83+
},
84+
{
85+
nodeType: "MethodDefinition",
86+
name: "execute",
87+
inspectArgs: (args) => this.inspectQuery("mysql2.execute", args),
88+
operationKind: "sql_op",
89+
},
90+
];
91+
}
92+
7593
wrap(hooks: Hooks) {
7694
const wrapConnection = (
7795
exports: any,
@@ -103,7 +121,11 @@ export class MySQL2 implements Wrapper {
103121
// For all versions of mysql2 newer than 3.0.0
104122
pkg
105123
.withVersion("^3.0.0")
106-
.onRequire((exports, pkgInfo) => wrapConnection(exports, pkgInfo, false));
124+
.onRequire((exports, pkgInfo) => wrapConnection(exports, pkgInfo, false))
125+
.addFileInstrumentation({
126+
path: "lib/connection.js",
127+
functions: this.getFunctionInstructions(),
128+
});
107129

108130
// For all versions of mysql2 newer than / equal 3.11.5
109131
// Reason: https://github.com/sidorares/node-mysql2/pull/3081
@@ -114,20 +136,7 @@ export class MySQL2 implements Wrapper {
114136
})
115137
.addFileInstrumentation({
116138
path: "lib/base/connection.js",
117-
functions: [
118-
{
119-
nodeType: "MethodDefinition",
120-
name: "query",
121-
inspectArgs: (args) => this.inspectQuery("mysql2.query", args),
122-
operationKind: "sql_op",
123-
},
124-
{
125-
nodeType: "MethodDefinition",
126-
name: "execute",
127-
inspectArgs: (args) => this.inspectQuery("mysql2.execute", args),
128-
operationKind: "sql_op",
129-
},
130-
],
139+
functions: this.getFunctionInstructions(),
131140
});
132141
}
133142
}

0 commit comments

Comments
 (0)