Support accessing module local vars

+ Add SQLite3 support

Author: timokoessler

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

@@ -230,3 +230,47 @@ fn get_insert_pos(body: &FunctionBody, is_constructor: bool) -> usize {
         0 // No super call found, insert at the beginning
     }
 }
+
+// Add a statement to the end of the body: __instrumentAccessLocalVariables('identifier', [var1, var2]);
+pub fn insert_access_local_var<'a>(
+    allocator: &'a Allocator,
+    builder: &'a AstBuilder,
+    identifier: &str,
+    var_names: &Vec<String>,
+    body: &mut OxcVec<'a, Statement<'a>>,
+) {
+    let mut instrument_args: OxcVec<'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,
+    )));
+
+    // [var1, var2]
+    let mut array_elements: OxcVec<'a, ArrayExpressionElement<'a>> =
+        builder.vec_with_capacity(var_names.len());
+    for name in var_names {
+        array_elements.push(ArrayExpressionElement::Identifier(
+            builder.alloc_identifier_reference(SPAN, allocator.alloc_str(name)),
+        ));
+    }
+
+    instrument_args.push(Argument::ArrayExpression(
+        builder.alloc_array_expression(SPAN, array_elements),
+    ));
+
+    // Build and add a call expression
+    let call_expr = builder.expression_call(
+        SPAN,
+        builder.expression_identifier(SPAN, "__instrumentAccessLocalVariables"),
+        NONE,
+        instrument_args,
+        false,
+    );
+
+    let stmt_expression = builder.statement_expression(SPAN, call_expr);
+
+    body.push(stmt_expression);
+}

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

@@ -14,6 +14,7 @@ const INSTRUMENT_IMPORT_SOURCE: &str = "@aikidosec/firewall/instrument/internals
 const INSTRUMENT_INSPECT_ARGS_METHOD_NAME: &str = "__instrumentInspectArgs";
 const INSTRUMENT_MODIFY_ARGS_METHOD_NAME: &str = "__instrumentModifyArgs";
 const INSTRUMENT_MODIFY_RETURN_VALUE_METHOD_NAME: &str = "__instrumentModifyReturnValue";
+const INSTRUMENT_ACCESS_LOCAL_VARS_METHOD_NAME: &str = "__instrumentAccessLocalVariables";
 
 type ImportMethodPredicate = fn(&FunctionInstructions) -> bool;
 
@@ -78,6 +79,28 @@ pub fn insert_import_statement<'a>(
             }
         }
 
+        if !file_instructions.access_local_variables.is_empty() {
+            binding_properties.push(
+                builder.binding_property(
+                    SPAN,
+                    builder.property_key_static_identifier(
+                        SPAN,
+                        INSTRUMENT_ACCESS_LOCAL_VARS_METHOD_NAME,
+                    ),
+                    builder.binding_pattern(
+                        builder.binding_pattern_kind_binding_identifier(
+                            SPAN,
+                            INSTRUMENT_ACCESS_LOCAL_VARS_METHOD_NAME,
+                        ),
+                        NONE,
+                        false,
+                    ),
+                    true,
+                    false,
+                ),
+            );
+        }
+
         if binding_properties.is_empty() {
             // If there are no methods to import, we can skip the require statement
             return;
@@ -131,6 +154,20 @@ pub fn insert_import_statement<'a>(
         }
     }
 
+    if !file_instructions.access_local_variables.is_empty() {
+        specifiers.push(
+            builder.import_declaration_specifier_import_specifier(
+                SPAN,
+                builder.module_export_name_identifier_name(
+                    SPAN,
+                    INSTRUMENT_ACCESS_LOCAL_VARS_METHOD_NAME,
+                ),
+                builder.binding_identifier(SPAN, INSTRUMENT_ACCESS_LOCAL_VARS_METHOD_NAME),
+                ImportOrExportKind::Value,
+            ),
+        );
+    }
+
     if specifiers.is_empty() {
         // If there are no methods to import, we can skip the import statement
         return;

instrumentation-wasm/src/js_transformer/instructions.rs

@@ -4,8 +4,10 @@ use serde::{Deserialize, Serialize};
 #[serde(rename_all = "camelCase")]
 pub struct FileInstructions {
     pub path: String,
+    pub identifier: String,
     pub version_range: String,
     pub functions: Vec<FunctionInstructions>,
+    pub access_local_variables: Vec<String>,
 }
 
 #[derive(Serialize, Deserialize)]

instrumentation-wasm/src/js_transformer/transformer.rs

@@ -7,7 +7,7 @@ use oxc_traverse::traverse_mut;
 
 use super::{
     helpers::{
-        insert_import_statement::insert_import_statement,
+        insert_code::insert_access_local_var, insert_import_statement::insert_import_statement,
         select_sourcetype_based_on_enum::select_sourcetype_based_on_enum,
     },
     instructions::FileInstructions,
@@ -74,6 +74,16 @@ pub fn transform_code_str(
         &file_instructions,
     );
 
+    if !file_instructions.access_local_variables.is_empty() {
+        insert_access_local_var(
+            &allocator,
+            &ast_builder,
+            &file_instructions.identifier,
+            &file_instructions.access_local_variables,
+            &mut program.body,
+        );
+    }
+
     // Todo: Update source map?
     let js = Codegen::new()
         .with_options(CodegenOptions {

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

@@ -41,6 +41,7 @@ t.test("Benchmark: Small code transformation", async (t) => {
       {
         path: "test.js",
         versionRange: "^1.0.0",
+        identifier: "testpkg.test.js.^1.0.0",
         functions: [
           {
             nodeType: "MethodDefinition",
@@ -52,6 +53,7 @@ t.test("Benchmark: Small code transformation", async (t) => {
             modifyArgumentsObject: false,
           },
         ],
+        accessLocalVariables: [],
       }
     );
   }
@@ -82,7 +84,8 @@ t.test("Benchmark: Large code transformation", async (t) => {
 
   for (let i = 0; i < iterations; i++) {
     transformCode("hono", "4.0.0", "dist/hono-base.js", honoSrc, "module", {
-      path: "dist/hono-base.js-----",
+      path: "dist/hono-base.js",
+      identifier: "hono.dist.hono-base.js.v4.0.0",
       versionRange: "^4.0.0",
       functions: [
         {
@@ -95,6 +98,7 @@ t.test("Benchmark: Large code transformation", async (t) => {
           modifyArgumentsObject: false,
         },
       ],
+      accessLocalVariables: [],
     });
   }