From 7f2f0983dfa74e6f47a3ee300f8d350abd0ef8eb Mon Sep 17 00:00:00 2001
From: chai2010 <chaishushan@gmail.com>
Date: Sun, 12 Mar 2023 23:11:32 +0800
Subject: [PATCH] =?UTF-8?q?=E5=A2=9E=E5=8A=A0=20os.Exit=20=E5=87=BD?=
 =?UTF-8?q?=E6=95=B0?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 internal/app/apputil/wabt_wat2wasm.go                  | 6 ++++++
 internal/waroot/_waroot/src/os/os.wa                   | 3 +++
 internal/waroot/_waroot/src/runtime/runtime_arduino.wa | 3 +++
 internal/waroot/_waroot/src/runtime/runtime_chrome.wa  | 3 +++
 internal/waroot/_waroot/src/runtime/runtime_wasi.wa    | 4 ++++
 main.go                                                | 6 +++++-
 6 files changed, 24 insertions(+), 1 deletion(-)

diff --git a/internal/app/apputil/wabt_wat2wasm.go b/internal/app/apputil/wabt_wat2wasm.go
index 1306ec7..ed59b59 100644
--- a/internal/app/apputil/wabt_wat2wasm.go
+++ b/internal/app/apputil/wabt_wat2wasm.go
@@ -15,6 +15,7 @@ import (
 	"sync"
 
 	"github.com/tetratelabs/wazero"
+	"github.com/tetratelabs/wazero/sys"
 
 	"wa-lang.org/wa/internal/app/waruntime"
 	"wa-lang.org/wa/internal/config"
@@ -78,6 +79,11 @@ func RunWasm(cfg *config.Config, filename string, wasmArgs ...string) (stdoutStd
 
 	_, err = r.InstantiateModule(ctx, code, conf)
 	if err != nil {
+		if exitErr, ok := err.(*sys.ExitError); ok {
+			if exitErr.ExitCode() == 0 {
+				return outputBuffer.Bytes(), nil
+			}
+		}
 		return outputBuffer.Bytes(), err
 	}
 
diff --git a/internal/waroot/_waroot/src/os/os.wa b/internal/waroot/_waroot/src/os/os.wa
index f188a27..2f01026 100644
--- a/internal/waroot/_waroot/src/os/os.wa
+++ b/internal/waroot/_waroot/src/os/os.wa
@@ -5,6 +5,9 @@ var Args []string
 #wa:linkname runtime.os_get_args
 func os_get_args() => []string
 
+#wa:linkname $runtime.procExit
+func Exit(code: i32)
+
 func init {
 	Args = os_get_args()
 }
diff --git a/internal/waroot/_waroot/src/runtime/runtime_arduino.wa b/internal/waroot/_waroot/src/runtime/runtime_arduino.wa
index 3a1ae89..ab4bdf6 100644
--- a/internal/waroot/_waroot/src/runtime/runtime_arduino.wa
+++ b/internal/waroot/_waroot/src/runtime/runtime_arduino.wa
@@ -19,6 +19,9 @@ func fdWrite(fd: i32, io: i32, iovs_len: i32, nwritten: i32) => (written: i32) {
 	return
 }
 
+#wa:linkname $runtime.procExit
+func procExit(code: i32) {}
+
 #wa:linkname $runtime.waPrintI32
 func waPrintI32(i: i32) {}
 
diff --git a/internal/waroot/_waroot/src/runtime/runtime_chrome.wa b/internal/waroot/_waroot/src/runtime/runtime_chrome.wa
index 1ee4184..6e801f8 100644
--- a/internal/waroot/_waroot/src/runtime/runtime_chrome.wa
+++ b/internal/waroot/_waroot/src/runtime/runtime_chrome.wa
@@ -17,6 +17,9 @@ func fdWrite(fd: i32, io: i32, iovs_len: i32, nwritten: i32) => (written: i32) {
 	return
 }
 
+#wa:linkname $runtime.procExit
+func procExit(code: i32) {}
+
 #wa:linkname $runtime.waPrintI32
 func waPrintI32(i: i32) {}
 
diff --git a/internal/waroot/_waroot/src/runtime/runtime_wasi.wa b/internal/waroot/_waroot/src/runtime/runtime_wasi.wa
index c2c8800..39a085c 100644
--- a/internal/waroot/_waroot/src/runtime/runtime_wasi.wa
+++ b/internal/waroot/_waroot/src/runtime/runtime_wasi.wa
@@ -14,6 +14,10 @@ func argsGet(result_argv: i32, result_argv_buf: i32) => (errno: i32)
 #wa:linkname $runtime.fdWrite
 func fdWrite(fd: i32, io: i32, iovs_len: i32, nwritten: i32) => (errno: i32)
 
+#wa:import wasi_snapshot_preview1 proc_exit
+#wa:linkname $runtime.procExit
+func procExit(code: i32)
+
 #wa:linkname puts
 func puts(ptr: i32, len: i32)
 
diff --git a/main.go b/main.go
index 813dd0d..25960cb 100644
--- a/main.go
+++ b/main.go
@@ -13,6 +13,8 @@ import (
 	"strings"
 	"time"
 
+	"github.com/tetratelabs/wazero/sys"
+
 	"wa-lang.org/wa/api"
 	"wa-lang.org/wa/internal/3rdparty/cli"
 	"wa-lang.org/wa/internal/app"
@@ -550,8 +552,10 @@ func cliRun(c *cli.Context) {
 		if len(stdoutStderr) > 0 {
 			fmt.Println(string(stdoutStderr))
 		}
+		if exitErr, ok := err.(*sys.ExitError); ok {
+			os.Exit(int(exitErr.ExitCode()))
+		}
 		fmt.Println(err)
-		os.Exit(1)
 	}
 	if len(stdoutStderr) > 0 {
 		fmt.Println(string(stdoutStderr))