corevm_engine/

engine.rs

use crate::{
	hash_encoded, FatalError, FsNode, HostCallError, HostCallHandler, InitError, InnerVm,
	InvokeArgs, KernelContext, Lookup, MemoryManager, OpenedFile, OuterVm, ProgramData, Reg,
	RunError, TouchError,
};
use alloc::vec::Vec;
use codec::{Compact, CompactLen, ConstEncodedLen, DecodeAll, Encode};
use corevm_host::{
	fs, AudioMode, CoreVmOutput, CoreVmPayload, ExecEnv, KernelFd, Outcome, RangeSet, VideoMode,
	VmOutput, VmState, PAGE_SIZE,
};
use jam_pvm_common::{ApiError, InvokeOutcome};
use jam_types::Hash;
use log::{debug, trace};
use polkakernel::{KernelState, Machine};
use polkavm::{MemoryMapBuilder, RETURN_TO_HOST};

/// An engine that drives execution of `refine` CoreVM service entry point.
pub struct Engine<O: OuterVm> {
	pub(crate) args: InvokeArgs,
	pub(crate) memory_man: MemoryManager<O>,
	pub(crate) video: Option<VideoMode>,
	pub(crate) audio: Option<AudioMode>,
	pub(crate) exec: ExecEnv,
	/// The number of video frames produced by the guest so far.
	pub(crate) num_video_frames: u64,
	/// In bytes.
	/// Total size of audio frames produced by the guest so far in bytes.
	pub(crate) num_audio_bytes: u64,
	exec_ref: fs::BlockRef,
	old_hash: Hash,
	host_call_handlers: Vec<HostCallHandler<O>>,
	/// The index of the host-call being handled.
	current_host_call: Option<u64>,
}

impl<O: OuterVm> Engine<O> {
	pub fn new(payload: CoreVmPayload, mut outer_vm: O) -> Result<Self, InitError> {
		let CoreVmPayload { gas, vm_state, exec_ref } = payload;
		let exec = outer_vm.read_file(&exec_ref)?;
		let exec = ExecEnv::decode_all(&mut &exec[..]).map_err(|_| InitError::ExecBlob)?;
		let max_exports = outer_vm.get_export_count();
		let auth_output_len = outer_vm.get_auth_output_len();
		let mut program_counter = vm_state.program_counter;
		let program = {
			let bytes = outer_vm.read_file(&exec.program)?;
			let corevm_blob = jam_program_blob_common::CoreVmProgramBlob::from_bytes(&bytes)
				.ok_or(InitError::ProgramBlob)?;
			polkavm::ProgramParts::from_bytes(corevm_blob.pvm_blob.into())?
		};
		let memory_map = MemoryMapBuilder::new(PAGE_SIZE as u32)
			.ro_data_size(program.ro_data_size)
			.rw_data_size(program.rw_data_size)
			.stack_size(program.stack_size)
			.build()
			.map_err(|_| InitError::ProgramBlob)?;
		let code_and_jump_table = program.code_and_jump_table.clone();
		let ro_data = program.ro_data.clone();
		let rw_data = program.rw_data.clone();
		let program_blob =
			polkavm::ProgramBlob::from_parts(program).map_err(|_| InitError::ProgramBlob)?;
		let host_call_handlers = Self::create_host_call_handlers(&program_blob);
		// Should include program hash and pages' hash.
		let old_hash = compute_state_hash(&vm_state);
		let mut args = InvokeArgs { regs: vm_state.regs, gas };
		let initial_program_run = program_counter == 0;
		let mut libc = false;
		if initial_program_run {
			debug!("This is initial program run");
			program_counter = program_blob
				.exports()
				.find(|export| {
					let name = export.symbol().as_bytes();
					if name == b"main" {
						debug!("Found `main` (no libc) entry point");
						return true;
					}
					if name == b"_pvm_start" {
						debug!("Found `_pvm_start` (libc) entry point");
						libc = true;
						return true;
					}
					false
				})
				.expect("Neither `main` nor `_pvm_start` entry point found")
				.program_counter()
				.0
				.into();
		}
		let inner_vm = outer_vm
			.machine(&code_and_jump_table[..], program_counter)
			.expect("Failed to load the code");
		let program_data = ProgramData::new(&memory_map, ro_data, rw_data);
		let kernel_state = KernelState {
			fds: vm_state
				.kernel
				.fds
				.into_iter()
				.map(|(fd, file)| {
					// TODO @ivan We probably want to open and read files lazily.
					let (block_ref, position) = (file.block_ref, file.position);
					let mut lookup = Lookup { outer_vm: &mut outer_vm };
					let node = fs::Node::open(&block_ref, &mut lookup)?;
					let node = match node {
						fs::Node::File(mut file) => {
							file.seek(position)?;
							FsNode::File(file)
						},
						fs::Node::Dir(dir) => FsNode::Dir { dir, position },
					};
					let opened_file = OpenedFile { block_ref, node };
					Ok((fd, opened_file))
				})
				.collect::<Result<_, fs::Error>>()?,
		};
		let mut memory_man = MemoryManager::new(
			outer_vm,
			inner_vm,
			program_data,
			vm_state.mapped_heap_pages,
			vm_state.resident_pages,
			max_exports as usize,
			auth_output_len as usize,
			kernel_state,
		)
		.expect("Failed to initialize memory pages");
		if initial_program_run {
			// Initialize the registers for the initial `invoke` call.
			let default_sp = memory_map.stack_address_high() as u64;
			if libc {
				let unused = fs::BlockRef { service_id: 0, hash: Default::default() };
				let mut context = KernelContext {
					memory_man: &mut memory_man,
					regs: &mut args.regs,
					root_dir: unused,
					error: None,
				};
				context
					.init(
						default_sp,
						RETURN_TO_HOST,
						exec.args.iter().map(|a| a.as_ref()),
						exec.env.iter().map(|a| a.as_ref()),
					)
					.expect("Failed to init arguments");
			} else {
				args.set_reg(Reg::RA, RETURN_TO_HOST);
				args.set_reg(Reg::SP, default_sp);
			}
		}
		Ok(Self {
			args,
			memory_man,
			video: vm_state.video,
			audio: vm_state.audio,
			old_hash,
			num_video_frames: 0,
			num_audio_bytes: 0,
			exec,
			exec_ref,
			host_call_handlers,
			current_host_call: vm_state.restart_host_call,
		})
	}

	pub fn run(mut self) -> Result<(CoreVmOutput, O), RunError> {
		// Restart the host-call that the engine was executing while ran out of output
		// space. Panic if the host-call needs to be restarted again.
		let mut restart_host_call = self.current_host_call.is_some();
		let (output, outer_vm) = loop {
			if let Some(index) = self.current_host_call.take() {
				use Outcome::*;
				if restart_host_call {
					debug!("Restarting host-call {index}");
				}
				let outcome = match self.handle_host_call_fault(index) {
					Ok(()) => {
						// Continue normal execution.
						None
					},
					Err(HostCallError::Outcome(
						outcome @ Panic | outcome @ Halt | outcome @ OutOfGas,
					)) => Some(outcome),
					Err(HostCallError::Outcome(outcome @ TimeLimitReached)) => {
						// NOTE We don't restart the host-call in this case.
						debug!("Time limit reached in host-call {index}");
						Some(outcome)
					},
					Err(HostCallError::Outcome(outcome @ OutputLimitReached)) => {
						if restart_host_call {
							return Err(FatalError::NotEnoughOutputSpace.into());
						}
						debug!("Output limit reached in host-call {index}");
						self.current_host_call = Some(index);
						Some(outcome)
					},
					Err(HostCallError::Outcome(outcome @ PageFault { page, num_pages })) => {
						if restart_host_call {
							debug!(
								"Hard page fault at {:#x?} while restarting host-call {index}",
								page..page + num_pages
							);
							// NOTE: This error might also occur if the builder haven't imported all
							// the necessary pages. Normally this shouldn't happen.
							return Err(FatalError::NotEnoughInputSpace.into());
						}
						debug!(
							"Hard page fault at {:#x?} in host-call {index}",
							page..page + num_pages
						);
						self.current_host_call = Some(index);
						Some(outcome)
					},
					Err(HostCallError::Jam(e)) => return Err(e.into()),
					Err(HostCallError::Fs(e)) => return Err(e.into()),
					Err(HostCallError::Fatal(e)) => return Err(e.into()),
				};
				if let Some(outcome) = outcome {
					break self.into_host_output(outcome)?;
				}
				restart_host_call = false;
			}
			let (outcome, gas, regs) =
				self.memory_man.inner_vm.invoke(self.args.gas, self.args.regs)?;
			trace!("Invoke outcome {:?}", DebugInvokeOutcome(outcome));
			self.args.gas = gas;
			self.args.regs = regs;
			match outcome {
				InvokeOutcome::Halt => {
					break self.into_host_output(Outcome::Halt)?;
				},
				InvokeOutcome::PageFault(address) => {
					match self.memory_man.touch(address) {
						Ok(..) => {},
						Err(TouchError::PageFault { page, num_pages }) => {
							trace!("Hard page fault at address {:#x}", page * PAGE_SIZE);
							// Suspend the execution when _either_ the program tries to
							// access the page that was not imported _or_
							// the max. no. of allocated pages is reached.
							break self.into_host_output(Outcome::PageFault { page, num_pages })?;
						},
						Err(TouchError::Jam(e)) => return Err(e.into()),
					}
				},
				InvokeOutcome::HostCallFault(index) => self.current_host_call = Some(index),
				InvokeOutcome::Panic => {
					break self.into_host_output(Outcome::Panic)?;
				},
				InvokeOutcome::OutOfGas => {
					break self.into_host_output(Outcome::OutOfGas)?;
				},
			}
		};
		debug!("Finished with outcome {:?}", output.vm_output.outcome);
		Ok((output, outer_vm))
	}

	fn handle_host_call_fault(&mut self, index: u64) -> Result<(), HostCallError> {
		let Some(handler) = self.host_call_handlers.get(index as usize) else {
			panic!("Unknown host call index: {index}");
		};
		handler(self)
	}

	fn into_host_output(self, outcome: Outcome) -> Result<(CoreVmOutput, O), ApiError> {
		let (
			outer_vm,
			inner_vm,
			mapped_heap_pages,
			resident_pages,
			touched_imported_pages,
			updated_pages,
			num_memory_pages,
			stream_len,
			kernel_state,
		) = self.memory_man.export()?;
		let program_counter = inner_vm.expunge().expect("Failed to close VM handle");
		let vm_state = VmState {
			mapped_heap_pages,
			resident_pages,
			regs: self.args.regs,
			program_counter,
			kernel: kernel_state,
			restart_host_call: self.current_host_call,
			video: self.video,
			audio: self.audio,
		};
		let new_hash = compute_state_hash(&vm_state);
		let work_output = CoreVmOutput {
			vm_output: VmOutput {
				remaining_gas: self.args.gas,
				outcome,
				num_memory_pages,
				stream_len,
			},
			vm_state,
			old_hash: self.old_hash,
			new_hash,
			touched_imported_pages,
			updated_pages,
			exec_ref: self.exec_ref,
		};
		Ok((work_output, outer_vm))
	}
}

/// This hash is used to check the validity of the state transition in JAM's `accumulate`.
pub fn compute_state_hash(vm: &VmState) -> Hash {
	hash_encoded((&vm.regs, vm.program_counter, &vm.mapped_heap_pages, &vm.resident_pages))
}

/// The hash of the initial state.
pub fn initial_state_hash() -> Hash {
	let initial_vm_state = VmState::initial();
	compute_state_hash(&initial_vm_state)
}

pub(crate) struct DebugInvokeOutcome(pub InvokeOutcome);

impl core::fmt::Debug for DebugInvokeOutcome {
	fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
		match self.0 {
			InvokeOutcome::Halt => f.write_str("Halt"),
			InvokeOutcome::PageFault(address) =>
				f.debug_tuple("PageFault").field(&format_args!("{address:#x}")).finish(),
			InvokeOutcome::HostCallFault(i) => f.debug_tuple("HostCallFault").field(&i).finish(),
			InvokeOutcome::Panic => f.write_str("Panic"),
			InvokeOutcome::OutOfGas => f.write_str("OutOfGas"),
		}
	}
}

pub(crate) fn get_work_output_len(
	mapped_heap_pages: &RangeSet,
	resident_pages: &RangeSet,
	num_imported_pages: usize,
	num_updated_pages: usize,
	num_opened_files: usize,
) -> usize {
	MAX_STATIC_WORK_OUTPUT_LEN +
		mapped_heap_pages.encoded_size() +
		resident_pages.encoded_size() +
		vec_map_encoded_len::<u64, Hash>(num_imported_pages) +
		vec_map_encoded_len::<u64, Hash>(num_updated_pages) +
		vec_map_encoded_len::<u32, KernelFd>(num_opened_files)
}

fn vec_map_encoded_len<K: ConstEncodedLen, V: ConstEncodedLen>(len: usize) -> usize {
	(K::max_encoded_len() + V::max_encoded_len()) * len + Compact::<u64>::compact_len(&(len as u64))
}

/// Maximum length of the work output in bytes without dynamically allocated components (e.g.
/// pages).
const MAX_STATIC_WORK_OUTPUT_LEN: usize = 269;

#[cfg(test)]
mod tests {
	use super::*;
	use corevm_host::{OutputStream, Range};
	use jam_types::{max_exports, max_imports, Encode, ServiceId, SignedGas};
	use rand::Rng;

	#[test]
	fn max_static_work_output_len_is_correct() {
		assert_eq!(MAX_STATIC_WORK_OUTPUT_LEN, encoded_work_output_len(max_exports()));
	}

	fn encoded_work_output_len(num_exports: u32) -> usize {
		let output = CoreVmOutput {
			vm_output: VmOutput {
				remaining_gas: SignedGas::MAX,
				outcome: Outcome::Halt,
				num_memory_pages: num_exports,
				stream_len: [u32::MAX; OutputStream::COUNT],
			},
			vm_state: VmState {
				regs: Default::default(),
				program_counter: Default::default(),
				mapped_heap_pages: Default::default(),
				resident_pages: Default::default(),
				kernel: Default::default(),
				restart_host_call: Some(u64::MAX),
				video: Some(
					VideoMode::new(1, 1, 1, corevm_host::VideoFrameFormat::Rgb888Indexed8).unwrap(),
				),
				audio: Some(AudioMode::new(1, 1, corevm_host::AudioSampleFormat::S16LE).unwrap()),
			},
			old_hash: Default::default(),
			new_hash: Default::default(),
			updated_pages: Default::default(),
			touched_imported_pages: Default::default(),
			exec_ref: fs::BlockRef { service_id: 0, hash: fs::Hash(Hash::default()) },
		};
		output.encode().len() -
			output.updated_pages.encode().len() -
			output.touched_imported_pages.encode().len() -
			output.vm_state.mapped_heap_pages.encode().len() -
			output.vm_state.resident_pages.encode().len() -
			output.vm_state.kernel.fds.encode().len()
	}

	#[test]
	fn get_work_output_len_works() {
		let mut rng = rand::rng();
		for num_imported_pages in [0, 1, 2, max_imports() as usize] {
			for num_updated_pages in [0, 1, 2, max_exports() as usize] {
				for num_opened_files in [0, 1, 2, 1000] {
					let mapped_heap_pages = random_range_set(&mut rng);
					let resident_pages = random_range_set(&mut rng);
					let expected_len = get_work_output_len(
						&mapped_heap_pages,
						&resident_pages,
						num_imported_pages,
						num_updated_pages,
						num_opened_files,
					);
					let actual_len = get_work_output_len_slow(
						&mapped_heap_pages,
						&resident_pages,
						num_imported_pages,
						num_updated_pages,
						num_opened_files,
					);
					assert_eq!(
						expected_len, actual_len,
						"expected = {expected_len}, \
                        actual = {actual_len}, \
                        num_imported_pages = {num_imported_pages}, \
                        num_updated_pages = {num_updated_pages}, \
                        num_opened_files = {num_opened_files}"
					);
				}
			}
		}
	}

	fn get_work_output_len_slow(
		mapped_heap_pages: &RangeSet,
		resident_pages: &RangeSet,
		num_imported_pages: usize,
		num_updated_pages: usize,
		num_opened_files: usize,
	) -> usize {
		let output = CoreVmOutput {
			vm_output: VmOutput {
				remaining_gas: SignedGas::MAX,
				outcome: Outcome::Halt,
				num_memory_pages: max_exports(),
				stream_len: [u32::MAX; OutputStream::COUNT],
			},
			vm_state: VmState {
				regs: Default::default(),
				program_counter: Default::default(),
				mapped_heap_pages: mapped_heap_pages.clone(),
				resident_pages: resident_pages.clone(),
				kernel: corevm_host::KernelState {
					fds: (0..num_opened_files)
						.map(|i| {
							(
								i as u32,
								KernelFd {
									block_ref: fs::BlockRef {
										service_id: ServiceId::MAX,
										hash: fs::Hash::default(),
									},
									position: u64::MAX,
								},
							)
						})
						.collect(),
				},
				restart_host_call: Some(u64::MAX),
				video: Some(
					VideoMode::new(1, 1, 1, corevm_host::VideoFrameFormat::Rgb888Indexed8).unwrap(),
				),
				audio: Some(AudioMode::new(1, 1, corevm_host::AudioSampleFormat::S16LE).unwrap()),
			},
			old_hash: Default::default(),
			new_hash: Default::default(),
			updated_pages: (0..num_updated_pages)
				.map(|i| (i as u64 * PAGE_SIZE, Hash::default()))
				.collect(),
			touched_imported_pages: (0..num_imported_pages)
				.map(|i| (i as u64 * PAGE_SIZE, Hash::default()))
				.collect(),
			exec_ref: fs::BlockRef { service_id: 0, hash: fs::Hash(Hash::default()) },
		};
		output.encode().len()
	}

	fn random_range_set<R: Rng>(rng: &mut R) -> RangeSet {
		let mut ranges = RangeSet::new();
		{
			let mut offset = 0;
			for _ in 0..10 {
				let start = rng.random_range(offset..=20);
				let end = rng.random_range(start..=20);
				offset = end;
				ranges.insert(Range::new(start, end));
			}
		}
		ranges
	}
}