From 81801e1c3375f2bcd1e096583954bf3bb41623ef Mon Sep 17 00:00:00 2001
From: Rui Guo <firemeteor.guo@gmail.com>
Date: Wed, 19 Aug 2009 22:25:24 +0800
Subject: [PATCH] Fix problem caused by resourcing a script when an callback is
 in function.

---
 src/lua.c | 83 ++++++++++++++++++++++++++++++++++++++++++++-------------------
 1 file changed, 58 insertions(+), 25 deletions(-)

diff --git a/src/lua.c b/src/lua.c
index ce8f2f1..2c73ea8 100644
--- a/src/lua.c
+++ b/src/lua.c
@@ -135,13 +135,35 @@ LuaFreeHandler(lua_State *L, struct lua_handler **lh)
 {
   if ((*lh)->type == LUA_HANDLER_TYPE_N)
     Free((*lh)->u.name)
-  else
+  else if (L)
     {
       LuaHUnRef(L, (*lh)->u.reference);
       (*lh)->u.reference = 0;
     }
   Free(*lh);
 }
+
+struct sfile {
+    lua_State *L;
+    ino_t inode;
+    int inuse;
+    struct sfile *next;
+};
+
+struct sfile *scripts = NULL;
+
+struct sfile *
+LuaGetSFile(lua_State *L)
+{
+  struct sfile *slist;
+  luaL_getmetatable(L, "screen");
+  lua_pushstring(L, "_sfile");
+  lua_rawget(L, -2);
+  slist = (struct sfile *)lua_touserdata(L, -1);
+  lua_pop(L, 2);
+  return slist;
+}
+
 /** Template {{{ */
 
 #define CHECK_TYPE(name, type) \
@@ -1061,8 +1083,15 @@ void
 script_input_fn(char *buf, int len, char *priv)
 {
   lua_handler lh = (lua_handler)priv;
-  lua_State *L = lh->L;
+  struct sfile *sf = (struct sfile *)lh->L;
+  lua_State *L = sf->L;
+  if (!L) {
+      free(sf);
+      LuaFreeHandler(L, &lh);
+      return;
+  }
 
+  lh->L = L;
   LuaPushHandler(lh);
   lua_pushstring(L, buf);
   if (lua_pcall(L, 1, 0, 0) == LUA_ERRRUN)
@@ -1072,6 +1101,7 @@ script_input_fn(char *buf, int len, char *priv)
           LuaShowErr(L);
 	}
     }
+  sf->inuse = 0;
   LuaFreeHandler(L, &lh);
 }
 
@@ -1082,9 +1112,12 @@ screen_input(lua_State *L)
   int n;
   const char * prompt = NULL, *prefill = NULL;
   lua_handler lh;
+  struct sfile *sf = LuaGetSFile(L);
 
   prompt = luaL_checkstring(L, 1);
   lh = LuaCheckHandler(L, 2, 1);
+  /* Hack! This is used to prevent from accessing unloaded script.*/
+  lh->L = (lua_State *)sf;
   if (!lh)
     luaL_error(L, "Out of Memory");
   
@@ -1092,6 +1125,7 @@ screen_input(lua_State *L)
     prefill = luaL_checkstring(L, 3);
 
 
+  sf->inuse = 1;
   Input((char *)prompt, 100, INP_COOKED, script_input_fn, (char *)lh, 0);
 
   if (!prefill)
@@ -1111,10 +1145,18 @@ static void
 sev_schedule_fn(struct event *ev, char *data)
 {
   lua_handler lh = (lua_handler)data;
-  lua_State *L = lh->L;
+  struct sfile *sf = (struct sfile *)lh->L;
+  lua_State *L = sf->L;
+  if (!L) {
+      free(sf);
+      LuaFreeHandler(L, &lh);
+      return;
+  }
+
   evdeq(ev);
   Free(ev);
 
+  lh->L = L;
   LuaPushHandler(lh);
   if (lua_pcall(L, 1, 0, 0) == LUA_ERRRUN)
     {
@@ -1123,6 +1165,7 @@ sev_schedule_fn(struct event *ev, char *data)
           LuaShowErr(L);
 	}
     }
+  sf->inuse = 0;
   LuaFreeHandler(L, &lh);
 }
 
@@ -1130,10 +1173,13 @@ static int
 screen_schedule(lua_State *L)
 {
   int timeout = luaL_checkinteger(L, 1);
-  lua_handler lh = LuaCheckHandler(L, 2, 1);
+  lua_handler lh;
+  struct sfile *sf = LuaGetSFile(L);
   struct event *ev;
   if (timeout <= 0)
     return 0;
+  lh = LuaCheckHandler(L, 2, 1);
+  lh->L = (lua_State *)sf;
 
   ev = (struct event *) calloc(1, sizeof(struct event));
   if (!ev)
@@ -1147,6 +1193,7 @@ screen_schedule(lua_State *L)
   ev->handler = sev_schedule_fn;
   evenq(ev);
   SetTimeout(ev, timeout * 1000);
+  sf->inuse = 1;
   return 0;
 }
 
@@ -1200,14 +1247,6 @@ prepare_weak_table(lua_State *L, const char *name, const char *mode)
   lua_pop(L, 1);
 }
 
-struct sfile {
-    lua_State *L;
-    ino_t inode;
-    struct sfile *next;
-};
-
-struct sfile *scripts = NULL;
-
 int LuaInit(void)
 {
   return 0;
@@ -1242,6 +1281,7 @@ LuaNewState(struct sfile *slist)
   lua_rawset(L, -3);
   lua_pop(L, 1);
   slist->L = L;
+  slist->inuse = 0;
   return L;
 }
 
@@ -1269,7 +1309,12 @@ LuaUnload(struct sfile *slist)
       }
       plist = &(*plist)->next;
   }
-  free(slist);
+
+  //Delay reclaiming the structure if it's still in use.
+  if (!slist->inuse)
+    free(slist);
+  else
+    slist->L = NULL;
 }
 
 /* FIXME: Think about this: will it affect the registered handlers?*/
@@ -1436,18 +1481,6 @@ LuaPushHTable(lua_State *L, int screen, const char * t)
   return lua_gettop(L);
 }
 
-struct sfile *
-LuaGetSFile(lua_State *L)
-{
-  struct sfile *slist;
-  luaL_getmetatable(L, "screen");
-  lua_pushstring(L, "_sfile");
-  lua_rawget(L, -1);
-  slist = (struct sfile *)lua_touserdata(L, -1);
-  lua_pop(L, 2);
-  return slist;
-}
-
 void
 LuaPushHandler(lua_handler lh)
 {
-- 
2.11.4.GIT