From 9311dda6ff5b2a15e0f5b4726a8ff0994e85600c Mon Sep 17 00:00:00 2001
From: Kyle Conroy <kyle@conroy.org>
Date: Fri, 21 Apr 2023 14:53:41 -0700
Subject: [PATCH] ext/wasm: Check exit code on returned error

As part of migrating plugins to the official Go toolchain, the wasmtime
call pattern needed to be updated. This should be backwards compatible
with TinyGo plugins.
---
 internal/ext/wasm/wasm.go | 44 +++++++++++++++++++++++++++++----------
 1 file changed, 33 insertions(+), 11 deletions(-)

diff --git a/internal/ext/wasm/wasm.go b/internal/ext/wasm/wasm.go
index 1d04ea176a..ee4b1841d7 100644
--- a/internal/ext/wasm/wasm.go
+++ b/internal/ext/wasm/wasm.go
@@ -258,25 +258,25 @@ func (r *Runner) Generate(ctx context.Context, req *plugin.CodeGenRequest) (*plu
 	store := wasmtime.NewStore(engine)
 	store.SetWasi(wasiConfig)
 
-	linkRegion := trace.StartRegion(ctx, "linker.Instantiate")
-	instance, err := linker.Instantiate(store, module)
+	linkRegion := trace.StartRegion(ctx, "linker.DefineModule")
+	err = linker.DefineModule(store, "", module)
 	linkRegion.End()
 	if err != nil {
 		return nil, fmt.Errorf("define wasi: %w", err)
 	}
 
 	// Run the function
+	fn, err := linker.GetDefault(store, "")
+	if err != nil {
+		return nil, fmt.Errorf("wasi: get default: %w", err)
+	}
+
 	callRegion := trace.StartRegion(ctx, "call _start")
-	nom := instance.GetExport(store, "_start").Func()
-	_, err = nom.Call(store)
+	_, err = fn.Call(store)
 	callRegion.End()
-	if err != nil {
-		// Print WASM stdout
-		stderrBlob, err := os.ReadFile(stderrPath)
-		if err == nil && len(stderrBlob) > 0 {
-			return nil, errors.New(string(stderrBlob))
-		}
-		return nil, fmt.Errorf("call: %w", err)
+
+	if cerr := checkError(err, stderrPath); cerr != nil {
+		return nil, cerr
 	}
 
 	// Print WASM stdout
@@ -284,6 +284,28 @@ func (r *Runner) Generate(ctx context.Context, req *plugin.CodeGenRequest) (*plu
 	if err != nil {
 		return nil, fmt.Errorf("read file: %w", err)
 	}
+
 	var resp plugin.CodeGenResponse
 	return &resp, resp.UnmarshalVT(stdoutBlob)
 }
+
+func checkError(err error, stderrPath string) error {
+	if err == nil {
+		return err
+	}
+
+	var wtError *wasmtime.Error
+	if errors.As(err, &wtError) {
+		if code, ok := wtError.ExitStatus(); ok {
+			if code == 0 {
+				return nil
+			}
+		}
+	}
+	// Print WASM stdout
+	stderrBlob, rferr := os.ReadFile(stderrPath)
+	if rferr == nil && len(stderrBlob) > 0 {
+		return errors.New(string(stderrBlob))
+	}
+	return fmt.Errorf("call: %w", err)
+}