From 74c8709d02584c95346fe200b226eaa1dd6d38d9 Mon Sep 17 00:00:00 2001
From: mqxf <mqxf@mvteam.dev>
Date: Wed, 27 Dec 2023 14:21:06 +0100
Subject: [PATCH] feat(core): Validate fns signature on code creation (#3559)

---
 core/src/code.rs          | 38 ++++++++++++++++++++++++++++++++++++--
 pallets/gear/src/tests.rs | 32 +++++++++++---------------------
 2 files changed, 47 insertions(+), 23 deletions(-)

diff --git a/core/src/code.rs b/core/src/code.rs
index 8fdb23ffcf8..e9eb03b0468 100644
--- a/core/src/code.rs
+++ b/core/src/code.rs
@@ -27,7 +27,9 @@ use alloc::{collections::BTreeSet, vec, vec::Vec};
 use gear_wasm_instrument::{
     parity_wasm::{
         self,
-        elements::{ExportEntry, GlobalEntry, GlobalType, InitExpr, Instruction, Internal, Module},
+        elements::{
+            ExportEntry, GlobalEntry, GlobalType, InitExpr, Instruction, Internal, Module, Type,
+        },
     },
     wasm_instrument::gas_metering::{ConstantCostRules, Rules},
     InstrumentationBuilder, STACK_END_EXPORT_NAME,
@@ -50,13 +52,36 @@ fn get_exports(
 ) -> Result<BTreeSet<DispatchKind>, CodeError> {
     let mut exports = BTreeSet::<DispatchKind>::new();
 
+    let funcs = module
+        .function_section()
+        .ok_or(CodeError::FunctionSectionNotFound)?
+        .entries();
+
+    let types = module
+        .type_section()
+        .ok_or(CodeError::TypeSectionNotFound)?
+        .types();
+
+    let import_count = module
+        .import_section()
+        .ok_or(CodeError::ImportSectionNotFound)?
+        .functions();
+
     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 {
+                // 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 let Some(kind) = DispatchKind::try_from_entry(entry.field()) {
                 exports.insert(kind);
             } else if !STATE_EXPORTS.contains(&entry.field()) && reject_unnecessary {
@@ -227,6 +252,15 @@ 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 function section of the wasm module is not present.
+    #[display(fmt = "Function section not found")]
+    FunctionSectionNotFound,
+    /// 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.
diff --git a/pallets/gear/src/tests.rs b/pallets/gear/src/tests.rs
index dda9bc2761f..af9b8b8c18c 100644
--- a/pallets/gear/src/tests.rs
+++ b/pallets/gear/src/tests.rs
@@ -64,7 +64,7 @@ use frame_support::{
 };
 use frame_system::pallet_prelude::BlockNumberFor;
 use gear_core::{
-    code::{self, Code},
+    code::{self, Code, CodeError},
     ids::{CodeId, MessageId, ProgramId},
     message::UserStoredMessage,
     pages::{PageNumber, PageU32Size, WasmPage},
@@ -73,7 +73,7 @@ use gear_core_backend::error::{
     TrapExplanation, UnrecoverableExecutionError, UnrecoverableExtError, UnrecoverableWaitError,
 };
 use gear_core_errors::*;
-use gear_wasm_instrument::STACK_END_EXPORT_NAME;
+use gear_wasm_instrument::{gas_metering::ConstantCostRules, STACK_END_EXPORT_NAME};
 use gstd::{collections::BTreeMap, errors::Error as GstdError};
 use pallet_gear_voucher::PrepaidCall;
 use sp_runtime::{traits::UniqueSaturatedInto, SaturatedConversion};
@@ -13729,7 +13729,6 @@ fn module_instantiation_error() {
 }
 
 #[test]
-#[ignore = "issue #3100"]
 fn wrong_entry_type() {
     let wat = r#"
     (module
@@ -13741,24 +13740,15 @@ fn wrong_entry_type() {
 
     init_logger();
     new_test_ext().execute_with(|| {
-        let pid = Gear::upload_program(
-            RuntimeOrigin::signed(USER_1),
-            ProgramCodeKind::Custom(wat).to_bytes(),
-            DEFAULT_SALT.to_vec(),
-            EMPTY_PAYLOAD.to_vec(),
-            50_000_000_000,
-            0,
-            false,
-        )
-        .map(|_| get_last_program_id())
-        .unwrap();
-        let mid = get_last_message_id();
-
-        run_to_next_block(None);
-
-        assert!(Gear::is_terminated(pid));
-        let err = get_last_event_error(mid);
-        assert!(err.starts_with(&ActorExecutionErrorReplyReason::Environment.to_string()));
+        assert_err!(
+            Code::try_new(
+                ProgramCodeKind::Custom(wat).to_bytes(),
+                1,
+                |_| ConstantCostRules::default(),
+                None
+            ),
+            CodeError::InvalidExportFnSignature
+        );
     });
 }