From d524c080c51bc2a83e92b1b67850273d3fd9f602 Mon Sep 17 00:00:00 2001
From: mqxf <mqxf@mvteam.dev>
Date: Fri, 1 Dec 2023 17:36:26 +0100
Subject: [PATCH] Check function signature

---
 core/src/code.rs | 45 ++++++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 44 insertions(+), 1 deletion(-)

diff --git a/core/src/code.rs b/core/src/code.rs
index 1561fb8c321..62fdc0258c8 100644
--- a/core/src/code.rs
+++ b/core/src/code.rs
@@ -40,6 +40,7 @@ use scale_info::{
     scale::{Decode, Encode},
     TypeInfo,
 };
+use gear_wasm_instrument::parity_wasm::elements::Type;
 
 /// Defines maximal permitted count of memory pages.
 pub const MAX_WASM_PAGE_COUNT: u16 = 512;
@@ -54,13 +55,18 @@ fn get_exports(
 ) -> Result<BTreeSet<DispatchKind>, CodeError> {
     let mut exports = BTreeSet::<DispatchKind>::new();
 
+    let mut raw_exports = Vec::new();
+
     for entry in module
         .export_section()
         .ok_or(CodeError::ExportSectionNotFound)?
         .entries()
         .iter()
     {
-        if let Internal::Function(_) = entry.internal() {
+        if let Internal::Function(i) = entry.internal() {
+            if reject_unnecessary {
+                raw_exports.push(*i);
+            }
             if let Some(kind) = DispatchKind::try_from_entry(entry.field()) {
                 exports.insert(kind);
             } else if !STATE_EXPORTS.contains(&entry.field()) && reject_unnecessary {
@@ -69,6 +75,19 @@ fn get_exports(
         }
     }
 
+    if reject_unnecessary {
+        let types = module.type_section()
+            .ok_or(CodeError::TypeSectionNotFound)?
+            .types();
+
+        for i in raw_exports {
+            let Type::Function(ref f) = types[i as usize];
+            if !f.params().is_empty() || !f.results().is_empty() {
+                return Err(CodeError::InvalidExportFnSignature);
+            }
+        }
+    }
+
     Ok(exports)
 }
 
@@ -238,6 +257,12 @@ pub enum CodeError {
     /// Gear protocol restriction for now.
     #[display(fmt = "Program cannot have mutable globals in export section")]
     MutGlobalExport,
+    /// The type section of the wasm module is not present.
+    #[display(fmt = "Type section not found")]
+    TypeSectionNotFound,
+    /// The signature of an exported function is invalid.
+    #[display(fmt = "Invalid function signature for exported function")]
+    InvalidExportFnSignature,
 }
 
 /// Contains instrumented binary code of a program and initial memory size from memory import.
@@ -795,4 +820,22 @@ mod tests {
         )
         .unwrap();
     }
+
+    #[test]
+    fn test_invalid_signature() {
+        const WAT: &str = r#"
+        (module
+            (import "env" "memory" (memory 1))
+            (export "handle" (func $handle))
+            (func $handle (param i32))
+        )
+        "#;
+
+        let original_code = wat2wasm(WAT);
+
+        assert_eq!(
+            Code::try_new(original_code, 1, |_| ConstantCostRules::default(), None),
+            Err(CodeError::InvalidExportFnSignature)
+        );
+    }
 }