Skip to content

Commit e974373

Browse files
timokoesslertimokoessler
committed
Use AstBuilder for imports
1 parent 6ebb39e commit e974373

File tree

7 files changed

+260
-48
lines changed

7 files changed

+260
-48
lines changed

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

Lines changed: 0 additions & 10 deletions
This file was deleted.

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

Lines changed: 177 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,177 @@
1+
use oxc_allocator::{Allocator, Vec};
2+
use oxc_ast::{
3+
ast::{
4+
Argument, BindingPatternKind, BindingProperty, ImportDeclarationSpecifier,
5+
ImportOrExportKind, Statement, VariableDeclarationKind,
6+
},
7+
AstBuilder, NONE,
8+
};
9+
use oxc_span::{SourceType, SPAN};
10+
11+
use crate::js_transformer::instructions::FileInstructions;
12+
13+
const INSTRUMENT_IMPORT_SOURCE: &str = "@aikidosec/firewall/instrument/internals";
14+
const INSTRUMENT_INSPECT_ARGS_METHOD_NAME: &str = "__instrumentInspectArgs";
15+
const INSTRUMENT_MODIFY_ARGS_METHOD_NAME: &str = "__instrumentModifyArgs";
16+
const INSTRUMENT_MODIFY_RETURN_VALUE_METHOD_NAME: &str = "__instrumentModifyReturnValue";
17+
18+
pub fn insert_import_statement<'a>(
19+
source_type: &SourceType,
20+
has_module_syntax: bool,
21+
allocator: &'a Allocator,
22+
builder: &'a AstBuilder,
23+
body: &mut Vec<'a, Statement<'a>>,
24+
file_instructions: &FileInstructions,
25+
) {
26+
// Common JS require() statement
27+
28+
if source_type.is_script() || source_type.is_unambiguous() && !has_module_syntax {
29+
let mut require_args: Vec<'a, Argument<'a>> = builder.vec_with_capacity(1);
30+
require_args.push(Argument::StringLiteral(builder.alloc_string_literal(
31+
SPAN,
32+
allocator.alloc_str(INSTRUMENT_IMPORT_SOURCE),
33+
None,
34+
)));
35+
36+
let mut binding_properties: Vec<'a, BindingProperty<'a>> = builder.vec_with_capacity(3);
37+
38+
if file_instructions.functions.iter().any(|f| f.inspect_args) {
39+
binding_properties.push(builder.binding_property(
40+
SPAN,
41+
builder.property_key_static_identifier(SPAN, INSTRUMENT_INSPECT_ARGS_METHOD_NAME),
42+
builder.binding_pattern(
43+
builder.binding_pattern_kind_binding_identifier(
44+
SPAN,
45+
INSTRUMENT_INSPECT_ARGS_METHOD_NAME,
46+
),
47+
NONE,
48+
false,
49+
),
50+
true,
51+
false,
52+
));
53+
}
54+
55+
if file_instructions.functions.iter().any(|f| f.modify_args) {
56+
binding_properties.push(builder.binding_property(
57+
SPAN,
58+
builder.property_key_static_identifier(SPAN, INSTRUMENT_MODIFY_ARGS_METHOD_NAME),
59+
builder.binding_pattern(
60+
builder.binding_pattern_kind_binding_identifier(
61+
SPAN,
62+
INSTRUMENT_MODIFY_ARGS_METHOD_NAME,
63+
),
64+
NONE,
65+
false,
66+
),
67+
true,
68+
false,
69+
));
70+
}
71+
72+
if file_instructions
73+
.functions
74+
.iter()
75+
.any(|f| f.modify_return_value)
76+
{
77+
binding_properties.push(builder.binding_property(
78+
SPAN,
79+
builder.property_key_static_identifier(
80+
SPAN,
81+
INSTRUMENT_MODIFY_RETURN_VALUE_METHOD_NAME,
82+
),
83+
builder.binding_pattern(
84+
builder.binding_pattern_kind_binding_identifier(
85+
SPAN,
86+
INSTRUMENT_MODIFY_RETURN_VALUE_METHOD_NAME,
87+
),
88+
NONE,
89+
false,
90+
),
91+
true,
92+
false,
93+
));
94+
}
95+
96+
let mut declarations = builder.vec_with_capacity(1);
97+
declarations.push(builder.variable_declarator(
98+
SPAN,
99+
VariableDeclarationKind::Const,
100+
builder.binding_pattern(
101+
BindingPatternKind::ObjectPattern(builder.alloc_object_pattern(
102+
SPAN,
103+
binding_properties,
104+
NONE,
105+
)),
106+
NONE,
107+
false,
108+
),
109+
Some(builder.expression_call(
110+
SPAN,
111+
builder.expression_identifier(SPAN, "require"),
112+
NONE,
113+
require_args,
114+
false,
115+
)),
116+
false,
117+
));
118+
119+
let var_declaration = Statement::VariableDeclaration(builder.alloc_variable_declaration(
120+
SPAN,
121+
VariableDeclarationKind::Const,
122+
declarations,
123+
false,
124+
));
125+
126+
body.insert(0, var_declaration);
127+
128+
return;
129+
}
130+
131+
let mut specifiers: Vec<'a, ImportDeclarationSpecifier<'a>> = builder.vec_with_capacity(3);
132+
133+
if file_instructions.functions.iter().any(|f| f.inspect_args) {
134+
specifiers.push(builder.import_declaration_specifier_import_specifier(
135+
SPAN,
136+
builder.module_export_name_identifier_name(SPAN, INSTRUMENT_INSPECT_ARGS_METHOD_NAME),
137+
builder.binding_identifier(SPAN, INSTRUMENT_INSPECT_ARGS_METHOD_NAME),
138+
ImportOrExportKind::Value,
139+
));
140+
}
141+
142+
if file_instructions.functions.iter().any(|f| f.modify_args) {
143+
specifiers.push(builder.import_declaration_specifier_import_specifier(
144+
SPAN,
145+
builder.module_export_name_identifier_name(SPAN, INSTRUMENT_MODIFY_ARGS_METHOD_NAME),
146+
builder.binding_identifier(SPAN, INSTRUMENT_MODIFY_ARGS_METHOD_NAME),
147+
ImportOrExportKind::Value,
148+
));
149+
}
150+
151+
if file_instructions
152+
.functions
153+
.iter()
154+
.any(|f| f.modify_return_value)
155+
{
156+
specifiers.push(builder.import_declaration_specifier_import_specifier(
157+
SPAN,
158+
builder.module_export_name_identifier_name(
159+
SPAN,
160+
INSTRUMENT_MODIFY_RETURN_VALUE_METHOD_NAME,
161+
),
162+
builder.binding_identifier(SPAN, INSTRUMENT_MODIFY_RETURN_VALUE_METHOD_NAME),
163+
ImportOrExportKind::Value,
164+
));
165+
}
166+
167+
let import_stmt = Statement::ImportDeclaration(builder.alloc_import_declaration(
168+
SPAN,
169+
Some(specifiers),
170+
builder.string_literal(SPAN, allocator.alloc_str(INSTRUMENT_IMPORT_SOURCE), None),
171+
None,
172+
NONE,
173+
ImportOrExportKind::Value,
174+
));
175+
176+
body.insert(0, import_stmt);
177+
}

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
pub mod get_arg_names;
2-
pub mod get_import_code_str;
32
pub mod get_name_str_for_member_expr;
43
pub mod insert_code;
4+
pub mod insert_import_statement;
55
pub mod parse_js_code_to_statements;
66
pub mod select_sourcetype_based_on_enum;
77
pub mod transform_return_statements;

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

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -7,13 +7,12 @@ use oxc_span::SPAN;
77

88
pub fn transform_return_statements<'a>(
99
allocator: &'a Allocator,
10+
builder: &AstBuilder<'a>,
1011
identifier: &str,
1112
body: &mut Box<'a, FunctionBody<'a>>,
1213
) {
1314
// Iterate through the statements in the function body
1415
for statement in &mut body.statements {
15-
let builder = AstBuilder::new(&allocator);
16-
1716
match statement {
1817
Statement::ReturnStatement(return_stmt) => {
1918
let mut instrument_args: Vec<'a, Argument<'a>> = builder.vec_with_capacity(2);

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

Lines changed: 12 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
use oxc_allocator::Allocator;
2+
use oxc_ast::AstBuilder;
23
use oxc_codegen::{Codegen, CodegenOptions};
34
use oxc_parser::Parser;
45
use oxc_semantic::SemanticBuilder;
@@ -7,7 +8,7 @@ use oxc_traverse::traverse_mut;
78

89
use super::{
910
helpers::{
10-
get_import_code_str::get_import_code_str,
11+
insert_import_statement::insert_import_statement,
1112
parse_js_code_to_statements::parse_js_code_to_statements,
1213
select_sourcetype_based_on_enum::select_sourcetype_based_on_enum,
1314
},
@@ -46,24 +47,25 @@ pub fn transform_code_str(
4647

4748
let (scopes, _nodes) = semantic.semantic.into_scoping_and_nodes();
4849

50+
let ast_builder = AstBuilder::new(&allocator);
51+
4952
let t = &mut Transformer {
5053
allocator: &allocator,
5154
file_instructions: &file_instructions,
5255
pkg_version: &pkg_version,
56+
ast_builder: &ast_builder,
5357
};
5458

5559
traverse_mut(t, &allocator, program, scopes);
5660

5761
// Add import / require statement
58-
program.body.insert(
59-
0,
60-
parse_js_code_to_statements(
61-
&allocator,
62-
get_import_code_str(&source_type, parser_result.module_record.has_module_syntax),
63-
SourceType::cjs(),
64-
)
65-
.pop()
66-
.unwrap(),
62+
insert_import_statement(
63+
&source_type,
64+
parser_result.module_record.has_module_syntax,
65+
&allocator,
66+
&ast_builder,
67+
&mut program.body,
68+
&file_instructions,
6769
);
6870

6971
// Todo: Update source map?

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

Lines changed: 13 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@ pub struct Transformer<'a> {
1515
pub allocator: &'a Allocator,
1616
pub file_instructions: &'a FileInstructions,
1717
pub pkg_version: &'a str,
18+
pub ast_builder: &'a oxc_ast::AstBuilder<'a>,
1819
}
1920

2021
impl<'a> Traverse<'a> for Transformer<'a> {
@@ -76,7 +77,12 @@ impl<'a> Traverse<'a> for Transformer<'a> {
7677
}
7778

7879
if instruction.modify_return_value {
79-
transform_return_statements(self.allocator, &instruction.identifier, body);
80+
transform_return_statements(
81+
self.allocator,
82+
self.ast_builder,
83+
&instruction.identifier,
84+
body,
85+
);
8086
}
8187
}
8288

@@ -152,7 +158,12 @@ impl<'a> Traverse<'a> for Transformer<'a> {
152158
}
153159

154160
if instruction.modify_return_value {
155-
transform_return_statements(self.allocator, &instruction.identifier, body);
161+
transform_return_statements(
162+
self.allocator,
163+
self.ast_builder,
164+
&instruction.identifier,
165+
body,
166+
);
156167
}
157168
}
158169
}

0 commit comments

Comments
 (0)