detect write op in view storage - refactor: removed has_format, chang…

…ed data structure for view_endpoints, impl default for WasmInfo and ReportCreator

framework/meta-lib/src/contract/sc_config/wasm_build.rs

@@ -1,4 +1,4 @@
-use std::{ffi::OsStr, fs, process::Command};
+use std::{collections::HashMap, ffi::OsStr, fs, process::Command};
 
 use super::ContractVariant;
 use crate::{
@@ -95,12 +95,8 @@ impl ContractVariant {
         let mut abi = ContractAbiJson::from(&self.abi);
         let mut view_endpoints = Vec::new();
         for endpoint in &abi.endpoints {
-            match endpoint.mutability {
-                crate::abi_json::EndpointMutabilityAbiJson::Readonly => {
-                    view_endpoints.push(&endpoint.name)
-                },
-                crate::abi_json::EndpointMutabilityAbiJson::Mutable => (),
-                crate::abi_json::EndpointMutabilityAbiJson::Pure => (),
+            if let crate::abi_json::EndpointMutabilityAbiJson::Readonly = endpoint.mutability {
+                view_endpoints.push(&endpoint.name)
             }
         }
         let build_info = core::mem::take(&mut abi.build_info).unwrap();
@@ -141,14 +137,10 @@ impl ContractVariant {
         let output_wasm_path = format!("{output_path}/{}", self.wasm_output_name(build_args));
 
         let abi = ContractAbiJson::from(&self.abi);
-        let mut view_endpoints = Vec::new();
+        let mut view_endpoints: HashMap<&str, usize> = HashMap::new();
         for endpoint in &abi.endpoints {
-            match endpoint.mutability {
-                crate::abi_json::EndpointMutabilityAbiJson::Readonly => {
-                    view_endpoints.push(endpoint.name.clone())
-                },
-                crate::abi_json::EndpointMutabilityAbiJson::Mutable => (),
-                crate::abi_json::EndpointMutabilityAbiJson::Pure => (),
+            if let crate::abi_json::EndpointMutabilityAbiJson::Readonly = endpoint.mutability {
+                view_endpoints.insert(&endpoint.name, 0);
             }
         }
 

framework/meta-lib/src/tools/report_creator.rs

@@ -8,3 +8,13 @@ pub struct ReportCreator {
 }
 
 impl ReportCreator {}
+
+impl Default for ReportCreator {
+    fn default() -> Self {
+        ReportCreator {
+            path: String::new(),
+            has_allocator: false,
+            has_panic: "none".to_string(),
+        }
+    }
+}

framework/meta-lib/src/tools/wasm_extractor.rs

@@ -12,26 +12,29 @@ use crate::ei::EIVersion;
 
 use super::report_creator::{ReportCreator, WITHOUT_MESSAGE, WITH_MESSAGE};
 
+type CallGraph = HashMap<usize, HashSet<usize>>;
+
 const PANIC_WITH_MESSAGE: &[u8; 16] = b"panic occurred: ";
 const PANIC_WITHOUT_MESSAGE: &[u8; 14] = b"panic occurred";
 const ERROR_FAIL_ALLOCATOR: &[u8; 27] = b"memory allocation forbidden";
 const MEMORY_GROW_OPCODE: u8 = 0x40;
 const WRITE_OP: [&str; 1] = ["mBufferStorageStore"];
 
+#[derive(Default)]
 pub struct WasmInfo {
     pub imports: Vec<String>,
     pub ei_check: bool,
     pub memory_grow_flag: bool,
-    pub has_format: bool,
     pub report: ReportCreator,
+    pub call_graph: CallGraph,
 }
 
 impl WasmInfo {
     pub fn extract_wasm_info(
         output_wasm_path: &str,
         extract_imports_enabled: bool,
         check_ei: &Option<EIVersion>,
-        view_endpoints: Vec<String>,
+        view_endpoints: HashMap<&str, usize>,
     ) -> Result<WasmInfo, BinaryReaderError> {
         let wasm_data = fs::read(output_wasm_path)
             .expect("error occured while extracting information from .wasm: file not found");
@@ -51,60 +54,67 @@ fn populate_wasm_info(
     wasm_data: Vec<u8>,
     extract_imports_enabled: bool,
     check_ei: &Option<EIVersion>,
-    view_endpoints: Vec<String>,
+    mut view_endpoints: HashMap<&str, usize>,
 ) -> Result<WasmInfo, BinaryReaderError> {
-    let mut imports = Vec::new();
-    let mut allocator_trigger = false;
-    let mut ei_check = false;
-    let mut memory_grow_flag = false;
-    let mut has_panic = "none";
-    let mut call_graph: HashMap<usize, Vec<usize>> = HashMap::new();
-    let mut views_data: HashMap<usize, String> = HashMap::new();
-    let mut write_functions: Vec<usize> = Vec::new();
-
-    let mut parser = Parser::new(0);
+    let mut wasm_info = WasmInfo::default();
+    let mut write_functions: HashSet<usize> = HashSet::new();
+
+    let parser = Parser::new(0);
     for payload in parser.parse_all(&wasm_data) {
         match payload? {
             Payload::ImportSection(import_section) => {
-                imports = extract_imports(
-                    import_section,
-                    extract_imports_enabled,
-                    &mut call_graph,
-                    &mut write_functions,
-                );
-                ei_check = is_ei_valid(&imports, check_ei);
+                write_functions =
+                    process_imports(import_section, extract_imports_enabled, &mut wasm_info);
+                wasm_info.ei_check = is_ei_valid(&wasm_info.imports, check_ei);
             },
             Payload::DataSection(data_section) => {
-                allocator_trigger = is_fail_allocator_triggered(data_section.clone());
+                wasm_info.report.has_allocator = is_fail_allocator_triggered(data_section.clone());
                 if is_panic_with_message_triggered(data_section.clone()) {
-                    has_panic = WITH_MESSAGE;
+                    wasm_info.report.has_panic = WITH_MESSAGE.to_owned();
                 } else if is_panic_without_message_triggered(data_section) {
-                    has_panic = WITHOUT_MESSAGE;
+                    wasm_info.report.has_panic = WITHOUT_MESSAGE.to_owned();
                 }
             },
             Payload::CodeSectionEntry(code_section) => {
-                memory_grow_flag = is_mem_grow(&code_section);
+                wasm_info.memory_grow_flag = is_mem_grow(&code_section);
+                create_call_graph(code_section, &mut wasm_info.call_graph);
             },
             Payload::ExportSection(export_section) => {
-                views_data = parse_export_section(export_section, &view_endpoints);
+                parse_export_section(export_section, &mut view_endpoints);
             },
             _ => (),
         }
     }
 
-    parser = Parser::new(0);
-    for payload in parser.parse_all(&wasm_data) {
-        if let Payload::CodeSectionEntry(body) = payload? {
-            create_call_graph(body, &mut call_graph);
-        }
-    }
+    detect_write_operations_in_views(&view_endpoints, &wasm_info.call_graph, &mut write_functions);
+
+    let report = ReportCreator {
+        path,
+        has_allocator: wasm_info.report.has_allocator,
+        has_panic: wasm_info.report.has_panic,
+    };
+
+    Ok(WasmInfo {
+        imports: wasm_info.imports,
+        ei_check: wasm_info.ei_check,
+        memory_grow_flag: wasm_info.memory_grow_flag,
+        call_graph: wasm_info.call_graph,
+        report,
+    })
+}
+
+fn detect_write_operations_in_views(
+    views_data: &HashMap<&str, usize>,
+    call_graph: &CallGraph,
+    write_functions: &mut HashSet<usize>,
+) {
     let mut visited: HashSet<usize> = HashSet::new();
-    for key in views_data.keys() {
-        mark_write(*key, &call_graph, &mut write_functions, &mut visited);
+    for index in views_data.values() {
+        mark_write(*index, call_graph, write_functions, &mut visited);
     }
 
-    for (index, name) in views_data {
-        if write_functions.contains(&index) {
+    for (name, index) in views_data {
+        if write_functions.contains(index) {
             println!(
                 "{} {}",
                 "Write storage operation in VIEW endpoint:"
@@ -115,42 +125,26 @@ fn populate_wasm_info(
             );
         }
     }
-
-    let report = ReportCreator {
-        path,
-        has_allocator: allocator_trigger,
-        has_panic: has_panic.to_string(),
-    };
-
-    Ok(WasmInfo {
-        imports,
-        ei_check,
-        memory_grow_flag,
-        has_format: true,
-        report,
-    })
 }
 
 fn parse_export_section(
     export_section: ExportSectionReader,
-    view_endpoints: &[String],
-) -> HashMap<usize, String> {
-    let mut views_data: HashMap<usize, String> = HashMap::new();
+    view_endpoints: &mut HashMap<&str, usize>,
+) {
     for export in export_section {
         let export = export.expect("Failed to read export section");
         if let wasmparser::ExternalKind::Func = export.kind {
-            if view_endpoints.contains(&export.name.to_string()) {
-                views_data.insert(export.index.try_into().unwrap(), export.name.to_string());
+            if let Some(endpoint_index) = view_endpoints.get_mut(export.name) {
+                *endpoint_index = export.index.try_into().unwrap();
             }
         }
     }
-    views_data
 }
 
 fn mark_write(
     func: usize,
-    call_graph: &HashMap<usize, Vec<usize>>,
-    write_functions: &mut Vec<usize>,
+    call_graph: &CallGraph,
+    write_functions: &mut HashSet<usize>,
     visited: &mut HashSet<usize>,
 ) {
     // Return early to prevent cycles.
@@ -163,27 +157,27 @@ fn mark_write(
     if let Some(callees) = call_graph.get(&func) {
         for &callee in callees {
             if write_functions.contains(&callee) {
-                write_functions.push(func);
+                write_functions.insert(func);
             } else {
                 mark_write(callee, call_graph, write_functions, visited);
                 if write_functions.contains(&callee) {
-                    write_functions.push(func);
+                    write_functions.insert(func);
                 }
             }
         }
     }
 }
 
-fn create_call_graph(body: FunctionBody, call_graph: &mut HashMap<usize, Vec<usize>>) {
+fn create_call_graph(body: FunctionBody, call_graph: &mut CallGraph) {
     let mut instructions_reader = body
         .get_operators_reader()
         .expect("Failed to get operators reader");
 
-    let mut call_functions = Vec::new();
+    let mut call_functions = HashSet::new();
     while let Ok(op) = instructions_reader.read() {
         if let Operator::Call { function_index } = op {
             let function_usize: usize = function_index.try_into().unwrap();
-            call_functions.push(function_usize);
+            call_functions.insert(function_usize);
         }
     }
 
@@ -239,28 +233,25 @@ fn is_panic_without_message_triggered(data_section: DataSectionReader) -> bool {
     false
 }
 
-pub fn extract_imports(
+pub fn process_imports(
     import_section: ImportSectionReader,
     import_extraction_enabled: bool,
-    call_graph: &mut HashMap<usize, Vec<usize>>,
-    write_functions: &mut Vec<usize>,
-) -> Vec<String> {
-    if !import_extraction_enabled {
-        return Vec::new();
-    }
-
-    let mut import_names = Vec::new();
+    wasm_info: &mut WasmInfo,
+) -> HashSet<usize> {
+    let mut write_functions = HashSet::new();
     for (index, import) in import_section.into_iter().flatten().enumerate() {
-        import_names.push(import.name.to_string());
-        call_graph.insert(index, vec![]);
+        if import_extraction_enabled {
+            wasm_info.imports.push(import.name.to_string());
+        }
+        wasm_info.call_graph.insert(index, HashSet::new());
         if WRITE_OP.contains(&import.name) {
-            write_functions.push(index);
+            write_functions.insert(index);
         }
     }
 
-    import_names.sort();
+    wasm_info.imports.sort();
 
-    import_names
+    write_functions
 }
 
 fn is_ei_valid(imports: &[String], check_ei: &Option<EIVersion>) -> bool {