Move ABI check to instantiation (#233)

* Move ABI check to instantiation

This is a follow-up to #197.

This fixes an extra case I was worried about even with that fix but didn't hit it in my initial examples so hoped it's not worth the complexity.

Now I've hit it in a different module: because ABI can have zeroes - e.g. 3.0 is encoded as `\00\00\03\00` in the data segment - linker might place it in the beginning of a data segment, in which case `wasm-opt` will trim those prefix \00 bytes and there won't be any single data segment containing the entire version, and the ABI check will still fail.

It's possible to make the `determine_spacetime_abi` account for this case and merge chunks from different data segments when they overlap with version, but I feel it's growing to the point where it doesn't justify the complexity.

Instead, I'm moving the check into Wasmer's instantiation, since we construct a Wasm instance right away anyway, and it takes care of zero-filling and loading all the data segments into memory.

This makes the check more reliable and simpler, as well as avoids double-parsing of Wasm, but it does mean that we'll need to copy it to other hosts in the unlikely event we want to support 2 Wasm hosts simultaneously in the future.

IMO in this case the tradeoff is worth it.

* Add comment explanation

* Fixup conflict resolution

Signed-off-by: Ingvar Stepanyan <me@rreverser.com>

---------

Signed-off-by: Ingvar Stepanyan <me@rreverser.com>
Signed-off-by: Tyler Cloutier <cloutiertyler@users.noreply.github.com>
Co-authored-by: Tyler Cloutier <cloutiertyler@users.noreply.github.com>

crates/core/src/host/wasm_common.rs

@@ -1,4 +1,3 @@
-pub mod abi;
 pub mod module_host_actor;
 
 use crate::error::{DBError, IndexError, NodesError};
@@ -226,7 +225,6 @@ impl FuncNames {
 #[error(transparent)]
 pub enum ModuleCreationError {
     WasmCompileError(anyhow::Error),
-    Abi(#[from] abi::AbiVersionError),
     Init(#[from] module_host_actor::InitializationError),
 }
 

crates/core/src/host/wasm_common/module_host_actor.rs

@@ -11,7 +11,7 @@ use bytes::Bytes;
 use nonempty::NonEmpty;
 use spacetimedb_lib::buffer::DecodeError;
 use spacetimedb_lib::identity::AuthCtx;
-use spacetimedb_lib::{bsatn, Address, IndexType, ModuleDef};
+use spacetimedb_lib::{bsatn, Address, IndexType, ModuleDef, VersionTuple};
 use spacetimedb_vm::expr::CrudExpr;
 
 use crate::client::ClientConnectionSender;
@@ -100,8 +100,20 @@ pub(crate) struct WasmModuleHostActor<T: WasmModule> {
     energy_monitor: Arc<dyn EnergyMonitor>,
 }
 
+#[derive(thiserror::Error, Debug)]
+pub enum AbiVersionError {
+    #[error("module doesn't indicate spacetime ABI version")]
+    NoVersion,
+    #[error("abi version is malformed somehow (out-of-bounds, etc)")]
+    Malformed,
+    #[error("abi version {got} is not supported (host implements {implement})")]
+    UnsupportedVersion { got: VersionTuple, implement: VersionTuple },
+}
+
 #[derive(thiserror::Error, Debug)]
 pub enum InitializationError {
+    #[error(transparent)]
+    Abi(#[from] AbiVersionError),
     #[error(transparent)]
     Validation(#[from] ValidationError),
     #[error("setup function returned an error: {0}")]

crates/core/src/host/wasmer/mod.rs

@@ -15,7 +15,7 @@ mod wasmer_module;
 use wasmer_module::WasmerModule;
 
 use super::scheduler::Scheduler;
-use super::wasm_common::{abi, module_host_actor::WasmModuleHostActor, ModuleCreationError};
+use super::wasm_common::{module_host_actor::WasmModuleHostActor, ModuleCreationError};
 use super::{EnergyMonitor, EnergyQuanta};
 
 pub fn make_actor(
@@ -47,15 +47,6 @@ pub fn make_actor(
 
     let module = Module::new(&engine, program_bytes).map_err(|e| ModuleCreationError::WasmCompileError(e.into()))?;
 
-    let abi = abi::determine_spacetime_abi(program_bytes)?;
-
-    if !WasmerModule::IMPLEMENTED_ABI.supports(abi) {
-        return Err(ModuleCreationError::Abi(abi::AbiVersionError::UnsupportedVersion {
-            implement: WasmerModule::IMPLEMENTED_ABI,
-            got: abi,
-        }));
-    }
-
     let module = WasmerModule::new(module, engine);
 
     WasmModuleHostActor::new(dbic, module_hash, module, scheduler, energy_monitor).map_err(Into::into)

crates/core/src/host/wasmer/wasmer_module.rs

@@ -1,13 +1,14 @@
 use super::wasm_instance_env::WasmInstanceEnv;
 use super::Mem;
 use crate::host::instance_env::InstanceEnv;
-use crate::host::wasm_common::module_host_actor::{DescribeError, InitializationError};
+use crate::host::wasm_common::module_host_actor::{AbiVersionError, DescribeError, InitializationError};
 use crate::host::wasm_common::*;
 use crate::host::{EnergyQuanta, Timestamp};
 use bytes::Bytes;
+use spacetimedb_lib::VersionTuple;
 use wasmer::{
     imports, AsStoreMut, Engine, ExternType, Function, FunctionEnv, Imports, Instance, Module, RuntimeError, Store,
-    TypedFunction,
+    TypedFunction, WasmPtr,
 };
 use wasmer_middlewares::metering as wasmer_metering;
 
@@ -45,7 +46,7 @@ impl WasmerModule {
         WasmerModule { module, engine }
     }
 
-    pub const IMPLEMENTED_ABI: abi::VersionTuple = abi::VersionTuple::new(4, 0);
+    pub const IMPLEMENTED_ABI: VersionTuple = VersionTuple::new(4, 0);
 
     fn imports(&self, store: &mut Store, env: &FunctionEnv<WasmInstanceEnv>) -> Imports {
         const _: () = assert!(WasmerModule::IMPLEMENTED_ABI.eq(spacetimedb_lib::MODULE_ABI_VERSION));
@@ -174,6 +175,43 @@ impl module_host_actor::WasmInstancePre for WasmerModule {
             .map_err(|err| InitializationError::Instantiation(err.into()))?;
 
         let mem = Mem::extract(&instance.exports).unwrap();
+
+        // We could (and did in the past) parse the ABI version manually before the instantiation,
+        // but it gets complicated in presence of wasm-opt optimisations which might split encoded
+        // versions like `[...other data...]\00\00\03\00[...other data...]` by zeroes
+        // into several segments, so there is no single data segment containing the entire version.
+        // Instead, it's more reliable to extract the version from an instantiated module
+        // when all the data segments are loaded into the flat memory at correct offsets.
+        let abi_version = instance
+            .exports
+            .get_global(STDB_ABI_SYM)
+            .map_err(|_| AbiVersionError::NoVersion)?;
+
+        let mut abi_version = match abi_version.get(&mut store) {
+            wasmer::Value::I32(x) => x as u32,
+            _ => return Err(AbiVersionError::Malformed.into()),
+        };
+
+        let abi_is_addr = instance.exports.get_global(STDB_ABI_IS_ADDR_SYM).is_ok();
+        if abi_is_addr {
+            abi_version = u32::from_le_bytes(
+                mem.read_bytes(&store, WasmPtr::new(abi_version), 4)
+                    .ok()
+                    .and_then(|bytes| bytes.try_into().ok())
+                    .ok_or(AbiVersionError::Malformed)?,
+            );
+        }
+
+        let abi_version = VersionTuple::from_u32(abi_version);
+
+        if !WasmerModule::IMPLEMENTED_ABI.supports(abi_version) {
+            return Err(AbiVersionError::UnsupportedVersion {
+                implement: WasmerModule::IMPLEMENTED_ABI,
+                got: abi_version,
+            }
+            .into());
+        }
+
         env.as_mut(&mut store).mem = Some(mem);
 
         // Note: this budget is just for initializers