Skip to content

Commit

Permalink
Only load module once when validating
Browse files Browse the repository at this point in the history
  • Loading branch information
pgherveou committed Mar 22, 2024
1 parent ea5f4e9 commit 94fe749
Show file tree
Hide file tree
Showing 5 changed files with 722 additions and 654 deletions.
15 changes: 11 additions & 4 deletions substrate/frame/contracts/src/benchmarking/sandbox.rs
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,8 @@
/// ! environment that provides the seal interface as imported functions.
use super::{code::WasmModule, Config};
use crate::wasm::{
AllowDeprecatedInterface, AllowUnstableInterface, Determinism, Environment, WasmBlob,
AllowDeprecatedInterface, AllowUnstableInterface, Determinism, Environment, LoadedModule,
LoadingMode, WasmBlob,
};
use sp_core::Get;
use wasmi::{errors::LinkerError, Func, Linker, StackLimits, Store};
Expand All @@ -42,12 +43,18 @@ impl<T: Config> From<&WasmModule<T>> for Sandbox {
/// Creates an instance from the supplied module.
/// Sets the execution engine fuel level to `u64::MAX`.
fn from(module: &WasmModule<T>) -> Self {
let (mut store, _memory, instance) = WasmBlob::<T>::instantiate::<EmptyEnv, _>(
let contract = LoadedModule::new::<T>(
&module.code,
Determinism::Relaxed,
Some(StackLimits::default()),
LoadingMode::Checked,
)
.expect("Failed to load Wasm module");

let (mut store, _memory, instance) = WasmBlob::<T>::instantiate::<EmptyEnv, _>(
contract,
(),
&<T>::Schedule::get(),
Determinism::Relaxed,
StackLimits::default(),
// We are testing with an empty environment anyways
AllowDeprecatedInterface::No,
)
Expand Down
32 changes: 31 additions & 1 deletion substrate/frame/contracts/src/tests.rs
Original file line number Diff line number Diff line change
Expand Up @@ -33,7 +33,7 @@ use crate::{
primitives::CodeUploadReturnValue,
storage::DeletionQueueManager,
tests::test_utils::{get_contract, get_contract_checked},
wasm::{Determinism, ReturnErrorCode as RuntimeReturnCode},
wasm::{Determinism, LoadingMode, ReturnErrorCode as RuntimeReturnCode},
weights::WeightInfo,
Array, BalanceOf, Code, CodeHash, CodeInfoOf, CollectEvents, Config, ContractInfo,
ContractInfoOf, DebugInfo, DefaultAddressGenerator, DeletionQueueCounter, Error, HoldReason,
Expand Down Expand Up @@ -1299,6 +1299,36 @@ fn delegate_call() {
});
}

#[test]
fn track_check_uncheck_module_call() {
let (wasm, code_hash) = compile_module::<Test>("dummy").unwrap();
ExtBuilder::default().build().execute_with(|| {
let _ = <Test as Config>::Currency::set_balance(&ALICE, 1_000_000);

Contracts::bare_upload_code(ALICE, wasm, None, Determinism::Enforced).unwrap();

Contracts::bare_instantiate(
ALICE,
1_000,
GAS_LIMIT,
None,
Code::Existing(code_hash),
vec![],
vec![],
DebugInfo::Skip,
CollectEvents::Skip,
)
.result
.unwrap()
.account_id;
});

assert_eq!(
crate::wasm::tracker::LOADED_MODULE.with(|stack| stack.borrow().clone()),
vec![LoadingMode::Checked, LoadingMode::Unchecked]
);
}

#[test]
fn transfer_expendable_cannot_kill_account() {
let (wasm, _code_hash) = compile_module::<Test>("dummy").unwrap();
Expand Down
43 changes: 25 additions & 18 deletions substrate/frame/contracts/src/wasm/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -21,26 +21,26 @@
mod prepare;
mod runtime;

#[cfg(test)]
pub use runtime::STABLE_API_COUNT;

#[cfg(doc)]
pub use crate::wasm::runtime::api_doc;

pub use crate::wasm::runtime::{
AllowDeprecatedInterface, AllowUnstableInterface, Environment, Runtime, RuntimeCosts,
};

#[cfg(test)]
pub use tests::MockExt;
pub use {
crate::wasm::{prepare::tracker, runtime::ReturnErrorCode},
runtime::STABLE_API_COUNT,
tests::MockExt,
};

#[cfg(test)]
pub use crate::wasm::runtime::ReturnErrorCode;
pub use crate::wasm::{
prepare::{LoadedModule, LoadingMode},
runtime::{
AllowDeprecatedInterface, AllowUnstableInterface, Environment, Runtime, RuntimeCosts,
},
};

use crate::{
exec::{ExecResult, Executable, ExportedFunction, Ext},
gas::{GasMeter, Token},
wasm::prepare::LoadedModule,
weights::WeightInfo,
AccountIdOf, BadOrigin, BalanceOf, CodeHash, CodeInfoOf, CodeVec, Config, Error, Event,
HoldReason, Pallet, PristineCode, Schedule, Weight, LOG_TARGET,
Expand Down Expand Up @@ -201,17 +201,14 @@ impl<T: Config> WasmBlob<T> {
/// When validating we pass `()` as `host_state`. Please note that such a dummy instance must
/// **never** be called/executed, since it will panic the executor.
pub fn instantiate<E, H>(
code: &[u8],
contract: LoadedModule,
host_state: H,
schedule: &Schedule<T>,
determinism: Determinism,
stack_limits: StackLimits,
allow_deprecated: AllowDeprecatedInterface,
) -> Result<(Store<H>, Memory, InstancePre), &'static str>
where
E: Environment<H>,
{
let contract = LoadedModule::new::<T>(&code, determinism, Some(stack_limits))?;
let mut store = Store::new(&contract.engine, host_state);
let mut linker = Linker::new(&contract.engine);
E::define(
Expand Down Expand Up @@ -358,12 +355,22 @@ impl<T: Config> Executable<T> for WasmBlob<T> {
// Instantiate the Wasm module to the engine.
let runtime = Runtime::new(ext, input_data);
let schedule = <T>::Schedule::get();

let contract = LoadedModule::new::<T>(
&code,
self.code_info.determinism,
Some(StackLimits::default()),
LoadingMode::Unchecked,
)
.map_err(|err| {
log::debug!(target: LOG_TARGET, "failed to create wasmi module: {err:?}");
Error::<T>::CodeRejected
})?;

let (mut store, memory, instance) = Self::instantiate::<crate::wasm::runtime::Env, _>(
code,
contract,
runtime,
&schedule,
self.code_info.determinism,
StackLimits::default(),
match function {
ExportedFunction::Call => AllowDeprecatedInterface::Yes,
ExportedFunction::Constructor => AllowDeprecatedInterface::No,
Expand Down
68 changes: 46 additions & 22 deletions substrate/frame/contracts/src/wasm/prepare.rs
Original file line number Diff line number Diff line change
Expand Up @@ -48,6 +48,20 @@ pub struct LoadedModule {
pub engine: Engine,
}

#[derive(PartialEq, Debug, Clone)]
pub enum LoadingMode {
Checked,
Unchecked,
}

#[cfg(test)]
pub mod tracker {
use sp_std::cell::RefCell;
thread_local! {
pub static LOADED_MODULE: RefCell<Vec<super::LoadingMode>> = RefCell::new(Vec::new());
}
}

impl LoadedModule {
/// Creates a new instance of `LoadedModule`.
///
Expand All @@ -57,6 +71,7 @@ impl LoadedModule {
code: &[u8],
determinism: Determinism,
stack_limits: Option<StackLimits>,
_mode: LoadingMode,
) -> Result<Self, &'static str> {
// NOTE: wasmi does not support unstable WebAssembly features. The module is implicitly
// checked for not having those ones when creating `wasmi::Module` below.
Expand All @@ -79,11 +94,16 @@ impl LoadedModule {
}

let engine = Engine::new(&config);

// TODO use Module::new_unchecked when validate_module is true once we are on [email protected]
let module = Module::new(&engine, code).map_err(|err| {
log::debug!(target: LOG_TARGET, "Module creation failed: {:?}", err);
"Can't load the module into wasmi!"
})?;

#[cfg(test)]
tracker::LOADED_MODULE.with(|t| t.borrow_mut().push(_mode));

// Return a `LoadedModule` instance with
// __valid__ module.
Ok(LoadedModule { module, engine })
Expand Down Expand Up @@ -229,24 +249,38 @@ where
E: Environment<()>,
T: Config,
{
(|| {
let module = (|| {
// We don't actually ever execute this instance so we can get away with a minimal stack
// which reduces the amount of memory that needs to be zeroed.
let stack_limits = Some(StackLimits::new(1, 1, 0).expect("initial <= max; qed"));

// We check that the module is generally valid,
// and does not have restricted WebAssembly features, here.
let contract_module = match *determinism {
Determinism::Relaxed =>
if let Ok(module) = LoadedModule::new::<T>(code, Determinism::Enforced, None) {
if let Ok(module) = LoadedModule::new::<T>(
code,
Determinism::Enforced,
stack_limits,
LoadingMode::Checked,
) {
*determinism = Determinism::Enforced;
module
} else {
LoadedModule::new::<T>(code, Determinism::Relaxed, None)?
LoadedModule::new::<T>(code, Determinism::Relaxed, None, LoadingMode::Checked)?
},
Determinism::Enforced => LoadedModule::new::<T>(code, Determinism::Enforced, None)?,
Determinism::Enforced => LoadedModule::new::<T>(
code,
Determinism::Enforced,
stack_limits,
LoadingMode::Checked,
)?,
};

// The we check that module satisfies constraints the pallet puts on contracts.
contract_module.scan_exports()?;
contract_module.scan_imports::<T>(schedule)?;
Ok(())
Ok(contract_module)
})()
.map_err(|msg: &str| {
log::debug!(target: LOG_TARGET, "New code rejected on validation: {}", msg);
Expand All @@ -257,22 +291,11 @@ where
//
// - It doesn't use any unknown imports.
// - It doesn't explode the wasmi bytecode generation.
//
// We don't actually ever execute this instance so we can get away with a minimal stack which
// reduces the amount of memory that needs to be zeroed.
let stack_limits = StackLimits::new(1, 1, 0).expect("initial <= max; qed");
WasmBlob::<T>::instantiate::<E, _>(
&code,
(),
schedule,
*determinism,
stack_limits,
AllowDeprecatedInterface::No,
)
.map_err(|err| {
log::debug!(target: LOG_TARGET, "{}", err);
(Error::<T>::CodeRejected.into(), "New code rejected on wasmi instantiation!")
})?;
WasmBlob::<T>::instantiate::<E, _>(module, (), schedule, AllowDeprecatedInterface::No)
.map_err(|err| {
log::debug!(target: LOG_TARGET, "{err}");
(Error::<T>::CodeRejected.into(), "New code rejected on wasmi instantiation!")
})?;

Ok(())
}
Expand Down Expand Up @@ -325,7 +348,8 @@ pub mod benchmarking {
owner: AccountIdOf<T>,
) -> Result<WasmBlob<T>, DispatchError> {
let determinism = Determinism::Enforced;
let contract_module = LoadedModule::new::<T>(&code, determinism, None)?;
let contract_module =
LoadedModule::new::<T>(&code, determinism, None, LoadingMode::Checked)?;
let _ = contract_module.scan_imports::<T>(schedule)?;
let code: CodeVec<T> = code.try_into().map_err(|_| <Error<T>>::CodeTooLarge)?;
let code_info = CodeInfo {
Expand Down
Loading

0 comments on commit 94fe749

Please sign in to comment.