Merge pull request #113 from OffchainLabs/stylus-compile-safety

Compile pricing for testnet

arbitrator/arbutil/src/format.rs

@@ -2,7 +2,10 @@
 // For license information, see https://github.com/OffchainLabs/nitro/blob/master/LICENSE
 
 use crate::color::Color;
-use std::{fmt::Display, time::Duration};
+use std::{
+    fmt::{Debug, Display},
+    time::Duration,
+};
 
 #[must_use]
 pub fn time(span: Duration) -> String {
@@ -37,3 +40,13 @@ where
 pub fn hex_fmt<T: AsRef<[u8]>>(data: T, f: &mut std::fmt::Formatter) -> std::fmt::Result {
     f.write_str(&hex::encode(data))
 }
+
+pub trait DebugBytes {
+    fn debug_bytes(self) -> Vec<u8>;
+}
+
+impl<T: Debug> DebugBytes for T {
+    fn debug_bytes(self) -> Vec<u8> {
+        format!("{:?}", self).as_bytes().to_vec()
+    }
+}

arbitrator/jit/src/gostack.rs

@@ -138,8 +138,16 @@ impl GoStack {
         self.read_u64() as *mut T
     }
 
+    /// TODO: replace `unbox` with a safe id-based API
     pub fn unbox<T>(&mut self) -> T {
-        unsafe { *Box::from_raw(self.read_ptr_mut()) }
+        let ptr: *mut T = self.read_ptr_mut();
+        unsafe { *Box::from_raw(ptr) }
+    }
+
+    /// TODO: replace `unbox_option` with a safe id-based API
+    pub fn unbox_option<T>(&mut self) -> Option<T> {
+        let ptr: *mut T = self.read_ptr_mut();
+        (!ptr.is_null()).then(|| unsafe { *Box::from_raw(ptr) })
     }
 
     pub fn write_u8(&mut self, x: u8) -> &mut Self {
@@ -243,6 +251,10 @@ impl GoStack {
         values
     }
 
+    pub fn read_bool32(&mut self) -> bool {
+        self.read_u32() != 0
+    }
+
     pub fn read_go_ptr(&mut self) -> u32 {
         self.read_u64().try_into().expect("go pointer doesn't fit")
     }

arbitrator/jit/src/machine.rs

@@ -123,6 +123,7 @@ pub fn create(opts: &Opts, env: WasmEnv) -> (Instance, FunctionEnv<WasmEnv>, Sto
             github!("arbos/programs.callUserWasmRustImpl") => func!(user::call_user_wasm),
             github!("arbos/programs.readRustVecLenImpl") => func!(user::read_rust_vec_len),
             github!("arbos/programs.rustVecIntoSliceImpl") => func!(user::rust_vec_into_slice),
+            github!("arbos/programs.rustMachineDropImpl") => func!(user::drop_machine),
             github!("arbos/programs.rustConfigImpl") => func!(user::rust_config_impl),
             github!("arbos/programs.rustEvmDataImpl") => func!(user::evm_data_impl),
 

arbitrator/jit/src/user/mod.rs

@@ -8,53 +8,73 @@ use crate::{
 };
 use arbutil::{
     evm::{user::UserOutcome, EvmData},
+    format::DebugBytes,
     heapify,
 };
 use prover::{
-    binary::WasmBinary,
     programs::{config::PricingParams, prelude::*},
+    Machine,
 };
 use std::mem;
 use stylus::native;
 
 mod evm_api;
 
-/// Compiles and instruments user wasm.
-/// go side: λ(wasm []byte, pageLimit, version u16, debug u32) (machine *Machine, footprint u32, err *Vec<u8>)
+/// Compiles and instruments a user wasm.
+///
+/// # Go side
+///
+/// The Go compiler expects the call to take the form
+///     λ(wasm []byte, pageLimit, version u16, debug u32) (module *Vec<u8>, info WasmInfo, err *Vec<u8>)
+///
+/// These values are placed on the stack as follows
+///     stack:  || wasm... || pageLimit | version | debug || mod ptr || info... || err ptr ||
+///     info:   || footprint | 2 pad | size ||
+///
 pub fn compile_user_wasm(env: WasmEnvMut, sp: u32) {
     let mut sp = GoStack::simple(sp, &env);
     let wasm = sp.read_go_slice_owned();
     let page_limit = sp.read_u16();
-    let compile = CompileConfig::version(sp.read_u16(), sp.read_u32() != 0);
+    let version = sp.read_u16();
+    let debug = sp.read_bool32();
+    let compile = CompileConfig::version(version, debug);
 
     macro_rules! error {
         ($error:expr) => {{
-            let error = $error.wrap_err("failed to compile");
-            let error = format!("{:?}", error).as_bytes().to_vec();
+            let error = $error.wrap_err("failed to compile").debug_bytes();
             sp.write_nullptr();
-            sp.skip_space(); // skip footprint
+            sp.skip_space(); // skip info
             sp.write_ptr(heapify(error));
             return;
         }};
     }
 
-    // ensure the wasm compiles during proving
-    let footprint = match WasmBinary::parse_user(&wasm, page_limit, &compile) {
-        Ok((.., pages)) => pages,
+    let (footprint, size) = match Machine::new_user_stub(&wasm, page_limit, version, debug) {
+        Ok((_, info)) => (info.footprint, info.size),
         Err(error) => error!(error),
     };
     let module = match native::module(&wasm, compile) {
         Ok(module) => module,
         Err(error) => error!(error),
     };
     sp.write_ptr(heapify(module));
-    sp.write_u16(footprint).skip_space();
+    sp.write_u16(footprint).skip_u16().write_u32(size); // wasm info
     sp.write_nullptr();
 }
 
 /// Links and executes a user wasm.
-/// λ(mach *Machine, calldata []byte, params *Configs, evmApi []byte, evmData: *EvmData, gas *u64, root *[32]byte)
-///     -> (status byte, out *Vec<u8>)
+///
+/// # Go side
+///
+/// The Go compiler expects the call to take the form
+///     λ(
+///           mach *Machine, calldata []byte, params *Configs, evmApi []byte, evmData: *EvmData,
+///           gas *u64, root *[32]byte
+///     ) -> (status byte, out *Vec<u8>)
+///
+/// These values are placed on the stack as follows
+///     || mach || calldata... || params || evmApi... || evmData || gas || root || status | 3 pad | out ptr ||
+///
 pub fn call_user_wasm(env: WasmEnvMut, sp: u32) -> MaybeEscape {
     let sp = &mut GoStack::simple(sp, &env);
     use UserOutcome::*;
@@ -91,31 +111,66 @@ pub fn call_user_wasm(env: WasmEnvMut, sp: u32) -> MaybeEscape {
 }
 
 /// Reads the length of a rust `Vec`
-/// go side: λ(vec *Vec<u8>) (len u32)
+///
+/// # Go side
+///
+/// The Go compiler expects the call to take the form
+///     λ(vec *Vec<u8>) (len u32)
+///
+/// These values are placed on the stack as follows
+///     || vec ptr || len u32 | pad 4 ||
+///
 pub fn read_rust_vec_len(env: WasmEnvMut, sp: u32) {
     let mut sp = GoStack::simple(sp, &env);
     let vec: &Vec<u8> = unsafe { &*sp.read_ptr() };
     sp.write_u32(vec.len() as u32);
 }
 
 /// Copies the contents of a rust `Vec` into a go slice, dropping it in the process
-/// go side: λ(vec *Vec<u8>, dest []byte)
+///
+/// # Go Side
+///
+/// The Go compiler expects the call to take the form
+///     λ(vec *Vec<u8>, dest []byte)
+///
+/// These values are placed on the stack as follows
+///     || vec ptr || dest... ||
+///
 pub fn rust_vec_into_slice(env: WasmEnvMut, sp: u32) {
     let mut sp = GoStack::simple(sp, &env);
-    let vec: Vec<u8> = unsafe { *Box::from_raw(sp.read_ptr_mut()) };
+    let vec: Vec<u8> = sp.unbox();
     let ptr: *mut u8 = sp.read_ptr_mut();
     sp.write_slice(ptr as u64, &vec);
     mem::drop(vec)
 }
 
+/// Drops module bytes. Note that in user-host this would be a `Machine`.
+///
+/// # Go side
+///
+/// The Go compiler expects the call to take the form
+///     λ(module *Vec<u8>)
+///
+pub fn drop_machine(env: WasmEnvMut, sp: u32) {
+    let mut sp = GoStack::simple(sp, &env);
+    if let Some(module) = sp.unbox_option::<Vec<u8>>() {
+        mem::drop(module);
+    }
+}
+
 /// Creates a `StylusConfig` from its component parts.
-/// go side: λ(version u16, maxDepth, inkPrice u32, debugMode: u32) *(CompileConfig, StylusConfig)
+///
+/// # Go side
+///
+/// The Go compiler expects the call to take the form
+///     λ(version u16, maxDepth, inkPrice u32, debugMode: u32) *(CompileConfig, StylusConfig)
+///
+/// The values are placed on the stack as follows
+///     || version | 2 garbage bytes | max_depth || ink_price | debugMode || result ptr ||
+///
 pub fn rust_config_impl(env: WasmEnvMut, sp: u32) {
     let mut sp = GoStack::simple(sp, &env);
 
-    // The Go compiler places these on the stack as follows
-    // | version | 2 garbage bytes | max_depth | ink_price | debugMode | result ptr |
-
     let config = StylusConfig {
         version: sp.read_u16(),
         max_depth: sp.skip_u16().read_u32(),
@@ -128,11 +183,20 @@ pub fn rust_config_impl(env: WasmEnvMut, sp: u32) {
 }
 
 /// Creates an `EvmData` from its component parts.
-/// go side: λ(
-///     blockBasefee *[32]byte, chainid u64, blockCoinbase *[20]byte, blockGasLimit,
-///     blockNumber, blockTimestamp u64, contractAddress, msgSender *[20]byte,
-///     msgValue, txGasPrice *[32]byte, txOrigin *[20]byte, reentrant u32,
-///) *EvmData
+///
+/// # Go side
+///
+/// The Go compiler expects the call to take the form
+///     λ(
+///         blockBasefee *[32]byte, chainid u64, blockCoinbase *[20]byte, blockGasLimit,
+///         blockNumber, blockTimestamp u64, contractAddress, msgSender *[20]byte,
+///         msgValue, txGasPrice *[32]byte, txOrigin *[20]byte, reentrant u32,
+///     ) -> *EvmData
+///
+/// These values are placed on the stack as follows
+///     || baseFee || chainid || coinbase || gas limit || block number || timestamp || address ||
+///     || sender || value || gas price || origin || reentrant | 4 pad || data ptr ||
+///
 pub fn evm_data_impl(env: WasmEnvMut, sp: u32) {
     let mut sp = GoStack::simple(sp, &env);
     let evm_data = EvmData {

arbitrator/prover/Cargo.toml

@@ -40,6 +40,7 @@ name = "prover"
 crate-type = ["staticlib", "lib"]
 
 [features]
-default = ["native", "singlepass_rayon"]
-native = ["dep:wasmer", "dep:wasmer-compiler-singlepass", "dep:rayon", "dep:brotli2"]
+default = ["native", "rayon", "singlepass_rayon"]
+native = ["dep:wasmer", "dep:wasmer-compiler-singlepass", "dep:brotli2"]
 singlepass_rayon = ["wasmer-compiler-singlepass?/rayon"]
+rayon = ["dep:rayon"]

arbitrator/prover/src/binary.rs

@@ -603,12 +603,54 @@ impl<'a> WasmBinary<'a> {
     ) -> Result<(WasmBinary<'a>, StylusData, u16)> {
         let mut bin = parse(wasm, Path::new("user"))?;
         let stylus_data = bin.instrument(compile)?;
-
         let pages = bin.memories.first().map(|m| m.initial).unwrap_or_default();
-        if pages > page_limit as u64 {
+        let debug = compile.debug.debug_funcs;
+
+        // ensure the wasm fits within the remaining amount of memory
+        if pages > page_limit.into() {
             let limit = page_limit.red();
             bail!("memory exceeds limit: {} > {limit}", pages.red());
         }
+
+        // not strictly necessary, but anti-DoS limits and extra checks in case of bugs
+        macro_rules! limit {
+            ($limit:expr, $count:expr, $name:expr) => {
+                if $count > $limit {
+                    bail!("too many wasm {}", $name);
+                }
+            };
+        }
+        limit!(1, bin.memories.len(), "memories");
+        limit!(100, bin.datas.len(), "datas");
+        limit!(100, bin.elements.len(), "elements");
+        limit!(1_000, bin.exports.len(), "exports");
+        limit!(1_000, bin.tables.len(), "tables");
+        limit!(10_000, bin.codes.len(), "functions");
+        limit!(10_000, bin.globals.len(), "globals");
+
+        let max_len = 500;
+        macro_rules! too_long {
+            ($name:expr, $len:expr) => {
+                bail!(
+                    "wasm {} too long: {} > {}",
+                    $name.red(),
+                    $len.red(),
+                    max_len.red()
+                )
+            };
+        }
+        if let Some((name, _)) = bin.exports.iter().find(|(name, _)| name.len() > max_len) {
+            too_long!("name", name.len())
+        }
+        if bin.names.module.len() > max_len {
+            too_long!("module name", bin.names.module.len())
+        }
+        if !debug && !bin.names.functions.is_empty() {
+            bail!("wasm custom names section not allowed")
+        }
+        if bin.start.is_some() {
+            bail!("wasm start functions not allowed");
+        }
         Ok((bin, stylus_data, pages as u16))
     }
 }