From 9dccd7533fd4f1c0962b43bdc67ddff7a4a59548 Mon Sep 17 00:00:00 2001
From: Jimmy Moore <jamesmoore@loopholelabs.io>
Date: Fri, 29 Sep 2023 19:38:39 +0100
Subject: [PATCH] Added error handling in go ext

Signed-off-by: Jimmy Moore <jamesmoore@loopholelabs.io>
---
 .../generator/golang/templates/guest.go.templ | 30 +++++++++++++++++--
 .../generator/golang/templates/host.go.templ  |  9 ++++++
 2 files changed, 37 insertions(+), 2 deletions(-)

diff --git a/extension/generator/golang/templates/guest.go.templ b/extension/generator/golang/templates/guest.go.templ
index 9e234fb1..4e9adc5e 100644
--- a/extension/generator/golang/templates/guest.go.templ
+++ b/extension/generator/golang/templates/guest.go.templ
@@ -53,10 +53,18 @@ func (d *_{{ $ifc.Name }}) {{ $fn.Name }}(params *{{ $fn.Params }}) ({{ $fn.Retu
   // Now make the call to the host.
 
   {{- if (IsInterface $schema $fn.Return) }}
+    readBuffer = nil
     v := ext_{{ $schema.Name }}_{{ $ifc.Name }}_{{ $fn.Name }}(d.instanceId, off, l)
     // IF the return type is an interface return ifc, which contains hidden instanceId.
 
-    // TODO: Handle error from host. In this case there'll be an error in the readBuffer
+    // Handle error from host. In this case there'll be an error in the readBuffer
+    if (readBuffer!=nil) {
+      val, err := polyglot.GetDecoder(readBuffer).Error()
+      if err!=nil {
+        panic(err)
+      }
+      return &_{{ $fn.Return}}{}, val
+    }
 
     ret := &_{{ $fn.Return }}{
       instanceId: v,
@@ -69,6 +77,11 @@ func (d *_{{ $ifc.Name }}) {{ $fn.Name }}(params *{{ $fn.Params }}) ({{ $fn.Retu
 
     ret := &{{ $fn.Return }}{}
     r, err := Decode{{ $fn.Return }}(ret, readBuffer)
+
+    if err!=nil {
+      return {{ $fn.Return }}{}, err
+    }
+
     return *r, err
   {{ end }}
 }
@@ -102,10 +115,18 @@ func {{ $fn.Name }}(params *{{ $fn.Params }}) ({{ $fn.Return }}, error) {
   // Now make the call to the host.
 
   {{- if (IsInterface $schema $fn.Return) }}
+    readBuffer = nil
     v := ext_{{ $schema.Name }}_{{ $fn.Name }}(0, off, l)
     // IF the return type is an interface return ifc, which contains hidden instanceId.
 
-    // TODO: Handle error from host. In this case there'll be an error in the readBuffer
+    // Handle error from host. In this case there'll be an error in the readBuffer
+    if (readBuffer!=nil) {
+      val, err := polyglot.GetDecoder(readBuffer).Error()
+      if err!=nil {
+        panic(err)
+      }
+      return nil, val
+    }
 
     ret := &_{{ $fn.Return }}{
       instanceId: v,
@@ -118,6 +139,11 @@ func {{ $fn.Name }}(params *{{ $fn.Params }}) ({{ $fn.Return }}, error) {
 
     ret := &{{ $fn.Return }}{}
     r, err := Decode{{ $fn.Return }}(ret, readBuffer)
+
+    if err!=nil {
+      return {{ $fn.Return }}{}, err
+    }
+
     return *r, err
   {{ end }}
 
diff --git a/extension/generator/golang/templates/host.go.templ b/extension/generator/golang/templates/host.go.templ
index 0f70208a..50c42fe3 100644
--- a/extension/generator/golang/templates/host.go.templ
+++ b/extension/generator/golang/templates/host.go.templ
@@ -100,12 +100,14 @@ func (h *{{ $schema.Name }}Host) host_ext_{{ $schema.Name }}_{{ $fn.Name}}(mem e
 	cd, err := Decode{{ $fn.Params }}(cd, data)
 	if err != nil {
 		hostError(mem, resize, err)
+    return
 	}
 
   // Call the implementation
 	r, err := h.impl.{{ $fn.Name }}(cd)
 	if err!=nil {
 		hostError(mem, resize, err)
+    return
 	}
 
 {{- if (IsInterface $schema $fn.Return) }}
@@ -127,10 +129,12 @@ func (h *{{ $schema.Name }}Host) host_ext_{{ $schema.Name }}_{{ $fn.Name}}(mem e
 
 	if err != nil {
 		hostError(mem, resize, err)
+    return
 	}
 
 	if !mem.Write(uint32(writeBuffer), b.Bytes()) {
 		hostError(mem, resize, err)
+    return
 	}
 
 {{ end }}
@@ -149,6 +153,7 @@ func (h *{{ $schema.Name }}Host) host_ext_{{ $schema.Name }}_{{ $ifc.Name }}_{{
 	h.instancesLock_{{ $ifc.Name }}.Unlock()
 	if !ok {
 		hostError(mem, resize, errors.New("Instance ID not found!"))
+    return
 	}
 
 	ptr := uint32(params[1])
@@ -159,11 +164,13 @@ func (h *{{ $schema.Name }}Host) host_ext_{{ $schema.Name }}_{{ $ifc.Name }}_{{
 	cd, err := Decode{{ $fn.Params }}(cd, data)
 	if err != nil {
 		hostError(mem, resize, err)
+    return
 	}
 
 	resp, err := r.{{ $fn.Name }}(cd)
 	if err != nil {
 		hostError(mem, resize, err)
+    return
 	}
 
 
@@ -186,10 +193,12 @@ func (h *{{ $schema.Name }}Host) host_ext_{{ $schema.Name }}_{{ $ifc.Name }}_{{
 
 	if err != nil {
 		hostError(mem, resize, err)
+    return
 	}
 
 	if !mem.Write(uint32(writeBuffer), b.Bytes()) {
 		hostError(mem, resize, err)
+    return
 	}
 
 {{ end }}