PVF: NaN canonicalization & deteriministic stack

client/executor/common/src/runtime_blob/runtime_blob.rs

@@ -81,6 +81,25 @@ impl RuntimeBlob {
 		export_mutable_globals(&mut self.raw_module, "exported_internal_global");
 	}
 
+	/// Run a pass that instrument this module so as to introduce a deterministic stack height limit.
+	///
+	/// It will introduce a global mutable counter. The instrumentation will increase the counter
+	/// according to the "cost" of the callee. If the cost exceeds the `stack_depth_limit` constant,
+	/// the instrumentation will trap. The counter will be decreased as soon as the the callee returns.
+	///
+	/// The stack cost of a function is computed based on how much locals there are and the maximum
+	/// depth of the wasm operand stack.
+	pub fn inject_stack_depth_metering(self, stack_depth_limit: u32) -> Result<Self, WasmError> {
+		let injected_module =
+			pwasm_utils::stack_height::inject_limiter(self.raw_module, stack_depth_limit).map_err(
+				|e| WasmError::Other(format!("cannot inject the stack limiter: {:?}", e)),
+			)?;
+
+		Ok(Self {
+			raw_module: injected_module,
+		})
+	}
+
 	/// Perform an instrumentation that makes sure that a specific function `entry_point` is exported
 	pub fn entry_point_exists(&self, entry_point: &str) -> bool {
 		self.raw_module.export_section().map(|e| {

client/executor/runtime-test/src/lib.rs

@@ -109,6 +109,12 @@ sp_core::wasm_export_functions! {
 
 	fn test_exhaust_heap() -> Vec<u8> { Vec::with_capacity(16777216) }
 
+	fn test_fp_f32add(a: [u8; 4], b: [u8; 4]) -> [u8; 4] {
+		let a = f32::from_le_bytes(a);
+		let b = f32::from_le_bytes(b);
+		f32::to_le_bytes(a + b)
+	}
+
 	fn test_panic() { panic!("test panic") }
 
 	fn test_conditional_panic(input: Vec<u8>) -> Vec<u8> {

client/executor/src/wasm_runtime.rs

@@ -329,6 +329,7 @@ pub fn create_wasm_runtime_with_code(
 					semantics: sc_executor_wasmtime::Semantics {
 						fast_instance_reuse: true,
 						stack_depth_metering: false,
+						canonicalize_nans: false,
 					},
 				},
 				host_functions,

client/executor/wasmtime/Cargo.toml

@@ -24,7 +24,12 @@ sp-wasm-interface = { version = "3.0.0", path = "../../../primitives/wasm-interf
 sp-runtime-interface = { version = "3.0.0", path = "../../../primitives/runtime-interface" }
 sp-core = { version = "3.0.0", path = "../../../primitives/core" }
 sp-allocator = { version = "3.0.0", path = "../../../primitives/allocator" }
-wasmtime = "0.27.0"
+pwasm-utils = { version = "0.18" }
+wasmtime = { version = "0.27.0", default-features = false, features = ["cache", "parallel-compilation"] }
 
 [dev-dependencies]
 assert_matches = "1.3.0"
+sc-runtime-test = { version = "2.0.0", path = "../runtime-test" }
+sp-io = { version = "3.0.0", path = "../../../primitives/io" }
+sp-wasm-interface = { version = "3.0.0", path = "../../../primitives/wasm-interface" }
+wat = "1.0"

client/executor/wasmtime/src/lib.rs

@@ -24,6 +24,9 @@ mod runtime;
 mod state_holder;
 mod util;
 
+#[cfg(test)]
+mod tests;
+
 pub use runtime::{
 	create_runtime, create_runtime_from_artifact, prepare_runtime_artifact, Config, Semantics,
 };

client/executor/wasmtime/src/runtime.rs

@@ -232,9 +232,39 @@ directory = \"{cache_dir}\"
 	Ok(())
 }
 
-fn common_config() -> wasmtime::Config {
+fn common_config(semantics: &Semantics) -> wasmtime::Config {
 	let mut config = wasmtime::Config::new();
 	config.cranelift_opt_level(wasmtime::OptLevel::SpeedAndSize);
+	config.cranelift_nan_canonicalization(semantics.canonicalize_nans);
+
+	if semantics.stack_depth_metering {
+		// With `stack_depth_metering` is enabled, the wasm code will instrumented with its own
+		// stack depth tracker.
+		//
+		// We do not ever want the wasmtime's stack overflow guard to be hit. The amount of data
+		// a stack frame of a compiled function consumes may depend on the wasmtime version and
+		// the architecture the host is running. This introduces some non-determinism and
+		// the `stack_depth_metering` option aims to fix exactly that.
+		//
+		// Hence 256 MiB for the stack. It seems unlikely that this will be reached even with specially
+		// crafted code.
+		config.max_wasm_stack(256 * 1024 * 1024).expect(
+			"this value is a literal and trivially is not 0;
+				we do not supply the async max stack;
+				according to the docs only Ok can be returned;
+				qed",
+		);
+	}
+
+	// Be clear and specific about the extensions we support. If an update brings new features
+	// they should be introduced here as well.
+	config.wasm_reference_types(false);
+	config.wasm_simd(false);
+	config.wasm_bulk_memory(false);
+	config.wasm_multi_value(false);
+	config.wasm_multi_memory(false);
+	config.wasm_module_linking(false);
+
 	config
 }
 
@@ -271,7 +301,20 @@ pub struct Semantics {
 	///
 	/// [stack_height]: https://github.com/paritytech/wasm-utils/blob/d9432baf/src/stack_height/mod.rs#L1-L50
 	pub stack_depth_metering: bool,
-	// Other things like nan canonicalization can be added here.
+
+	/// Controls whether wasmtime should compile floating point in a way that doesn't allow for
+	/// non-determinism.
+	///
+	/// By default, the wasm spec allows some local non-determinism wrt. certain floating point
+	/// operations. Specifically, those operations that are not defined to operate on bits (e.g. fneg)
+	/// can produce NaN values. The exact bit pattern for those is not specified and may depend
+	/// on the particular machine that executes wasmtime generated JITed machine code. That is
+	/// a source of non-deterministic values.
+	///
+	/// The classical runtime environment for Substrate allowed it and punted this on the runtime
+	/// developers. For PVFs, we want to ensure that execution is deterministic though. Therefore,
+	/// for PVF execution this flag is meant to be turned on.
+	pub canonicalize_nans: bool,
 }
 
 pub struct Config {
@@ -355,7 +398,7 @@ unsafe fn do_create_runtime(
 	host_functions: Vec<&'static dyn Function>,
 ) -> std::result::Result<WasmtimeRuntime, WasmError> {
 	// Create the engine, store and finally the module from the given code.
-	let mut wasmtime_config = common_config();
+	let mut wasmtime_config = common_config(&config.semantics);
 	if let Some(ref cache_path) = config.cache_path {
 		if let Err(reason) = setup_wasmtime_caching(cache_path, &mut wasmtime_config) {
 			log::warn!(
@@ -369,8 +412,8 @@ unsafe fn do_create_runtime(
 		.map_err(|e| WasmError::Other(format!("cannot create the engine for runtime: {}", e)))?;
 
 	let (module, snapshot_data) = match code_supply_mode {
-		CodeSupplyMode::Verbatim { mut blob } => {
-			instrument(&mut blob, &config.semantics);
+		CodeSupplyMode::Verbatim { blob } => {
+			let blob = instrument(blob, &config.semantics)?;
 
 			if config.semantics.fast_instance_reuse {
 				let data_segments_snapshot = DataSegmentsSnapshot::take(&blob).map_err(|e| {
@@ -412,25 +455,32 @@ unsafe fn do_create_runtime(
 	})
 }
 
-fn instrument(blob: &mut RuntimeBlob, semantics: &Semantics) {
+fn instrument(
+	mut blob: RuntimeBlob,
+	semantics: &Semantics,
+) -> std::result::Result<RuntimeBlob, WasmError> {
+	if semantics.stack_depth_metering {
+		const STACK_DEPTH_LIMIT: u32 = 65536;
+		blob = blob.inject_stack_depth_metering(STACK_DEPTH_LIMIT)?;
+	}
+
+	// If enabled, this should happen after all other passes that may introduce global variables.
 	if semantics.fast_instance_reuse {
 		blob.expose_mutable_globals();
 	}
 
-	if semantics.stack_depth_metering {
-		// TODO: implement deterministic stack metering https://github.com/paritytech/substrate/issues/8393
-	}
+	Ok(blob)
 }
 
 /// Takes a [`RuntimeBlob`] and precompiles it returning the serialized result of compilation. It
 /// can then be used for calling [`create_runtime`] avoiding long compilation times.
 pub fn prepare_runtime_artifact(
-	mut blob: RuntimeBlob,
+	blob: RuntimeBlob,
 	semantics: &Semantics,
 ) -> std::result::Result<Vec<u8>, WasmError> {
-	instrument(&mut blob, semantics);
+	let blob = instrument(blob, semantics)?;
 
-	let engine = Engine::new(&common_config())
+	let engine = Engine::new(&common_config(semantics))
 		.map_err(|e| WasmError::Other(format!("cannot create the engine: {}", e)))?;
 
 	engine