corevm-engine 0.1.28

use crate::hash_encoded;
use alloc::{borrow::Cow, vec::Vec};
use codec::{Decode, DecodeAll, Encode};
use corevm_host::{
	fs, CoreVmInstruction, CoreVmOutput, PageInfo, PageNum, Range, RangeSet, StorageKey, VmSpec,
	VmState,
};
use jam_types::{
	AccumulateItem, Balance, Hash, Memo, SegmentTreeRoot, ServiceId, SignedGas, Slot, UnsignedGas,
	WorkError, WorkOutput,
};
use log::debug;

mod messages;

const ALL_STORAGE_KEYS_WITHOUT_PARAMS: &[StorageKey] = {
	use StorageKey::*;
	&[
		Gas,
		StateHash,
		VmSpec,
		VideoMode,
		AudioMode,
		ExecEnvRef,
		Owner,
		StoredPages,
		IncomingServiceMessages,
		OutgoingServiceMessages,
	]
};

const _: () = assert!(ALL_STORAGE_KEYS_WITHOUT_PARAMS.len() == StorageKey::COUNT_WITHOUT_PARAMS);

/// Storage keys that should be removed on VM reset.
const REMOVE_ON_RESET_KEYS: &[StorageKey] = {
	use StorageKey::*;
	&[VmSpec, VideoMode, AudioMode, StoredPages, IncomingServiceMessages, OutgoingServiceMessages]
};

/// Storage keys that should _not_ be removed on VM reset.
const DO_NOT_REMOVE_ON_RESET_KEYS: &[StorageKey] = {
	use StorageKey::*;
	&[Gas, ExecEnvRef, StateHash, Owner]
};

const _: () = assert!(
	REMOVE_ON_RESET_KEYS.len() + DO_NOT_REMOVE_ON_RESET_KEYS.len() ==
		StorageKey::COUNT_WITHOUT_PARAMS
);

/// Operations that [`AccumulateEngine`] uses during accumulation.
///
/// Each operation corresponds to a JAM host-call.
pub trait AccumulateOps {
	/// Implemenatation-specific error type.
	type Error: core::fmt::Debug;

	/// Read value under the specified key from the storage of the specified service.
	fn get(&self, id: ServiceId, key: &StorageKey) -> Option<Cow<'_, [u8]>>;

	/// Write value under the specified key to the storage.
	fn set(&mut self, key: StorageKey, value: Cow<'_, [u8]>) -> Result<(), Self::Error>;

	/// Remove value under the specified key from the storage of the service being accumulated.
	///
	/// Returns `true` if the value was present in the storage before the removal.
	fn remove(&mut self, key: &StorageKey) -> bool;

	/// Same as [`set`](Self::set) but encodes the provided value using JAM codec.
	fn set_typed(&mut self, key: StorageKey, value: &impl Encode) -> Result<(), Self::Error> {
		value.using_encoded(|bytes| self.set(key, bytes.into()))
	}

	/// Same as [`get`](Self::get) but decodes the retrieved value using JAM codec.
	fn get_typed<T: Decode>(
		&self,
		id: ServiceId,
		key: &StorageKey,
	) -> Result<Option<T>, codec::Error> {
		self.get(id, key).map(|bytes| T::decode_all(&mut bytes.as_ref())).transpose()
	}

	/// Get min. memo gas of the specified service.
	fn min_memo_gas(&self, service_id: ServiceId) -> Option<UnsignedGas>;

	/// Transfer data and/or funds to another service.
	fn transfer(
		&self,
		destination: ServiceId,
		amount: Balance,
		gas_limit: UnsignedGas,
		memo: &Memo,
	) -> Result<(), Self::Error>;

	/// Same as [`transfer`](Self::transfer) but encodes the provided memo using JAM codec.
	///
	/// Panics if the encoded value doesn't fit into [`Memo`].
	fn transfer_typed(
		&self,
		destination: ServiceId,
		amount: Balance,
		gas_limit: UnsignedGas,
		memo: &impl Encode,
	) -> Result<(), Self::Error> {
		let mut raw_memo = Memo::zero();
		memo.encode_to(&mut SliceOutput(&mut raw_memo[..]));
		self.transfer(destination, amount, gas_limit, &raw_memo)
	}

	/// Reset the privileged services.
	fn bless<'a>(
		&mut self,
		manager: ServiceId,
		assigner: ServiceId,
		designator: ServiceId,
		registrar: ServiceId,
		always_acc: impl IntoIterator<Item = &'a (ServiceId, UnsignedGas)>,
	);

	/// Zombify the service.
	fn zombify(&mut self, ejector: ServiceId);
}

impl<A: AccumulateOps + ?Sized> AccumulateOps for &mut A {
	type Error = A::Error;

	fn get(&self, service_id: ServiceId, key: &StorageKey) -> Option<Cow<'_, [u8]>> {
		AccumulateOps::get(*self, service_id, key)
	}

	fn set(&mut self, key: StorageKey, value: Cow<'_, [u8]>) -> Result<(), Self::Error> {
		AccumulateOps::set(*self, key, value)
	}

	fn remove(&mut self, key: &StorageKey) -> bool {
		AccumulateOps::remove(*self, key)
	}

	fn min_memo_gas(&self, service_id: ServiceId) -> Option<UnsignedGas> {
		AccumulateOps::min_memo_gas(*self, service_id)
	}

	fn transfer(
		&self,
		destination: ServiceId,
		amount: Balance,
		gas_limit: UnsignedGas,
		memo: &Memo,
	) -> Result<(), Self::Error> {
		AccumulateOps::transfer(*self, destination, amount, gas_limit, memo)
	}

	fn bless<'a>(
		&mut self,
		manager: ServiceId,
		assigner: ServiceId,
		designator: ServiceId,
		registrar: ServiceId,
		always_acc: impl IntoIterator<Item = &'a (ServiceId, UnsignedGas)>,
	) {
		AccumulateOps::bless(*self, manager, assigner, designator, registrar, always_acc);
	}

	fn zombify(&mut self, ejector: ServiceId) {
		AccumulateOps::zombify(*self, ejector);
	}
}

/// An engine that drives execution of `accumulate` and `on_transfer` CoreVM service entry points.
pub struct AccumulateEngine<A: AccumulateOps> {
	ops: A,
	/// The current slot.
	slot: Slot,
	/// Id of the service that is being accumulated.
	service_id: ServiceId,
}

impl<A: AccumulateOps> AccumulateEngine<A> {
	/// Create new engine with the specified accumulation API implementation.
	pub fn new(ops: A, slot: Slot, service_id: ServiceId) -> Self {
		Self { ops, slot, service_id }
	}

	/// Accumulate the specified items.
	///
	/// Returns accumulation result for each item.
	pub fn run(
		&mut self,
		items: &[AccumulateItem],
	) -> Vec<Result<(), AccumulationError<A::Error>>> {
		items
			.iter()
			.map(|item| match item {
				AccumulateItem::WorkItem(item) =>
					self.accumulate(item.result.as_ref(), item.exports_root),
				AccumulateItem::Transfer(item) => self.transfer(item.source, item.memo),
			})
			.collect()
	}

	/// Get the underlying accumulation API implementation.
	pub fn into_inner(self) -> A {
		self.ops
	}

	fn accumulate(
		&mut self,
		result: Result<&WorkOutput, &WorkError>,
		exports_root: SegmentTreeRoot,
	) -> Result<(), AccumulationError<A::Error>> {
		use AccumulationError::*;
		let output = result.map_err(|e| Work(e.clone()))?;
		let output = CoreVmOutput::decode_all(&mut &output[..])?;

		if self.get_typed::<fs::BlockRef>(&StorageKey::ExecEnvRef)? != Some(output.exec_ref) {
			return Err(ExecEnvRefMismatch);
		}

		// Update VM state hash.
		if self.get_typed::<Hash>(&StorageKey::StateHash)? != Some(output.old_hash) {
			return Err(PriorStateMismatch);
		}
		self.set_typed(StorageKey::StateHash, &output.new_hash)?;
		let mut stored_pages =
			self.get_typed::<RangeSet>(&StorageKey::StoredPages)?.unwrap_or_default();
		// Check imported pages' hashes.
		for (page, our_hash) in output.touched_imported_pages.iter() {
			let key = StorageKey::PageInfo(*page);
			let their_hash = self.get_typed::<PageInfo>(&key)?.ok_or(NoSuchPage(*page))?.hash;
			if &their_hash != our_hash {
				return Err(PageHashMismatch);
			}
		}
		// Update memory pages' metadata.
		let mut export_index = 0;
		for (page, hash) in output.updated_pages.iter() {
			let key = StorageKey::PageInfo(*page);
			if hash == &[0; 32] {
				// The page was freed.
				self.get_typed::<PageInfo>(&key)?.ok_or(NoSuchPage(*page))?;
				stored_pages.remove(&Range::new(page.0, page.0 + 1));
				self.ops.remove(&key);
			} else {
				stored_pages.insert(Range::new(page.0, page.0 + 1));
				let info = PageInfo { hash: *hash, exports_root, export_index };
				self.set_typed(key, &info)?;
				export_index += 1;
			}
		}
		self.set_typed(StorageKey::StoredPages, &stored_pages)?;
		// Update video mode.
		if let Some(ref mode) = output.vm_state.video {
			self.set_typed(StorageKey::VideoMode, mode)?;
		}
		// Update audio mode.
		if let Some(ref mode) = output.vm_state.audio {
			self.set_typed(StorageKey::AudioMode, mode)?;
		}
		self.remove_processed_incoming_service_messages(&output.processed_service_messages)?;
		self.send_new_outgoing_service_messages_and_remove_old(output.outgoing_messages)?;
		// Store VM output specification.
		let spec = VmSpec { exports_root, output: output.vm_output, state: output.vm_state };
		self.set_typed(StorageKey::VmSpec, &spec)?;

		Ok(())
	}

	fn transfer(
		&mut self,
		source: ServiceId,
		memo: Memo,
	) -> Result<(), AccumulationError<A::Error>> {
		use AccumulationError::*;
		let instr = CoreVmInstruction::decode(&mut &memo[..])?;
		let owner = self.get_typed::<ServiceId>(&StorageKey::Owner)?;
		if let Some(owner) = owner {
			if owner != source {
				return Err(NotTheOwner(owner, source));
			}
		}
		match instr {
			CoreVmInstruction::Reset { gas, exec_ref } => {
				if owner.is_none() {
					debug!("Setting owner to {source:x}");
					self.set_typed(StorageKey::Owner, &source)?;
				}
				self.reset(gas, exec_ref)?;
			},
			CoreVmInstruction::SetOwner(owner) => {
				debug!("Setting owner to {owner:x}");
				self.set_typed(StorageKey::Owner, &owner)?;
			},
			CoreVmInstruction::PushServiceMessage(message) => {
				self.push_incoming_service_message(message)?;
			},
			CoreVmInstruction::Destroy { ejector } => {
				self.destroy()?;
				self.ops.bless(ejector, ejector, ejector, ejector, []);
				self.ops.zombify(ejector);
			},
		}
		Ok(())
	}

	/// Reset the VM.
	///
	/// Sets `ExecEnv` and `Gas`; resets everything else.
	fn reset(
		&mut self,
		gas: SignedGas,
		exec_ref: fs::BlockRef,
	) -> Result<(), AccumulationError<A::Error>> {
		debug!("Resetting VM: gas = {gas}, exec = {exec_ref}");
		let state_hash = hash_encoded(VmState::initial());
		self.set_typed(StorageKey::Gas, &gas)?;
		self.set_typed(StorageKey::ExecEnvRef, &exec_ref)?;
		self.set_typed(StorageKey::StateHash, &state_hash)?;
		self.remove_memory_pages()?;
		self.remove_all_incoming_service_messages()?;
		self.remove_all_outgoing_service_messages()?;
		// We should remove every key except the ones that we've overwritten above and except
		// Owner.
		for key in REMOVE_ON_RESET_KEYS {
			self.ops.remove(key);
		}
		Ok(())
	}

	/// Destroy the VM.
	///
	/// Removes everything from the storage.
	fn destroy(&mut self) -> Result<(), AccumulationError<A::Error>> {
		self.remove_memory_pages()?;
		self.remove_all_incoming_service_messages()?;
		self.remove_all_outgoing_service_messages()?;
		for key in ALL_STORAGE_KEYS_WITHOUT_PARAMS.iter() {
			self.ops.remove(key);
		}
		Ok(())
	}

	fn remove_memory_pages(&mut self) -> Result<(), AccumulationError<A::Error>> {
		let stored_pages =
			self.get_typed::<RangeSet>(&StorageKey::StoredPages)?.unwrap_or_default();
		for range in stored_pages.as_slice().iter() {
			for page in range.start..range.end {
				self.ops.remove(&StorageKey::PageInfo(PageNum(page)));
			}
		}
		Ok(())
	}

	fn set_typed(
		&mut self,
		key: StorageKey,
		value: &impl Encode,
	) -> Result<(), AccumulationError<A::Error>> {
		self.ops.set_typed(key, value).map_err(AccumulationError::Api)
	}

	fn get_typed<T: Decode>(&self, key: &StorageKey) -> Result<Option<T>, codec::Error> {
		self.ops.get_typed::<T>(self.service_id, key)
	}
}

/// Accumulation error.
#[derive(Debug, thiserror::Error)]
pub enum AccumulationError<E: core::fmt::Debug> {
	#[error("Accumulation API error: {0:?}")]
	Api(E),
	#[error("Work error: {0:?}")]
	Work(WorkError),
	#[error("JAM codec error: {0:?}")]
	Codec(codec::Error),
	#[error("Prior VM state mismatch")]
	PriorStateMismatch,
	#[error("Page {0} not found in storage")]
	NoSuchPage(PageNum),
	#[error("Page hash mismatch")]
	PageHashMismatch,
	#[error("ExecEnvRef mismatch")]
	ExecEnvRefMismatch,
	#[error("Processed message doesn't match the message in the queue")]
	MessageMismatch,
	#[error("Invalid message")]
	InvalidMessage,
	#[error("Unknown service {0:x}")]
	UnknownService(ServiceId),
	#[error("Transfer source is not the owner of the service: actual owner {0:x}, transfer source {1:x}")]
	NotTheOwner(ServiceId, ServiceId),
}

impl<E: core::fmt::Debug> From<codec::Error> for AccumulationError<E> {
	fn from(e: codec::Error) -> Self {
		Self::Codec(e)
	}
}

struct SliceOutput<'a>(&'a mut [u8]);

impl codec::Output for SliceOutput<'_> {
	fn write(&mut self, src: &[u8]) {
		let (dst, rest) = core::mem::take(&mut self.0).split_at_mut(src.len());
		dst.copy_from_slice(src);
		self.0 = rest;
	}
}