From d145b73eba13d4b2bb095ab10ae5fcc576c054bd Mon Sep 17 00:00:00 2001 From: BiancaIalangi Date: Fri, 13 Sep 2024 19:11:46 +0300 Subject: [PATCH] detect write op in view storage - refactor: removed has_format, changed data structure for view_endpoints, impl default for WasmInfo and ReportCreator --- .../src/contract/sc_config/wasm_build.rs | 20 +-- .../meta-lib/src/tools/report_creator.rs | 10 ++ .../meta-lib/src/tools/wasm_extractor.rs | 143 ++++++++---------- 3 files changed, 83 insertions(+), 90 deletions(-) diff --git a/framework/meta-lib/src/contract/sc_config/wasm_build.rs b/framework/meta-lib/src/contract/sc_config/wasm_build.rs index e7a9bec5ed..4311bd0979 100644 --- a/framework/meta-lib/src/contract/sc_config/wasm_build.rs +++ b/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); } } diff --git a/framework/meta-lib/src/tools/report_creator.rs b/framework/meta-lib/src/tools/report_creator.rs index 37af10fa4e..60de81dc31 100644 --- a/framework/meta-lib/src/tools/report_creator.rs +++ b/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(), + } + } +} diff --git a/framework/meta-lib/src/tools/wasm_extractor.rs b/framework/meta-lib/src/tools/wasm_extractor.rs index 4cd80507d4..63dad6b3ad 100644 --- a/framework/meta-lib/src/tools/wasm_extractor.rs +++ b/framework/meta-lib/src/tools/wasm_extractor.rs @@ -12,18 +12,21 @@ use crate::ei::EIVersion; use super::report_creator::{ReportCreator, WITHOUT_MESSAGE, WITH_MESSAGE}; +type CallGraph = HashMap>; + 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, pub ei_check: bool, pub memory_grow_flag: bool, - pub has_format: bool, pub report: ReportCreator, + pub call_graph: CallGraph, } impl WasmInfo { @@ -31,7 +34,7 @@ impl WasmInfo { output_wasm_path: &str, extract_imports_enabled: bool, check_ei: &Option, - view_endpoints: Vec, + view_endpoints: HashMap<&str, usize>, ) -> Result { 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, extract_imports_enabled: bool, check_ei: &Option, - view_endpoints: Vec, + mut view_endpoints: HashMap<&str, usize>, ) -> Result { - 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> = HashMap::new(); - let mut views_data: HashMap = HashMap::new(); - let mut write_functions: Vec = Vec::new(); - - let mut parser = Parser::new(0); + let mut wasm_info = WasmInfo::default(); + let mut write_functions: HashSet = 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, +) { let mut visited: HashSet = 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 { - let mut views_data: HashMap = 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>, - write_functions: &mut Vec, + call_graph: &CallGraph, + write_functions: &mut HashSet, visited: &mut HashSet, ) { // 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>) { +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>, - write_functions: &mut Vec, -) -> Vec { - if !import_extraction_enabled { - return Vec::new(); - } - - let mut import_names = Vec::new(); + wasm_info: &mut WasmInfo, +) -> HashSet { + 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) -> bool {