From ba03c665728b8567321c7bd1895b3832cba76ccb Mon Sep 17 00:00:00 2001
From: rrai35 <ritesh.rai@aexp.com>
Date: Fri, 22 Sep 2023 10:27:33 -0400
Subject: [PATCH] feat(hostcallDuringLoad): Allow updating the context with
 InvokeContext which allows hostcalls during module load

---
 engines/wazero/wazero.go      |  9 ++++++++-
 engines/wazero/wazero_test.go | 13 +++++++++++++
 2 files changed, 21 insertions(+), 1 deletion(-)

diff --git a/engines/wazero/wazero.go b/engines/wazero/wazero.go
index 7b9dd72..f734bb2 100644
--- a/engines/wazero/wazero.go
+++ b/engines/wazero/wazero.go
@@ -239,7 +239,7 @@ func (w *wapcHost) hostCall(ctx context.Context, m api.Module, stack []uint64) {
 	namespace := requireReadString(mem, "namespace", nsPtr, nsLen)
 	operation := requireReadString(mem, "operation", cmdPtr, cmdLen)
 	payload := requireRead(mem, "payload", payloadPtr, payloadLen)
-
+	
 	if ic.hostResp, ic.hostErr = w.callHandler(ctx, binding, namespace, operation, payload); ic.hostErr != nil {
 		stack[0] = 0 // false: error (assumed to be logged already?)
 	} else {
@@ -465,3 +465,10 @@ func requireRead(mem api.Memory, fieldName string, offset, byteCount uint32) []b
 	}
 	return buf
 }
+
+// AddInvokeContext allows the use of invocation context during module load. This allows host calls from the wapc guest's main()
+func AddInvokeContext(ctx context.Context) context.Context {
+	ic := invokeContext{}
+	ctx = newInvokeContext(ctx, &ic)
+	return ctx
+}
diff --git a/engines/wazero/wazero_test.go b/engines/wazero/wazero_test.go
index 648a4a6..6a1869f 100644
--- a/engines/wazero/wazero_test.go
+++ b/engines/wazero/wazero_test.go
@@ -33,6 +33,16 @@ func TestMain(m *testing.M) {
 	os.Exit(m.Run())
 }
 
+// TestAddInvokeContext ensures we can read the example wasm prior to running unit tests.
+func TestAddInvokeContext(t *testing.T) {
+	ctx := context.Background()
+	ctx = AddInvokeContext(ctx)
+	ic := fromInvokeContext(ctx)
+	if ic == nil {
+		t.Errorf("Error adding InvocationContext to context ")
+	}
+}
+
 // TestModule_UnwrapRuntime ensures the Unwrap returns the correct Runtime interface
 func TestModule_UnwrapRuntime(t *testing.T) {
 	m, err := EngineWithRuntime(DefaultRuntime).New(testCtx, wapc.NoOpHostCallHandler, guest, mc)
@@ -141,4 +151,7 @@ func TestEngineWithRuntime(t *testing.T) {
 			t.Errorf("Unexpected error, got %v, expected %v", err, expectedErr)
 		}
 	})
+
 }
+
+