Skip to content

Commit 96a85e8

Browse files
timokoesslertimokoessler
committed
Refactor builtin interceptors (broken WIP)
1 parent d930443 commit 96a85e8

File tree

11 files changed

+108
-172
lines changed

11 files changed

+108
-172
lines changed

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -132,7 +132,7 @@ impl<'a> Traverse<'a> for Transformer<'a> {
132132

133133
if instruction.inspect_args {
134134
let source_text: &'a str = self.allocator.alloc_str(&format!(
135-
"__instrumentInspectArgs('{}', false, arguments);",
135+
"__instrumentInspectArgs('{}', arguments);",
136136
instruction.identifier
137137
));
138138

‎library/agent/hooks/BuiltinModule.ts‎

Lines changed: 0 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,7 @@
1-
import { BuiltinInstrumentationInstruction } from "./instrumentation/types";
21
import { RequireInterceptor } from "./RequireInterceptor";
32

43
export class BuiltinModule {
54
private requireInterceptors: RequireInterceptor[] = [];
6-
private instrumentationInstructions:
7-
| BuiltinInstrumentationInstruction
8-
| undefined;
95

106
constructor(private readonly name: string) {
117
if (!this.name) {
@@ -24,14 +20,4 @@ export class BuiltinModule {
2420
getRequireInterceptors() {
2521
return this.requireInterceptors;
2622
}
27-
28-
setInstrumentationInstruction(
29-
instruction: BuiltinInstrumentationInstruction
30-
) {
31-
this.instrumentationInstructions = instruction;
32-
}
33-
34-
getInstrumentationInstruction() {
35-
return this.instrumentationInstructions;
36-
}
3723
}

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

Lines changed: 17 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ import * as t from "tap";
22
import { generateBuildinShim } from "./builtinShim";
33

44
t.test("generate fs/promises shim", async (t) => {
5+
// Todo update
56
const shim = generateBuildinShim("fs/promises", "fs/promises", [
67
{
78
name: "readFile",
@@ -10,22 +11,30 @@ t.test("generate fs/promises shim", async (t) => {
1011
modifyReturnValue: false,
1112
},
1213
]);
14+
if (!shim) {
15+
t.fail("shim is undefined");
16+
return;
17+
}
1318

1419
t.match(
15-
shim?.replace(/\s+/g, " "),
20+
shim.replace(/\s+/g, " "),
1621
`const orig = process.getBuiltinModule("fs/promises");
1722
const { __instrumentInspectArgs } = require('@aikidosec/firewall/instrument/internals');
1823
19-
exports.readFile = function() {
20-
__instrumentInspectArgs("fs/promises.readFile", true, arguments);
24+
orig.readFile = function() {
25+
__instrumentInspectArgs("fs/promises.readFile", arguments);
2126
return orig.readFile(...arguments);
2227
};
28+
29+
exports = orig;
2330
`.replace(/\s+/g, " ")
2431
);
2532

26-
t.match(shim, "exports.rename = orig.rename;");
27-
t.match(shim, "exports.rmdir = orig.rmdir;");
28-
t.match(shim, "exports.rm = orig.rm;");
29-
t.match(shim, "exports.stat = orig.stat;");
30-
t.match(shim, "exports.symlink = orig.symlink;");
33+
const modifiedShim = shim.replace(
34+
"@aikidosec/firewall/instrument/internals",
35+
"./injectedFunctions"
36+
);
37+
38+
let modifiedExports = eval(modifiedShim);
39+
console.log(modifiedExports);
3140
});

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

Lines changed: 10 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -1,40 +1,22 @@
1-
import { getExportsForBuiltin } from "./getExportsForBuiltin";
2-
import { BuiltinInstrumentationInstructionJSON } from "./types";
3-
41
export function generateBuildinShim(
52
builtinName: string,
63
builtinNameWithoutPrefix: string,
7-
instructions: BuiltinInstrumentationInstructionJSON["functions"]
4+
format: "commonjs" | "module"
85
): string | undefined {
9-
const exports = getExportsForBuiltin(builtinName);
6+
// Todo fix broken with ESM
7+
// Todo access unwrapped methods inside our own library to prevent infinite recursion
108

11-
// Filter out non-existing exports
12-
const methods = instructions.filter((m) => exports.has(m.name));
9+
if (format === "module") {
10+
return `const orig = process.getBuiltinModule(${JSON.stringify(builtinName)});
11+
import { __wrapBuiltinExports } from '@aikidosec/firewall/instrument/internals';
1312
14-
if (methods.length === 0) {
15-
// No methods to instrument, return undefined
16-
return;
13+
export default __wrapBuiltinExports("${builtinNameWithoutPrefix}", orig);
14+
`;
1715
}
1816

19-
const unpatchedExports = Array.from(exports).filter(
20-
(m: any) => !methods.some((method) => method.name === m.name)
21-
);
22-
23-
// Todos: Copy over properties of patched methods?
24-
// Todo check default export !!!
2517
return `const orig = process.getBuiltinModule(${JSON.stringify(builtinName)});
26-
const { __instrumentInspectArgs } = require('@aikidosec/firewall/instrument/internals');
27-
28-
${methods
29-
.map(
30-
(method) => `
31-
exports.${method.name} = function() {
32-
${method.inspectArgs ? `__instrumentInspectArgs("${builtinNameWithoutPrefix}.${method.name}", true, arguments);` : ""}
33-
return orig.${method.name}(...arguments);
34-
};`
35-
)
36-
.join("\n")}
18+
const { __wrapBuiltinExports } = require('@aikidosec/firewall/instrument/internals');
3719
38-
${unpatchedExports.map((e) => `exports.${e} = orig.${e};`).join("\n")}
20+
module.exports = __wrapBuiltinExports("${builtinNameWithoutPrefix}", orig);
3921
`;
4022
}

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

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,6 @@
11
import { onModuleLoad } from "./loadHook";
22
import * as mod from "node:module";
33
import type { RegisterHookFunction } from "./types";
4-
import { Hooks } from "../Hooks";
54

65
export function registerNodeHooks() {
76
if (!("registerHooks" in mod) || typeof mod.registerHooks !== "function") {

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

Lines changed: 38 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -1,31 +1,52 @@
11
import { getInstance } from "../../AgentSingleton";
2-
import { getBuiltinCallbacks, getPackageCallbacks } from "./instructions";
2+
import { WrapPackageInfo } from "../WrapPackageInfo";
3+
import { getBuiltinInterceptors, getPackageCallbacks } from "./instructions";
34

4-
export function __instrumentInspectArgs(
5-
id: string,
6-
isBuiltin: boolean,
7-
args: unknown[]
8-
): void {
5+
export function __instrumentInspectArgs(id: string, args: unknown[]): void {
96
const agent = getInstance();
107
if (!agent) {
118
return;
129
}
1310

14-
if (isBuiltin) {
15-
const callbacks = getBuiltinCallbacks(id);
16-
17-
if (typeof callbacks.inspectArgs === "function") {
18-
// Todo support subject?
19-
callbacks.inspectArgs(args, agent, undefined);
20-
}
21-
22-
return;
23-
}
24-
2511
const cbFuncs = getPackageCallbacks(id);
2612

2713
if (typeof cbFuncs.inspectArgs === "function") {
2814
// Todo support subject?
2915
cbFuncs.inspectArgs(args, agent, undefined);
3016
}
3117
}
18+
19+
export function __wrapBuiltinExports(id: string, exports: unknown): unknown {
20+
const agent = getInstance();
21+
if (!agent) {
22+
return exports;
23+
}
24+
25+
const interceptors = getBuiltinInterceptors(id);
26+
27+
if (interceptors.length === 0) {
28+
return exports;
29+
}
30+
31+
const pkgInfo: WrapPackageInfo = {
32+
name: id,
33+
type: "builtin",
34+
};
35+
36+
// Todo check if cache is needed
37+
for (const interceptor of interceptors) {
38+
try {
39+
const returnVal = interceptor(exports, pkgInfo);
40+
// If the interceptor returns a value, we want to use this value as the new exports
41+
if (typeof returnVal !== "undefined") {
42+
exports = returnVal;
43+
}
44+
} catch (error) {
45+
if (error instanceof Error) {
46+
getInstance()?.onFailedToWrapModule(pkgInfo.name, error);
47+
}
48+
}
49+
}
50+
51+
return exports;
52+
}

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

Lines changed: 12 additions & 52 deletions
Original file line numberDiff line numberDiff line change
@@ -1,19 +1,18 @@
11
import { Package } from "../Package";
22
import { BuiltinModule } from "../BuiltinModule";
33
import {
4-
BuiltinInstrumentationInstructionJSON,
54
IntereptorFunctionsObj,
65
PackageFileInstrumentationInstructionJSON,
76
} from "./types";
87
import { satisfiesVersion } from "../../../helpers/satisfiesVersion";
8+
import { RequireInterceptor } from "../RequireInterceptor";
99

1010
// Keys are package / builtin names
1111
let packages = new Map<string, PackageFileInstrumentationInstructionJSON[]>();
12-
let builtins = new Map<string, BuiltinInstrumentationInstructionJSON>();
1312

1413
// Stores the callbacks for the instrumented functions of builtin modules
1514
// Identifier for builtin is: moduleName.functionName
16-
let builtinCallbacks = new Map<string, IntereptorFunctionsObj>();
15+
let builtinRequireInterceptors = new Map<string, RequireInterceptor[]>();
1716
// Stores the callbacks for the instrumented functions of package files
1817
// Identifier for package file is: packageName.relativePath.functionName.matchingVersion
1918
let packageFileCallbacks = new Map<string, IntereptorFunctionsObj>();
@@ -65,58 +64,17 @@ export function setPackagesToInstrument(_packages: Package[]) {
6564

6665
export function setBuiltinsToInstrument(builtinModules: BuiltinModule[]) {
6766
// Clear the previous builtins
68-
builtins = new Map();
69-
builtinCallbacks = new Map();
67+
builtinRequireInterceptors = new Map();
7068

7169
for (const builtin of builtinModules) {
72-
const instructions = builtin.getInstrumentationInstruction();
73-
74-
if (
75-
!instructions ||
76-
!instructions.functions ||
77-
instructions.functions.length === 0
78-
) {
79-
continue;
80-
}
70+
const interceptors = builtin.getRequireInterceptors();
8171

82-
// Check if function is included twice
83-
const functionNames = new Set<string>();
84-
for (const f of instructions.functions) {
85-
if (functionNames.has(f.name)) {
86-
throw new Error(
87-
`Function ${f.name} is included twice in the instrumentation instructions for ${builtin.getName()}`
88-
);
89-
}
90-
functionNames.add(f.name);
72+
if (interceptors.length > 0) {
73+
builtinRequireInterceptors.set(builtin.getName(), interceptors);
9174
}
92-
93-
const functions = instructions.functions.map((f) => {
94-
builtinCallbacks.set(`${builtin.getName()}.${f.name}`, {
95-
inspectArgs: f.inspectArgs,
96-
modifyArgs: f.modifyArgs,
97-
modifyReturnValue: f.modifyReturnValue,
98-
});
99-
100-
return {
101-
name: f.name,
102-
inspectArgs: !!f.inspectArgs,
103-
modifyArgs: !!f.modifyArgs,
104-
modifyReturnValue: !!f.modifyReturnValue,
105-
};
106-
});
107-
108-
builtins.set(builtin.getName(), {
109-
functions,
110-
});
11175
}
11276
}
11377

114-
export function getBuiltinInstrumentationInstructions(
115-
name: string
116-
): BuiltinInstrumentationInstructionJSON | undefined {
117-
return builtins.get(name);
118-
}
119-
12078
export function shouldPatchPackage(name: string): boolean {
12179
return packages.has(name);
12280
}
@@ -139,8 +97,10 @@ export function getPackageCallbacks(
13997
return packageFileCallbacks.get(identifier) || {};
14098
}
14199

142-
export function getBuiltinCallbacks(
143-
identifier: string
144-
): IntereptorFunctionsObj {
145-
return builtinCallbacks.get(identifier) || {};
100+
export function getBuiltinInterceptors(name: string): RequireInterceptor[] {
101+
return builtinRequireInterceptors.get(name) || [];
102+
}
103+
104+
export function shouldPatchBuiltin(name: string): boolean {
105+
return builtinRequireInterceptors.has(name);
146106
}

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

Lines changed: 14 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -6,8 +6,8 @@ import { getPackageVersionFromPath } from "./getPackageVersionFromPath";
66
import { transformCode } from "./codeTransformation";
77
import { generateBuildinShim } from "./builtinShim";
88
import {
9-
getBuiltinInstrumentationInstructions,
109
getPackageFileInstrumentationInstructions,
10+
shouldPatchBuiltin,
1111
shouldPatchPackage,
1212
} from "./instructions";
1313
import { removeNodePrefix } from "../../../helpers/removeNodePrefix";
@@ -41,7 +41,7 @@ export function onModuleLoad(
4141

4242
// For Node.js builtin modules
4343
if (isBuiltin) {
44-
return patchBuiltin(path, previousLoadResult, isModifiedBuiltin);
44+
return patchBuiltin(path, previousLoadResult, context, isModifiedBuiltin);
4545
}
4646

4747
return patchPackage(path, previousLoadResult);
@@ -123,6 +123,7 @@ function patchPackage(
123123
function patchBuiltin(
124124
builtinName: string,
125125
previousLoadResult: ReturnType<LoadFunction>,
126+
context: Parameters<LoadFunction>[1],
126127
isAlreadyModified: boolean
127128
) {
128129
if (isAlreadyModified) {
@@ -132,29 +133,31 @@ function patchBuiltin(
132133

133134
const builtinNameWithoutPrefix = removeNodePrefix(builtinName);
134135

135-
const builtin = getBuiltinInstrumentationInstructions(
136-
builtinNameWithoutPrefix
137-
);
136+
const builtin = shouldPatchBuiltin(builtinNameWithoutPrefix);
138137
if (!builtin) {
139138
return previousLoadResult;
140139
}
141140

142-
if (!builtin.functions || builtin.functions.length === 0) {
143-
// We don't want to modify this module
144-
return previousLoadResult;
145-
}
141+
const isCJSRequire =
142+
(Array.isArray(context.conditions) &&
143+
context.conditions.includes("require")) ||
144+
("has" in context.conditions &&
145+
typeof context.conditions.has === "function" &&
146+
context.conditions.has("require"));
147+
148+
const format = isCJSRequire ? "commonjs" : "module";
146149

147150
const shim = generateBuildinShim(
148151
builtinName,
149152
builtinNameWithoutPrefix,
150-
builtin.functions
153+
format
151154
);
152155
if (!shim) {
153156
return previousLoadResult;
154157
}
155158

156159
return {
157-
format: "commonjs",
160+
format: format,
158161
shortCircuit: previousLoadResult.shortCircuit,
159162
source: shim,
160163
};

0 commit comments

Comments
 (0)