From 4eb614d18311b3b0184dc5c8a42c2845aa10b0ce Mon Sep 17 00:00:00 2001
From: chvancooten <chvancooten@gmail.com>
Date: Sun, 13 Oct 2024 22:02:55 +0200
Subject: [PATCH] Add panic handler for COFF execution

---
 src/coff/coff_windows.go | 13 ++++++++++++-
 1 file changed, 12 insertions(+), 1 deletion(-)

diff --git a/src/coff/coff_windows.go b/src/coff/coff_windows.go
index cc2709c..f62f4d8 100644
--- a/src/coff/coff_windows.go
+++ b/src/coff/coff_windows.go
@@ -3,11 +3,13 @@ package coff
 import (
 	_ "embed"
 	"fmt"
-	"github.com/praetorian-inc/goffloader/src/lighthouse"
+	"runtime/debug"
 	"strings"
 	"syscall"
 	"unsafe"
 
+	"github.com/praetorian-inc/goffloader/src/lighthouse"
+
 	"github.com/RIscRIpt/pecoff"
 	"github.com/RIscRIpt/pecoff/binutil"
 	"github.com/RIscRIpt/pecoff/windef"
@@ -327,6 +329,15 @@ func LoadWithMethod(coffBytes []byte, argBytes []byte, method string) (string, e
 func invokeMethod(methodName string, argBytes []byte, parsedCoff *pecoff.File, sectionMap map[string]CoffSection, outChannel chan<- interface{}) {
 	defer close(outChannel)
 
+	// Catch unexpected panics and propagate them to the output channel
+	// This prevents the host program from terminating unexpectedly
+	defer func() {
+		if r := recover(); r != nil {
+			errorMsg := fmt.Sprintf("Panic occurred when executing COFF: %v\n%s", r, debug.Stack())
+			outChannel <- errorMsg
+		}
+	}()
+
 	// Call the entry point
 	for _, symbol := range parsedCoff.Symbols {
 		if symbol.NameString() == methodName {