Implement return value modification

Author: timokoessler

instrumentation-wasm/Cargo.lock

@@ -10,13 +10,13 @@ name = "node_code_instrumentation"
 crate-type = ["cdylib", "rlib"]
 
 [dependencies]
-oxc_allocator = "0.72.0"
-oxc_ast = "0.72.0"
-oxc_codegen = "0.72.0"
-oxc_parser = "0.72.0"
-oxc_semantic = "0.72.0"
-oxc_span = "0.72.0"
-oxc_traverse = "0.72.0"
+oxc_allocator = "0.72.2"
+oxc_ast = "0.72.2"
+oxc_codegen = "0.72.2"
+oxc_parser = "0.72.2"
+oxc_semantic = "0.72.2"
+oxc_span = "0.72.2"
+oxc_traverse = "0.72.2"
 serde = "1.0.219"
 serde_json = "1.0.140"
 wasm-bindgen = "0.2.100"

instrumentation-wasm/Cargo.toml

@@ -3,8 +3,8 @@ use oxc_span::SourceType;
 pub fn get_import_code_str(source_type: &SourceType, has_module_syntax: bool) -> &'static str {
     // Todo maybe import only needed functions, not sure if it is worth it
     if source_type.is_script() || source_type.is_unambiguous() && !has_module_syntax {
-        return "const { __instrumentInspectArgs, __instrumentModifyArgs } = require('@aikidosec/firewall/instrument/internals');";
+        return "const { __instrumentInspectArgs, __instrumentModifyArgs, __instrumentModifyReturnValue } = require('@aikidosec/firewall/instrument/internals');";
     }
 
-    "import { __instrumentInspectArgs, __instrumentModifyArgs } from '@aikidosec/firewall/instrument/internals';"
+    "import { __instrumentInspectArgs, __instrumentModifyArgs, __instrumentModifyReturnValue } from '@aikidosec/firewall/instrument/internals';"
 }

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

@@ -4,3 +4,4 @@ pub mod get_name_str_for_member_expr;
 pub mod insert_code;
 pub mod parse_js_code_to_statements;
 pub mod select_sourcetype_based_on_enum;
+pub mod transform_return_statements;

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

@@ -3,6 +3,7 @@ use oxc_ast::ast::Statement;
 use oxc_parser::{ParseOptions, Parser};
 use oxc_span::SourceType;
 
+// Todo replace with AstBuilder in the future, most likely better performance
 pub fn parse_js_code_to_statements<'a>(
     allocator: &'a Allocator,
     source_text: &'a str,

instrumentation-wasm/src/js_transformer/helpers/parse_js_code_to_statements.rs

@@ -0,0 +1,51 @@
+use oxc_allocator::{Allocator, Box, Vec};
+use oxc_ast::{
+    ast::{Argument, FunctionBody, Statement},
+    AstBuilder, NONE,
+};
+use oxc_span::SPAN;
+
+pub fn transform_return_statements<'a>(
+    allocator: &'a Allocator,
+    identifier: &str,
+    body: &mut Box<'a, FunctionBody<'a>>,
+) {
+    // Iterate through the statements in the function body
+    for statement in &mut body.statements {
+        let builder = AstBuilder::new(&allocator);
+
+        match statement {
+            Statement::ReturnStatement(return_stmt) => {
+                let mut instrument_args: Vec<'a, Argument<'a>> = builder.vec_with_capacity(2);
+
+                // Add the identifier to the arguments
+                instrument_args.push(Argument::StringLiteral(builder.alloc_string_literal(
+                    SPAN,
+                    allocator.alloc_str(identifier),
+                    None,
+                )));
+                // Add original return value as the second argument
+                let arg_expr = return_stmt
+                    .argument
+                    .take()
+                    .unwrap_or_else(|| builder.expression_identifier(SPAN, "undefined"));
+                instrument_args.push(arg_expr.into());
+
+                let new_call_expr = builder.expression_call(
+                    SPAN,
+                    builder.expression_identifier(SPAN, "__instrumentModifyReturnValue"),
+                    NONE,
+                    instrument_args,
+                    false,
+                );
+
+                // Replace the return statement with a call to the identifier
+                return_stmt.argument = Some(new_call_expr);
+            }
+            _ => {
+                // Ignore
+                continue;
+            }
+        }
+    }
+}

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

@@ -7,6 +7,7 @@ use super::helpers::{
     get_arg_names::get_method_arg_names,
     get_name_str_for_member_expr::get_name_str_for_member_expr,
     insert_code::{insert_inspect_args, insert_modify_args},
+    transform_return_statements::transform_return_statements,
 };
 use super::instructions::FileInstructions;
 
@@ -74,7 +75,9 @@ impl<'a> Traverse<'a> for Transformer<'a> {
             );
         }
 
-        // Todo support return value modification
+        if instruction.modify_return_value {
+            transform_return_statements(self.allocator, &instruction.identifier, body);
+        }
     }
 
     fn enter_assignment_expression(
@@ -148,6 +151,8 @@ impl<'a> Traverse<'a> for Transformer<'a> {
             );
         }
 
-        // Todo support return value modification
+        if instruction.modify_return_value {
+            transform_return_statements(self.allocator, &instruction.identifier, body);
+        }
     }
 }