PVF: NaN canonicalization & deteriministic stack (#9069)

* NaN canonicalization

* Introduce a simple stack depth metering

* Be explicit about the wasm features we enable

* Pull the latest latast fix for the pwasm-utils crate

* Disable `wasm_threads` as well.

* Factor out deterministic stack params

* Add more docs

* Remove redundant dep

* Refine comments

* Typo

Co-authored-by: Andronik Ordian <write@reusable.software>

Co-authored-by: Andronik Ordian <write@reusable.software>

substrate/client/executor/wasmtime/src/lib.rs

@@ -24,6 +24,10 @@ 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,
+	DeterministicStackLimit,
 };

substrate/client/executor/wasmtime/src/runtime.rs

@@ -232,10 +232,75 @@ directory = \"{cache_dir}\"
 	Ok(())
 }
 
-fn common_config() -> wasmtime::Config {
+fn common_config(semantics: &Semantics) -> std::result::Result<wasmtime::Config, WasmError> {
 	let mut config = wasmtime::Config::new();
 	config.cranelift_opt_level(wasmtime::OptLevel::SpeedAndSize);
-	config
+	config.cranelift_nan_canonicalization(semantics.canonicalize_nans);
+
+	if let Some(DeterministicStackLimit {
+		native_stack_max, ..
+	}) = semantics.deterministic_stack_limit
+	{
+		config
+			.max_wasm_stack(native_stack_max as usize)
+			.map_err(|e| WasmError::Other(format!("cannot set max wasm stack: {}", e)))?;
+	}
+
+	// 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.wasm_threads(false);
+
+	Ok(config)
+}
+
+/// Knobs for deterministic stack height limiting.
+///
+/// The WebAssembly standard defines a call/value stack but it doesn't say anything about its
+/// size except that it has to be finite. The implementations are free to choose their own notion
+/// of limit: some may count the number of calls or values, others would rely on the host machine
+/// stack and trap on reaching a guard page.
+///
+/// This obviously is a source of non-determinism during execution. This feature can be used
+/// to instrument the code so that it will count the depth of execution in some deterministic
+/// way (the machine stack limit should be so high that the deterministic limit always triggers
+/// first).
+///
+/// The deterministic stack height limiting feature allows to instrument the code so that it will
+/// count the number of items that may be on the stack. This counting will only act as an rough
+/// estimate of the actual stack limit in wasmtime. This is because wasmtime measures it's stack
+/// usage in bytes.
+///
+/// The actual number of bytes consumed by a function is not trivial to compute  without going through
+/// full compilation. Therefore, it's expected that `native_stack_max` is grealy overestimated and
+/// thus never reached in practice. The stack overflow check introduced by the instrumentation and
+/// that relies on the logical item count should be reached first.
+///
+/// See [here][stack_height] for more details of the instrumentation
+///
+/// [stack_height]: https://github.com/paritytech/wasm-utils/blob/d9432baf/src/stack_height/mod.rs#L1-L50
+pub struct DeterministicStackLimit {
+	/// A number of logical "values" that can be pushed on the wasm stack. A trap will be triggered
+	/// if exceeded.
+	///
+	/// A logical value is a local, an argument or a value pushed on operand stack.
+	pub logical_max: u32,
+	/// The maximum number of bytes for stack used by wasmtime JITed code.
+	///
+	/// It's not specified how much bytes will be consumed by a stack frame for a given wasm function
+	/// after translation into machine code. It is also not quite trivial.
+	///
+	/// Therefore, this number should be choosen conservatively. It must be so large so that it can
+	/// fit the [`logical_max`] logical values on the stack, according to the current instrumentation
+	/// algorithm.
+	///
+	/// This value cannot be 0.
+	pub native_stack_max: u32,
 }
 
 pub struct Semantics {
@@ -254,24 +319,30 @@ pub struct Semantics {
 	/// is used.
 	pub fast_instance_reuse: bool,
 
-	/// The WebAssembly standard defines a call/value stack but it doesn't say anything about its
-	/// size except that it has to be finite. The implementations are free to choose their own notion
-	/// of limit: some may count the number of calls or values, others would rely on the host machine
-	/// stack and trap on reaching a guard page.
+	/// Specifiying `Some` will enable deterministic stack height. That is, all executor invocations
+	/// will reach stack overflow at the exactly same point across different wasmtime versions and
+	/// architectures.
 	///
-	/// This obviously is a source of non-determinism during execution. This feature can be used
-	/// to instrument the code so that it will count the depth of execution in some deterministic
-	/// way (the machine stack limit should be so high that the deterministic limit always triggers
-	/// first).
-	///
-	/// See [here][stack_height] for more details of the instrumentation
+	/// This is achieved by a combination of running an instrumentation pass on input code and
+	/// configuring wasmtime accordingly.
 	///
 	/// Since this feature depends on instrumentation, it can be set only if [`CodeSupplyMode::Verbatim`]
 	/// is used.
+	pub deterministic_stack_limit: Option<DeterministicStackLimit>,
+
+	/// Controls whether wasmtime should compile floating point in a way that doesn't allow for
+	/// non-determinism.
 	///
-	/// [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.
+	/// 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 +426,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 +440,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 +483,31 @@ 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 let Some(DeterministicStackLimit { logical_max, .. }) = semantics.deterministic_stack_limit {
+		blob = blob.inject_stack_depth_metering(logical_max)?;
+	}
+
+	// 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

substrate/client/executor/wasmtime/src/tests.rs

@@ -0,0 +1,173 @@
+// This file is part of Substrate.
+
+// Copyright (C) 2021 Parity Technologies (UK) Ltd.
+// SPDX-License-Identifier: GPL-3.0-or-later WITH Classpath-exception-2.0
+
+// This program is free software: you can redistribute it and/or modify
+// it under the terms of the GNU General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+
+// This program is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+// GNU General Public License for more details.
+
+// You should have received a copy of the GNU General Public License
+// along with this program. If not, see <https://www.gnu.org/licenses/>.
+
+use sc_executor_common::{
+	runtime_blob::RuntimeBlob,
+	wasm_runtime::WasmModule,
+};
+use sc_runtime_test::wasm_binary_unwrap;
+use codec::{Encode as _, Decode as _};
+use std::sync::Arc;
+
+type HostFunctions = sp_io::SubstrateHostFunctions;
+
+struct RuntimeBuilder {
+	code: Option<&'static str>,
+	fast_instance_reuse: bool,
+	canonicalize_nans: bool,
+	deterministic_stack: bool,
+	heap_pages: u32,
+}
+
+impl RuntimeBuilder {
+	/// Returns a new builder that won't use the fast instance reuse mechanism, but instead will
+	/// create a new runtime instance each time.
+	fn new_on_demand() -> Self {
+		Self {
+			code: None,
+			fast_instance_reuse: false,
+			canonicalize_nans: false,
+			deterministic_stack: false,
+			heap_pages: 1024,
+		}
+	}
+
+	fn use_wat(&mut self, code: &'static str) {
+		self.code = Some(code);
+	}
+
+	fn canonicalize_nans(&mut self, canonicalize_nans: bool) {
+		self.canonicalize_nans = canonicalize_nans;
+	}
+
+	fn deterministic_stack(&mut self, deterministic_stack: bool) {
+		self.deterministic_stack = deterministic_stack;
+	}
+
+	fn build(self) -> Arc<dyn WasmModule> {
+		let blob = {
+			let wasm: Vec<u8>;
+
+			let wasm = match self.code {
+				None => wasm_binary_unwrap(),
+				Some(wat) => {
+					wasm = wat::parse_str(wat).unwrap();
+					&wasm
+				}
+			};
+
+			RuntimeBlob::uncompress_if_needed(&wasm)
+				.expect("failed to create a runtime blob out of test runtime")
+		};
+
+		let rt = crate::create_runtime(
+			blob,
+			crate::Config {
+				heap_pages: self.heap_pages,
+				allow_missing_func_imports: true,
+				cache_path: None,
+				semantics: crate::Semantics {
+					fast_instance_reuse: self.fast_instance_reuse,
+					deterministic_stack_limit:
+						match self.deterministic_stack {
+							true => Some(crate::DeterministicStackLimit {
+								logical_max: 65536,
+								native_stack_max: 256 * 1024 * 1024,
+							}),
+							false => None,
+						},
+					canonicalize_nans: self.canonicalize_nans,
+				},
+			},
+			{
+				use sp_wasm_interface::HostFunctions as _;
+				HostFunctions::host_functions()
+			}
+		)
+		.expect("cannot create runtime");
+
+		Arc::new(rt) as Arc<dyn WasmModule>
+	}
+}
+
+#[test]
+fn test_nan_canonicalization() {
+	let runtime = {
+		let mut builder = RuntimeBuilder::new_on_demand();
+		builder.canonicalize_nans(true);
+		builder.build()
+	};
+
+	let instance = runtime
+		.new_instance()
+		.expect("failed to instantiate a runtime");
+
+	/// A NaN with canonical payload bits.
+	const CANONICAL_NAN_BITS: u32 = 0x7fc00000;
+	/// A NaN value with an abitrary payload.
+	const ARBITRARY_NAN_BITS: u32 = 0x7f812345;
+
+	// This test works like this: we essentially do
+	//
+	//     a + b
+	//
+	// where
+	//
+	//     * a is a nan with arbitrary bits in its payload
+	//     * b is 1.
+	//
+	// according to the wasm spec, if one of the inputs to the operation is a non-canonical NaN
+	// then the value be a NaN with non-deterministic payload bits.
+	//
+	// However, with the `canonicalize_nans` option turned on above, we expect that the output will
+	// be a canonical NaN.
+	//
+	// We exterpolate the results of this tests so that we assume that all intermediate computations
+	// that involve floats are sanitized and cannot produce a non-deterministic NaN.
+
+	let params = (u32::to_le_bytes(ARBITRARY_NAN_BITS), u32::to_le_bytes(1)).encode();
+	let res = {
+		let raw_result = instance.call_export(
+			"test_fp_f32add",
+			&params,
+		).unwrap();
+		u32::from_le_bytes(<[u8; 4]>::decode(&mut &raw_result[..]).unwrap())
+	};
+	assert_eq!(res, CANONICAL_NAN_BITS);
+}
+
+#[test]
+fn test_stack_depth_reaching() {
+	const TEST_GUARD_PAGE_SKIP: &str = include_str!("test-guard-page-skip.wat");
+
+	let runtime = {
+		let mut builder = RuntimeBuilder::new_on_demand();
+		builder.use_wat(TEST_GUARD_PAGE_SKIP);
+		builder.deterministic_stack(true);
+		builder.build()
+	};
+	let instance = runtime
+		.new_instance()
+		.expect("failed to instantiate a runtime");
+
+	let err = instance.call_export("test-many-locals", &[]).unwrap_err();
+
+	assert!(
+		format!("{:?}", err).starts_with("Other(\"Wasm execution trapped: wasm trap: unreachable")
+	);
+}