Create a function to avoid duplication of code

nervosnetwork · Jan 20, 2025 · 74fd84b · 74fd84b
1 parent ac0c2b0
commit 74fd84b
Show file tree

Hide file tree

Showing 3 changed files with 66 additions and 52 deletions.
diff --git a/script/src/syscalls/exec_v2.rs b/script/src/syscalls/exec_v2.rs
@@ -1,6 +1,6 @@
+use crate::syscalls::utils::validate_offset_length;
 use crate::syscalls::{
-    Place, Source, EXEC, INDEX_OUT_OF_BOUND, SLICE_OUT_OF_BOUND, SOURCE_ENTRY_MASK,
-    SOURCE_GROUP_FLAG,
+    Place, Source, EXEC, INDEX_OUT_OF_BOUND, SOURCE_ENTRY_MASK, SOURCE_GROUP_FLAG,
 };
 use crate::types::{DataLocation, DataPieceId, ExecV2Args, Message, TxData, VmId};
 use ckb_traits::{CellDataProvider, ExtensionProvider, HeaderProvider};
@@ -74,33 +74,16 @@ where
         let offset = bounds >> 32;
         let length = bounds as u32 as u64;
 
-        // We are fetching the actual cell here for some in-place validation
-        let mut sc = self
-            .snapshot2_context
-            .lock()
-            .map_err(|e| VMError::Unexpected(e.to_string()))?;
-        let (_, full_length) = match sc.load_data(&data_piece_id, 0, 0) {
-            Ok(val) => val,
-            Err(VMError::SnapshotDataLoadError) => {
-                // This comes from TxData results in an out of bound error, to
-                // mimic current behavior, we would return INDEX_OUT_OF_BOUND error.
-                machine.set_register(A0, Mac::REG::from_u8(INDEX_OUT_OF_BOUND));
-                return Ok(true);
-            }
-            Err(e) => return Err(e),
-        };
-        if offset >= full_length {
-            machine.set_register(A0, Mac::REG::from_u8(SLICE_OUT_OF_BOUND));
+        // We are fetching the actual cell here for some in-place validation.
+        if !validate_offset_length(
+            machine,
+            self.snapshot2_context.clone(),
+            &data_piece_id,
+            offset,
+            length,
+        )? {
             return Ok(true);
         }
-        if length > 0 {
-            // Both offset and length are <= u32::MAX, so offset.checked_add(length) will be always a Some.
-            let end = offset.checked_add(length).ok_or(VMError::MemOutOfBound)?;
-            if end > full_length {
-                machine.set_register(A0, Mac::REG::from_u8(SLICE_OUT_OF_BOUND));
-                return Ok(true);
-            }
-        }
 
         let argc = machine.registers()[A4].to_u64();
         let argv = machine.registers()[A5].to_u64();

diff --git a/script/src/syscalls/spawn.rs b/script/src/syscalls/spawn.rs
@@ -1,5 +1,6 @@
+use crate::syscalls::utils::validate_offset_length;
 use crate::syscalls::{
-    Source, INDEX_OUT_OF_BOUND, SLICE_OUT_OF_BOUND, SOURCE_ENTRY_MASK, SOURCE_GROUP_FLAG, SPAWN,
+    Source, INDEX_OUT_OF_BOUND, SOURCE_ENTRY_MASK, SOURCE_GROUP_FLAG, SPAWN,
     SPAWN_EXTRA_CYCLES_BASE, SPAWN_YIELD_CYCLES_BASE,
 };
 use crate::types::{DataLocation, DataPieceId, Fd, Message, SpawnArgs, TxData, VmId};
@@ -108,32 +109,17 @@ where
             }
         }
 
-        // We are fetching the actual cell here for some in-place validation
-        let mut sc = self
-            .snapshot2_context
-            .lock()
-            .map_err(|e| VMError::Unexpected(e.to_string()))?;
-        let (_, full_length) = match sc.load_data(&data_piece_id, 0, 0) {
-            Ok(val) => val,
-            Err(VMError::SnapshotDataLoadError) => {
-                // This comes from TxData results in an out of bound error, to
-                // mimic current behavior, we would return INDEX_OUT_OF_BOUND error.
-                machine.set_register(A0, Mac::REG::from_u8(INDEX_OUT_OF_BOUND));
-                return Ok(true);
-            }
-            Err(e) => return Err(e),
-        };
-        if offset >= full_length {
-            machine.set_register(A0, Mac::REG::from_u8(SLICE_OUT_OF_BOUND));
+        // We are fetching the actual cell here for some in-place validation.
+        if !validate_offset_length(
+            machine,
+            self.snapshot2_context.clone(),
+            &data_piece_id,
+            offset,
+            length,
+        )? {
             return Ok(true);
         }
-        if length > 0 {
-            let end = offset.checked_add(length).ok_or(VMError::MemOutOfBound)?;
-            if end > full_length {
-                machine.set_register(A0, Mac::REG::from_u8(SLICE_OUT_OF_BOUND));
-                return Ok(true);
-            }
-        }
+
         machine.add_cycles_no_checking(SPAWN_EXTRA_CYCLES_BASE)?;
         machine.add_cycles_no_checking(SPAWN_YIELD_CYCLES_BASE)?;
         self.message_box

diff --git a/script/src/syscalls/utils.rs b/script/src/syscalls/utils.rs
@@ -1,9 +1,16 @@
 use byteorder::{ByteOrder, LittleEndian};
+use ckb_traits::{CellDataProvider, ExtensionProvider, HeaderProvider};
 use ckb_vm::{
     registers::{A0, A1, A2},
+    snapshot2::Snapshot2Context,
     Error as VMError, Memory, Register, SupportMachine,
 };
 use std::cmp;
+use std::sync::{Arc, Mutex};
+
+use crate::syscalls::{INDEX_OUT_OF_BOUND, SLICE_OUT_OF_BOUND};
+use crate::types::TxData;
+use crate::DataPieceId;
 
 pub fn store_data<Mac: SupportMachine>(machine: &mut Mac, data: &[u8]) -> Result<u64, VMError> {
     let addr = machine.registers()[A0].to_u64();
@@ -28,3 +35,41 @@ pub fn store_u64<Mac: SupportMachine>(machine: &mut Mac, v: u64) -> Result<u64,
     LittleEndian::write_u64(&mut buffer, v);
     store_data(machine, &buffer)
 }
+
+pub fn validate_offset_length<Mac, DL>(
+    machine: &mut Mac,
+    snapshot2_context: Arc<Mutex<Snapshot2Context<DataPieceId, TxData<DL>>>>,
+    data_piece_id: &DataPieceId,
+    offset: u64,
+    length: u64,
+) -> Result<bool, VMError>
+where
+    Mac: SupportMachine,
+    DL: CellDataProvider + HeaderProvider + ExtensionProvider + Send + Sync + Clone + 'static,
+{
+    let mut sc = snapshot2_context
+        .lock()
+        .map_err(|e| VMError::Unexpected(e.to_string()))?;
+    let (_, full_length) = match sc.load_data(&data_piece_id, 0, 0) {
+        Ok(val) => val,
+        Err(VMError::SnapshotDataLoadError) => {
+            // This comes from TxData results in an out of bound error, to
+            // mimic current behavior, we would return INDEX_OUT_OF_BOUND error.
+            machine.set_register(A0, Mac::REG::from_u8(INDEX_OUT_OF_BOUND));
+            return Ok(false);
+        }
+        Err(e) => return Err(e),
+    };
+    if offset >= full_length {
+        machine.set_register(A0, Mac::REG::from_u8(SLICE_OUT_OF_BOUND));
+        return Ok(false);
+    }
+    if length > 0 {
+        let end = offset.checked_add(length).ok_or(VMError::MemOutOfBound)?;
+        if end > full_length {
+            machine.set_register(A0, Mac::REG::from_u8(SLICE_OUT_OF_BOUND));
+            return Ok(false);
+        }
+    }
+    Ok(true)
+}