Merge pull request #4785 from libraries/args_reader

Improving spawn and exec syscalls

script/Cargo.toml

@@ -22,7 +22,7 @@ ckb-traits = { path = "../traits", version = "= 0.121.0-pre" }
 byteorder = "1.3.1"
 ckb-types = { path = "../util/types", version = "= 0.121.0-pre" }
 ckb-hash = { path = "../util/hash", version = "= 0.121.0-pre" }
-ckb-vm = { version = "= 0.24.12", default-features = false }
+ckb-vm = { version = "= 0.24.13", default-features = false }
 faster-hex = "0.6"
 ckb-logger = { path = "../util/logger", version = "= 0.121.0-pre", optional = true }
 serde = { version = "1.0", features = ["derive"] }

script/src/scheduler.rs

@@ -1,15 +1,15 @@
 use crate::cost_model::transferred_byte_cycles;
 use crate::syscalls::{
-    INVALID_FD, MAX_FDS_CREATED, MAX_VMS_SPAWNED, OTHER_END_CLOSED, SPAWN_EXTRA_CYCLES_BASE,
-    SUCCESS, WAIT_FAILURE,
+    EXEC_LOAD_ELF_V2_CYCLES_BASE, INVALID_FD, MAX_FDS_CREATED, MAX_VMS_SPAWNED, OTHER_END_CLOSED,
+    SPAWN_EXTRA_CYCLES_BASE, SUCCESS, WAIT_FAILURE,
 };
 use crate::types::MachineContext;
 use crate::verify::TransactionScriptsSyscallsGenerator;
 use crate::ScriptVersion;
 
 use crate::types::{
-    CoreMachineType, DataPieceId, Fd, FdArgs, FullSuspendedState, Machine, Message, ReadState,
-    RunMode, TxData, VmId, VmState, WriteState, FIRST_FD_SLOT, FIRST_VM_ID,
+    CoreMachineType, DataLocation, DataPieceId, Fd, FdArgs, FullSuspendedState, Machine, Message,
+    ReadState, RunMode, TxData, VmId, VmState, WriteState, FIRST_FD_SLOT, FIRST_VM_ID,
 };
 use ckb_traits::{CellDataProvider, ExtensionProvider, HeaderProvider};
 use ckb_types::core::Cycle;
@@ -22,7 +22,7 @@ use ckb_vm::{
     memory::Memory,
     registers::A0,
     snapshot2::Snapshot2,
-    Error, Register,
+    Error, FlattenedArgsReader, Register,
 };
 use std::collections::{BTreeMap, HashMap};
 use std::sync::{Arc, Mutex};
@@ -208,7 +208,14 @@ where
         if self.states.is_empty() {
             // Booting phase, we will need to initialize the first VM.
             assert_eq!(
-                self.boot_vm(&DataPieceId::Program, 0, u64::MAX, &[])?,
+                self.boot_vm(
+                    &DataLocation {
+                        data_piece_id: DataPieceId::Program,
+                        offset: 0,
+                        length: u64::MAX,
+                    },
+                    None
+                )?,
                 ROOT_VM_ID
             );
         }
@@ -340,6 +347,39 @@ where
         let messages: Vec<Message> = self.message_box.lock().expect("lock").drain(..).collect();
         for message in messages {
             match message {
+                Message::ExecV2(vm_id, args) => {
+                    let (old_context, old_machine) = self
+                        .instantiated
+                        .get_mut(&vm_id)
+                        .ok_or_else(|| Error::Unexpected("Unable to find VM Id".to_string()))?;
+                    old_machine
+                        .machine
+                        .add_cycles_no_checking(EXEC_LOAD_ELF_V2_CYCLES_BASE)?;
+                    let old_cycles = old_machine.machine.cycles();
+                    let max_cycles = old_machine.machine.max_cycles();
+                    let program = {
+                        let mut sc = old_context.snapshot2_context().lock().expect("lock");
+                        sc.load_data(
+                            &args.location.data_piece_id,
+                            args.location.offset,
+                            args.location.length,
+                        )?
+                        .0
+                    };
+                    let (context, mut new_machine) = self.create_dummy_vm(&vm_id)?;
+                    new_machine.set_max_cycles(max_cycles);
+                    new_machine.machine.add_cycles_no_checking(old_cycles)?;
+                    self.load_vm_program(
+                        &context,
+                        &mut new_machine,
+                        &args.location,
+                        program,
+                        Some((vm_id, args.argc, args.argv)),
+                    )?;
+                    // The insert operation removes the old vm instance and adds the new vm instance.
+                    debug_assert!(self.instantiated.contains_key(&vm_id));
+                    self.instantiated.insert(vm_id, (context, new_machine));
+                }
                 Message::Spawn(vm_id, args) => {
                     // All fds must belong to the correct owner
                     if args.fds.iter().any(|fd| self.fds.get(fd) != Some(&vm_id)) {
@@ -353,7 +393,7 @@ where
                         continue;
                     }
                     let spawned_vm_id =
-                        self.boot_vm(&args.data_piece_id, args.offset, args.length, &args.argv)?;
+                        self.boot_vm(&args.location, Some((vm_id, args.argc, args.argv)))?;
                     // Move passed fds from spawner to spawnee
                     for fd in &args.fds {
                         self.fds.insert(*fd, spawned_vm_id);
@@ -746,11 +786,17 @@ where
     /// Boot a vm by given program and args.
     pub fn boot_vm(
         &mut self,
-        data_piece_id: &DataPieceId,
-        offset: u64,
-        length: u64,
-        args: &[Bytes],
+        location: &DataLocation,
+        args: Option<(u64, u64, u64)>,
     ) -> Result<VmId, Error> {
+        let id = self.next_vm_id;
+        self.next_vm_id += 1;
+        let (context, mut machine) = self.create_dummy_vm(&id)?;
+        let (program, _) = {
+            let mut sc = context.snapshot2_context().lock().expect("lock");
+            sc.load_data(&location.data_piece_id, location.offset, location.length)?
+        };
+        self.load_vm_program(&context, &mut machine, location, program, args)?;
         // Newly booted VM will be instantiated by default
         while self.instantiated.len() >= MAX_INSTANTIATED_VMS {
             // Instantiated is a BTreeMap, first_entry will maintain key order
@@ -761,26 +807,43 @@ where
                 .key();
             self.suspend_vm(&id)?;
         }
-
-        let id = self.next_vm_id;
-        self.next_vm_id += 1;
-        let (context, mut machine) = self.create_dummy_vm(&id)?;
-        {
-            let mut sc = context.snapshot2_context().lock().expect("lock");
-            let (program, _) = sc.load_data(data_piece_id, offset, length)?;
-            let metadata = parse_elf::<u64>(&program, machine.machine.version())?;
-            let bytes = machine.load_program_with_metadata(&program, &metadata, args)?;
-            sc.mark_program(&mut machine.machine, &metadata, data_piece_id, offset)?;
-            machine
-                .machine
-                .add_cycles_no_checking(transferred_byte_cycles(bytes))?;
-        }
         self.instantiated.insert(id, (context, machine));
         self.states.insert(id, VmState::Runnable);
 
         Ok(id)
     }
 
+    // Load the program into an empty vm.
+    fn load_vm_program(
+        &mut self,
+        context: &MachineContext<DL>,
+        machine: &mut Machine,
+        location: &DataLocation,
+        program: Bytes,
+        args: Option<(u64, u64, u64)>,
+    ) -> Result<u64, Error> {
+        let metadata = parse_elf::<u64>(&program, machine.machine.version())?;
+        let bytes = match args {
+            Some((vm_id, argc, argv)) => {
+                let (_, machine_from) = self.ensure_get_instantiated(&vm_id)?;
+                let argv = FlattenedArgsReader::new(machine_from.machine.memory_mut(), argc, argv);
+                machine.load_program_with_metadata(&program, &metadata, argv)?
+            }
+            None => machine.load_program_with_metadata(&program, &metadata, vec![].into_iter())?,
+        };
+        let mut sc = context.snapshot2_context().lock().expect("lock");
+        sc.mark_program(
+            &mut machine.machine,
+            &metadata,
+            &location.data_piece_id,
+            location.offset,
+        )?;
+        machine
+            .machine
+            .add_cycles_no_checking(transferred_byte_cycles(bytes))?;
+        Ok(bytes)
+    }
+
     // Create a new VM instance with syscalls attached
     fn create_dummy_vm(&self, id: &VmId) -> Result<(MachineContext<DL>, Machine), Error> {
         // The code here looks slightly weird, since I don't want to copy over all syscall

script/src/syscalls/exec.rs

@@ -1,5 +1,4 @@
 use crate::cost_model::transferred_byte_cycles;
-use crate::syscalls::utils::load_c_string;
 use crate::syscalls::{
     Place, Source, SourceEntry, EXEC, INDEX_OUT_OF_BOUND, MAX_ARGV_LENGTH, SLICE_OUT_OF_BOUND,
     WRONG_FORMAT,
@@ -9,6 +8,7 @@ use ckb_traits::CellDataProvider;
 use ckb_types::core::cell::{CellMeta, ResolvedTransaction};
 use ckb_types::core::error::ARGV_TOO_LONG_TEXT;
 use ckb_types::packed::{Bytes as PackedBytes, BytesVec};
+use ckb_vm::memory::load_c_string_byte_by_byte;
 use ckb_vm::Memory;
 use ckb_vm::{
     registers::{A0, A1, A2, A3, A4, A5, A7},
@@ -24,7 +24,6 @@ pub struct Exec<DL> {
     outputs: Arc<Vec<CellMeta>>,
     group_inputs: Indices,
     group_outputs: Indices,
-    load_elf_base_fee: u64,
 }
 
 impl<DL: CellDataProvider> Exec<DL> {
@@ -34,15 +33,13 @@ impl<DL: CellDataProvider> Exec<DL> {
         outputs: Arc<Vec<CellMeta>>,
         group_inputs: Indices,
         group_outputs: Indices,
-        load_elf_base_fee: u64,
     ) -> Exec<DL> {
         Exec {
             data_loader,
             rtx,
             outputs,
             group_inputs,
             group_outputs,
-            load_elf_base_fee,
         }
     }
 
@@ -154,6 +151,7 @@ impl<Mac: SupportMachine, DL: CellDataProvider + Send + Sync> Syscalls<Mac> for
         let data = if length == 0 {
             data.slice(offset..data_size)
         } else {
+            // 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 > data_size {
                 machine.set_register(A0, Mac::REG::from_u8(SLICE_OUT_OF_BOUND));
@@ -166,12 +164,8 @@ impl<Mac: SupportMachine, DL: CellDataProvider + Send + Sync> Syscalls<Mac> for
         let mut argv = Vec::new();
         let mut argv_length: u64 = 0;
         for _ in 0..argc {
-            let target_addr = machine
-                .memory_mut()
-                .load64(&Mac::REG::from_u64(addr))?
-                .to_u64();
-
-            let cstr = load_c_string(machine, target_addr)?;
+            let target_addr = machine.memory_mut().load64(&Mac::REG::from_u64(addr))?;
+            let cstr = load_c_string_byte_by_byte(machine.memory_mut(), &target_addr)?;
             let cstr_len = cstr.len();
             argv.push(cstr);
 
@@ -191,7 +185,6 @@ impl<Mac: SupportMachine, DL: CellDataProvider + Send + Sync> Syscalls<Mac> for
         machine.reset(max_cycles);
         machine.set_cycles(cycles);
 
-        machine.add_cycles_no_checking(self.load_elf_base_fee)?;
         match machine.load_elf(&data, true) {
             Ok(size) => {
                 machine.add_cycles_no_checking(transferred_byte_cycles(size))?;
@@ -203,7 +196,7 @@ impl<Mac: SupportMachine, DL: CellDataProvider + Send + Sync> Syscalls<Mac> for
         }
 
         match machine.initialize_stack(
-            &argv,
+            argv.into_iter().map(Ok),
             (RISCV_MAX_MEMORY - DEFAULT_STACK_SIZE) as u64,
             DEFAULT_STACK_SIZE as u64,
         ) {

script/src/syscalls/exec_v2.rs

@@ -0,0 +1,65 @@
+use crate::syscalls::{EXEC, INDEX_OUT_OF_BOUND};
+use crate::types::{DataLocation, DataPieceId, ExecV2Args, Message, VmId};
+use ckb_vm::{
+    registers::{A0, A1, A2, A3, A4, A5, A7},
+    Error as VMError, Register, SupportMachine, Syscalls,
+};
+use std::sync::{Arc, Mutex};
+
+pub struct ExecV2 {
+    id: VmId,
+    message_box: Arc<Mutex<Vec<Message>>>,
+}
+
+impl ExecV2 {
+    pub fn new(id: VmId, message_box: Arc<Mutex<Vec<Message>>>) -> ExecV2 {
+        ExecV2 { id, message_box }
+    }
+}
+
+impl<Mac> Syscalls<Mac> for ExecV2
+where
+    Mac: SupportMachine,
+{
+    fn initialize(&mut self, _machine: &mut Mac) -> Result<(), VMError> {
+        Ok(())
+    }
+
+    fn ecall(&mut self, machine: &mut Mac) -> Result<bool, VMError> {
+        if machine.registers()[A7].to_u64() != EXEC {
+            return Ok(false);
+        }
+        let index = machine.registers()[A0].to_u64();
+        let source = machine.registers()[A1].to_u64();
+        let place = machine.registers()[A2].to_u64();
+        let data_piece_id = match DataPieceId::try_from((source, index, place)) {
+            Ok(id) => id,
+            Err(_) => {
+                machine.set_register(A0, Mac::REG::from_u8(INDEX_OUT_OF_BOUND));
+                return Ok(true);
+            }
+        };
+        let bounds = machine.registers()[A3].to_u64();
+        let offset = bounds >> 32;
+        let length = bounds as u32 as u64;
+
+        let argc = machine.registers()[A4].to_u64();
+        let argv = machine.registers()[A5].to_u64();
+        self.message_box
+            .lock()
+            .map_err(|e| VMError::Unexpected(e.to_string()))?
+            .push(Message::ExecV2(
+                self.id,
+                ExecV2Args {
+                    location: DataLocation {
+                        data_piece_id,
+                        offset,
+                        length,
+                    },
+                    argc,
+                    argv,
+                },
+            ));
+        Err(VMError::Yield)
+    }
+}

script/src/syscalls/mod.rs

@@ -2,6 +2,7 @@ mod close;
 mod current_cycles;
 mod debugger;
 mod exec;
+mod exec_v2;
 mod inherited_fd;
 mod load_block_extension;
 mod load_cell;
@@ -31,6 +32,7 @@ pub use self::close::Close;
 pub use self::current_cycles::CurrentCycles;
 pub use self::debugger::Debugger;
 pub use self::exec::Exec;
+pub use self::exec_v2::ExecV2;
 pub use self::inherited_fd::InheritedFd;
 pub use self::load_block_extension::LoadBlockExtension;
 pub use self::load_cell::LoadCell;