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

core/src/code.rs

@@ -24,33 +24,50 @@ use crate::{
     pages::{PageNumber, PageU32Size, WasmPage},
 };
 use alloc::{collections::BTreeSet, vec, vec::Vec};
-use gear_wasm_instrument::{
-    parity_wasm::{
-        self,
-        elements::{
-            ExportEntry, GlobalEntry, GlobalType, InitExpr, Instruction, Internal, Module, Type,
-        },
+use gear_wasm_instrument::{parity_wasm::{
+    self,
+    elements::{
+        ExportEntry, GlobalEntry, GlobalType, InitExpr, Instruction, Internal, Module, Type,
     },
-    wasm_instrument::gas_metering::{ConstantCostRules, Rules},
-    InstrumentationBuilder, STACK_END_EXPORT_NAME,
-};
+}, wasm_instrument::gas_metering::{ConstantCostRules, Rules}, InstrumentationBuilder, STACK_END_EXPORT_NAME, SyscallName};
 use scale_info::{
     scale::{Decode, Encode},
     TypeInfo,
 };
+use gear_wasm_instrument::parity_wasm::elements::{External, ImportCountType, ValueType};
+use crate::costs::RuntimeCosts::PayProgramRent;
 
 /// 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"];
+
+/// Name of exports required on chain (only 1 of these is required).
+pub const REQUIRED_EXPORTS: [&str; 2] = ["init", "handle"];
 
-/// Parse function exports from wasm module into [`DispatchKind`].
 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,
-    reject_unnecessary: bool,
-) -> Result<BTreeSet<DispatchKind>, CodeError> {
-    let mut exports = BTreeSet::<DispatchKind>::new();
+    config: &TryNewCodeConfig,
+) -> Result<(), CodeError> {
+    //return Ok(());
 
     let funcs = module
         .function_section()
@@ -62,35 +79,65 @@ 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()) && config.check_exports {
+                    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::all().collect::<Vec<_>>();
+        for import in imports {
+            if let External::Function(i) = import.external() {
+                let Type::Function(types) = &types[*i as usize];
+                // We can likely improve this by adding some helper function in SyscallName
+                let syscall = syscalls.iter().find(|s| s.to_str() == import.field()).ok_or(CodeError::UnknownImport)?;
+                let signature = syscall.signature();
+
+                let params = signature.params().iter().copied().map(Into::<ValueType>::into).collect::<Vec<_>>();
+                if &params != types.params() {
+                    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 +308,12 @@ 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,
 }
 
 /// Contains instrumented binary code of a program and initial memory size from memory import.
@@ -330,6 +383,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 +402,7 @@ impl Default for TryNewCodeConfig {
             stack_height: None,
             export_stack_height: false,
             check_exports: true,
+            check_imports: true,
             check_and_canonize_stack_end: true,
             check_mut_global_exports: true,
             check_start_section: true,
@@ -417,12 +473,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/wasm_project.rs

@@ -16,11 +16,7 @@
 // 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 crate::{
-    crate_info::CrateInfo,
-    optimize::{self, OptType, Optimizer},
-    smart_fs,
-};
+use crate::{crate_info::CrateInfo, optimize::{self, OptType, Optimizer}, smart_fs};
 use anyhow::{Context, Result};
 use chrono::offset::Local as ChronoLocal;
 use gmeta::MetadataRepr;