feat(wasm-builder): Check all gear requirements for code at compile…

… time (#3649)

Co-authored-by: StackOverflowExcept1on <[email protected]>

core/src/code.rs

@@ -28,11 +28,12 @@ use gear_wasm_instrument::{
     parity_wasm::{
         self,
         elements::{
-            ExportEntry, GlobalEntry, GlobalType, InitExpr, Instruction, Internal, Module, Type,
+            ExportEntry, External, GlobalEntry, GlobalType, ImportCountType, InitExpr, Instruction,
+            Internal, Module, Type, ValueType,
         },
     },
     wasm_instrument::gas_metering::{ConstantCostRules, Rules},
-    InstrumentationBuilder, STACK_END_EXPORT_NAME,
+    InstrumentationBuilder, SyscallName, STACK_END_EXPORT_NAME,
 };
 use scale_info::{
     scale::{Decode, Encode},
@@ -42,16 +43,40 @@ use scale_info::{
 /// Defines maximal permitted count of memory pages.
 pub const MAX_WASM_PAGE_COUNT: u16 = 512;
 
-/// Name of exports allowed on chain except execution kinds.
-pub const STATE_EXPORTS: [&str; 2] = ["state", "metahash"];
+/// Name of exports allowed on chain.
+pub const ALLOWED_EXPORTS: [&str; 6] = [
+    "init",
+    "handle",
+    "handle_reply",
+    "handle_signal",
+    "state",
+    "metahash",
+];
 
-/// Parse function exports from wasm module into [`DispatchKind`].
-fn get_exports(
-    module: &Module,
-    reject_unnecessary: bool,
-) -> Result<BTreeSet<DispatchKind>, CodeError> {
-    let mut exports = BTreeSet::<DispatchKind>::new();
+/// Name of exports required on chain (only 1 of these is required).
+pub const REQUIRED_EXPORTS: [&str; 2] = ["init", "handle"];
+
+fn get_exports(module: &Module) -> BTreeSet<DispatchKind> {
+    let mut entries = BTreeSet::new();
+
+    for entry in module
+        .export_section()
+        .expect("Exports section has been checked for already")
+        .entries()
+        .iter()
+    {
+        if let Internal::Function(_) = entry.internal() {
+            if let Some(entry) = DispatchKind::try_from_entry(entry.field()) {
+                entries.insert(entry);
+            }
+        }
+    }
 
+    entries
+}
+
+/// Parse function exports from wasm module into [`DispatchKind`].
+fn check_code(module: &Module, config: &TryNewCodeConfig) -> Result<(), CodeError> {
     let funcs = module
         .function_section()
         .ok_or(CodeError::FunctionSectionNotFound)?
@@ -62,35 +87,76 @@ fn get_exports(
         .ok_or(CodeError::TypeSectionNotFound)?
         .types();
 
-    let import_count = module
+    let import_count = module.import_count(ImportCountType::Function);
+
+    let imports = module
         .import_section()
         .ok_or(CodeError::ImportSectionNotFound)?
-        .functions();
+        .entries();
 
-    for entry in module
+    let exports = module
         .export_section()
         .ok_or(CodeError::ExportSectionNotFound)?
-        .entries()
-        .iter()
-    {
-        if let Internal::Function(i) = entry.internal() {
-            if reject_unnecessary {
+        .entries();
+
+    if config.check_exports {
+        let mut entry = false;
+        for export in exports {
+            if let Internal::Function(i) = export.internal() {
                 // Index access into arrays cannot panic unless the Module structure is invalid
                 let type_id = funcs[*i as usize - import_count].type_ref();
                 let Type::Function(ref f) = types[type_id as usize];
                 if !f.params().is_empty() || !f.results().is_empty() {
                     return Err(CodeError::InvalidExportFnSignature);
                 }
+                if !ALLOWED_EXPORTS.contains(&export.field()) {
+                    return Err(CodeError::NonGearExportFnFound);
+                }
+                if REQUIRED_EXPORTS.contains(&export.field()) {
+                    entry = true;
+                }
             }
-            if let Some(kind) = DispatchKind::try_from_entry(entry.field()) {
-                exports.insert(kind);
-            } else if !STATE_EXPORTS.contains(&entry.field()) && reject_unnecessary {
-                return Err(CodeError::NonGearExportFnFound);
+        }
+
+        if !entry {
+            return Err(CodeError::RequiredExportFnNotFound);
+        }
+    }
+
+    if config.check_imports {
+        let syscalls = SyscallName::instrumentable_map();
+        let mut seen = BTreeSet::new();
+        for import in imports {
+            if let External::Function(i) = import.external() {
+                let Type::Function(types) = &types[*i as usize];
+                let syscall = syscalls
+                    .get(import.field())
+                    .ok_or(CodeError::UnknownImport)?;
+
+                if !seen.insert(*syscall) {
+                    return Err(CodeError::DuplicateImport);
+                }
+
+                let signature = syscall.signature();
+
+                let params = signature
+                    .params()
+                    .iter()
+                    .copied()
+                    .map(Into::<ValueType>::into);
+                if !params.eq(types.params().iter().copied()) {
+                    return Err(CodeError::InvalidImportFnSignature);
+                }
+
+                let results = signature.results().unwrap_or(&[]);
+                if results != types.results() {
+                    return Err(CodeError::InvalidImportFnSignature);
+                }
             }
         }
     }
 
-    Ok(exports)
+    Ok(())
 }
 
 fn get_export_entry<'a>(module: &'a Module, name: &str) -> Option<&'a ExportEntry> {
@@ -261,6 +327,15 @@ pub enum CodeError {
     /// The signature of an exported function is invalid.
     #[display(fmt = "Invalid function signature for exported function")]
     InvalidExportFnSignature,
+    /// An imported function was not recognized.
+    #[display(fmt = "Unknown function name in import section")]
+    UnknownImport,
+    /// The signature of an imported function is invalid.
+    #[display(fmt = "Invalid function signature for imported function")]
+    InvalidImportFnSignature,
+    /// An imported was declared multiple times.
+    #[display(fmt = "An import was declared multiple times")]
+    DuplicateImport,
 }
 
 /// Contains instrumented binary code of a program and initial memory size from memory import.
@@ -330,6 +405,8 @@ pub struct TryNewCodeConfig {
     pub export_stack_height: bool,
     /// Check exports (wasm contains init or handle exports)
     pub check_exports: bool,
+    /// Check imports (check that all imports are valid syscalls with correct signature)
+    pub check_imports: bool,
     /// Check and canonize stack end
     pub check_and_canonize_stack_end: bool,
     /// Check mutable global exports
@@ -347,6 +424,7 @@ impl Default for TryNewCodeConfig {
             stack_height: None,
             export_stack_height: false,
             check_exports: true,
+            check_imports: false,
             check_and_canonize_stack_end: true,
             check_mut_global_exports: true,
             check_start_section: true,
@@ -417,12 +495,9 @@ impl Code {
             return Err(CodeError::InvalidStaticPageCount);
         }
 
-        let exports = get_exports(&module, config.check_exports)?;
-        if config.check_exports
-            && !(exports.contains(&DispatchKind::Init) || exports.contains(&DispatchKind::Handle))
-        {
-            return Err(CodeError::RequiredExportFnNotFound);
-        }
+        check_code(&module, &config)?;
+
+        let exports = get_exports(&module);
 
         let mut instrumentation_builder = InstrumentationBuilder::new("env");
 

utils/wasm-builder/src/optimize.rs

@@ -131,11 +131,14 @@ impl Optimizer {
             .map_err(BuilderError::CodeCheckFailed)?,
             // validate wasm code
             // see `pallet_gear::pallet::Pallet::upload_program(...)`
-            OptType::Opt => Code::try_new(
+            OptType::Opt => Code::try_new_mock_with_rules(
                 original_code,
-                1,
                 |_| CustomConstantCostRules::default(),
-                None,
+                TryNewCodeConfig {
+                    version: 1,
+                    check_imports: true,
+                    ..Default::default()
+                },
             )
             .map(|_| ())
             .map_err(BuilderError::CodeCheckFailed)?,

utils/wasm-instrument/src/syscalls.rs

@@ -19,7 +19,12 @@
 //! Gear syscalls for programs execution signatures.
 
 use crate::parity_wasm::elements::{FunctionType, ValueType};
-use alloc::{borrow::ToOwned, collections::BTreeSet, vec::Vec};
+use alloc::{
+    borrow::ToOwned,
+    collections::{BTreeMap, BTreeSet},
+    string::{String, ToString},
+    vec::Vec,
+};
 use core::iter;
 use enum_iterator::{self, Sequence};
 pub use pointers::*;
@@ -233,10 +238,20 @@ impl SyscallName {
             Self::ReserveGas,
             Self::UnreserveGas,
             Self::Random,
+            Self::SystemReserveGas,
+            Self::EnvVars,
         ]
         .into()
     }
 
+    /// Returns map of all syscall string values to syscall names.
+    pub fn instrumentable_map() -> BTreeMap<String, SyscallName> {
+        Self::instrumentable()
+            .into_iter()
+            .map(|n| (n.to_str().to_string(), n))
+            .collect()
+    }
+
     /// Returns signature for syscall by name.
     pub fn signature(self) -> SyscallSignature {
         use RegularParamType::*;