diff --git a/src/elf.rs b/src/elf.rs new file mode 100644 index 00000000..971dd7c8 --- /dev/null +++ b/src/elf.rs @@ -0,0 +1,206 @@ +// This module maps the data structure of different versions of goblin to the +// same internal structure. +use crate::machine::VERSION1; +use crate::memory::{round_page_down, round_page_up, FLAG_EXECUTABLE, FLAG_FREEZED}; +use crate::{Error, Register}; +use bytes::Bytes; +use scroll::Pread; +use std::ops::Range; + +// Even for different versions of goblin, their values must be consistent. +pub use goblin_v023::elf::program_header::{PF_R, PF_W, PF_X, PT_LOAD}; +pub use goblin_v023::elf::section_header::SHF_EXECINSTR; + +/// Converts goblin's ELF flags into RISC-V flags +pub fn convert_flags(p_flags: u32, allow_freeze_writable: bool) -> Result { + let readable = p_flags & PF_R != 0; + let writable = p_flags & PF_W != 0; + let executable = p_flags & PF_X != 0; + if !readable { + return Err(Error::ElfSegmentUnreadable); + } + if writable && executable { + return Err(Error::ElfSegmentWritableAndExecutable); + } + if executable { + Ok(FLAG_EXECUTABLE | FLAG_FREEZED) + } else if writable && !allow_freeze_writable { + Ok(0) + } else { + Ok(FLAG_FREEZED) + } +} + +/// Same as goblin::elf::ProgramHeader. +pub struct ProgramHeader { + pub p_type: u32, + pub p_flags: u32, + pub p_offset: u64, + pub p_vaddr: u64, + pub p_paddr: u64, + pub p_filesz: u64, + pub p_memsz: u64, + pub p_align: u64, +} + +impl ProgramHeader { + pub fn from_v0(header: &goblin_v023::elf::ProgramHeader) -> Self { + Self { + p_type: header.p_type, + p_flags: header.p_flags, + p_offset: header.p_offset, + p_vaddr: header.p_vaddr, + p_paddr: header.p_paddr, + p_filesz: header.p_filesz, + p_memsz: header.p_memsz, + p_align: header.p_align, + } + } + + pub fn from_v1(header: &goblin_v040::elf::ProgramHeader) -> Self { + Self { + p_type: header.p_type, + p_flags: header.p_flags, + p_offset: header.p_offset, + p_vaddr: header.p_vaddr, + p_paddr: header.p_paddr, + p_filesz: header.p_filesz, + p_memsz: header.p_memsz, + p_align: header.p_align, + } + } +} + +/// Same as goblin::elf::SectionHeader. +pub struct SectionHeader { + pub sh_name: usize, + pub sh_type: u32, + pub sh_flags: u64, + pub sh_addr: u64, + pub sh_offset: u64, + pub sh_size: u64, + pub sh_link: u32, + pub sh_info: u32, + pub sh_addralign: u64, + pub sh_entsize: u64, +} + +impl SectionHeader { + pub fn from_v0(header: &goblin_v023::elf::SectionHeader) -> Self { + Self { + sh_name: header.sh_name, + sh_type: header.sh_type, + sh_flags: header.sh_flags, + sh_addr: header.sh_addr, + sh_offset: header.sh_offset, + sh_size: header.sh_size, + sh_link: header.sh_link, + sh_info: header.sh_info, + sh_addralign: header.sh_addralign, + sh_entsize: header.sh_entsize, + } + } + + pub fn from_v1(header: &goblin_v040::elf::SectionHeader) -> Self { + Self { + sh_name: header.sh_name, + sh_type: header.sh_type, + sh_flags: header.sh_flags, + sh_addr: header.sh_addr, + sh_offset: header.sh_offset, + sh_size: header.sh_size, + sh_link: header.sh_link, + sh_info: header.sh_info, + sh_addralign: header.sh_addralign, + sh_entsize: header.sh_entsize, + } + } +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct LoadingAction { + pub addr: u64, + pub size: u64, + pub flags: u8, + pub source: Range, + pub offset_from_addr: u64, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct ProgramMetadata { + pub actions: Vec, + pub entry: u64, +} + +pub fn parse_elf(program: &Bytes, version: u32) -> Result { + // We did not use Elf::parse here to avoid triggering potential bugs in goblin. + // * https://github.com/nervosnetwork/ckb-vm/issues/143 + let (entry, program_headers): (u64, Vec) = if version < VERSION1 { + use goblin_v023::container::Ctx; + use goblin_v023::elf::{program_header::ProgramHeader as GoblinProgramHeader, Header}; + let header = program.pread::
(0)?; + let container = header.container().map_err(|_e| Error::ElfBits)?; + let endianness = header.endianness().map_err(|_e| Error::ElfBits)?; + if R::BITS != if container.is_big() { 64 } else { 32 } { + return Err(Error::ElfBits); + } + let ctx = Ctx::new(container, endianness); + let program_headers = GoblinProgramHeader::parse( + program, + header.e_phoff as usize, + header.e_phnum as usize, + ctx, + )? + .iter() + .map(ProgramHeader::from_v0) + .collect(); + (header.e_entry, program_headers) + } else { + use goblin_v040::container::Ctx; + use goblin_v040::elf::{program_header::ProgramHeader as GoblinProgramHeader, Header}; + let header = program.pread::
(0)?; + let container = header.container().map_err(|_e| Error::ElfBits)?; + let endianness = header.endianness().map_err(|_e| Error::ElfBits)?; + if R::BITS != if container.is_big() { 64 } else { 32 } { + return Err(Error::ElfBits); + } + let ctx = Ctx::new(container, endianness); + let program_headers = GoblinProgramHeader::parse( + program, + header.e_phoff as usize, + header.e_phnum as usize, + ctx, + )? + .iter() + .map(ProgramHeader::from_v1) + .collect(); + (header.e_entry, program_headers) + }; + let mut bytes: u64 = 0; + let mut actions = vec![]; + for program_header in program_headers { + if program_header.p_type == PT_LOAD { + let aligned_start = round_page_down(program_header.p_vaddr); + let padding_start = program_header.p_vaddr.wrapping_sub(aligned_start); + let size = round_page_up(program_header.p_memsz.wrapping_add(padding_start)); + let slice_start = program_header.p_offset; + let slice_end = program_header + .p_offset + .wrapping_add(program_header.p_filesz); + if slice_start > slice_end || slice_end > program.len() as u64 { + return Err(Error::ElfSegmentAddrOrSizeError); + } + actions.push(LoadingAction { + addr: aligned_start, + size, + flags: convert_flags(program_header.p_flags, version < VERSION1)?, + source: slice_start..slice_end, + offset_from_addr: padding_start, + }); + bytes = bytes.checked_add(slice_end - slice_start).ok_or_else(|| { + Error::Unexpected(String::from("The bytes count overflowed on loading elf")) + })?; + } + } + Ok(ProgramMetadata { actions, entry }) +} diff --git a/src/lib.rs b/src/lib.rs index 103bca61..73a9f33e 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -5,11 +5,13 @@ pub mod bits; pub mod cost_model; pub mod debugger; pub mod decoder; +pub mod elf; pub mod error; pub mod instructions; pub mod machine; pub mod memory; pub mod snapshot; +pub mod snapshot2; pub mod syscalls; pub use bytes; diff --git a/src/machine/asm/mod.rs b/src/machine/asm/mod.rs index de556944..b062c9ab 100644 --- a/src/machine/asm/mod.rs +++ b/src/machine/asm/mod.rs @@ -16,6 +16,7 @@ use std::os::raw::c_uchar; use crate::{ decoder::{build_decoder, Decoder}, + elf::ProgramMetadata, instructions::{ blank_instruction, execute_instruction, extract_opcode, instruction_length, is_basic_block_end_instruction, @@ -423,6 +424,10 @@ impl SupportMachine for Box { self.max_cycles } + fn set_max_cycles(&mut self, max_cycles: u64) { + self.max_cycles = max_cycles; + } + fn reset(&mut self, max_cycles: u64) { self.registers = [0; RISCV_GENERAL_REGISTER_NUMBER]; self.pc = 0; @@ -481,6 +486,16 @@ impl AsmMachine { self.machine.load_program(program, args) } + pub fn load_program_with_metadata( + &mut self, + program: &Bytes, + metadata: &ProgramMetadata, + args: &[Bytes], + ) -> Result { + self.machine + .load_program_with_metadata(program, metadata, args) + } + pub fn run(&mut self) -> Result { if self.machine.isa() & ISA_MOP != 0 && self.machine.version() == VERSION0 { return Err(Error::InvalidVersion); diff --git a/src/machine/elf_adaptor.rs b/src/machine/elf_adaptor.rs deleted file mode 100644 index 99e9e3e0..00000000 --- a/src/machine/elf_adaptor.rs +++ /dev/null @@ -1,114 +0,0 @@ -// This module maps the data structure of different versions of goblin to the -// same internal structure. -use crate::memory::{FLAG_EXECUTABLE, FLAG_FREEZED}; -use crate::Error; - -// Even for different versions of goblin, their values must be consistent. -pub use goblin_v023::elf::program_header::{PF_R, PF_W, PF_X, PT_LOAD}; -pub use goblin_v023::elf::section_header::SHF_EXECINSTR; - -/// Converts goblin's ELF flags into RISC-V flags -pub fn convert_flags(p_flags: u32, allow_freeze_writable: bool) -> Result { - let readable = p_flags & PF_R != 0; - let writable = p_flags & PF_W != 0; - let executable = p_flags & PF_X != 0; - if !readable { - return Err(Error::ElfSegmentUnreadable); - } - if writable && executable { - return Err(Error::ElfSegmentWritableAndExecutable); - } - if executable { - Ok(FLAG_EXECUTABLE | FLAG_FREEZED) - } else if writable && !allow_freeze_writable { - Ok(0) - } else { - Ok(FLAG_FREEZED) - } -} - -/// Same as goblin::elf::ProgramHeader. -pub struct ProgramHeader { - pub p_type: u32, - pub p_flags: u32, - pub p_offset: u64, - pub p_vaddr: u64, - pub p_paddr: u64, - pub p_filesz: u64, - pub p_memsz: u64, - pub p_align: u64, -} - -impl ProgramHeader { - pub fn from_v0(header: &goblin_v023::elf::ProgramHeader) -> Self { - Self { - p_type: header.p_type, - p_flags: header.p_flags, - p_offset: header.p_offset, - p_vaddr: header.p_vaddr, - p_paddr: header.p_paddr, - p_filesz: header.p_filesz, - p_memsz: header.p_memsz, - p_align: header.p_align, - } - } - - pub fn from_v1(header: &goblin_v040::elf::ProgramHeader) -> Self { - Self { - p_type: header.p_type, - p_flags: header.p_flags, - p_offset: header.p_offset, - p_vaddr: header.p_vaddr, - p_paddr: header.p_paddr, - p_filesz: header.p_filesz, - p_memsz: header.p_memsz, - p_align: header.p_align, - } - } -} - -/// Same as goblin::elf::SectionHeader. -pub struct SectionHeader { - pub sh_name: usize, - pub sh_type: u32, - pub sh_flags: u64, - pub sh_addr: u64, - pub sh_offset: u64, - pub sh_size: u64, - pub sh_link: u32, - pub sh_info: u32, - pub sh_addralign: u64, - pub sh_entsize: u64, -} - -impl SectionHeader { - pub fn from_v0(header: &goblin_v023::elf::SectionHeader) -> Self { - Self { - sh_name: header.sh_name, - sh_type: header.sh_type, - sh_flags: header.sh_flags, - sh_addr: header.sh_addr, - sh_offset: header.sh_offset, - sh_size: header.sh_size, - sh_link: header.sh_link, - sh_info: header.sh_info, - sh_addralign: header.sh_addralign, - sh_entsize: header.sh_entsize, - } - } - - pub fn from_v1(header: &goblin_v040::elf::SectionHeader) -> Self { - Self { - sh_name: header.sh_name, - sh_type: header.sh_type, - sh_flags: header.sh_flags, - sh_addr: header.sh_addr, - sh_offset: header.sh_offset, - sh_size: header.sh_size, - sh_link: header.sh_link, - sh_info: header.sh_info, - sh_addralign: header.sh_addralign, - sh_entsize: header.sh_entsize, - } - } -} diff --git a/src/machine/mod.rs b/src/machine/mod.rs index c3623bf0..cd1141f1 100644 --- a/src/machine/mod.rs +++ b/src/machine/mod.rs @@ -1,6 +1,5 @@ #[cfg(has_asm)] pub mod asm; -pub mod elf_adaptor; pub mod trace; use std::fmt::{self, Display}; @@ -8,12 +7,12 @@ use std::sync::atomic::{AtomicU8, Ordering}; use std::sync::Arc; use bytes::Bytes; -use scroll::Pread; use super::debugger::Debugger; use super::decoder::{build_decoder, Decoder}; +use super::elf::{parse_elf, LoadingAction, ProgramMetadata}; use super::instructions::{execute, Instruction, Register}; -use super::memory::{round_page_down, round_page_up, Memory}; +use super::memory::Memory; use super::syscalls::Syscalls; use super::{ registers::{A0, A7, REGISTER_ABI_NAMES, SP}, @@ -70,6 +69,7 @@ pub trait SupportMachine: CoreMachine { fn cycles(&self) -> u64; fn set_cycles(&mut self, cycles: u64); fn max_cycles(&self) -> u64; + fn set_max_cycles(&mut self, cycles: u64); fn running(&self) -> bool; fn set_running(&mut self, running: bool); @@ -101,84 +101,8 @@ pub trait SupportMachine: CoreMachine { fn load_elf_inner(&mut self, program: &Bytes, update_pc: bool) -> Result { let version = self.version(); - // We did not use Elf::parse here to avoid triggering potential bugs in goblin. - // * https://github.com/nervosnetwork/ckb-vm/issues/143 - let (e_entry, program_headers): (u64, Vec) = - if version < VERSION1 { - use goblin_v023::container::Ctx; - use goblin_v023::elf::{program_header::ProgramHeader, Header}; - let header = program.pread::
(0)?; - let container = header.container().map_err(|_e| Error::ElfBits)?; - let endianness = header.endianness().map_err(|_e| Error::ElfBits)?; - if Self::REG::BITS != if container.is_big() { 64 } else { 32 } { - return Err(Error::ElfBits); - } - let ctx = Ctx::new(container, endianness); - let program_headers = ProgramHeader::parse( - program, - header.e_phoff as usize, - header.e_phnum as usize, - ctx, - )? - .iter() - .map(elf_adaptor::ProgramHeader::from_v0) - .collect(); - (header.e_entry, program_headers) - } else { - use goblin_v040::container::Ctx; - use goblin_v040::elf::{program_header::ProgramHeader, Header}; - let header = program.pread::
(0)?; - let container = header.container().map_err(|_e| Error::ElfBits)?; - let endianness = header.endianness().map_err(|_e| Error::ElfBits)?; - if Self::REG::BITS != if container.is_big() { 64 } else { 32 } { - return Err(Error::ElfBits); - } - let ctx = Ctx::new(container, endianness); - let program_headers = ProgramHeader::parse( - program, - header.e_phoff as usize, - header.e_phnum as usize, - ctx, - )? - .iter() - .map(elf_adaptor::ProgramHeader::from_v1) - .collect(); - (header.e_entry, program_headers) - }; - let mut bytes: u64 = 0; - for program_header in program_headers { - if program_header.p_type == elf_adaptor::PT_LOAD { - let aligned_start = round_page_down(program_header.p_vaddr); - let padding_start = program_header.p_vaddr.wrapping_sub(aligned_start); - let size = round_page_up(program_header.p_memsz.wrapping_add(padding_start)); - let slice_start = program_header.p_offset; - let slice_end = program_header - .p_offset - .wrapping_add(program_header.p_filesz); - if slice_start > slice_end || slice_end > program.len() as u64 { - return Err(Error::ElfSegmentAddrOrSizeError); - } - self.memory_mut().init_pages( - aligned_start, - size, - elf_adaptor::convert_flags(program_header.p_flags, version < VERSION1)?, - Some(program.slice(slice_start as usize..slice_end as usize)), - padding_start, - )?; - if version < VERSION1 { - self.memory_mut() - .store_byte(aligned_start, padding_start, 0)?; - } - bytes = bytes.checked_add(slice_end - slice_start).ok_or_else(|| { - Error::Unexpected(String::from("The bytes count overflowed on loading elf")) - })?; - } - } - if update_pc { - self.update_pc(Self::REG::from_u64(e_entry)); - self.commit_pc(); - } - Ok(bytes) + let metadata = parse_elf::(program, version)?; + self.load_binary(program, &metadata, update_pc) } fn load_elf(&mut self, program: &Bytes, update_pc: bool) -> Result { @@ -195,6 +119,56 @@ pub trait SupportMachine: CoreMachine { self.load_elf_inner(program, update_pc) } + fn load_binary_inner( + &mut self, + program: &Bytes, + metadata: &ProgramMetadata, + update_pc: bool, + ) -> Result { + let version = self.version(); + let mut bytes: u64 = 0; + for action in &metadata.actions { + let LoadingAction { + addr, + size, + flags, + source, + offset_from_addr, + } = action; + + self.memory_mut().init_pages( + *addr, + *size, + *flags, + Some(program.slice(source.start as usize..source.end as usize)), + *offset_from_addr, + )?; + if version < VERSION1 { + self.memory_mut().store_byte(*addr, *offset_from_addr, 0)?; + } + bytes = bytes + .checked_add(source.end - source.start) + .ok_or_else(|| { + Error::Unexpected(String::from("The bytes count overflowed on loading elf")) + })?; + } + if update_pc { + self.update_pc(Self::REG::from_u64(metadata.entry)); + self.commit_pc(); + } + Ok(bytes) + } + + fn load_binary( + &mut self, + program: &Bytes, + metadata: &ProgramMetadata, + update_pc: bool, + ) -> Result { + // Similar to load_elf, this provides a way to adjust the behavior of load_binary_inner + self.load_binary_inner(program, metadata, update_pc) + } + fn initialize_stack( &mut self, args: &[Bytes], @@ -351,6 +325,10 @@ impl> SupportMachine for DefaultCoreMachine SupportMachine for DefaultMachine { self.inner.max_cycles() } + fn set_max_cycles(&mut self, max_cycles: u64) { + self.inner.set_max_cycles(max_cycles) + } + fn reset(&mut self, max_cycles: u64) { self.inner_mut().reset(max_cycles); } @@ -566,6 +548,32 @@ impl Display for DefaultMachine { impl DefaultMachine { pub fn load_program(&mut self, program: &Bytes, args: &[Bytes]) -> Result { let elf_bytes = self.load_elf(program, true)?; + let stack_bytes = self.initialize(args)?; + let bytes = elf_bytes.checked_add(stack_bytes).ok_or_else(|| { + Error::Unexpected(String::from( + "The bytes count overflowed on loading program", + )) + })?; + Ok(bytes) + } + + pub fn load_program_with_metadata( + &mut self, + program: &Bytes, + metadata: &ProgramMetadata, + args: &[Bytes], + ) -> Result { + let elf_bytes = self.load_binary(program, metadata, true)?; + let stack_bytes = self.initialize(args)?; + let bytes = elf_bytes.checked_add(stack_bytes).ok_or_else(|| { + Error::Unexpected(String::from( + "The bytes count overflowed on loading program", + )) + })?; + Ok(bytes) + } + + fn initialize(&mut self, args: &[Bytes]) -> Result { for syscall in &mut self.syscalls { syscall.initialize(&mut self.inner)?; } @@ -580,12 +588,7 @@ impl DefaultMachine { if self.inner.version() >= VERSION1 { debug_assert!(self.registers()[SP].to_u64() % 16 == 0); } - let bytes = elf_bytes.checked_add(stack_bytes).ok_or_else(|| { - Error::Unexpected(String::from( - "The bytes count overflowed on loading program", - )) - })?; - Ok(bytes) + Ok(stack_bytes) } pub fn take_inner(self) -> Inner { diff --git a/src/machine/trace.rs b/src/machine/trace.rs index 7eff3f9a..be25817a 100644 --- a/src/machine/trace.rs +++ b/src/machine/trace.rs @@ -1,6 +1,7 @@ use super::{ super::{ decoder::build_decoder, + elf::ProgramMetadata, instructions::{ execute_with_thread, extract_opcode, handle_invalid_op, instruction_length, is_basic_block_end_instruction, Instruction, Register, Thread, ThreadFactory, @@ -116,6 +117,16 @@ impl TraceMachine { self.machine.load_program(program, args) } + pub fn load_program_with_metadata( + &mut self, + program: &Bytes, + metadata: &ProgramMetadata, + args: &[Bytes], + ) -> Result { + self.machine + .load_program_with_metadata(program, metadata, args) + } + pub fn run(&mut self) -> Result { let mut decoder = build_decoder::(self.isa(), self.version()); self.machine.set_running(true); diff --git a/src/snapshot2.rs b/src/snapshot2.rs new file mode 100644 index 00000000..dc271048 --- /dev/null +++ b/src/snapshot2.rs @@ -0,0 +1,265 @@ +use crate::{ + bits::roundup, + elf::{LoadingAction, ProgramMetadata}, + machine::SupportMachine, + memory::{Memory, FLAG_DIRTY}, + Error, Register, RISCV_GENERAL_REGISTER_NUMBER, RISCV_PAGESIZE, +}; +use bytes::Bytes; +use serde::{Deserialize, Serialize}; +use std::cmp::min; +use std::collections::HashMap; + +const PAGE_SIZE: u64 = RISCV_PAGESIZE as u64; + +/// DataSource represents data source that can stay stable and possibly +/// immutable for the entire lifecycle duration of a VM instance. One example +/// can be the enclosing transaction when using CKB-VM in CKB's environment, +/// no matter where and when we run the CKB smart contract, the enclosing +/// transaction will always be the same down to every last byte. As a result, +/// we can leverage DataSource for snapshot optimizations: data that is already +/// locatable in the DataSource will not need to be included in the snapshot +/// again, all we need is an id to locate it, together with a pair of +/// offset / length to cut in to the correct slices. +pub trait DataSource { + fn load_data(&self, id: &I, offset: u64, length: u64) -> Result; +} + +#[derive(Clone, Debug)] +pub struct Snapshot2Context> { + // page index -> (id, offset, flag) + pages: HashMap, + data_source: D, +} + +impl + Default> Default for Snapshot2Context { + fn default() -> Self { + Self::new(D::default()) + } +} + +impl> Snapshot2Context { + pub fn new(data_source: D) -> Self { + Self { + pages: HashMap::default(), + data_source, + } + } + + /// Resume a previously suspended machine from snapshot. + pub fn resume( + &mut self, + machine: &mut M, + snapshot: &Snapshot2, + ) -> Result<(), Error> { + if machine.version() != snapshot.version { + return Err(Error::InvalidVersion); + } + // A resume basically means we reside in a new context + self.pages.clear(); + for (i, v) in snapshot.registers.iter().enumerate() { + machine.set_register(i, M::REG::from_u64(*v)); + } + machine.update_pc(M::REG::from_u64(snapshot.pc)); + machine.commit_pc(); + machine.set_cycles(snapshot.cycles); + machine.set_max_cycles(snapshot.max_cycles); + for (address, flag, id, offset, length) in &snapshot.pages_from_source { + if address % PAGE_SIZE != 0 { + return Err(Error::MemPageUnalignedAccess); + } + let data = self.data_source().load_data(id, *offset, *length)?; + if data.len() as u64 % PAGE_SIZE != 0 { + return Err(Error::MemPageUnalignedAccess); + } + machine.memory_mut().store_bytes(*address, &data)?; + for i in 0..(data.len() as u64 / PAGE_SIZE) { + let page = address / PAGE_SIZE + i; + machine.memory_mut().set_flag(page, *flag)?; + } + self.track_pages(machine, *address, data.len() as u64, id, *offset)?; + } + for (address, flag, content) in &snapshot.dirty_pages { + if address % PAGE_SIZE != 0 { + return Err(Error::MemPageUnalignedAccess); + } + if content.len() as u64 % PAGE_SIZE != 0 { + return Err(Error::MemPageUnalignedAccess); + } + machine.memory_mut().store_bytes(*address, content)?; + for i in 0..(content.len() as u64 / PAGE_SIZE) { + let page = address / PAGE_SIZE + i; + machine.memory_mut().set_flag(page, *flag)?; + } + } + machine + .memory_mut() + .set_lr(&M::REG::from_u64(snapshot.load_reservation_address)); + Ok(()) + } + + pub fn data_source(&self) -> &D { + &self.data_source + } + + pub fn data_source_mut(&mut self) -> &mut D { + &mut self.data_source + } + + /// Similar to Memory::store_bytes, but this method also tracks memory + /// pages whose entire content comes from DataSource + pub fn store_bytes( + &mut self, + machine: &mut M, + addr: u64, + id: &I, + offset: u64, + length: u64, + ) -> Result<(), Error> { + let data = self.data_source.load_data(id, offset, length)?; + machine.memory_mut().store_bytes(addr, &data)?; + self.track_pages(machine, addr, data.len() as u64, id, offset) + } + + /// Due to the design of ckb-vm right now, load_program function does not + /// belong to SupportMachine yet. For Snapshot2Context to track memory pages + /// from program in DataSource, we have to use the following steps now: + /// + /// 1. use elf::parse_elf to generate ProgramMetadata + /// 2. use DefaultMachine::load_program_with_metadata to load the program + /// 3. Pass ProgramMetadata to this method so we can track memory pages from + /// program, so as to further reduce the size of the generated snapshot. + /// + /// One can also use the original DefaultMachine::load_program, and parse the + /// ELF a second time to extract metadata for this method. However the above + /// listed process saves us the time to parse the ELF again. + pub fn mark_program( + &mut self, + machine: &mut M, + metadata: &ProgramMetadata, + id: &I, + offset: u64, + ) -> Result<(), Error> { + for action in &metadata.actions { + self.init_pages(machine, action, id, offset)?; + } + Ok(()) + } + + /// Create a snapshot for the passed machine. + pub fn make_snapshot(&self, machine: &mut M) -> Result, Error> { + let mut dirty_pages: Vec<(u64, u8, Vec)> = vec![]; + for i in 0..machine.memory().memory_pages() as u64 { + if self.pages.contains_key(&i) { + continue; + } + let flag = machine.memory_mut().fetch_flag(i)?; + if flag & FLAG_DIRTY == 0 { + continue; + } + let address = i * PAGE_SIZE; + let mut data: Vec = machine.memory_mut().load_bytes(address, PAGE_SIZE)?.into(); + if let Some(last) = dirty_pages.last_mut() { + if last.0 + last.2.len() as u64 == address && last.1 == flag { + last.2.append(&mut data); + } + } + if !data.is_empty() { + dirty_pages.push((address, flag, data)); + } + } + let mut pages_from_source: Vec<(u64, u8, I, u64, u64)> = vec![]; + let mut pages: Vec = self.pages.keys().copied().collect(); + pages.sort_unstable(); + for page in pages { + let address = page * PAGE_SIZE; + let (id, offset, flag) = &self.pages[&page]; + let mut appended_to_last = false; + if let Some((last_address, last_flag, last_id, last_offset, last_length)) = + pages_from_source.last_mut() + { + if *last_address + *last_length == address + && *last_flag == *flag + && *last_id == *id + && *last_offset + *last_length == *offset + { + *last_length += PAGE_SIZE; + appended_to_last = true; + } + } + if !appended_to_last { + pages_from_source.push((address, *flag, id.clone(), *offset, PAGE_SIZE)); + } + } + let mut registers = [0u64; RISCV_GENERAL_REGISTER_NUMBER]; + for (i, v) in machine.registers().iter().enumerate() { + registers[i] = v.to_u64(); + } + Ok(Snapshot2 { + pages_from_source, + dirty_pages, + version: machine.version(), + registers, + pc: machine.pc().to_u64(), + cycles: machine.cycles(), + max_cycles: machine.max_cycles(), + load_reservation_address: machine.memory().lr().to_u64(), + }) + } + + fn init_pages( + &mut self, + machine: &mut M, + action: &LoadingAction, + id: &I, + offset: u64, + ) -> Result<(), Error> { + let start = action.addr + action.offset_from_addr; + let length = min( + action.source.end - action.source.start, + action.size - action.offset_from_addr, + ); + self.track_pages(machine, start, length, id, offset + action.source.start) + } + + fn track_pages( + &mut self, + machine: &mut M, + start: u64, + mut length: u64, + id: &I, + mut offset: u64, + ) -> Result<(), Error> { + let mut aligned_start = roundup(start, PAGE_SIZE); + let aligned_bytes = aligned_start - start; + if length < aligned_bytes { + return Ok(()); + } + offset += aligned_bytes; + length -= aligned_bytes; + while length >= PAGE_SIZE { + let page = aligned_start / PAGE_SIZE; + machine.memory_mut().clear_flag(page, FLAG_DIRTY)?; + let flag = machine.memory_mut().fetch_flag(page)?; + self.pages.insert(page, (id.clone(), offset, flag)); + aligned_start += PAGE_SIZE; + length -= PAGE_SIZE; + offset += PAGE_SIZE; + } + Ok(()) + } +} + +#[derive(Clone, Debug, Deserialize, Serialize)] +pub struct Snapshot2 { + // (address, flag, id, source offset, source length) + pub pages_from_source: Vec<(u64, u8, I, u64, u64)>, + // (address, flag, content) + pub dirty_pages: Vec<(u64, u8, Vec)>, + pub version: u32, + pub registers: [u64; RISCV_GENERAL_REGISTER_NUMBER], + pub pc: u64, + pub cycles: u64, + pub max_cycles: u64, + pub load_reservation_address: u64, +} diff --git a/tests/programs/_build_all_native.sh b/tests/programs/_build_all_native.sh index 1987653a..07005aea 100755 --- a/tests/programs/_build_all_native.sh +++ b/tests/programs/_build_all_native.sh @@ -1,5 +1,6 @@ set -ex +riscv64-unknown-elf-gcc -o resume2_load_data resume2_load_data.c riscv64-unknown-elf-gcc -o alloc_many alloc_many.c riscv64-unknown-elf-as -o amo_check_write.o amo_check_write.S && riscv64-unknown-elf-ld -T amo_check_write.lds -o amo_check_write amo_check_write.o && rm amo_check_write.o riscv64-unknown-elf-as -o amo_compare.o amo_compare.S && riscv64-unknown-elf-ld -T amo_compare.lds -o amo_compare amo_compare.o && rm amo_compare.o diff --git a/tests/programs/resume2_load_data b/tests/programs/resume2_load_data new file mode 100755 index 00000000..d69b2811 Binary files /dev/null and b/tests/programs/resume2_load_data differ diff --git a/tests/programs/resume2_load_data.c b/tests/programs/resume2_load_data.c new file mode 100644 index 00000000..8fdf42cd --- /dev/null +++ b/tests/programs/resume2_load_data.c @@ -0,0 +1,40 @@ +#include + +static inline long __internal_syscall(long n, long _a0, long _a1, long _a2, + long _a3, long _a4, long _a5) { + register long a0 asm("a0") = _a0; + register long a1 asm("a1") = _a1; + register long a2 asm("a2") = _a2; + register long a3 asm("a3") = _a3; + register long a4 asm("a4") = _a4; + register long a5 asm("a5") = _a5; + +#ifdef __riscv_32e + register long syscall_id asm("t0") = n; +#else + register long syscall_id asm("a7") = n; +#endif + + asm volatile("scall" + : "+r"(a0) + : "r"(a1), "r"(a2), "r"(a3), "r"(a4), "r"(a5), "r"(syscall_id)); + return a0; +} + +#define syscall(n, a, b, c, d, e, f) \ + __internal_syscall(n, (long)(a), (long)(b), (long)(c), (long)(d), (long)(e), \ + (long)(f)) + +#define SIZE (16 * 4096) + +int main() { + volatile char buf[SIZE] = { 0 }; + syscall(1111, (void *) buf, SIZE, 0, 0, 0, 0); + + char ret = 0; + for (int i = 0; i < SIZE; i++) { + ret += buf[i]; + } + + return (int) ret; +} diff --git a/tests/test_a_extension.rs b/tests/test_a_extension.rs index 99073492..6e60dbbf 100644 --- a/tests/test_a_extension.rs +++ b/tests/test_a_extension.rs @@ -1,4 +1,6 @@ -use ckb_vm::{CoreMachine, Error, Memory}; +use ckb_vm::Error; +#[cfg(has_asm)] +use ckb_vm::{CoreMachine, Memory}; pub mod machine_build; #[test] diff --git a/tests/test_resume2.rs b/tests/test_resume2.rs new file mode 100644 index 00000000..96717507 --- /dev/null +++ b/tests/test_resume2.rs @@ -0,0 +1,622 @@ +#![cfg(has_asm)] +pub mod machine_build; +use bytes::Bytes; +use ckb_vm::cost_model::constant_cycles; +use ckb_vm::elf::parse_elf; +use ckb_vm::machine::asm::{AsmCoreMachine, AsmMachine}; +use ckb_vm::machine::trace::TraceMachine; +use ckb_vm::machine::{ + CoreMachine, DefaultCoreMachine, DefaultMachine, SupportMachine, VERSION0, VERSION1, VERSION2, +}; +use ckb_vm::memory::{sparse::SparseMemory, wxorx::WXorXMemory}; +use ckb_vm::registers::{A0, A1, A7}; +use ckb_vm::snapshot2::{DataSource, Snapshot2, Snapshot2Context}; +use ckb_vm::{DefaultMachineBuilder, Error, Register, Syscalls, ISA_A, ISA_IMC}; +use std::collections::HashMap; +use std::fs::File; +use std::io::Read; +use std::sync::{Arc, Mutex}; + +#[test] +fn test_resume2_interpreter_with_trace_2_asm() { + resume_interpreter_with_trace_2_asm_inner(VERSION1, 8126917); + resume_interpreter_with_trace_2_asm_inner(VERSION0, 8126917); +} + +#[test] +fn test_resume2_interpreter_2_asm() { + resume_interpreter_2_asm(VERSION1, 8126917); + resume_interpreter_2_asm(VERSION0, 8126917); +} + +#[test] +fn test_resume2_interpreter_2_interpreter() { + resume_interpreter_2_interpreter(VERSION1, 8126917); + resume_interpreter_2_interpreter(VERSION0, 8126917); +} + +#[test] +fn test_resume2_asm_2_interpreter() { + resume_asm_2_interpreter(VERSION1, 8126917); + resume_asm_2_interpreter(VERSION0, 8126917); +} + +#[test] +fn test_resume2_asm_2_asm_2_asm() { + resume_asm_2_asm_2_asm(VERSION1, 8126917); + resume_asm_2_asm_2_asm(VERSION0, 8126917); +} + +#[test] +fn test_resume2_asm_2_asm() { + resume_asm_2_asm(VERSION1, 8126917); + resume_asm_2_asm(VERSION0, 8126917); +} + +#[test] +fn test_resume2_secp256k1_asm_2_interpreter_2_asm() { + let data_source = load_program("benches/data/secp256k1_bench"); + + let version = VERSION1; + let except_cycles = 613073; + + let mut machine1 = MachineTy::Asm.build(data_source.clone(), version); + machine1.set_max_cycles(100000); + machine1 + .load_program(&vec![ + "secp256k1_bench".into(), + "033f8cf9c4d51a33206a6c1c6b27d2cc5129daa19dbd1fc148d395284f6b26411f".into(), + "304402203679d909f43f073c7c1dcf8468a485090589079ee834e6eed92fea9b09b06a2402201e46f1075afa18f306715e7db87493e7b7e779569aa13c64ab3d09980b3560a3".into(), + "foo".into(), + "bar".into(), + ]) + .unwrap(); + let result1 = machine1.run(); + assert_eq!(result1.unwrap_err(), Error::CyclesExceeded); + let snapshot1 = machine1.snapshot().unwrap(); + assert!(!snapshot1.pages_from_source.is_empty()); + + let mut machine2 = MachineTy::Interpreter.build(data_source.clone(), version); + machine2.resume(snapshot1).unwrap(); + + assert_eq!(machine1.cycles(), machine2.cycles()); + assert_eq!(machine1.full_registers(), machine2.full_registers()); + #[cfg(not(feature = "enable-chaos-mode-by-default"))] + assert_eq!(machine1.full_memory(), machine2.full_memory()); + + machine2.set_max_cycles(100000 + 200000); + let result2 = machine2.run(); + assert_eq!(result2.unwrap_err(), Error::CyclesExceeded); + let snapshot2 = machine2.snapshot().unwrap(); + assert!(!snapshot2.pages_from_source.is_empty()); + + let mut machine3 = MachineTy::Asm.build(data_source, version); + machine3.resume(snapshot2).unwrap(); + + assert_eq!(machine2.cycles(), machine3.cycles()); + assert_eq!(machine2.full_registers(), machine3.full_registers()); + #[cfg(not(feature = "enable-chaos-mode-by-default"))] + assert_eq!(machine2.full_memory(), machine3.full_memory()); + + machine3.set_max_cycles(100000 + 200000 + 400000); + let result3 = machine3.run(); + let cycles3 = machine3.cycles(); + assert_eq!(result3.unwrap(), 0); + assert_eq!(cycles3, except_cycles); +} + +#[test] +fn test_resume2_load_data_asm_2_interpreter() { + let data_source = load_program("tests/programs/resume2_load_data"); + + let version = VERSION1; + let except_cycles = 1476715; + + let mut machine1 = MachineTy::Asm.build(data_source.clone(), version); + machine1.set_max_cycles(300000); + machine1 + .load_program(&vec!["resume2_load_data".into()]) + .unwrap(); + let result1 = machine1.run(); + assert!(result1.is_err()); + assert_eq!(result1.unwrap_err(), Error::CyclesExceeded); + let snapshot = machine1.snapshot().unwrap(); + assert!(!snapshot.pages_from_source.is_empty()); + + let mut machine2 = MachineTy::Interpreter.build(data_source, version); + machine2.resume(snapshot).unwrap(); + + assert_eq!(machine1.cycles(), machine2.cycles()); + assert_eq!(machine1.full_registers(), machine2.full_registers()); + #[cfg(not(feature = "enable-chaos-mode-by-default"))] + assert_eq!(machine1.full_memory(), machine2.full_memory()); + + machine2.set_max_cycles(except_cycles + 10); + + let result2 = machine2.run(); + let cycles2 = machine2.cycles(); + assert_eq!(result2.unwrap(), 0); + assert_eq!(cycles2, except_cycles); +} + +#[test] +fn test_resume2_load_data_interpreter_2_asm() { + let data_source = load_program("tests/programs/resume2_load_data"); + + let version = VERSION1; + let except_cycles = 1476715; + + let mut machine1 = MachineTy::Interpreter.build(data_source.clone(), version); + machine1.set_max_cycles(300000); + machine1 + .load_program(&vec!["resume2_load_data".into()]) + .unwrap(); + let result1 = machine1.run(); + assert!(result1.is_err()); + assert_eq!(result1.unwrap_err(), Error::CyclesExceeded); + let snapshot = machine1.snapshot().unwrap(); + assert!(!snapshot.pages_from_source.is_empty()); + + let mut machine2 = MachineTy::Asm.build(data_source, version); + machine2.resume(snapshot).unwrap(); + + assert_eq!(machine1.cycles(), machine2.cycles()); + assert_eq!(machine1.full_registers(), machine2.full_registers()); + #[cfg(not(feature = "enable-chaos-mode-by-default"))] + assert_eq!(machine1.full_memory(), machine2.full_memory()); + + machine2.set_max_cycles(except_cycles + 10); + + let result2 = machine2.run(); + let cycles2 = machine2.cycles(); + assert_eq!(result2.unwrap(), 0); + assert_eq!(cycles2, except_cycles); +} + +pub fn resume_asm_2_asm(version: u32, except_cycles: u64) { + let data_source = load_program("tests/programs/alloc_many"); + + // The cycles required for complete execution is 4194622 + let mut machine1 = MachineTy::Asm.build(data_source.clone(), version); + machine1.set_max_cycles(except_cycles - 30); + machine1.load_program(&vec!["alloc_many".into()]).unwrap(); + let result1 = machine1.run(); + assert!(result1.is_err()); + assert_eq!(result1.unwrap_err(), Error::CyclesExceeded); + let snapshot = machine1.snapshot().unwrap(); + + let mut machine2 = MachineTy::Asm.build(data_source, version); + machine2.resume(snapshot).unwrap(); + machine2.set_max_cycles(except_cycles + 10); + let result2 = machine2.run(); + let cycles2 = machine2.cycles(); + assert_eq!(result2.unwrap(), 0); + assert_eq!(cycles2, except_cycles); +} + +pub fn resume_asm_2_asm_2_asm(version: u32, except_cycles: u64) { + let data_source = load_program("tests/programs/alloc_many"); + + let mut machine1 = MachineTy::Asm.build(data_source.clone(), version); + machine1.set_max_cycles(1000000); + machine1.load_program(&vec!["alloc_many".into()]).unwrap(); + let result1 = machine1.run(); + assert!(result1.is_err()); + assert_eq!(result1.unwrap_err(), Error::CyclesExceeded); + let snapshot1 = machine1.snapshot().unwrap(); + + let mut machine2 = MachineTy::Asm.build(data_source.clone(), version); + machine2.resume(snapshot1).unwrap(); + machine2.set_max_cycles(1000000 + 4000000); + let result2 = machine2.run(); + assert!(result2.is_err()); + assert_eq!(result2.unwrap_err(), Error::CyclesExceeded); + let snapshot2 = machine2.snapshot().unwrap(); + + let mut machine3 = MachineTy::Asm.build(data_source, version); + machine3.resume(snapshot2).unwrap(); + machine3.set_max_cycles(1000000 + 4000000 + 4000000); + let result3 = machine3.run(); + let cycles3 = machine3.cycles(); + assert_eq!(result3.unwrap(), 0); + assert_eq!(cycles3, except_cycles); +} + +pub fn resume_asm_2_interpreter(version: u32, except_cycles: u64) { + let data_source = load_program("tests/programs/alloc_many"); + + let mut machine1 = MachineTy::Asm.build(data_source.clone(), version); + machine1.set_max_cycles(except_cycles - 30); + machine1.load_program(&vec!["alloc_many".into()]).unwrap(); + let result1 = machine1.run(); + assert!(result1.is_err()); + assert_eq!(result1.unwrap_err(), Error::CyclesExceeded); + let snapshot = machine1.snapshot().unwrap(); + + let mut machine2 = MachineTy::Interpreter.build(data_source, version); + machine2.resume(snapshot).unwrap(); + machine2.set_max_cycles(except_cycles + 10); + + let result2 = machine2.run(); + let cycles2 = machine2.cycles(); + assert_eq!(result2.unwrap(), 0); + assert_eq!(cycles2, except_cycles); +} + +pub fn resume_interpreter_2_interpreter(version: u32, except_cycles: u64) { + let data_source = load_program("tests/programs/alloc_many"); + + let mut machine1 = MachineTy::Interpreter.build(data_source.clone(), version); + machine1.set_max_cycles(except_cycles - 30); + machine1.load_program(&vec!["alloc_many".into()]).unwrap(); + let result1 = machine1.run(); + assert!(result1.is_err()); + assert_eq!(result1.unwrap_err(), Error::CyclesExceeded); + let snapshot = machine1.snapshot().unwrap(); + + let mut machine2 = MachineTy::Interpreter.build(data_source, version); + machine2.resume(snapshot).unwrap(); + machine2.set_max_cycles(except_cycles + 10); + let result2 = machine2.run(); + let cycles2 = machine2.cycles(); + assert_eq!(result2.unwrap(), 0); + assert_eq!(cycles2, except_cycles); +} + +pub fn resume_interpreter_2_asm(version: u32, except_cycles: u64) { + let data_source = load_program("tests/programs/alloc_many"); + + let mut machine1 = MachineTy::Interpreter.build(data_source.clone(), version); + machine1.set_max_cycles(except_cycles - 30); + machine1.load_program(&vec!["alloc_many".into()]).unwrap(); + let result1 = machine1.run(); + assert!(result1.is_err()); + assert_eq!(result1.unwrap_err(), Error::CyclesExceeded); + let snapshot = machine1.snapshot().unwrap(); + + let mut machine2 = MachineTy::Asm.build(data_source, version); + machine2.resume(snapshot).unwrap(); + machine2.set_max_cycles(except_cycles); + let result2 = machine2.run(); + let cycles2 = machine2.cycles(); + assert_eq!(result2.unwrap(), 0); + assert_eq!(cycles2, except_cycles); +} + +pub fn resume_interpreter_with_trace_2_asm_inner(version: u32, except_cycles: u64) { + let data_source = load_program("tests/programs/alloc_many"); + + let mut machine1 = MachineTy::InterpreterWithTrace.build(data_source.clone(), version); + machine1.set_max_cycles(except_cycles - 30); + machine1.load_program(&vec!["alloc_many".into()]).unwrap(); + let result1 = machine1.run(); + assert!(result1.is_err()); + assert_eq!(result1.unwrap_err(), Error::CyclesExceeded); + let snapshot = machine1.snapshot().unwrap(); + + let mut machine2 = MachineTy::Asm.build(data_source, version); + machine2.resume(snapshot).unwrap(); + machine2.set_max_cycles(except_cycles); + let result2 = machine2.run(); + let cycles2 = machine2.cycles(); + assert_eq!(result2.unwrap(), 0); + assert_eq!(cycles2, except_cycles); +} + +fn load_program(name: &str) -> TestSource { + let mut file = File::open(name).unwrap(); + let mut buffer = Vec::new(); + file.read_to_end(&mut buffer).unwrap(); + let program = buffer.into(); + + let data = vec![7; 16 * 4096]; + let mut m = HashMap::default(); + m.insert(DATA_ID, data.into()); + m.insert(PROGRAM_ID, program); + + TestSource(m) +} + +const PROGRAM_ID: u64 = 0x1234; +const DATA_ID: u64 = 0x2000; + +#[derive(Clone)] +struct TestSource(HashMap); + +impl DataSource for TestSource { + fn load_data(&self, id: &u64, offset: u64, length: u64) -> Result { + match self.0.get(id) { + Some(data) => Ok(data.slice( + offset as usize..(if length > 0 { + (offset + length) as usize + } else { + data.len() + }), + )), + None => Err(Error::Unexpected(format!( + "Id {} is missing in source!", + id + ))), + } + } +} + +struct InsertDataSyscall(Arc>>); + +impl Syscalls for InsertDataSyscall { + fn initialize(&mut self, _machine: &mut Mac) -> Result<(), Error> { + Ok(()) + } + + fn ecall(&mut self, machine: &mut Mac) -> Result { + let code = &machine.registers()[A7]; + if code.to_i32() != 1111 { + return Ok(false); + } + let addr = machine.registers()[A0].to_u64(); + let size = machine.registers()[A1].to_u64(); + + self.0 + .lock() + .unwrap() + .store_bytes(machine, addr, &DATA_ID, 0, size)?; + + machine.add_cycles_no_checking(100000)?; + + machine.set_register(A0, Mac::REG::from_u64(0)); + Ok(true) + } +} + +enum MachineTy { + Asm, + Interpreter, + InterpreterWithTrace, +} + +impl MachineTy { + fn build(self, data_source: TestSource, version: u32) -> Machine { + match self { + MachineTy::Asm => { + let context = Arc::new(Mutex::new(Snapshot2Context::new(data_source))); + let asm_core1 = AsmCoreMachine::new(ISA_IMC | ISA_A, version, 0); + let core1 = DefaultMachineBuilder::>::new(asm_core1) + .instruction_cycle_func(Box::new(constant_cycles)) + .syscall(Box::new(InsertDataSyscall(context.clone()))) + .build(); + Machine::Asm(AsmMachine::new(core1), context) + } + MachineTy::Interpreter => { + let context = Arc::new(Mutex::new(Snapshot2Context::new(data_source))); + let core_machine1 = DefaultCoreMachine::>>::new( + ISA_IMC | ISA_A, + version, + 0, + ); + Machine::Interpreter( + DefaultMachineBuilder::>>>::new( + core_machine1, + ) + .instruction_cycle_func(Box::new(constant_cycles)) + .syscall(Box::new(InsertDataSyscall(context.clone()))) + .build(), + context, + ) + } + MachineTy::InterpreterWithTrace => { + let context = Arc::new(Mutex::new(Snapshot2Context::new(data_source))); + let core_machine1 = DefaultCoreMachine::>>::new( + ISA_IMC | ISA_A, + version, + 0, + ); + Machine::InterpreterWithTrace( + TraceMachine::new( + DefaultMachineBuilder::< + DefaultCoreMachine>>, + >::new(core_machine1) + .instruction_cycle_func(Box::new(constant_cycles)) + .syscall(Box::new(InsertDataSyscall(context.clone()))) + .build(), + ), + context, + ) + } + } + } +} + +enum Machine { + Asm(AsmMachine, Arc>>), + Interpreter( + DefaultMachine>>>, + Arc>>, + ), + InterpreterWithTrace( + TraceMachine>>>, + Arc>>, + ), +} + +impl Machine { + fn load_program(&mut self, args: &[Bytes]) -> Result { + use Machine::*; + match self { + Asm(inner, context) => { + let program = context + .lock() + .unwrap() + .data_source() + .load_data(&PROGRAM_ID, 0, 0) + .unwrap(); + let metadata = parse_elf::(&program, inner.machine.version())?; + let bytes = inner.load_program_with_metadata(&program, &metadata, args)?; + context.lock().unwrap().mark_program( + inner.machine.inner_mut(), + &metadata, + &PROGRAM_ID, + 0, + )?; + Ok(bytes) + } + Interpreter(inner, context) => { + let program = context + .lock() + .unwrap() + .data_source() + .load_data(&PROGRAM_ID, 0, 0) + .unwrap(); + let metadata = parse_elf::(&program, inner.version())?; + let bytes = inner.load_program_with_metadata(&program, &metadata, args)?; + context.lock().unwrap().mark_program( + inner.inner_mut(), + &metadata, + &PROGRAM_ID, + 0, + )?; + Ok(bytes) + } + InterpreterWithTrace(inner, context) => { + let program = context + .lock() + .unwrap() + .data_source() + .load_data(&PROGRAM_ID, 0, 0) + .unwrap(); + let metadata = parse_elf::(&program, inner.machine.version())?; + let bytes = inner.load_program_with_metadata(&program, &metadata, args)?; + context.lock().unwrap().mark_program( + inner.machine.inner_mut(), + &metadata, + &PROGRAM_ID, + 0, + )?; + Ok(bytes) + } + } + } + + fn run(&mut self) -> Result { + use Machine::*; + match self { + Asm(inner, _) => inner.run(), + Interpreter(inner, _) => inner.run(), + InterpreterWithTrace(inner, _) => inner.run(), + } + } + + fn set_max_cycles(&mut self, cycles: u64) { + use Machine::*; + match self { + Asm(inner, _) => inner.machine.set_max_cycles(cycles), + Interpreter(inner, _) => inner.set_max_cycles(cycles), + InterpreterWithTrace(inner, _) => inner.machine.set_max_cycles(cycles), + } + } + + fn cycles(&self) -> u64 { + use Machine::*; + match self { + Asm(inner, _) => inner.machine.cycles(), + Interpreter(inner, _) => inner.cycles(), + InterpreterWithTrace(inner, _) => inner.machine.cycles(), + } + } + + #[cfg(not(feature = "enable-chaos-mode-by-default"))] + fn full_memory(&mut self) -> Result { + use ckb_vm::{Memory, RISCV_MAX_MEMORY}; + use Machine::*; + match self { + Asm(inner, _) => inner + .machine + .memory_mut() + .load_bytes(0, RISCV_MAX_MEMORY as u64), + Interpreter(inner, _) => inner.memory_mut().load_bytes(0, RISCV_MAX_MEMORY as u64), + InterpreterWithTrace(inner, _) => inner + .machine + .memory_mut() + .load_bytes(0, RISCV_MAX_MEMORY as u64), + } + } + + fn full_registers(&self) -> [u64; 33] { + use Machine::*; + let mut regs = [0u64; 33]; + match self { + Asm(inner, _) => { + regs[0..32].copy_from_slice(inner.machine.registers()); + regs[32] = *inner.machine.pc(); + } + Interpreter(inner, _) => { + regs[0..32].copy_from_slice(inner.registers()); + regs[32] = *inner.pc(); + } + InterpreterWithTrace(inner, _) => { + regs[0..32].copy_from_slice(inner.machine.registers()); + regs[32] = *inner.machine.pc(); + } + }; + regs + } + + fn snapshot(&mut self) -> Result, Error> { + use Machine::*; + match self { + Asm(inner, context) => { + let context = context.lock().unwrap(); + Ok(context.make_snapshot(inner.machine.inner_mut())?) + } + Interpreter(inner, context) => { + let context = context.lock().unwrap(); + Ok(context.make_snapshot(inner.inner_mut())?) + } + InterpreterWithTrace(inner, context) => { + let context = context.lock().unwrap(); + Ok(context.make_snapshot(inner.machine.inner_mut())?) + } + } + } + + fn resume(&mut self, snap: Snapshot2) -> Result<(), Error> { + use Machine::*; + match self { + Asm(inner, context) => { + context + .lock() + .unwrap() + .resume(inner.machine.inner_mut(), &snap)?; + } + Interpreter(inner, context) => { + context.lock().unwrap().resume(inner.inner_mut(), &snap)?; + } + InterpreterWithTrace(inner, context) => { + context + .lock() + .unwrap() + .resume(inner.machine.inner_mut(), &snap)?; + } + }; + Ok(()) + } +} + +#[test] +pub fn test_sc_after_snapshot2() { + let data_source = load_program("tests/programs/sc_after_snapshot"); + + let mut machine1 = MachineTy::Interpreter.build(data_source.clone(), VERSION2); + machine1.set_max_cycles(5); + machine1.load_program(&vec!["main".into()]).unwrap(); + let result1 = machine1.run(); + assert!(result1.is_err()); + assert_eq!(result1.unwrap_err(), Error::CyclesExceeded); + let snapshot = machine1.snapshot().unwrap(); + + let mut machine2 = MachineTy::Interpreter.build(data_source, VERSION2); + machine2.resume(snapshot).unwrap(); + machine2.set_max_cycles(20); + let result2 = machine2.run(); + assert!(result2.is_ok()); + assert_eq!(result2.unwrap(), 0); +}