wavm: fix function export bug. (#200)

Signed-off-by: chaoqin-li1123 <[email protected]>

Author: chaoqin-li1123

src/wavm/wavm.cc

@@ -443,8 +443,10 @@ void getFunctionWavm(WasmVm *vm, std::string_view function_name,
     return;
   }
   if (!checkFunctionType(f, inferStdFunctionType(function))) {
+    *function = nullptr;
     wavm->fail(FailState::UnableToInitializeCode,
                "Bad function signature for: " + std::string(function_name));
+    return;
   }
   *function = [wavm, f, function_name](ContextBase *context, Args... args) -> R {
     WasmUntaggedValue values[] = {args...};
@@ -473,8 +475,10 @@ void getFunctionWavm(WasmVm *vm, std::string_view function_name,
     return;
   }
   if (!checkFunctionType(f, inferStdFunctionType(function))) {
+    *function = nullptr;
     wavm->fail(FailState::UnableToInitializeCode,
                "Bad function signature for: " + std::string(function_name));
+    return;
   }
   *function = [wavm, f, function_name](ContextBase *context, Args... args) {
     WasmUntaggedValue values[] = {args...};

test/runtime_test.cc

@@ -171,6 +171,35 @@ TEST_P(TestVM, StraceLogLevel) {
   EXPECT_NE(integration->trace_message_, "");
 }
 
+TEST_P(TestVM, BadExportFunction) {
+  auto source = readTestWasmFile("callback.wasm");
+  ASSERT_TRUE(vm_->load(source, {}, {}));
+
+  TestContext context;
+  vm_->registerCallback(
+      "env", "callback", &callback,
+      &ConvertFunctionWordToUint32<decltype(callback), callback>::convertFunctionWordToUint32);
+  vm_->registerCallback(
+      "env", "callback2", &callback2,
+      &ConvertFunctionWordToUint32<decltype(callback2), callback2>::convertFunctionWordToUint32);
+  ASSERT_TRUE(vm_->link(""));
+
+  WasmCallVoid<0> run;
+  vm_->getFunction("non-existent", &run);
+  EXPECT_TRUE(run == nullptr);
+
+  WasmCallWord<2> bad_signature_run;
+  vm_->getFunction("run", &bad_signature_run);
+  EXPECT_TRUE(bad_signature_run == nullptr);
+
+  vm_->getFunction("run", &run);
+  EXPECT_TRUE(run != nullptr);
+  for (auto i = 0; i < 100; i++) {
+    run(&context);
+  }
+  ASSERT_EQ(context.counter, 100);
+}
+
 TEST_P(TestVM, Callback) {
   auto source = readTestWasmFile("callback.wasm");
   ASSERT_TRUE(vm_->load(source, {}, {}));