fix(RUN-907): Allow export of imports

rs/embedders/src/wasm_utils/validation.rs

@@ -833,11 +833,14 @@ fn validate_export_section(
     max_sum_exported_function_name_lengths: usize,
 ) -> Result<(), WasmValidationError> {
     if !module.exports.is_empty() {
-        let num_imported_functions = module
+        let imported_function_types: Vec<_> = module
             .imports
             .iter()
-            .filter(|i| matches!(i.ty, TypeRef::Func(_)))
-            .count();
+            .filter_map(|i| match i.ty {
+                TypeRef::Func(type_index) => Some(&module.types[type_index as usize]),
+                _ => None,
+            })
+            .collect();
 
         let mut seen_funcs: HashSet<&str> = HashSet::new();
         let valid_exported_functions = get_valid_exported_functions();
@@ -897,23 +900,18 @@ fn validate_export_section(
                     // module, so we need to subtract the number of imported functions to get the
                     // correct index from the general function space.
                     let fn_index = fn_index as usize;
-                    let import_count = num_imported_functions;
-                    if fn_index < import_count {
-                        return Err(WasmValidationError::InvalidFunctionIndex {
-                            index: fn_index,
-                            import_count,
-                        });
-                    }
-                    let actual_fn_index = fn_index - import_count;
-                    let type_index = module.functions[actual_fn_index] as usize;
-                    let func_ty = if let CompositeType::Func(func_ty) =
-                        &module.types[type_index].composite_type
-                    {
-                        func_ty
+                    let import_count = imported_function_types.len();
+                    let composite_type = if fn_index < import_count {
+                        &imported_function_types[fn_index].composite_type
                     } else {
+                        let actual_fn_index = fn_index - import_count;
+                        let type_index = module.functions[actual_fn_index] as usize;
+                        &module.types[type_index].composite_type
+                    };
+                    let CompositeType::Func(func_ty) = composite_type else {
                         return Err(WasmValidationError::InvalidExportSection(format!(
                             "Function export doesn't have a function type. Type found: {:?}",
-                            &module.types[type_index]
+                            composite_type
                         )));
                     };
                     validate_function_signature(valid_signature, export.name, func_ty)?;

rs/embedders/tests/misc_tests.rs

@@ -5,9 +5,16 @@ use ic_embedders::{
     wasm_utils::{decoding::decode_wasm, validate_and_instrument_for_testing},
     WasmtimeEmbedder,
 };
+use ic_interfaces::execution_environment::HypervisorError;
 use ic_logger::replica_logger::no_op_logger;
+use ic_test_utilities::wasmtime_instance::WasmtimeInstanceBuilder;
+use ic_test_utilities_time::mock_time;
+use ic_types::{
+    methods::{FuncRef, WasmMethod},
+    Cycles, PrincipalId,
+};
 use ic_wasm_transform::Module;
-use ic_wasm_types::BinaryEncodedWasm;
+use ic_wasm_types::{BinaryEncodedWasm, WasmValidationError};
 use std::sync::Arc;
 use wasmparser::ExternalKind;
 
@@ -152,3 +159,91 @@ fn test_decode_large_compressed_module_with_tweaked_size() {
     contents[n - 4..n].copy_from_slice(&100u32.to_le_bytes());
     decode_wasm(Arc::new(contents)).unwrap();
 }
+
+fn run_go_export(wat: &str) -> Result<(), HypervisorError> {
+    const LARGE_INSTRUCTION_LIMIT: u64 = 1_000_000_000_000;
+
+    let mut instance = WasmtimeInstanceBuilder::new()
+        .with_wat(wat)
+        .with_api_type(ic_system_api::ApiType::update(
+            mock_time(),
+            vec![],
+            Cycles::from(0_u128),
+            PrincipalId::new_user_test_id(0),
+            0.into(),
+        ))
+        .with_num_instructions(LARGE_INSTRUCTION_LIMIT.into())
+        .try_build()
+        .map_err(|(err, _)| err)?;
+
+    instance
+        .run(FuncRef::Method(WasmMethod::Update("go".to_string())))
+        .unwrap();
+    Ok(())
+}
+
+/// Test that we can handle a module that exports one of its imports.
+#[test]
+fn direct_export_of_import() {
+    run_go_export(
+        r#"
+		(module
+			(func $reply (export "canister_update go") (import "ic0" "msg_reply"))
+		)
+	"#,
+    )
+    .unwrap();
+}
+
+/// Test that we can handle a module that exports one of its imports when there
+/// are other imports.
+#[test]
+fn direct_export_of_one_import_from_many() {
+    run_go_export(
+        r#"
+		(module
+            (import "ic0" "call_cycles_add" (func $ic0_call_cycles_add (param $amount i64)))
+            (import "ic0" "canister_cycle_balance" (func $ic0_canister_cycle_balance (result i64)))
+			(func $reply (export "canister_update go") (import "ic0" "msg_reply"))
+            (import "ic0" "msg_cycles_accept" (func $ic0_msg_cycles_accept (param $amount i64) (result i64)))
+		)
+	"#,
+    ).unwrap();
+}
+
+#[test]
+fn direct_export_of_import_fails_with_wrong_type() {
+    let err = run_go_export(
+        r#"
+		(module
+            (import "ic0" "call_cycles_add" (func $ic0_call_cycles_add (param $amount i64)))
+            (import "ic0" "msg_reply" (func $reply))
+			(func (export "canister_update go") (import "ic0" "canister_cycle_balance") (result i64))
+            (import "ic0" "msg_cycles_accept" (func $ic0_msg_cycles_accept (param $amount i64) (result i64)))
+            (func $other (call $reply))
+		)
+	"#,
+    ).unwrap_err();
+    assert_eq!(
+        err,
+        HypervisorError::InvalidWasm(WasmValidationError::InvalidFunctionSignature(
+            "Expected return type [] for 'canister_update go', got [I64].".to_string()
+        ))
+    );
+}
+
+/// A module should be allowed to export a direct import which doesn't have type
+/// () -> () as long as it doesn't have a `canister_` prefix in the export name.
+#[test]
+fn direct_export_of_import_without_unit_type() {
+    run_go_export(
+        r#"
+		(module
+            (import "ic0" "call_cycles_add" (func $ic0_call_cycles_add (param $amount i64)))
+			(func (export "foo") (import "ic0" "canister_cycle_balance") (result i64))
+            (import "ic0" "msg_cycles_accept" (func $ic0_msg_cycles_accept (param $amount i64) (result i64)))
+            (func (export "canister_update go"))
+		)
+	"#,
+    ).unwrap();
+}

rs/embedders/tests/validation.rs

@@ -904,18 +904,24 @@ fn can_validate_module_cycles_related_imports() {
 }
 
 #[test]
-fn can_validate_valid_export_section_with_invalid_function_index() {
-    let wasm = BinaryEncodedWasm::new(
-        include_bytes!("instrumentation-test-data/export_section_invalid_function_index.wasm")
-            .to_vec(),
-    );
-    assert_matches!(
-        validate_wasm_binary(&wasm, &EmbeddersConfig::default()),
-        Err(WasmValidationError::InvalidFunctionIndex {
-            index: 0,
-            import_count: 1
-        })
-    );
+fn can_validate_export_section_exporting_import() {
+    let wasm = wat2wasm(
+        r#"
+        (module
+            (type (;0;) (func))
+            (import "env" "memory" (memory (;0;) 529))
+            (import "env" "table" (table (;0;) 33 33 funcref))
+            (import "ic0" "msg_reply" (func (;0;) (type 0)))
+            (func (;1;) (type 0))
+            (func (;2;) (type 0))
+            (func (;3;) (type 0))
+            (export "canister_heartbeat" (func 1))
+            (export "canister_pre_upgrade" (func 2))
+            (export "canister_post_upgrade" (func 0)))
+        "#,
+    )
+    .unwrap();
+    validate_wasm_binary(&wasm, &EmbeddersConfig::default()).unwrap();
 }
 
 #[test]

rs/types/wasm_types/src/errors.rs

@@ -44,8 +44,6 @@ pub enum WasmValidationError {
     TooManyFunctions { defined: usize, allowed: usize },
     /// Module contains too many custom sections.
     TooManyCustomSections { defined: usize, allowed: usize },
-    /// Module defines an invalid index for a local function.
-    InvalidFunctionIndex { index: usize, import_count: usize },
     /// A function was too complex.
     FunctionComplexityTooHigh {
         index: usize,
@@ -107,14 +105,6 @@ impl std::fmt::Display for WasmValidationError {
                 "Wasm module defined {} custom sections which exceeds the maximum number allowed {}.",
                 defined, allowed
             ),
-            Self::InvalidFunctionIndex {
-                index,
-                import_count,
-            } => write!(
-                f,
-                "Function has index {} but should start from {}.",
-                index, import_count
-            ),
             Self::FunctionComplexityTooHigh{ index, complexity, allowed } => write!(
                 f,
                 "Wasm module contains a function at index {} with complexity {} which exceeds the maximum complexity allowed {}",