Skip to content

Commit 0da9a90

Browse files
timokoesslertimokoessler
committed
Add unit tests, fix multiple bugs
1 parent 6b0bdfe commit 0da9a90

File tree

10 files changed

+371
-106
lines changed

10 files changed

+371
-106
lines changed

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

Lines changed: 28 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -14,12 +14,7 @@ use super::{
1414
instructions::FileInstructions,
1515
};
1616

17-
pub fn transform_code_str(
18-
code: &str,
19-
module_name: &str,
20-
instructions_json: &str,
21-
src_type: i32,
22-
) -> String {
17+
pub fn transform_code_str(code: &str, instructions_json: &str, src_type: i32) -> String {
2318
let allocator = Allocator::default();
2419

2520
let file_instructions: FileInstructions =
@@ -54,23 +49,38 @@ pub fn transform_code_str(
5449
file_instructions: &file_instructions,
5550
//current_function_identifier: None,
5651
//modify_return_value: false,
57-
module_name,
5852
//sub_function_counter: 0,
5953
};
6054

6155
traverse_mut(t, &allocator, program, symbols, scopes);
6256

63-
// Add import statement
64-
program.body.insert(
65-
0,
66-
parse_js_code_to_statements(
67-
&allocator,
68-
"import { __instrumentInspectArgs } from '@aikidosec/firewall/instrument/internals';",
69-
SourceType::mjs(),
70-
)
71-
.pop()
72-
.unwrap(),
73-
);
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+
}
7484

7585
// Todo: Update source map?
7686
let js = Codegen::new()
@@ -94,7 +104,6 @@ struct Transformer<'a> {
94104
//current_function_identifier: Option<String>, // Only set if we want to instrument the current function
95105
//sub_function_counter: i32, // Counter to keep track of how many sub functions we are in
96106
//modify_return_value: bool,
97-
module_name: &'a str,
98107
}
99108

100109
impl<'a> Traverse<'a> for Transformer<'a> {

‎instrumentation-wasm/src/wasm_bindings/mod.rs‎

Lines changed: 2 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -3,11 +3,6 @@ use wasm_bindgen::prelude::*;
33
use crate::js_transformer::transformer::transform_code_str;
44

55
#[wasm_bindgen]
6-
pub fn wasm_transform_code_str(
7-
code: &str,
8-
module_name: &str,
9-
instructions_json: &str,
10-
source_type: i32,
11-
) -> String {
12-
return transform_code_str(code, module_name, instructions_json, source_type);
6+
pub fn wasm_transform_code_str(code: &str, instructions_json: &str, source_type: i32) -> String {
7+
return transform_code_str(code, instructions_json, source_type);
138
}

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

Lines changed: 44 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -3,53 +3,59 @@ import { generateBuildinShim } from "./builtinShim";
33
import { setBuiltinsToInstrument } from "./instructions";
44
import { BuiltinModule } from "../BuiltinModule";
55
import { createTestAgent } from "../../../helpers/createTestAgent";
6+
import { getMajorNodeVersion } from "../../../helpers/getNodeVersion";
67

7-
t.test("Wrap fs/promises", (t) => {
8-
const agent = createTestAgent();
8+
t.test(
9+
"Wrap fs/promises",
10+
{
11+
skip: getMajorNodeVersion() < 20 ? "Node.js 20+ required" : false,
12+
},
13+
(t) => {
14+
const agent = createTestAgent();
915

10-
const fsPromises = new BuiltinModule("fs/promises");
16+
const fsPromises = new BuiltinModule("fs/promises");
1117

12-
fsPromises.onRequire((exports, pkgInfo) => {
13-
exports.test = 42;
14-
exports.__aikido_is_awesome = true;
15-
});
18+
fsPromises.onRequire((exports, pkgInfo) => {
19+
exports.test = 42;
20+
exports.__aikidoIsAwesome = true;
21+
});
1622

17-
setBuiltinsToInstrument([fsPromises]);
23+
setBuiltinsToInstrument([fsPromises]);
1824

19-
const shim = generateBuildinShim("fs/promises", "fs/promises", true);
20-
if (!shim) {
21-
t.fail("shim is undefined");
22-
return;
23-
}
25+
const shim = generateBuildinShim("fs/promises", "fs/promises", true);
26+
if (!shim) {
27+
t.fail("shim is undefined");
28+
return;
29+
}
2430

25-
t.match(
26-
shim.replace(/\s+/g, " "),
27-
`const orig = process.getBuiltinModule("fs/promises");
31+
t.match(
32+
shim.replace(/\s+/g, " "),
33+
`const orig = process.getBuiltinModule("fs/promises");
2834
const { __wrapBuiltinExports } = require('@aikidosec/firewall/instrument/internals');
2935
module.exports = __wrapBuiltinExports("fs/promises", orig);
3036
`.replace(/\s+/g, " ")
31-
);
37+
);
3238

33-
const modifiedShim = shim.replace(
34-
"@aikidosec/firewall/instrument/internals",
35-
"./injectedFunctions"
36-
);
39+
const modifiedShim = shim.replace(
40+
"@aikidosec/firewall/instrument/internals",
41+
"./injectedFunctions"
42+
);
3743

38-
let modifiedExports = eval(modifiedShim);
39-
t.equal(modifiedExports.test, 42);
40-
t.equal(modifiedExports.__aikido_is_awesome, true);
41-
t.ok(typeof modifiedExports.readFile === "function");
42-
t.ok(typeof modifiedExports.writeFile === "function");
44+
const modifiedExports = eval(modifiedShim);
45+
t.equal(modifiedExports.test, 42);
46+
t.equal(modifiedExports.__aikidoIsAwesome, true);
47+
t.ok(typeof modifiedExports.readFile === "function");
48+
t.ok(typeof modifiedExports.writeFile === "function");
4349

44-
const shimESM = generateBuildinShim("fs/promises", "fs/promises", false);
45-
if (!shimESM) {
46-
t.fail("shim is undefined");
47-
return;
48-
}
50+
const shimESM = generateBuildinShim("fs/promises", "fs/promises", false);
51+
if (!shimESM) {
52+
t.fail("shim is undefined");
53+
return;
54+
}
4955

50-
t.match(
51-
shimESM.replace(/\s+/g, " "),
52-
`const orig = process.getBuiltinModule("fs/promises");
56+
t.match(
57+
shimESM.replace(/\s+/g, " "),
58+
`const orig = process.getBuiltinModule("fs/promises");
5359
const { __wrapBuiltinExports } = require('@aikidosec/firewall/instrument/internals');
5460
5561
const wrapped = __wrapBuiltinExports("fs/promises", orig);
@@ -58,7 +64,8 @@ t.test("Wrap fs/promises", (t) => {
5864
exports.default = wrapped.default;
5965
exports.access = wrapped.access;
6066
`.replace(/\s+/g, " ")
61-
);
67+
);
6268

63-
t.end();
64-
});
69+
t.end();
70+
}
71+
);

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

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -8,13 +8,11 @@ import { getSourceType } from "./getSourceType";
88
export function transformCode(
99
path: string,
1010
code: string,
11-
moduleName: string,
1211
isESM: boolean,
1312
fileInstructions: PackageFileInstrumentationInstructionJSON
1413
): string {
1514
const result = wasm_transform_code_str(
1615
code,
17-
moduleName,
1816
JSON.stringify(fileInstructions),
1917
getSourceType(path, isESM)
2018
);

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

Lines changed: 43 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -1,40 +1,47 @@
11
import * as t from "tap";
22
import { getExportsForBuiltin } from "./getExportsForBuiltin";
3+
import { getMajorNodeVersion } from "../../../helpers/getNodeVersion";
34

4-
t.test("getExportsForBuiltin works", async (t) => {
5-
t.match(
6-
getExportsForBuiltin("fs"),
7-
new Set([
8-
"default",
9-
"appendFile",
10-
"appendFileSync",
11-
"access",
12-
"accessSync",
13-
"chown",
14-
"chownSync",
15-
"chmod",
16-
"chmodSync",
17-
"close",
18-
"closeSync",
19-
"copyFile",
20-
"copyFileSync",
21-
"writeFile",
22-
"writeFileSync",
23-
"constants",
24-
"promises",
25-
"F_OK",
26-
// ...
27-
])
28-
);
5+
t.test(
6+
"getExportsForBuiltin works",
7+
{
8+
skip: getMajorNodeVersion() < 20 ? "Node.js 20+ required" : false,
9+
},
10+
async (t) => {
11+
t.match(
12+
getExportsForBuiltin("fs"),
13+
new Set([
14+
"default",
15+
"appendFile",
16+
"appendFileSync",
17+
"access",
18+
"accessSync",
19+
"chown",
20+
"chownSync",
21+
"chmod",
22+
"chmodSync",
23+
"close",
24+
"closeSync",
25+
"copyFile",
26+
"copyFileSync",
27+
"writeFile",
28+
"writeFileSync",
29+
"constants",
30+
"promises",
31+
"F_OK",
32+
// ...
33+
])
34+
);
2935

30-
t.match(
31-
getExportsForBuiltin("net"),
32-
new Set([
33-
"default",
34-
"createServer",
35-
"createConnection",
36-
"isIP",
37-
// ...
38-
])
39-
);
40-
});
36+
t.match(
37+
getExportsForBuiltin("net"),
38+
new Set([
39+
"default",
40+
"createServer",
41+
"createConnection",
42+
"isIP",
43+
// ...
44+
])
45+
);
46+
}
47+
);

0 commit comments

Comments
 (0)