From f1b7052154d746899d8bf3462dc78aebb9799d54 Mon Sep 17 00:00:00 2001 From: Jiyuan Zheng Date: Thu, 6 Jun 2024 11:53:18 +0800 Subject: [PATCH] feat(poc): extension system (#27) * Every Extension is represented by a trait * Extension trait has corresponding Dispatchables (enum generic on extension impl) * dummy implementations for trait like extension_core and extension_fungibles * An ExtensionExecutor which supports initializing extension impls tuple, dispatching polkavm host calls * A PermController safeguards extensions according to invocation source --- Cargo.lock | 10 ++ Cargo.toml | 10 +- poc/executor/src/lib.rs | 9 +- poc/extensions/Cargo.toml | 14 ++ poc/extensions/src/dispatchable.rs | 7 + poc/extensions/src/error.rs | 10 ++ poc/extensions/src/extension_core.rs | 48 ++++++ poc/extensions/src/extension_fungibles.rs | 44 +++++ poc/extensions/src/extension_id.rs | 5 + poc/extensions/src/guest.rs | 10 ++ poc/extensions/src/lib.rs | 194 ++++++++++++++++++++++ poc/extensions/src/macros.rs | 20 +++ poc/extensions/src/perm_controller.rs | 18 ++ poc/hosts/pass-custom-type/src/main.rs | 4 +- poc/runtime/src/xcq.rs | 2 +- 15 files changed, 398 insertions(+), 7 deletions(-) create mode 100644 poc/extensions/Cargo.toml create mode 100644 poc/extensions/src/dispatchable.rs create mode 100644 poc/extensions/src/error.rs create mode 100644 poc/extensions/src/extension_core.rs create mode 100644 poc/extensions/src/extension_fungibles.rs create mode 100644 poc/extensions/src/extension_id.rs create mode 100644 poc/extensions/src/guest.rs create mode 100644 poc/extensions/src/lib.rs create mode 100644 poc/extensions/src/macros.rs create mode 100644 poc/extensions/src/perm_controller.rs diff --git a/Cargo.lock b/Cargo.lock index f6b5175..f7a8c7c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2521,6 +2521,16 @@ dependencies = [ "polkavm", ] +[[package]] +name = "poc-extension" +version = "0.1.0" +dependencies = [ + "impl-trait-for-tuples", + "parity-scale-codec", + "poc-executor", + "scale-info", +] + [[package]] name = "poc-host-pass-custom-type" version = "0.1.0" diff --git a/Cargo.toml b/Cargo.toml index 4cd4fd4..387a449 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -11,6 +11,7 @@ members = [ "poc/hosts/*", "poc/runtime", "poc/executor", + "poc/extensions", "xcq-api", "xcq-executor", @@ -51,8 +52,13 @@ polkavm = { path = "vendor/polkavm/crates/polkavm", default-features = false } sp-api = { version = "29.0.0", default-features = false } # nostd -parity-scale-codec = { version = "3.6.12", default-features = false, features = ["derive", "max-encoded-len"] } -scale-info = { version = "2.11.3", default-features = false, features = ["derive"] } +parity-scale-codec = { version = "3.6.12", default-features = false, features = [ + "derive", + "max-encoded-len", +] } +scale-info = { version = "2.11.3", default-features = false, features = [ + "derive", +] } tracing = { version = "0.1.40", default-features = false } # std diff --git a/poc/executor/src/lib.rs b/poc/executor/src/lib.rs index 2cc99cd..07fbf87 100644 --- a/poc/executor/src/lib.rs +++ b/poc/executor/src/lib.rs @@ -52,7 +52,12 @@ impl XcqExecutor { } } - pub fn execute(&mut self, raw_blob: &[u8], input: &[u8]) -> Result, XcqExecutorError> { + pub fn execute( + &mut self, + raw_blob: &[u8], + method: impl AsRef<[u8]>, + input: &[u8], + ) -> Result, XcqExecutorError> { let blob = ProgramBlob::parse(raw_blob.into())?; let module = Module::from_blob(&self.engine, &Default::default(), blob)?; let instance_pre = self.linker.instantiate_pre(&module)?; @@ -71,7 +76,7 @@ impl XcqExecutor { 0 }; - let res = instance.call_typed::<(u32, u32), u64>(&mut self.context, "main", (input_ptr, input.len() as u32))?; + let res = instance.call_typed::<(u32, u32), u64>(&mut self.context, method, (input_ptr, input.len() as u32))?; let res_ptr = (res >> 32) as u32; let res_size = (res & 0xffffffff) as u32; let result = instance diff --git a/poc/extensions/Cargo.toml b/poc/extensions/Cargo.toml new file mode 100644 index 0000000..1333cfc --- /dev/null +++ b/poc/extensions/Cargo.toml @@ -0,0 +1,14 @@ +[package] +name = "poc-extension" +version = "0.1.0" +edition = "2021" + +[dependencies] +parity-scale-codec = { version = "3.6.12", default-features = false } +scale-info = { version = "2.6.0", default-features = false } +poc-executor = { path = "../executor", default-features = false } +impl-trait-for-tuples = "0.2.2" + +[features] +default = ["std"] +std = ["parity-scale-codec/std", "scale-info/std", "poc-executor/std"] diff --git a/poc/extensions/src/dispatchable.rs b/poc/extensions/src/dispatchable.rs new file mode 100644 index 0000000..a66758c --- /dev/null +++ b/poc/extensions/src/dispatchable.rs @@ -0,0 +1,7 @@ +pub trait Dispatchable { + fn dispatch(self) -> Result, DispatchError>; +} + +pub enum DispatchError { + PhantomData, +} diff --git a/poc/extensions/src/error.rs b/poc/extensions/src/error.rs new file mode 100644 index 0000000..4e87e71 --- /dev/null +++ b/poc/extensions/src/error.rs @@ -0,0 +1,10 @@ +// TODO: contain source error +use crate::DispatchError; +use parity_scale_codec::Error as CodeCError; +pub enum ExtensionError { + PermissionError, + PolkavmError, + DecodeError(CodeCError), + DispatchError(DispatchError), + UnsupportedExtension, +} diff --git a/poc/extensions/src/extension_core.rs b/poc/extensions/src/extension_core.rs new file mode 100644 index 0000000..692695d --- /dev/null +++ b/poc/extensions/src/extension_core.rs @@ -0,0 +1,48 @@ +use crate::{DispatchError, Dispatchable}; +use crate::{ExtensionId, ExtensionIdTy}; +use parity_scale_codec::{Decode, Encode}; + +pub trait ExtensionCore { + type Config: Config; + fn some_host_function( + args: ::ArgsOfSomeHostFunction, + ) -> ::ResultOfSomeHostFunction; +} + +pub trait Config { + type ArgsOfSomeHostFunction: Decode; + type ResultOfSomeHostFunction: Encode; +} + +// #[extension(ExtensionCore)] +// type Call; + +mod generated_by_extension_decl { + use super::*; + + #[derive(Decode)] + pub enum ExtensionCoreCall { + SomeHostFunction { + args: ::ArgsOfSomeHostFunction, + }, + } + + impl Dispatchable for ExtensionCoreCall { + fn dispatch(self) -> Result, DispatchError> { + match self { + Self::SomeHostFunction { args } => Ok(Impl::some_host_function(args).encode()), + } + } + } + + impl ExtensionId for ExtensionCoreCall { + const EXTENSION_ID: ExtensionIdTy = 0u64; + } + + // TODO: remove this when formalized + #[allow(dead_code)] + pub type Call = ExtensionCoreCall; +} + +#[allow(unused_imports)] +pub use generated_by_extension_decl::*; diff --git a/poc/extensions/src/extension_fungibles.rs b/poc/extensions/src/extension_fungibles.rs new file mode 100644 index 0000000..9dbfbb8 --- /dev/null +++ b/poc/extensions/src/extension_fungibles.rs @@ -0,0 +1,44 @@ +use crate::{DispatchError, Dispatchable}; +use crate::{ExtensionId, ExtensionIdTy}; +use core::marker::PhantomData; +use parity_scale_codec::{Decode, Encode}; + +pub trait ExtensionFungibles { + fn free_balance_of(who: [u8; 32]) -> u32; + fn reserved_balance_of(who: [u8; 32]) -> u32; +} + +// #[extension(ExtensionFungibles)] +// type Call; + +mod generated_by_extension_decl { + + use super::*; + #[derive(Decode)] + pub enum ExtensionFungiblesCall { + FreeBalanceOf { who: [u8; 32] }, + ReservedBalanceOf { who: [u8; 32] }, + _Marker(PhantomData), + } + + impl Dispatchable for ExtensionFungiblesCall { + fn dispatch(self) -> Result, DispatchError> { + match self { + Self::FreeBalanceOf { who } => Ok(Impl::free_balance_of(who).encode()), + Self::ReservedBalanceOf { who } => Ok(Impl::reserved_balance_of(who).encode()), + Self::_Marker(_) => Err(DispatchError::PhantomData), + } + } + } + + impl ExtensionId for ExtensionFungiblesCall { + const EXTENSION_ID: ExtensionIdTy = 1u64; + } + + // TODO: remove this when formalized + #[allow(dead_code)] + pub type Call = ExtensionFungiblesCall; +} + +#[allow(unused_imports)] +pub use generated_by_extension_decl::*; diff --git a/poc/extensions/src/extension_id.rs b/poc/extensions/src/extension_id.rs new file mode 100644 index 0000000..095b5df --- /dev/null +++ b/poc/extensions/src/extension_id.rs @@ -0,0 +1,5 @@ +pub type ExtensionIdTy = u64; + +pub trait ExtensionId { + const EXTENSION_ID: ExtensionIdTy; +} diff --git a/poc/extensions/src/guest.rs b/poc/extensions/src/guest.rs new file mode 100644 index 0000000..f05c850 --- /dev/null +++ b/poc/extensions/src/guest.rs @@ -0,0 +1,10 @@ +pub trait Guest { + fn program(&self) -> &[u8]; +} + +pub type Method = String; + +pub trait Input { + fn method(&self) -> Method; + fn args(&self) -> &[u8]; +} diff --git a/poc/extensions/src/lib.rs b/poc/extensions/src/lib.rs new file mode 100644 index 0000000..489a0f3 --- /dev/null +++ b/poc/extensions/src/lib.rs @@ -0,0 +1,194 @@ +#![cfg_attr(not(feature = "std"), no_std)] +use core::marker::PhantomData; + +use parity_scale_codec::Decode; +use poc_executor::{XcqExecutor, XcqExecutorContext}; +#[cfg(not(feature = "std"))] +use scale_info::prelude::{format, string::String}; +pub type XcqResponse = Vec; +pub type XcqError = String; +pub type XcqResult = Result; + +mod dispatchable; +pub use dispatchable::{DispatchError, Dispatchable}; +mod extension_id; +pub use extension_id::{ExtensionId, ExtensionIdTy}; +mod error; +pub use error::ExtensionError; +mod macros; + +mod extension_core; +mod extension_fungibles; + +mod perm_controller; +pub use perm_controller::{InvokeSource, PermController}; + +mod guest; +pub use guest::{Guest, Input, Method}; + +// alias trait +trait Extension: Dispatchable + ExtensionId + Decode {} +impl Extension for T where T: Dispatchable + ExtensionId + Decode {} + +trait ExtensionTuple { + fn dispatch(extension_id: ExtensionIdTy, data: &[u8]) -> Result, ExtensionError>; +} + +struct Context { + invoke_source: InvokeSource, + _marker: PhantomData<(E, P)>, +} + +impl Context { + pub fn new(invoke_source: InvokeSource) -> Self { + Self { + invoke_source, + _marker: PhantomData, + } + } +} + +impl XcqExecutorContext for Context { + fn register_host_functions(&mut self, linker: &mut poc_executor::Linker) { + let invoke_source = self.invoke_source; + linker + .func_wrap( + "_", + move |mut caller: poc_executor::Caller<_>, + extension_id: u64, + call_ptr: u32, + call_len: u32, + res_ptr: u32| + -> u32 { + // useful closure to handle early return + let mut func_with_result = || -> Result { + let call_bytes = caller + .read_memory_into_vec(call_ptr, call_len) + .map_err(|_| ExtensionError::PolkavmError)?; + if P::is_allowed(extension_id, &call_bytes, invoke_source) { + return Err(ExtensionError::PermissionError); + } + let res_bytes = E::dispatch(extension_id, &call_bytes)?; + caller + .write_memory(res_ptr, &res_bytes[..]) + .map_err(|_| ExtensionError::PolkavmError)?; + Ok(res_bytes.len() as u32) + }; + func_with_result().unwrap_or(0) + }, + ) + .unwrap(); + } +} + +struct ExtensionsExecutor { + executor: XcqExecutor>, +} +impl ExtensionsExecutor { + #[allow(dead_code)] + pub fn new(source: InvokeSource) -> Self { + let context = Context::::new(source); + let executor = XcqExecutor::new(Default::default(), context); + Self { executor } + } + // In PoC, guest and input are opaque to the runtime + // In SDK, we can make them has type + #[allow(dead_code)] + fn execute_method(&mut self, guest: G, input: I) -> XcqResult { + self.executor + .execute(guest.program(), input.method(), input.args()) + .map_err(|e| format!("{:?}", e)) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::extension_core::ExtensionCore; + use crate::extension_fungibles::ExtensionFungibles; + use parity_scale_codec::{Decode, Encode}; + + // extension_core impls + pub struct ExtensionCoreImpl; + + #[derive(Encode, Decode)] + pub struct ArgsImpl { + pub a: u32, + pub b: u32, + } + + pub struct ConfigImpl; + impl extension_core::Config for ConfigImpl { + // this associated type is generated by the macro + type ArgsOfSomeHostFunction = ArgsImpl; + type ResultOfSomeHostFunction = u32; + } + + impl ExtensionCore for ExtensionCoreImpl { + type Config = ConfigImpl; + fn some_host_function( + args: ::ArgsOfSomeHostFunction, + ) -> ::ResultOfSomeHostFunction { + args.a + args.b + } + } + + // extension_fungibles impls + pub struct ExtensionFungiblesImpl; + + impl ExtensionFungibles for ExtensionFungiblesImpl { + fn free_balance_of(_who: [u8; 32]) -> u32 { + 100 + } + fn reserved_balance_of(_who: [u8; 32]) -> u32 { + 42 + } + } + + type Extensions = ( + extension_core::Call, + extension_fungibles::Call, + ); + + // guest impls + pub struct GuestImpl { + pub program: Vec, + } + + impl Guest for GuestImpl { + fn program(&self) -> &[u8] { + &self.program + } + } + + pub struct InputImpl { + pub method: Method, + pub args: Vec, + } + + impl Input for InputImpl { + fn method(&self) -> Method { + self.method.clone() + } + fn args(&self) -> &[u8] { + &self.args + } + } + + // TODO: refine the test + #[test] + fn extensions_executor_fails() { + let mut executor = ExtensionsExecutor::::new(InvokeSource::RuntimeAPI); + let guest = GuestImpl { + program: vec![0, 1, 2, 3], + }; + let input = InputImpl { + method: "main".to_string(), + args: vec![0, 1, 2, 3], + }; + let res = executor.execute_method(guest, input); + assert!(res.is_err()) + } + + // TODO: add success test +} diff --git a/poc/extensions/src/macros.rs b/poc/extensions/src/macros.rs new file mode 100644 index 0000000..c743d35 --- /dev/null +++ b/poc/extensions/src/macros.rs @@ -0,0 +1,20 @@ +use crate::Extension; +use crate::ExtensionError; +use crate::ExtensionIdTy; +use crate::ExtensionTuple; + +// Use the macro to implement ExtensionTuple for tuples of different lengths +#[impl_trait_for_tuples::impl_for_tuples(10)] +#[tuple_types_custom_trait_bound(Extension)] +impl ExtensionTuple for Tuple { + fn dispatch(extension_id: ExtensionIdTy, mut call: &[u8]) -> Result, ExtensionError> { + for_tuples!( + #( + if extension_id == Tuple::EXTENSION_ID { + return Tuple::decode(&mut call).map_err(ExtensionError::DecodeError)?.dispatch().map_err(ExtensionError::DispatchError); + } + )* + ); + Err(ExtensionError::UnsupportedExtension) + } +} diff --git a/poc/extensions/src/perm_controller.rs b/poc/extensions/src/perm_controller.rs new file mode 100644 index 0000000..0d4a62e --- /dev/null +++ b/poc/extensions/src/perm_controller.rs @@ -0,0 +1,18 @@ +use crate::ExtensionIdTy; +pub trait PermController { + fn is_allowed(extension_id: ExtensionIdTy, call: &[u8], source: InvokeSource) -> bool; +} + +impl PermController for () { + fn is_allowed(_extension_id: ExtensionIdTy, _call: &[u8], _context: InvokeSource) -> bool { + true + } +} + +#[derive(Copy, Clone)] +pub enum InvokeSource { + RuntimeAPI, + XCM, + Extrinsic, + Runtime, +} diff --git a/poc/hosts/pass-custom-type/src/main.rs b/poc/hosts/pass-custom-type/src/main.rs index 0596df4..3518a23 100644 --- a/poc/hosts/pass-custom-type/src/main.rs +++ b/poc/hosts/pass-custom-type/src/main.rs @@ -54,9 +54,9 @@ fn main() { let config = Config::from_env().unwrap(); let mut executor: poc_executor::XcqExecutor = poc_executor::XcqExecutor::new(config, HostFunctions); - let res = executor.execute(raw_blob, &[0u8]).unwrap(); + let res = executor.execute(raw_blob, "main", &[0u8]).unwrap(); println!("Result: {:?}", res); - let res = executor.execute(raw_blob, &[1u8, 40u8]).unwrap(); + let res = executor.execute(raw_blob, "main", &[1u8, 40u8]).unwrap(); println!("Result: {:?}", res); } diff --git a/poc/runtime/src/xcq.rs b/poc/runtime/src/xcq.rs index 05fbbe6..04490ab 100644 --- a/poc/runtime/src/xcq.rs +++ b/poc/runtime/src/xcq.rs @@ -51,7 +51,7 @@ impl poc_executor::XcqExecutorContext for HostFunctions { pub fn execute_query(query: Vec, input: Vec) -> XcqResult { let mut executor = poc_executor::XcqExecutor::new(Default::default(), HostFunctions); - executor.execute(&query, &input).map_err(|e| format!("{:?}", e)) + executor.execute(&query, "main", &input).map_err(|e| format!("{:?}", e)) } #[cfg(test)]