From 233144b19d16a7b8bec74de78642ead0d5ab532d Mon Sep 17 00:00:00 2001
From: Paul Kulchenko <paul@kulchenko.com>
Date: Thu, 14 Apr 2022 18:13:53 -0700
Subject: [PATCH] Fix memory deallocation while yielding in redbean. (#384)

The yielded coroutine was removed from the stack too early,
leaving it not being anchored, which led to memory freed prematurely.
---
 third_party/lua/luacallwithtrace.c | 5 ++++-
 tool/net/luacheck.h                | 1 -
 tool/net/redbean.c                 | 4 ++--
 3 files changed, 6 insertions(+), 4 deletions(-)

diff --git a/third_party/lua/luacallwithtrace.c b/third_party/lua/luacallwithtrace.c
index 6217a120f..53947299e 100644
--- a/third_party/lua/luacallwithtrace.c
+++ b/third_party/lua/luacallwithtrace.c
@@ -42,7 +42,10 @@ int LuaCallWithTrace(lua_State *L, int nargs, int nres, lua_State *C) {
   lua_xmove(L, C, 1 + nargs);
   // resume the coroutine thus executing the function
   status = lua_resume(C, L, nargs, &nresults);
-  lua_remove(L, 1);  // remove coroutine (still) at the bottom
+  // remove coroutine (still) at the bottom, but only if not yielding
+  // keep it when yielding to anchor, so it's not GC-collected
+  // it's going to be removed at the beggining of the request handling
+  if (!canyield) lua_remove(L, 1);
   if (status != LUA_OK && status != LUA_YIELD) {
     // move the error message
     lua_xmove(C, L, 1);
diff --git a/tool/net/luacheck.h b/tool/net/luacheck.h
index ea0a8f44e..c5506387a 100644
--- a/tool/net/luacheck.h
+++ b/tool/net/luacheck.h
@@ -13,7 +13,6 @@ COSMOPOLITAN_C_START_
       char *s = LuaFormatStack(L);                \
       WARNF("lua stack should be empty!\n%s", s); \
       free(s);                                    \
-      lua_settop(L, 0);                           \
     }                                             \
   } while (0)
 
diff --git a/tool/net/redbean.c b/tool/net/redbean.c
index 806313083..f6b35e3ef 100644
--- a/tool/net/redbean.c
+++ b/tool/net/redbean.c
@@ -2406,6 +2406,7 @@ static int LuaCallWithYield(lua_State *L) {
   // the second set of headers is not going to be sent
   lua_State *co = lua_newthread(L);
   if ((status = LuaCallWithTrace(L, 0, 0, co)) == LUA_YIELD) {
+    CHECK_GT(lua_gettop(L), 0);  // make sure that coroutine is anchored
     YL = co;
     generator = YieldGenerator;
     if (!isyielding) isyielding = 1;
@@ -3002,9 +3003,9 @@ static char *LuaOnHttpRequest(void) {
   lua_State *L = GL;
   effectivepath.p = url.path.p;
   effectivepath.n = url.path.n;
+  lua_settop(L, 0);  // clear Lua stack, as it needs to start fresh
   lua_getglobal(L, "OnHttpRequest");
   if (LuaCallWithYield(L) == LUA_OK) {
-    AssertLuaStackIsEmpty(L);
     return CommitOutput(GetLuaResponse());
   } else {
     LogLuaError("OnHttpRequest", lua_tostring(L, -1));
@@ -3012,7 +3013,6 @@ static char *LuaOnHttpRequest(void) {
         500, "Internal Server Error",
         ShouldServeCrashReportDetails() ? lua_tostring(L, -1) : NULL);
     lua_pop(L, 1);  // pop error
-    AssertLuaStackIsEmpty(L);
     return error;
   }
 }