From cbf83e6ce8500b7c8e764e8e8b3295c644e7c87c Mon Sep 17 00:00:00 2001 From: Maciej Kot Date: Fri, 17 May 2024 00:56:04 +0000 Subject: [PATCH] Feat: RUN-979: Support for wasm64 in native stable64 --- .../src/wasm_utils/instrumentation.rs | 29 +++-- .../src/wasm_utils/system_api_replacements.rs | 20 +++- rs/embedders/src/wasm_utils/validation.rs | 42 +++++-- .../tests/wasmtime_random_memory_writes.rs | 107 ++++++++++++++++++ .../src/hypervisor/tests.rs | 36 ++++++ .../execution_environment/src/lib.rs | 5 + 6 files changed, 214 insertions(+), 25 deletions(-) diff --git a/rs/embedders/src/wasm_utils/instrumentation.rs b/rs/embedders/src/wasm_utils/instrumentation.rs index 7d3b3be85e7..56fb857aa7d 100644 --- a/rs/embedders/src/wasm_utils/instrumentation.rs +++ b/rs/embedders/src/wasm_utils/instrumentation.rs @@ -142,6 +142,16 @@ pub(crate) enum WasmMemoryType { Wasm64, } +pub(crate) fn main_memory_type(module: &Module<'_>) -> WasmMemoryType { + let mut mem_type = WasmMemoryType::Wasm32; + if let Some(mem) = module.memories.first() { + if mem.memory64 { + mem_type = WasmMemoryType::Wasm64; + } + } + mem_type +} + // The indices of injected function imports. pub(crate) enum InjectedImports { OutOfInstructions = 0, @@ -908,12 +918,7 @@ pub(super) fn instrument( subnet_type: SubnetType, dirty_page_overhead: NumInstructions, ) -> Result { - let mut main_memory_type = WasmMemoryType::Wasm32; - if let Some(mem) = module.memories.first() { - if mem.memory64 { - main_memory_type = WasmMemoryType::Wasm64; - } - } + let main_memory_type = main_memory_type(&module); let stable_memory_index; let mut module = inject_helper_functions(module, wasm_native_stable_memory, main_memory_type); module = export_table(module); @@ -1012,6 +1017,7 @@ pub(super) fn instrument( special_indices, subnet_type, dirty_page_overhead, + main_memory_type, ) } @@ -1095,6 +1101,7 @@ fn replace_system_api_functions( special_indices: SpecialIndices, subnet_type: SubnetType, dirty_page_overhead: NumInstructions, + main_memory_type: WasmMemoryType, ) { let api_indexes = calculate_api_indexes(module); let number_of_func_imports = module @@ -1106,9 +1113,12 @@ fn replace_system_api_functions( // Collect a single map of all the function indexes that need to be // replaced. let mut func_index_replacements = BTreeMap::new(); - for (api, (ty, body)) in - replacement_functions(special_indices, subnet_type, dirty_page_overhead) - { + for (api, (ty, body)) in replacement_functions( + special_indices, + subnet_type, + dirty_page_overhead, + main_memory_type, + ) { if let Some(old_index) = api_indexes.get(&api) { let type_idx = add_func_type(module, ty); let new_index = (number_of_func_imports + module.functions.len()) as u32; @@ -1825,6 +1835,7 @@ fn get_data( offset_expr, } => match offset_expr { Operator::I32Const { value } => *value as usize, + Operator::I64Const { value } => *value as usize, _ => return Some(Err(WasmInstrumentationError::WasmDeserializeError(WasmError::new( "complex initialization expressions for data segments are not supported!".into() )))), diff --git a/rs/embedders/src/wasm_utils/system_api_replacements.rs b/rs/embedders/src/wasm_utils/system_api_replacements.rs index 849177f0552..395d4e8feb4 100644 --- a/rs/embedders/src/wasm_utils/system_api_replacements.rs +++ b/rs/embedders/src/wasm_utils/system_api_replacements.rs @@ -12,8 +12,9 @@ //! use crate::{ - wasm_utils::instrumentation::InjectedImports, - wasmtime_embedder::system_api_complexity::overhead_native, InternalErrorCode, + wasm_utils::instrumentation::{InjectedImports, WasmMemoryType}, + wasmtime_embedder::system_api_complexity::overhead_native, + InternalErrorCode, }; use ic_interfaces::execution_environment::StableMemoryApi; use ic_registry_subnet_type::SubnetType; @@ -31,6 +32,7 @@ pub(super) fn replacement_functions( special_indices: SpecialIndices, subnet_type: SubnetType, dirty_page_overhead: NumInstructions, + main_memory_type: WasmMemoryType, ) -> Vec<(SystemApiFunc, (FuncType, Body<'static>))> { let count_clean_pages_fn_index = special_indices.count_clean_pages_fn.unwrap(); let dirty_pages_counter_index = special_indices.dirty_pages_counter_ix.unwrap(); @@ -41,6 +43,12 @@ pub(super) fn replacement_functions( use Operator::*; let page_size_shift = PAGE_SIZE.trailing_zeros() as i32; let stable_memory_bytemap_index = stable_memory_index + 1; + + let cast_to_heap_addr_type = match main_memory_type { + WasmMemoryType::Wasm32 => I32WrapI64, + WasmMemoryType::Wasm64 => Nop, + }; + vec![ ( SystemApiFunc::StableSize, @@ -731,10 +739,10 @@ pub(super) fn replacement_functions( }, Else, LocalGet { local_index: DST }, - I32WrapI64, + cast_to_heap_addr_type.clone(), LocalGet { local_index: SRC }, LocalGet { local_index: LEN }, - I32WrapI64, + cast_to_heap_addr_type.clone(), MemoryCopy { dst_mem: 0, src_mem: stable_memory_index, @@ -1205,9 +1213,9 @@ pub(super) fn replacement_functions( // copy memory contents LocalGet { local_index: DST }, LocalGet { local_index: SRC }, - I32WrapI64, + cast_to_heap_addr_type.clone(), LocalGet { local_index: LEN }, - I32WrapI64, + cast_to_heap_addr_type, MemoryCopy { dst_mem: stable_memory_index, src_mem: 0, diff --git a/rs/embedders/src/wasm_utils/validation.rs b/rs/embedders/src/wasm_utils/validation.rs index 10c56129bd8..b70e18a8c76 100644 --- a/rs/embedders/src/wasm_utils/validation.rs +++ b/rs/embedders/src/wasm_utils/validation.rs @@ -20,7 +20,8 @@ use crate::wasmtime_embedder::{ }; use crate::{ wasm_utils::instrumentation::{ - ACCESSED_PAGES_COUNTER_GLOBAL_NAME, DIRTY_PAGES_COUNTER_GLOBAL_NAME, + main_memory_type, WasmMemoryType, ACCESSED_PAGES_COUNTER_GLOBAL_NAME, + DIRTY_PAGES_COUNTER_GLOBAL_NAME, }, MAX_WASM_STACK_SIZE, MIN_GUARD_REGION_SIZE, }; @@ -954,24 +955,45 @@ fn validate_export_section( // expression. Required because of OP. See also: // instrumentation.rs fn validate_data_section(module: &Module) -> Result<(), WasmValidationError> { - fn validate_segment(s: &DataSegment) -> Result<(), WasmValidationError> { - match &s.kind { - DataSegmentKind::Passive => Ok(()), - DataSegmentKind::Active { - memory_index: _, - offset_expr, - } => match offset_expr { + fn validate_segment( + s: &DataSegment, + mem_type: WasmMemoryType, + ) -> Result<(), WasmValidationError> { + match (&s.kind, mem_type) { + (DataSegmentKind::Passive, _) => Ok(()), + ( + DataSegmentKind::Active { + memory_index: _, + offset_expr, + }, + WasmMemoryType::Wasm32, + ) => match offset_expr { Operator::I32Const { .. } => Ok(()), _ => Err(WasmValidationError::InvalidDataSection(format!( - "Invalid offset expression in data segment: {:?}", + "Invalid offset expression in data segment for 32bit memory: {:?}", + offset_expr + ))), + }, + ( + DataSegmentKind::Active { + memory_index: _, + offset_expr, + }, + WasmMemoryType::Wasm64, + ) => match offset_expr { + Operator::I64Const { .. } => Ok(()), + _ => Err(WasmValidationError::InvalidDataSection(format!( + "Invalid offset expression in data segment for 64bit memory: {:?}", offset_expr ))), }, } } + let mem_type = main_memory_type(module); + for d in &module.data { - validate_segment(d)?; + validate_segment(d, mem_type)?; } Ok(()) } diff --git a/rs/embedders/tests/wasmtime_random_memory_writes.rs b/rs/embedders/tests/wasmtime_random_memory_writes.rs index 871f4ab5e56..4a8467da8a0 100644 --- a/rs/embedders/tests/wasmtime_random_memory_writes.rs +++ b/rs/embedders/tests/wasmtime_random_memory_writes.rs @@ -259,6 +259,57 @@ fn make_module_wat_for_api_calls(heap_size: usize) -> String { ) } +fn make_module64_wat_for_api_calls(heap_size: usize) -> String { + format!( + r#" + (module + (import "ic0" "msg_reply" (func $msg_reply)) + (import "ic0" "msg_reply_data_append" + (func $msg_reply_data_append (param i32) (param i32))) + (import "ic0" "msg_arg_data_copy" + (func $ic0_msg_arg_data_copy (param i32) (param i32) (param i32))) + (import "ic0" "msg_arg_data_size" + (func $ic0_msg_arg_data_size (result i32))) + (import "ic0" "msg_caller_copy" + (func $ic0_msg_caller_copy (param i32) (param i32) (param i32))) + (import "ic0" "msg_caller_size" + (func $ic0_msg_caller_size (result i32))) + (import "ic0" "canister_self_copy" + (func $ic0_canister_self_copy (param i32) (param i32) (param i32))) + (import "ic0" "canister_self_size" + (func $ic0_canister_self_size (result i32))) + + (import "ic0" "canister_cycle_balance128" + (func $ic0_canister_cycle_balance128 (param i32))) + + (import "ic0" "stable_grow" + (func $ic0_stable_grow (param $pages i32) (result i32))) + (import "ic0" "stable64_read" + (func $ic0_stable64_read (param $dst i64) (param $offset i64) (param $size i64))) + (import "ic0" "stable64_write" + (func $ic0_stable64_write (param $offset i64) (param $src i64) (param $size i64))) + + (func $touch_heap_with_api_calls + (call $ic0_msg_caller_copy (i32.const 4096) (i32.const 0) (call $ic0_msg_caller_size)) + (call $ic0_msg_arg_data_copy (i32.const 12288) (i32.const 0) (call $ic0_msg_arg_data_size)) + (call $ic0_canister_self_copy (i32.const 20480) (i32.const 0) (call $ic0_canister_self_size)) + (call $ic0_canister_cycle_balance128 (i32.const 36864)) + + (; Write some data to page 10 using stable_read, by first copying 4 + bytes from the second page to stable memory, then copying back ;) + (drop (call $ic0_stable_grow (i32.const 1))) + (call $ic0_stable64_write (i64.const 0) (i64.const 4096) (i64.const 4)) + (call $ic0_stable64_read (i64.const 40960) (i64.const 0) (i64.const 4)) + ) + + (memory $memory i64 {HEAP_SIZE}) + (export "memory" (memory $memory)) + (export "canister_update touch_heap_with_api_calls" (func $touch_heap_with_api_calls)) + )"#, + HEAP_SIZE = heap_size + ) +} + fn make_module_wat_with_write_fun(heap_size: usize, write_fun: &str) -> String { format!( r#" @@ -1385,6 +1436,62 @@ mod tests { }); } + #[test] + fn touch_heap64_with_api_calls() { + with_test_replica_logger(|log| { + let wat = make_module64_wat_for_api_calls(TEST_NUM_PAGES); + let wasm = wat2wasm(&wat).unwrap(); + let mut config = EmbeddersConfig::default(); + config.feature_flags.wasm64 = FlagStatus::Enabled; + let embedder = WasmtimeEmbedder::new(config, log); + let (embedder_cache, result) = compile(&embedder, &wasm); + result.unwrap(); + + let mut dirty_pages: BTreeSet = BTreeSet::new(); + + let payload = vec![0, 1, 2, 3, 4, 5, 6, 7]; + + let api = test_api_for_update( + no_op_logger(), + None, + payload, + SubnetType::Application, + MAX_NUM_INSTRUCTIONS, + ); + let instruction_limit = api.slice_instruction_limit(); + let mut instance = embedder + .new_instance( + canister_test_id(1), + &embedder_cache, + None, + &Memory::new(PageMap::new_for_testing(), NumWasmPages::from(0)), + &Memory::new(PageMap::new_for_testing(), NumWasmPages::from(0)), + ModificationTracking::Track, + Some(api), + ) + .map_err(|r| r.0) + .expect("Failed to create instance"); + instance.set_instruction_counter(i64::try_from(instruction_limit.get()).unwrap()); + + let result = instance + .run(FuncRef::Method(WasmMethod::Update( + "touch_heap_with_api_calls".to_string(), + ))) + .expect("call to touch_heap_with_api_calls failed"); + dirty_pages.extend(result.dirty_pages.iter().map(|x| x.get())); + + let mut expected_dirty_pages: BTreeSet = BTreeSet::new(); + expected_dirty_pages.insert(1); // caller_copy + expected_dirty_pages.insert(3); // data_copy + expected_dirty_pages.insert(5); // canister_self_copy + expected_dirty_pages.insert(9); // canister_cycle_balance128 + expected_dirty_pages.insert(9); // msg_cycles_available128 + expected_dirty_pages.insert(10); // stable_read + + assert_eq!(expected_dirty_pages, dirty_pages); + }); + } + #[test] fn wasmtime_random_memory_writes_ignore_dirty_pages() { // The seed value will always be the same for a particular version of diff --git a/rs/execution_environment/src/hypervisor/tests.rs b/rs/execution_environment/src/hypervisor/tests.rs index 70be81f1755..160d09d5924 100644 --- a/rs/execution_environment/src/hypervisor/tests.rs +++ b/rs/execution_environment/src/hypervisor/tests.rs @@ -1438,6 +1438,42 @@ fn ic0_msg_reject_works() { assert_eq!(WasmResult::Reject("panic!".to_string()), result); } +#[test] +fn wasm64_active_data_segments() { + let mut test = ExecutionTestBuilder::new().with_wasm64().build(); + let wat = r#" + (module + (import "ic0" "msg_reply" (func $msg_reply)) + (func (export "canister_update test") + (if (i64.ne + (i64.load8_u (i64.const 0)) + (i64.const 112) ;; p + ) + (then (unreachable)) + ) + (if (i64.ne + (i64.load8_u (i64.const 1)) + (i64.const 0) + ) + (then (unreachable)) + ) + (if (i64.ne + (i64.load8_u (i64.const 2)) + (i64.const 97) ;; a + ) + (then (unreachable)) + ) + (call $msg_reply) + ) + (memory i64 1 1) + (data (i64.const 0) "p") + (data (i64.const 2) "a") + )"#; + let canister_id = test.canister_from_wat(wat).unwrap(); + let result = test.ingress(canister_id, "test", vec![]).unwrap(); + assert_eq!(WasmResult::Reply(vec![]), result); +} + #[test] fn ic0_msg_caller_size_works_in_reply_callback() { let mut test = ExecutionTestBuilder::new().build(); diff --git a/rs/test_utilities/execution_environment/src/lib.rs b/rs/test_utilities/execution_environment/src/lib.rs index c474c4f117e..62aa01114b8 100644 --- a/rs/test_utilities/execution_environment/src/lib.rs +++ b/rs/test_utilities/execution_environment/src/lib.rs @@ -1929,6 +1929,11 @@ impl ExecutionTestBuilder { self } + pub fn with_wasm64(mut self) -> Self { + self.execution_config.embedders_config.feature_flags.wasm64 = FlagStatus::Enabled; + self + } + pub fn with_metering_type(mut self, metering_type: MeteringType) -> Self { self.execution_config.embedders_config.metering_type = metering_type; self