corevm_engine/

engine.rs

use crate::{
	hash_encoded, AppendOutputError, Arg, InnerVm, InvokeArgs, MemoryManager, OuterVm,
	PageFaultError, ProgramData, Reg,
};
use codec::{Compact, CompactLen};
use corevm_host::{
	CoreVmOutput, CoreVmPayload, Outcome, OutputStream, RangeSet, VideoMode, VmOutput, VmState,
	PAGE_SIZE,
};
use jam_pvm_common::{ApiError, InvokeOutcome};
use jam_types::{Hash, ServiceId};
use log::{debug, trace};
use polkavm::{MemoryMapBuilder, RETURN_TO_HOST};

pub struct Engine<O: OuterVm> {
	outer_vm: O,
	args: InvokeArgs,
	memory_man: MemoryManager<O::InnerVm>,
	video: Option<VideoMode>,
	frame_number: u64,
	old_hash: Hash,
	code_hash: Hash,
	code_host: ServiceId,
}

impl<O: OuterVm> Engine<O> {
	pub fn new(payload: CoreVmPayload, mut outer_vm: O) -> Result<Self, ApiError> {
		let CoreVmPayload { gas, vm_state, program: guest_code_ref } = payload;
		let max_exports = outer_vm.get_export_count();
		let auth_output_len = outer_vm.get_auth_output_len();
		let program_counter = vm_state.program_counter;
		let program = {
			let bytes = outer_vm.read_file(&guest_code_ref).expect("Failed to fetch program");
			let corevm_blob = jam_program_blob::CoreVmProgramBlob::from_bytes(&bytes)
				.expect("Failed to parse CoreVM blob");
			polkavm::ProgramParts::from_bytes(corevm_blob.pvm_blob.into())
				.expect("Failed to parse PolkaVM blob")
		};
		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()
			.expect("Failed to build memory map");
		// Should include program hash and pages' hash.
		let old_hash = compute_state_hash(&vm_state);
		let mut args = InvokeArgs { regs: vm_state.regs, gas };
		if program_counter == 0 {
			// Initialize the registers for the initial `invoke` call.
			args.set_reg(Reg::RA, RETURN_TO_HOST);
			args.set_reg(Reg::SP, memory_map.stack_address_high() as u64);
		}
		let code = &program.code_and_jump_table[..];
		let inner_vm = outer_vm.machine(code, program_counter).expect("Failed to load the code");
		let program_data = ProgramData::new(&memory_map, program.ro_data, program.rw_data);
		let memory_man = MemoryManager::new(
			inner_vm,
			&mut outer_vm,
			program_data,
			vm_state.mapped_heap_pages,
			vm_state.resident_pages,
			max_exports as usize,
			auth_output_len as usize,
		)
		.expect("Failed to initialize memory pages");
		Ok(Self {
			outer_vm,
			args,
			memory_man,
			video: None,
			frame_number: vm_state.frame_number,
			old_hash,
			code_hash: guest_code_ref.hash,
			code_host: guest_code_ref.service_id,
		})
	}

	pub fn run(mut self) -> Result<(CoreVmOutput, O), ApiError> {
		let (output, outer_vm) = loop {
			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(PageFaultError::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(PageFaultError::ApiError(e)) => return Err(e),
					}
				},
				InvokeOutcome::HostCallFault(index) => match self.handle_host_call_fault(index) {
					Ok(..) => {},
					Err(AppendOutputError::OutputLimitReached) => {
						trace!("Output limit reached");
						break self.into_host_output(Outcome::OutputLimitReached)?;
					},
					Err(AppendOutputError::PageFault { page, num_pages }) => {
						trace!(
							"Hard page fault at address {:#x}-{:#x}",
							page * PAGE_SIZE,
							(page + num_pages) * PAGE_SIZE
						);
						break self.into_host_output(Outcome::PageFault { page, num_pages })?;
					},
					Err(AppendOutputError::ApiError(e)) => return Err(e),
				},
				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<(), AppendOutputError> {
		let args = &mut self.args;
		match index {
			index_for_corevm_call!(gas) => {
				args.set_return_value(args.gas);
			},
			index_for_corevm_call!(yield_console_data) => {
				let mut i = 0;
				let stream: u64 = Arg::get(args, &mut i);
				let inner_src: u64 = Arg::get(args, &mut i);
				let length: u64 = Arg::get(args, &mut i);
				let stream = match stream {
					1 => OutputStream::Stdout,
					2 => OutputStream::Stderr,
					other => panic!("Invalid output stream specified: {other}"),
				};
				trace!("Call yield_console_data({stream:?}, {inner_src:#x}, {length})");
				args.set_return_value(1_u64);
				self.memory_man.append_output(stream, inner_src, length)?;
				args.set_return_value(0_u64);
			},
			index_for_corevm_call!(yield_video_frame) => {
				let mut i = 0;
				let inner_src: u64 = Arg::get(args, &mut i);
				let length: u64 = Arg::get(args, &mut i);
				trace!("Call yield_video_frame({inner_src:#x}, {length})");
				args.set_return_value(1_u64);
				self.memory_man.append_output(OutputStream::Video, inner_src, length)?;
				args.set_return_value(0_u64);
				self.frame_number += 1;
			},
			index_for_corevm_call!(alloc) => {
				let mut i = 0;
				let size: u64 = Arg::get(args, &mut i);
				let (address, _size) = match self.memory_man.alloc(size) {
					Some((address, size)) => (address, size),
					None => {
						debug!("Failed to map guest memory block of size {}", size);
						(0, 0)
					},
				};
				trace!("Call alloc({size}) = {address:#x}");
				args.set_return_value(address);
			},
			index_for_corevm_call!(free) => {
				let mut i = 0;
				let address: u64 = Arg::get(args, &mut i);
				let size: u64 = Arg::get(args, &mut i);
				self.memory_man.dealloc(address, size)?;
				trace!("Call free({address:#x}, {size})");
				args.set_return_value(0_u64);
			},
			index_for_corevm_call!(video_mode) => {
				let mut i = 0;
				let width: u64 = Arg::get(args, &mut i);
				let height: u64 = Arg::get(args, &mut i);
				let refresh_rate: u64 = Arg::get(args, &mut i);
				let format: u64 = Arg::get(args, &mut i);
				trace!("Call video_mode({width}, {height}, {refresh_rate}, {format})");
				self.video = Some(VideoMode {
					width: width as u32,
					height: height as u32,
					mode: format as u32,
					refresh_rate: refresh_rate as u16,
				});
				args.set_return_value(0_u64);
			},
			index => {
				panic!("Unknown host call index: {}", index);
			},
		}
		Ok(())
	}

	fn into_host_output(mut self, outcome: Outcome) -> Result<(CoreVmOutput, O), ApiError> {
		let (
			mapped_heap_pages,
			resident_pages,
			touched_imported_pages,
			updated_pages,
			num_memory_pages,
			stream_len,
		) = self.memory_man.export(&mut self.outer_vm)?;
		let program_counter =
			self.memory_man.inner_vm.expunge().expect("Failed to close VM handle");
		let vm_state = VmState {
			mapped_heap_pages,
			resident_pages,
			regs: self.args.regs,
			program_counter,
			frame_number: self.frame_number,
		};
		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,
			video: self.video,
			guest_code_hash: self.code_hash,
			guest_code_host: self.code_host,
		};
		Ok((work_output, self.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.frame_number,
		&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"),
		}
	}
}

// TODO @ivan Handle imports and exports by name, not by index.
#[rustfmt::skip]
macro_rules! index_for_corevm_call {
	(gas) => {0};
	(alloc) => {1};
	(free) => {2};
	(yield_console_data) => {3};
	(yield_video_frame) => {4};
	(video_mode) => {5};
}

use index_for_corevm_call;

pub fn get_work_output_len(
	mapped_heap_pages: &RangeSet,
	resident_pages: &RangeSet,
	num_imported_pages: usize,
	num_updated_pages: usize,
) -> usize {
	let imported_pages_encoded_len = core::mem::size_of::<(u64, Hash)>() * num_imported_pages +
		Compact::<u64>::compact_len(&(num_imported_pages as u64));
	let updated_pages_encoded_len = core::mem::size_of::<(u64, Hash)>() * num_updated_pages +
		Compact::<u64>::compact_len(&(num_updated_pages as u64));
	MAX_STATIC_WORK_OUTPUT_LEN +
		mapped_heap_pages.encoded_len() +
		resident_pages.encoded_len() +
		imported_pages_encoded_len +
		updated_pages_encoded_len
}

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

/// The maximum total size of all output blobs in a work-report, in octets.
pub const MAX_TOTAL_OUTPUT_BLOB_SIZE: usize = 48 * 1024;

#[cfg(test)]
mod tests {
	use super::*;
	use jam_types::{max_exports, Encode, SignedGas};

	#[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 {
				frame_number: Default::default(),
				regs: Default::default(),
				program_counter: Default::default(),
				mapped_heap_pages: Default::default(),
				resident_pages: Default::default(),
			},
			old_hash: Default::default(),
			new_hash: Default::default(),
			updated_pages: Default::default(),
			touched_imported_pages: Default::default(),
			video: Some(VideoMode {
				mode: Default::default(),
				width: Default::default(),
				height: Default::default(),
				refresh_rate: Default::default(),
			}),
			guest_code_hash: Default::default(),
			guest_code_host: Default::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()
	}
}