Disallow calling random lifecycle reducers

Author: coolreader18

crates/core/src/host/module_host.rs

@@ -1,3 +1,4 @@
+use super::wasm_common::{IDENTITY_CONNECTED_DUNDER, IDENTITY_DISCONNECTED_DUNDER};
 use super::{ArgsTuple, InvalidReducerArguments, ReducerArgs, ReducerCallResult, ReducerId, Timestamp};
 use crate::client::{ClientActorId, ClientConnectionSender};
 use crate::database_logger::LogLevel;
@@ -677,9 +678,9 @@ impl ModuleHost {
                 None,
                 None,
                 if connected {
-                    "__identity_connected__"
+                    IDENTITY_CONNECTED_DUNDER
                 } else {
-                    "__identity_disconnected__"
+                    IDENTITY_DISCONNECTED_DUNDER
                 },
                 ReducerArgs::Nullary,
             )
@@ -735,6 +736,9 @@ impl ModuleHost {
         reducer_name: &str,
         args: ReducerArgs,
     ) -> Result<ReducerCallResult, ReducerCallError> {
+        if reducer_name.starts_with("__") && reducer_name.ends_with("__") {
+            return Err(ReducerCallError::NoSuchReducer);
+        }
         let res = self
             .call_reducer_inner(
                 caller_identity,

crates/core/src/host/wasm_common.rs

@@ -21,6 +21,8 @@ pub const SETUP_DUNDER: &str = "__setup__";
 pub const INIT_DUNDER: &str = "__init__";
 /// the reducer with this name is invoked when updating the database
 pub const UPDATE_DUNDER: &str = "__update__";
+pub const IDENTITY_CONNECTED_DUNDER: &str = "__identity_connected__";
+pub const IDENTITY_DISCONNECTED_DUNDER: &str = "__identity_disconnected__";
 
 #[derive(Debug, Copy, Clone, PartialEq, Eq)]
 #[allow(unused)]

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

@@ -5,7 +5,7 @@ use std::time::Duration;
 
 use spacetimedb_lib::buffer::DecodeError;
 use spacetimedb_lib::identity::AuthCtx;
-use spacetimedb_lib::{bsatn, Address, ModuleDef, TableDesc};
+use spacetimedb_lib::{bsatn, Address, ModuleDef, ModuleValidationError, TableDesc};
 use spacetimedb_vm::expr::CrudExpr;
 
 use super::instrumentation::CallTimes;
@@ -90,6 +90,8 @@ pub(crate) struct WasmModuleHostActor<T: WasmModule> {
 pub enum InitializationError {
     #[error(transparent)]
     Validation(#[from] ValidationError),
+    #[error(transparent)]
+    ModuleValidation(#[from] ModuleValidationError),
     #[error("setup function returned an error: {0}")]
     Setup(Box<str>),
     #[error("wasm trap while calling {func:?}")]
@@ -153,7 +155,8 @@ impl<T: WasmModule> WasmModuleHostActor<T> {
         )?;
 
         let desc = instance.extract_descriptions()?;
-        let desc = bsatn::from_slice(&desc).map_err(DescribeError::Decode)?;
+        let desc: ModuleDef = bsatn::from_slice(&desc).map_err(DescribeError::Decode)?;
+        desc.validate_reducers()?;
         let ModuleDef {
             mut typespace,
             mut tables,
@@ -173,7 +176,10 @@ impl<T: WasmModule> WasmModuleHostActor<T> {
             tables
                 .into_iter()
                 .map(|x| (x.schema.table_name.clone(), EntityDef::Table(x))),
-            reducers.iter().map(|x| (x.name.clone(), EntityDef::Reducer(x.clone()))),
+            reducers
+                .iter()
+                .filter(|r| !(r.name.starts_with("__") && r.name.ends_with("__")))
+                .map(|x| (x.name.clone(), EntityDef::Reducer(x.clone()))),
         )
         .collect();
         let reducers = ReducersMap(reducers.into_iter().map(|x| (x.name.clone(), x)).collect());

crates/lib/src/lib.rs

@@ -187,3 +187,33 @@ pub struct TypeAlias {
     pub name: String,
     pub ty: sats::AlgebraicTypeRef,
 }
+
+impl ModuleDef {
+    pub fn validate_reducers(&self) -> Result<(), ModuleValidationError> {
+        for reducer in &self.reducers {
+            match &*reducer.name {
+                "__init__" => {}
+                "__identity_connected__" | "__identity_disconnected__" | "__update__" | "__migrate__" => {
+                    if !reducer.args.is_empty() {
+                        return Err(ModuleValidationError::InvalidLifecycleReducer {
+                            reducer: reducer.name.clone(),
+                        });
+                    }
+                }
+                name if name.starts_with("__") && name.ends_with("__") => {
+                    return Err(ModuleValidationError::UnknownDunderscore)
+                }
+                _ => {}
+            }
+        }
+        Ok(())
+    }
+}
+
+#[derive(thiserror::Error, Debug)]
+pub enum ModuleValidationError {
+    #[error("lifecycle reducer {reducer:?} has invalid signature")]
+    InvalidLifecycleReducer { reducer: Box<str> },
+    #[error("reducers with double-underscores at the start and end of their names are not allowed")]
+    UnknownDunderscore,
+}