Use correct memory pages and size

definitions/src/lib.rs

@@ -8,6 +8,7 @@ pub const RISCV_PAGESIZE: usize = 1 << RISCV_PAGE_SHIFTS;
 pub const RISCV_GENERAL_REGISTER_NUMBER: usize = 32;
 // 4 MB
 pub const RISCV_MAX_MEMORY: usize = 4 << 20;
+pub const RISCV_MAX_MEMORY_V0: usize = 4 << 20;
 // 1 MB
 pub const DEFAULT_STACK_SIZE: usize = 1 << 20;
 pub const RISCV_PAGES: usize = RISCV_MAX_MEMORY / RISCV_PAGESIZE;

src/decoder.rs

@@ -96,8 +96,9 @@ impl Decoder {
     }
 
     pub fn decode_raw<M: Memory>(&mut self, memory: &mut M, pc: u64) -> Result<Instruction, Error> {
-        // since we are using RISCV_MAX_MEMORY as the default key in the instruction cache, have to check out of bound error first
-        if pc as usize >= RISCV_MAX_MEMORY {
+        // since we are using RISCV_MAX_MEMORY as the default key in the instruction cache, have to check out of bound
+        // error first.
+        if pc as usize >= memory.memory_size() {
             return Err(Error::MemOutOfBound);
         }
         let instruction_cache_key = {

src/instructions/common.rs

@@ -1,6 +1,6 @@
 use super::super::machine::Machine;
 use super::super::memory::Memory;
-use super::super::RISCV_MAX_MEMORY;
+use super::super::RISCV_MAX_MEMORY_V0;
 use super::register::Register;
 use super::utils::update_register;
 use super::{Error, RegisterIndex, SImmediate, UImmediate};
@@ -85,7 +85,8 @@ fn check_load_boundary<R: Register>(version0: bool, address: &R, bytes: u64) ->
     if version0 {
         let address = address.to_u64();
         let end = address.checked_add(bytes).ok_or(Error::MemOutOfBound)?;
-        if end == RISCV_MAX_MEMORY as u64 {
+        // The memory size of CKB-VM V0 is 4M.
+        if end == RISCV_MAX_MEMORY_V0 as u64 {
             return Err(Error::MemOutOfBound);
         }
     }

src/lib.rs

@@ -31,8 +31,8 @@ pub use bytes::Bytes;
 
 pub use ckb_vm_definitions::{
     registers, DEFAULT_STACK_SIZE, ISA_A, ISA_B, ISA_IMC, ISA_MOP, MEMORY_FRAMES, MEMORY_FRAMESIZE,
-    MEMORY_FRAME_SHIFTS, RISCV_GENERAL_REGISTER_NUMBER, RISCV_MAX_MEMORY, RISCV_PAGES,
-    RISCV_PAGESIZE, RISCV_PAGE_SHIFTS,
+    MEMORY_FRAME_SHIFTS, RISCV_GENERAL_REGISTER_NUMBER, RISCV_MAX_MEMORY, RISCV_MAX_MEMORY_V0,
+    RISCV_PAGES, RISCV_PAGESIZE, RISCV_PAGE_SHIFTS,
 };
 
 pub use error::Error;

src/machine/asm/mod.rs

@@ -28,7 +28,7 @@ use crate::{
         FLAG_EXECUTABLE, FLAG_FREEZED, FLAG_WRITABLE, FLAG_WXORX_BIT,
     },
     CoreMachine, DefaultMachine, Error, Machine, Memory, SupportMachine, MEMORY_FRAME_SHIFTS,
-    RISCV_MAX_MEMORY, RISCV_PAGES, RISCV_PAGESIZE,
+    RISCV_PAGES, RISCV_PAGESIZE,
 };
 
 impl CoreMachine for Box<AsmCoreMachine> {
@@ -115,7 +115,7 @@ fn check_memory_writable(
 ) -> Result<(), Error> {
     debug_assert!(size == 1 || size == 2 || size == 4 || size == 8);
     let page = addr >> RISCV_PAGE_SHIFTS;
-    if page as usize >= RISCV_PAGES {
+    if page as usize >= machine.memory_pages() {
         return Err(Error::MemOutOfBound);
     }
     check_permission(machine, page, FLAG_WRITABLE)?;
@@ -126,7 +126,7 @@ fn check_memory_writable(
     let page_offset = addr as usize % RISCV_PAGESIZE;
     if page_offset + size > RISCV_PAGESIZE {
         let page = page + 1;
-        if page as usize >= RISCV_PAGES {
+        if page as usize >= machine.memory_pages() {
             return Err(Error::MemOutOfBound);
         } else {
             check_permission(machine, page, FLAG_WRITABLE)?;
@@ -146,7 +146,7 @@ fn check_memory_executable(
     debug_assert!(size == 2 || size == 4);
 
     let page = addr >> RISCV_PAGE_SHIFTS;
-    if page as usize >= RISCV_PAGES {
+    if page as usize >= machine.memory_pages() {
         return Err(Error::MemOutOfBound);
     }
     check_permission(machine, page, FLAG_EXECUTABLE)?;
@@ -156,7 +156,7 @@ fn check_memory_executable(
     let page_offset = addr as usize % RISCV_PAGESIZE;
     if page_offset + size > RISCV_PAGESIZE {
         let page = page + 1;
-        if page as usize >= RISCV_PAGES {
+        if page as usize >= machine.memory_pages() {
             return Err(Error::MemOutOfBound);
         } else {
             check_permission(machine, page, FLAG_EXECUTABLE)?;
@@ -174,7 +174,7 @@ fn check_memory_inited(
 ) -> Result<(), Error> {
     debug_assert!(size == 1 || size == 2 || size == 4 || size == 8);
     let page = addr >> RISCV_PAGE_SHIFTS;
-    if page as usize >= RISCV_PAGES {
+    if page as usize >= machine.memory_pages() {
         return Err(Error::MemOutOfBound);
     }
     check_memory(machine, page);
@@ -183,7 +183,7 @@ fn check_memory_inited(
     let page_offset = addr as usize % RISCV_PAGESIZE;
     if page_offset + size > RISCV_PAGESIZE {
         let page = page + 1;
-        if page as usize >= RISCV_PAGES {
+        if page as usize >= machine.memory_pages() {
             return Err(Error::MemOutOfBound);
         } else {
             check_memory(machine, page);
@@ -349,9 +349,10 @@ impl Memory for Box<AsmCoreMachine> {
         if round_page_down(addr) != addr || round_page_up(size) != size {
             return Err(Error::MemPageUnalignedAccess);
         }
-        if addr > RISCV_MAX_MEMORY as u64
-            || size > RISCV_MAX_MEMORY as u64
-            || addr + size > RISCV_MAX_MEMORY as u64
+        let memory_size = self.memory_size() as u64;
+        if addr > memory_size
+            || size > memory_size as u64
+            || addr + size > memory_size as u64
             || offset_from_addr > size
         {
             return Err(Error::MemOutOfBound);
@@ -381,15 +382,15 @@ impl Memory for Box<AsmCoreMachine> {
     }
 
     fn fetch_flag(&mut self, page: u64) -> Result<u8, Error> {
-        if page < RISCV_PAGES as u64 {
+        if page < self.memory_pages() as u64 {
             Ok(self.flags[page as usize])
         } else {
             Err(Error::MemOutOfBound)
         }
     }
 
     fn set_flag(&mut self, page: u64, flag: u8) -> Result<(), Error> {
-        if page < RISCV_PAGES as u64 {
+        if page < self.memory_pages() as u64 {
             self.flags[page as usize] |= flag;
             // Clear last write page cache
             self.last_write_page = u64::max_value();
@@ -400,7 +401,7 @@ impl Memory for Box<AsmCoreMachine> {
     }
 
     fn clear_flag(&mut self, page: u64, flag: u8) -> Result<(), Error> {
-        if page < RISCV_PAGES as u64 {
+        if page < self.memory_pages() as u64 {
             self.flags[page as usize] &= !flag;
             // Clear last write page cache
             self.last_write_page = u64::max_value();

src/memory/mod.rs

@@ -44,6 +44,9 @@ pub trait Memory {
     fn set_flag(&mut self, page: u64, flag: u8) -> Result<(), Error>;
     fn clear_flag(&mut self, page: u64, flag: u8) -> Result<(), Error>;
     fn memory_size(&self) -> usize;
+    fn memory_pages(&self) -> usize {
+        self.memory_size() >> RISCV_PAGE_SHIFTS
+    }
 
     // This is in fact just memset
     fn store_byte(&mut self, addr: u64, size: u64, value: u8) -> Result<(), Error>;

src/snapshot.rs

@@ -44,7 +44,7 @@ pub fn make_snapshot<T: CoreMachine>(machine: &mut T) -> Result<Snapshot, Error>
         snap.registers[i] = v.to_u64();
     }
 
-    for i in 0..(machine.memory().memory_size() >> RISCV_PAGE_SHIFTS) {
+    for i in 0..machine.memory().memory_pages() {
         let flag = machine.memory_mut().fetch_flag(i as u64)?;
         if flag & FLAG_DIRTY != 0 {
             let addr_from = i << RISCV_PAGE_SHIFTS;

src/snapshot2.rs

@@ -3,7 +3,7 @@ use crate::{
     elf::{LoadingAction, ProgramMetadata},
     machine::SupportMachine,
     memory::{Memory, FLAG_DIRTY},
-    Error, Register, RISCV_GENERAL_REGISTER_NUMBER, RISCV_PAGESIZE, RISCV_PAGE_SHIFTS,
+    Error, Register, RISCV_GENERAL_REGISTER_NUMBER, RISCV_PAGESIZE,
 };
 use bytes::Bytes;
 use serde::{Deserialize, Serialize};
@@ -146,7 +146,7 @@ impl<I: Clone + PartialEq, D: DataSource<I>> Snapshot2Context<I, D> {
     /// Create a snapshot for the passed machine.
     pub fn make_snapshot<M: SupportMachine>(&self, machine: &mut M) -> Result<Snapshot2<I>, Error> {
         let mut dirty_pages: Vec<(u64, u8, Vec<u8>)> = vec![];
-        for i in 0..(machine.memory().memory_size() as u64 >> RISCV_PAGE_SHIFTS) {
+        for i in 0..machine.memory().memory_pages() as u64 {
             if self.pages.contains_key(&i) {
                 continue;
             }