add RestoreStackFunc method for preventing stack leaks

2 files changed, 21 insertions, 2 deletions

diff --git a/lua.go b/lua.go
index 874dac8..9e7b6ec 100644
--- a/lua.go
+++ b/lua.go
@@ -122,3 +122,21 @@ func (l *Lua) PushTableItem (key string) {
 	C.lua_pushstring(l.l, ckey);
 	C.lua_gettable(l.l, -2);
 }
+
+func (l *Lua) RestoreStackFunc() func () {
+	before := l.StackLen()
+	return func () {
+		after := l.StackLen()
+		diff := after - before
+		if diff == 0 {
+			return
+		} else if diff < 0 {
+			log.Fatalf(
+				"too many stack pops: len before: %d, after: %d\n",
+				before,
+				after,
+			)
+		}
+		C.lua_settop(l.l, C.int(before))
+	}
+}
diff --git a/worker.go b/worker.go
index 3c64d68..6d659da 100644
--- a/worker.go
+++ b/worker.go
@@ -57,6 +57,7 @@ func (w *Worker) Start (filename string) error {
 func (w *Worker) Listen () {
 	for {
 		r := <- w.read
+		resStack := w.lua.RestoreStackFunc()
 		handlerRef := w.routes[r.Route]
 		w.lua.PushFromRef(handlerRef)
 		w.lua.PushString(r.Request.Method)
@@ -70,7 +71,7 @@ func (w *Worker) Listen () {
 		
 		body, err := io.ReadAll(r.Request.Body)
 		if err != nil {
-			w.lua.Pop(4);
+			resStack()
 			r.result <- &WorkerResponse {
 				Code: 500,
 				Headers: make(map[string]string),
@@ -101,7 +102,7 @@ func (w *Worker) Listen () {
 			k := w.lua.ToString(-1)
 			headers[k] = v
 		}
-		w.lua.Pop(2)
+		resStack()
 
 		r.result <- &WorkerResponse {
 			Code: int(code),