From b973132c2b8e9e342c5bcd753522fba8f36a0674 Mon Sep 17 00:00:00 2001 From: Jiyuan Zheng Date: Fri, 23 Aug 2024 12:53:31 +0800 Subject: [PATCH] fix: sbrk must be called twice (#41) --- xcq-api/procedural/src/program/expand/mod.rs | 16 ++++++---------- xcq-executor/src/lib.rs | 17 +++++++++++------ xcq-extension/src/lib.rs | 5 ++++- .../tests/with_associated_types_works.rs | 2 +- xcq-test-runner/src/main.rs | 10 +++++----- 5 files changed, 27 insertions(+), 23 deletions(-) diff --git a/xcq-api/procedural/src/program/expand/mod.rs b/xcq-api/procedural/src/program/expand/mod.rs index 1eddf8f..7c2287f 100644 --- a/xcq-api/procedural/src/program/expand/mod.rs +++ b/xcq-api/procedural/src/program/expand/mod.rs @@ -62,22 +62,19 @@ fn pass_byte_to_host() -> TokenStream2 { // TODO check res type to determine the appropriate serializing method quote! { let res_bytes = res.to_le_bytes(); - let ptr = polkavm_derive::sbrk(res_bytes.len()); - if ptr.is_null(){ + let res_ptr = polkavm_derive::sbrk(0); + let end_ptr = polkavm_derive::sbrk(res_bytes.len()); + if end_ptr.is_null(){ return 0; } unsafe { - core::ptr::copy_nonoverlapping(res_bytes.as_ptr(),ptr,res_bytes.len()); + core::ptr::copy_nonoverlapping(res_bytes.as_ptr(),res_ptr,res_bytes.len()); } - (res_bytes.len() as u64) << 32 | (ptr as u64) + (res_bytes.len() as u64) << 32 | (res_ptr as u64) } } fn generate_main(entrypoint: &EntrypointDef) -> Result { - let move_to_stack = quote! { - let arg_bytes = unsafe {alloc::vec::Vec::from_raw_parts(ptr as *mut u8, size as usize, size as usize)}; - let mut arg_ptr = arg_bytes.as_ptr() as u32; - }; // Construct call_data let mut get_call_data = TokenStream2::new(); for (arg_type_index, arg_type) in entrypoint.arg_types.iter().enumerate() { @@ -141,8 +138,7 @@ fn generate_main(entrypoint: &EntrypointDef) -> Result { let main = quote! { #[polkavm_derive::polkavm_export] - extern "C" fn main(ptr: u32, size:u32) -> u64 { - #move_to_stack + extern "C" fn main(mut arg_ptr: u32, size:u32) -> u64 { #get_call_data #call_entrypoint #pass_bytes_back diff --git a/xcq-executor/src/lib.rs b/xcq-executor/src/lib.rs index 77f82c1..bf14375 100644 --- a/xcq-executor/src/lib.rs +++ b/xcq-executor/src/lib.rs @@ -63,15 +63,20 @@ impl XcqExecutor { let instance_pre = self.linker.instantiate_pre(&module)?; let instance = instance_pre.instantiate()?; - // Args are passed via guest's heap let input_ptr = if !input.is_empty() { - let ptr = instance - .sbrk(input.len() as u32)? - .expect("sbrk must be able to allocate memory here"); + // First sbrk call to get the start of the heap + let start_ptr = instance.sbrk(0)?.expect("should not fail because we don't allocate"); + // Second sbrk call to check the allocation doesn't exceed the heap limit + if instance.sbrk(input.len() as u32)?.is_none() { + return Err(XcqExecutorError::ExecutionError(polkavm::ExecutionError::Error( + polkavm::Error::from("cannot srk enough memory"), + ))); + }; + // Args are passed via guest's heap instance - .write_memory(ptr, input) + .write_memory(start_ptr, input) .map_err(|e| XcqExecutorError::ExecutionError(polkavm::ExecutionError::Trap(e)))?; - ptr + start_ptr } else { 0 }; diff --git a/xcq-extension/src/lib.rs b/xcq-extension/src/lib.rs index 6fb2ea8..418a033 100644 --- a/xcq-extension/src/lib.rs +++ b/xcq-extension/src/lib.rs @@ -72,7 +72,10 @@ impl XcqExecutorContext for Context let res_bytes = E::dispatch(extension_id, &call_bytes)?; tracing::debug!("(host call): res_bytes: {:?}", res_bytes); let res_bytes_len = res_bytes.len(); - let res_ptr = caller.sbrk(res_bytes_len as u32).ok_or(ExtensionError::PolkavmError)?; + let res_ptr = caller.sbrk(0).ok_or(ExtensionError::PolkavmError)?; + if caller.sbrk(res_bytes_len as u32).is_none() { + return Err(ExtensionError::PolkavmError); + } caller .write_memory(res_ptr, &res_bytes) .map_err(|_| ExtensionError::PolkavmError)?; diff --git a/xcq-extension/tests/with_associated_types_works.rs b/xcq-extension/tests/with_associated_types_works.rs index 937379f..04b9933 100644 --- a/xcq-extension/tests/with_associated_types_works.rs +++ b/xcq-extension/tests/with_associated_types_works.rs @@ -122,7 +122,7 @@ fn call_core_works() { } #[test] fn multi_calls_works() { - let blob = include_bytes!("../../output/poc-guest-balance-sum-percent.polkavm"); + let blob = include_bytes!("../../output/poc-guest-sum-balance-percent.polkavm"); let mut executor = ExtensionsExecutor::::new(InvokeSource::RuntimeAPI); let guest = GuestImpl { program: blob.to_vec() }; let mut input_data = extension_fungibles::EXTENSION_ID.encode(); diff --git a/xcq-test-runner/src/main.rs b/xcq-test-runner/src/main.rs index 1b550a8..87efa6c 100644 --- a/xcq-test-runner/src/main.rs +++ b/xcq-test-runner/src/main.rs @@ -42,11 +42,11 @@ fn main() { }; input_data.extend_from_slice(&method1_encoded); input_data.extend_from_slice(&method2.encode()); - input_data.extend_from_slice(&xcq_extension_fungibles::EXTENSION_ID.encode()); - let method3 = FungiblesMethod::TotalSupply { asset: 1 }; - let method3_encoded = method3.encode(); - input_data.extend_from_slice(&[method3_encoded.len() as u8]); - input_data.extend_from_slice(&method3_encoded); + // input_data.extend_from_slice(&xcq_extension_fungibles::EXTENSION_ID.encode()); + // let method3 = FungiblesMethod::TotalSupply { asset: 1 }; + // let method3_encoded = method3.encode(); + // input_data.extend_from_slice(&[method3_encoded.len() as u8]); + // input_data.extend_from_slice(&method3_encoded); tracing::info!("Input data: {:?}", input_data); let input = InputImpl { method: "main".to_string(),