corevm-engine 0.1.28

use crate::{
	hash_encoded, BoundedInput, BoundedWorkOutput, FatalError, FsNode, HostCallError,
	HostCallHandler, InitError, InnerVm, InvokeArgs, KernelContext, Lookup, MemoryMap, OpenedFile,
	OuterVm, OutputBuffers, Reg, RunError, TouchError,
};
use alloc::vec::Vec;
use codec::DecodeAll;
use corevm_codec::video;
use corevm_host::{
	fs, CoreVmExtrinsics, CoreVmOutput, CoreVmPayload, ExecEnv, Outcome, OutputStream, PageNum,
	Range, VmOutput, PAGE_SIZE,
};
use jam_pvm_common::InvokeOutcome;
use jam_types::{max_exports, VecSet};
use log::{debug, trace};
use polkakernel::{KernelState, Machine};
use polkavm::{ArcBytes, RETURN_TO_HOST};

/// Parameters derived from work package.
pub struct WorkPackageParams {
	/// Total encoded size of the package.
	pub encoded_size: u32,
	/// The size of the authorizer output.
	pub auth_output_size: u32,
	/// The number of exports.
	pub export_count: u16,
}

/// This structure records which portions of the input were read/accessed by the guest.
#[derive(Debug)]
pub struct InputStats {
	/// Imported memory pages that were accessed by the guest.
	pub accessed_imported_pages: VecSet<PageNum>,
	/// The number of inter-service messages read by the guest.
	pub num_service_messages_read: u32,
	/// The number of input chunks read.
	pub num_input_chunks_read: u32,
}

/// An engine that drives execution of `refine` CoreVM service entry point.
pub struct Engine<O: OuterVm> {
	pub(crate) args: InvokeArgs,
	pub(crate) outer_vm: O,
	pub(crate) inner_vm: O::InnerVm,
	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,
	/// The current timestamp within the time slot, measured as the no. of milliseconds since the
	/// start of the slot.
	///
	/// Incremented when a video frame or audio samples are yielded by the program. Never goes
	/// backwards.
	pub(crate) timestamp: u64,
	pub(crate) memory_map: MemoryMap,
	/// Should be equal to `WorkItem::export_count`.
	pub(crate) export_count: u16,
	/// Guest program's output streams.
	pub(crate) output: OutputBuffers,
	pub(crate) ro_data: ArcBytes,
	pub(crate) rw_data: ArcBytes,
	pub(crate) work_output: BoundedWorkOutput,
	pub(crate) host_call_handlers: Vec<HostCallHandler<O>>,
	/// Empty when logs are disabled.
	pub(crate) host_call_names: Vec<&'static str>,
	pub(crate) input: BoundedInput,
	/// Video output encoder.
	pub(crate) video_encoder: Option<video::Encoder>,
}

impl<O: OuterVm> Engine<O> {
	pub fn new(
		payload: CoreVmPayload,
		extrinsics: CoreVmExtrinsics,
		work_package_params: WorkPackageParams,
		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 export_count = work_package_params.export_count;
		let auth_output_len = work_package_params.auth_output_size;
		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 polkavm_memory_map = polkavm::MemoryMapBuilder::new(PAGE_SIZE)
			.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, host_call_names) = Self::create_host_call_handlers(&program_blob);
		// Should include program hash and pages' hash.
		let old_hash = hash_encoded(&vm_state);
		let 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 input = BoundedInput::new(work_package_params.encoded_size, extrinsics);
		let inner_vm = outer_vm
			.machine(&code_and_jump_table[..], program_counter)
			.expect("Failed to load the code");
		let memory_map = MemoryMap::new(&polkavm_memory_map);
		let kernel_state = KernelState {
			fds: vm_state
				.kernel
				.fds
				.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 heap_address_range = memory_map.heap_range();
		let heap_page_range =
			Range::new(heap_address_range.start / PAGE_SIZE, heap_address_range.end / PAGE_SIZE);
		// Sanity checks.
		assert!(u32::from(export_count) <= max_exports());
		assert!(
			ro_data.len() as u32 <=
				memory_map.ro_data_range().end - memory_map.ro_data_range().start
		);
		assert!(
			rw_data.len() as u32 <=
				memory_map.rw_data_range().end - memory_map.rw_data_range().start
		);
		let mut engine = Self {
			// Pre-populate output with the data that we already have.
			work_output: BoundedWorkOutput::try_from(
				CoreVmOutput {
					vm_output: VmOutput {
						remaining_gas: 0,
						outcome: Outcome::Halt,
						num_memory_pages: u32::from(export_count),
						stream_len: [u32::MAX; OutputStream::COUNT],
					},
					vm_state,
					old_hash,
					new_hash: Default::default(),
					touched_imported_pages: Default::default(),
					updated_pages: Default::default(),
					exec_ref,
					outgoing_messages: Default::default(),
					processed_service_messages: Default::default(),
				},
				kernel_state,
				auth_output_len,
				heap_page_range,
			)
			.expect("Empty work output shouldn't exceed max_report_elective_data"),
			args,
			host_call_handlers,
			host_call_names,
			num_video_frames: 0,
			num_audio_bytes: 0,
			timestamp: 0,
			outer_vm,
			inner_vm,
			memory_map,
			ro_data,
			rw_data,
			export_count,
			output: Default::default(),
			exec,
			input,
			video_encoder: None,
		};
		if initial_program_run {
			// Initialize the registers for the initial `invoke` call.
			let default_sp = polkavm_memory_map.stack_address_high() as u64;
			if libc {
				let args = core::mem::take(&mut engine.exec.args);
				let env = core::mem::take(&mut engine.exec.env);
				let mut context = KernelContext { engine: &mut engine, error: None };
				context
					.init(
						default_sp,
						RETURN_TO_HOST,
						args.iter().map(|a| a.as_ref()),
						env.iter().map(|a| a.as_ref()),
					)
					.expect("Failed to init arguments");
				engine.exec.args = args;
				engine.exec.env = env;
			} else {
				engine.args.set_reg(Reg::RA, RETURN_TO_HOST);
				engine.args.set_reg(Reg::SP, default_sp);
			}
		}
		Ok(engine)
	}

	pub fn run(mut self) -> Result<(CoreVmOutput, O, InputStats), 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.work_output.vm_state.restart_host_call.is_some();
		let (output, outer_vm, input_stats) = loop {
			if let Some(index) = self.work_output.take_restart_host_call() {
				use Outcome::*;
				if restart_host_call {
					debug!("Restarting host-call {:?}", self.pretty_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 {:?}",
							self.pretty_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 {:?}",
							self.pretty_host_call(index)
						);
						self.work_output.set_restart_host_call(index);
						Some(outcome)
					},
					Err(HostCallError::Outcome(outcome @ InputLimitReached)) => {
						if restart_host_call {
							return Err(FatalError::NotEnoughInputSpace.into());
						}
						self.work_output.set_restart_host_call(index);
						Some(outcome)
					},
					Err(HostCallError::Outcome(outcome @ WaitingForInput(..))) => {
						self.work_output.set_restart_host_call(index);
						Some(outcome)
					},
					Err(HostCallError::Outcome(outcome @ PageFault { page, num_pages })) => {
						if restart_host_call {
							debug!(
								"Hard page fault at {:?} while restarting host-call {:?}",
								self.memory_map.pretty_page_range(page.0..page.0 + num_pages),
								self.pretty_host_call(index)
							);
							// 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 {:?}",
							self.memory_map.pretty_page_range(page.0..page.0 + num_pages),
							self.pretty_host_call(index)
						);
						self.work_output.set_restart_host_call(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.suspend(outcome)?;
				}
				restart_host_call = false;
			}
			let (outcome, gas, regs) = self.inner_vm.invoke(self.args.gas, self.args.regs)?;
			trace!("Invoke outcome {:?}", self.pretty_invoke_outcome(outcome));
			self.args.gas = gas;
			self.args.regs = regs;
			match outcome {
				InvokeOutcome::Halt => {
					break self.suspend(Outcome::Halt)?;
				},
				InvokeOutcome::PageFault(address) => {
					match self.touch(PageNum::from_address(address as u32)) {
						Ok(..) => {},
						Err(TouchError::PageFault { page, num_pages }) => {
							trace!("Hard page fault at {:?}", self.memory_map.pretty_page(page.0));
							// 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.suspend(Outcome::PageFault { page, num_pages })?;
						},
						Err(TouchError::Jam(e)) => return Err(e.into()),
						Err(TouchError::InvalidMemoryAccess) =>
							return Err(RunError::Fatal(FatalError::InvalidMemoryAccess)),
						Err(TouchError::InputLimitReached) =>
							break self.suspend(Outcome::InputLimitReached)?,
					}
				},
				InvokeOutcome::HostCallFault(index) =>
					self.work_output.set_restart_host_call(index),
				InvokeOutcome::Panic => {
					break self.suspend(Outcome::Panic)?;
				},
				InvokeOutcome::OutOfGas => {
					break self.suspend(Outcome::OutOfGas)?;
				},
			}
		};
		debug!("Finished with outcome {:?}", output.vm_output.outcome);
		Ok((output, outer_vm, input_stats))
	}

	fn handle_host_call_fault(&mut self, index: u64) -> Result<(), HostCallError> {
		let handler =
			self.host_call_handlers.get(index as usize).ok_or(FatalError::UnknownHostCall)?;
		handler(self)
	}
}