From 7733fd76571e5391835fc59b3c838f8877222bdc Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Eduardo=20Leegwater=20Sim=C3=B5es?= Date: Wed, 1 Nov 2023 17:05:20 +0100 Subject: [PATCH] piecrust: distinguish 32bit and 64bit memories Two different trees are created - one for 32-bit and another for 64-bit - and ensure they're properly distinguished. --- piecrust/src/instance.rs | 4 +-- piecrust/src/store.rs | 2 +- piecrust/src/store/memory.rs | 43 ++++++++++++++++++++--------- piecrust/src/store/session.rs | 5 ++-- piecrust/src/store/tree.rs | 51 ++++++++++++++++++++++++++++++----- 5 files changed, 81 insertions(+), 24 deletions(-) diff --git a/piecrust/src/instance.rs b/piecrust/src/instance.rs index 5aa6408b..e99adf1f 100644 --- a/piecrust/src/instance.rs +++ b/piecrust/src/instance.rs @@ -13,7 +13,7 @@ use piecrust_uplink::{ContractId, Event, ARGBUF_LEN}; use crate::contract::WrappedContract; use crate::imports::Imports; use crate::session::Session; -use crate::store::{Memory, MAX_MEM_SIZE}; +use crate::store::Memory; use crate::Error; pub struct WrappedInstance { @@ -169,7 +169,7 @@ impl WrappedInstance { _ => return Err(Error::InvalidArgumentBuffer), }; - if arg_buf_ofs + ARGBUF_LEN >= MAX_MEM_SIZE { + if arg_buf_ofs + ARGBUF_LEN >= memory.len() { return Err(Error::InvalidArgumentBuffer); } diff --git a/piecrust/src/store.rs b/piecrust/src/store.rs index 08fa5447..69ef53a7 100644 --- a/piecrust/src/store.rs +++ b/piecrust/src/store.rs @@ -22,7 +22,7 @@ use std::{fs, io, thread}; pub use bytecode::Bytecode; use dusk_wasmtime::Engine; -pub use memory::{Memory, MAX_MEM_SIZE, PAGE_SIZE}; +pub use memory::{Memory, PAGE_SIZE}; pub use metadata::Metadata; pub use module::Module; use piecrust_uplink::ContractId; diff --git a/piecrust/src/store/memory.rs b/piecrust/src/store/memory.rs index ce549a49..abe63b9c 100644 --- a/piecrust/src/store/memory.rs +++ b/piecrust/src/store/memory.rs @@ -16,18 +16,15 @@ use crumbles::{LocateFile, Mmap}; use dusk_wasmtime::LinearMemory; pub const PAGE_SIZE: usize = 0x10000; -const WASM_MAX_PAGES: u32 = 0x4000000; -const MIN_PAGES: usize = 4; -const MIN_MEM_SIZE: usize = MIN_PAGES * PAGE_SIZE; -const MAX_PAGES: usize = WASM_MAX_PAGES as usize; - -pub const MAX_MEM_SIZE: usize = MAX_PAGES * PAGE_SIZE; +const WASM32_MAX_PAGES: usize = 0x10000; +const WASM64_MAX_PAGES: usize = 0x4000000; pub struct MemoryInner { pub mmap: Mmap, pub current_len: usize, pub is_new: bool, + is_64: bool, ref_count: AtomicUsize, } @@ -63,32 +60,54 @@ pub struct Memory { } impl Memory { - pub fn new() -> io::Result { + pub fn new(is_64: bool) -> io::Result { + let max_pages = if is_64 { + WASM64_MAX_PAGES + } else { + WASM32_MAX_PAGES + }; + Ok(Self { inner: Box::leak(Box::new(MemoryInner { - mmap: Mmap::new(MAX_PAGES, PAGE_SIZE)?, - current_len: MIN_MEM_SIZE, + mmap: Mmap::new(max_pages, PAGE_SIZE)?, + current_len: 0, is_new: true, + is_64, ref_count: AtomicUsize::new(1), })), }) } - pub fn from_files(file_locator: FL, len: usize) -> io::Result + pub fn from_files( + is_64: bool, + file_locator: FL, + len: usize, + ) -> io::Result where FL: 'static + LocateFile, { + let max_pages = if is_64 { + WASM64_MAX_PAGES + } else { + WASM32_MAX_PAGES + }; + Ok(Self { inner: Box::leak(Box::new(MemoryInner { mmap: unsafe { - Mmap::with_files(MAX_PAGES, PAGE_SIZE, file_locator)? + Mmap::with_files(max_pages, PAGE_SIZE, file_locator)? }, current_len: len, is_new: false, + is_64, ref_count: AtomicUsize::new(1), })), }) } + + pub fn is_64(&self) -> bool { + self.inner.is_64 + } } /// This implementation of clone is dangerous, and must be accompanied by the @@ -143,7 +162,7 @@ unsafe impl LinearMemory for Memory { } fn maximum_byte_size(&self) -> Option { - Some(MAX_MEM_SIZE) + Some(self.inner.len()) } fn grow_to(&mut self, new_size: usize) -> Result<(), dusk_wasmtime::Error> { diff --git a/piecrust/src/store/session.rs b/piecrust/src/store/session.rs index 2212828f..2689516d 100644 --- a/piecrust/src/store/session.rs +++ b/piecrust/src/store/session.rs @@ -182,6 +182,7 @@ impl ContractSession { let memory_path = memory_path.clone(); Memory::from_files( + module.is_64(), move |page_index: usize| { match page_indices .contains(&page_index) @@ -199,7 +200,7 @@ impl ContractSession { elem.len, )? } - None => Memory::new()?, + None => Memory::new(module.is_64())?, }; let contract = entry @@ -250,10 +251,10 @@ impl ContractSession { metadata: ContractMetadata, metadata_bytes: B, ) -> io::Result<()> { - let memory = Memory::new()?; let bytecode = Bytecode::new(bytecode)?; let module = Module::new(&self.engine, objectcode)?; let metadata = Metadata::new(metadata_bytes, metadata)?; + let memory = Memory::new(module.is_64())?; // If the position is already filled in the tree, the contract cannot be // inserted. diff --git a/piecrust/src/store/tree.rs b/piecrust/src/store/tree.rs index aff4656e..bb55ed3b 100644 --- a/piecrust/src/store/tree.rs +++ b/piecrust/src/store/tree.rs @@ -15,16 +15,53 @@ use rkyv::{Archive, Deserialize, Serialize}; use crate::store::memory::Memory; -// There are max `2^26` pages in a memory -const P_HEIGHT: usize = 13; -const P_ARITY: usize = 4; +// There are max `2^16` pages in a 32-bit memory +const P32_HEIGHT: usize = 8; +const P32_ARITY: usize = 4; -pub type PageTree = dusk_merkle::Tree; +type PageTree32 = dusk_merkle::Tree; + +// There are max `2^26` pages in a 64-bit memory +const P64_HEIGHT: usize = 13; +const P64_ARITY: usize = 4; + +type PageTree64 = dusk_merkle::Tree; // This means we have max `2^32` contracts const C_HEIGHT: usize = 32; const C_ARITY: usize = 2; +#[derive(Debug, Clone, Archive, Deserialize, Serialize)] +#[archive_attr(derive(CheckBytes))] +pub enum PageTree { + Wasm32(PageTree32), + Wasm64(PageTree64), +} + +impl PageTree { + pub fn new(is_64: bool) -> Self { + if is_64 { + Self::Wasm64(PageTree64::new()) + } else { + Self::Wasm32(PageTree32::new()) + } + } + + pub fn insert(&mut self, position: u64, item: impl Into) { + match self { + Self::Wasm32(tree) => tree.insert(position, item), + Self::Wasm64(tree) => tree.insert(position, item), + } + } + + pub fn root(&self) -> Ref { + match self { + Self::Wasm32(tree) => tree.root(), + Self::Wasm64(tree) => tree.root(), + } + } +} + pub type Tree = dusk_merkle::Tree; #[derive(Debug, Clone, Archive, Deserialize, Serialize)] @@ -57,7 +94,7 @@ impl ContractIndex { self.contracts.insert( contract, ContractIndexElement { - tree: PageTree::new(), + tree: PageTree::new(memory.is_64()), len: 0, page_indices: BTreeSet::new(), }, @@ -155,10 +192,10 @@ impl dusk_merkle::Aggregate for Hash { } } -impl dusk_merkle::Aggregate for Hash { +impl dusk_merkle::Aggregate for Hash { const EMPTY_SUBTREE: Self = Hash([0; blake3::OUT_LEN]); - fn aggregate(items: [&Self; P_ARITY]) -> Self { + fn aggregate(items: [&Self; P32_ARITY]) -> Self { let mut hasher = Hasher::new(); for item in items { hasher.update(item.as_bytes());