Closing methods should not interfere with returning values

A closing method cannot be called in its own stack slot, as there may
be returning values in the stack after that slot, and the call would
corrupt those values. Instead, the closing method must be copied to the
top of the stack to be called.

Moreover, even when a function returns no value, its return istruction
still has to have its position (which will set the stack top) after
the local variables, otherwise a closing method might corrupt another
not-yet-called closing method.

4 files changed, 98 insertions, 42 deletions

diff --git a/lfunc.c b/lfunc.c
index 4f9362f3..15874f88 100644
--- a/lfunc.c
+++ b/lfunc.c
@@ -98,27 +98,29 @@ UpVal *luaF_findupval (lua_State *L, StkId level) {
 
 
 static void callclose (lua_State *L, void *ud) {
-  luaD_callnoyield(L, cast(StkId, ud), 0);
+  UNUSED(ud);
+  luaD_callnoyield(L, L->top - 2, 0);
 }
 
 
 /*
-** Prepare closing method with its argument for object at
-** index 'func' in the stack. Assume there is an error message
-** (or nil) just below the object.
+** Prepare closing method plus its argument for object 'obj' with
+** error message 'err'. (This function assumes EXTRA_STACK.)
 */
-static int prepclosingmethod (lua_State *L, StkId func) {
-  if (ttisfunction(s2v(func))) {  /* object to-be-closed is a function? */
-    setobjs2s(L, func + 1, func - 1);  /* push error msg. as argument */
+static int prepclosingmethod (lua_State *L, TValue *obj, TValue *err) {
+  StkId top = L->top;
+  if (ttisfunction(obj)) {  /* object to-be-closed is a function? */
+    setobj2s(L, top, obj);  /* push function */
+    setobj2s(L, top + 1, err);  /* push error msg. as argument */
   }
   else {  /* try '__close' metamethod */
-    const TValue *tm = luaT_gettmbyobj(L, s2v(func), TM_CLOSE);
+    const TValue *tm = luaT_gettmbyobj(L, obj, TM_CLOSE);
     if (ttisnil(tm))  /* no metamethod? */
       return 0;  /* nothing to call */
-    setobjs2s(L, func + 1, func);  /* 'self' is the argument */
-    setobj2s(L, func, tm);  /* will call metamethod */
+    setobj2s(L, top, tm);  /* will call metamethod... */
+    setobj2s(L, top + 1, obj);  /* with 'self' as the argument */
   }
-  L->top = func + 2;  /* add function and argument */
+  L->top = top + 2;  /* add function and argument */
   return 1;
 }
 
@@ -129,22 +131,24 @@ static int prepclosingmethod (lua_State *L, StkId func) {
 ** will be handled there. Otherwise, a previous error already
 ** activated original protected call, and so the call to the
 ** closing method must be protected here.
+** If status is OK, the call to the closing method will be pushed
+** at the top of the stack. Otherwise, values are pushed after
+** the 'level' of the upvalue being closed, as everything after
+** that won't be used again.
 */
 static int closeupval (lua_State *L, TValue *uv, StkId level, int status) {
-  StkId func = level + 1;  /* save slot for old error message */
-  if (unlikely(status != LUA_OK))  /* was there an error? */
-    luaD_seterrorobj(L, status, level);  /* save error message */
-  else
-    setnilvalue(s2v(level));  /* no error message */
-  setobj2s(L, func, uv);  /* put object on top of error message */
-  if (!prepclosingmethod(L, func))
-    return status;  /* nothing to call */
-  if (likely(status == LUA_OK))  /* not in "error mode"? */
-    callclose(L, func);  /* call closing method */
-  else {  /* already inside error handler; cannot raise another error */
-    int newstatus = luaD_pcall(L, callclose, func, savestack(L, level), 0);
-    if (newstatus != LUA_OK)  /* error when closing? */
-      status = newstatus;  /* this will be the new error */
+  if (likely(status == LUA_OK)) {
+    if (prepclosingmethod(L, uv, &G(L)->nilvalue))  /* something to call? */
+      callclose(L, NULL);  /* call closing method */
+  }
+  else {  /* there was an error */
+    /* save error message and set stack top to 'level + 1' */
+    luaD_seterrorobj(L, status, level);
+    if (prepclosingmethod(L, uv, s2v(level))) {  /* something to call? */
+      int newstatus = luaD_pcall(L, callclose, NULL, savestack(L, level), 0);
+      if (newstatus != LUA_OK)  /* another error when closing? */
+        status = newstatus;  /* this will be the new error */
+    }
   }
   return status;
 }
@@ -169,12 +173,10 @@ static void trynewtbcupval (lua_State *L, void *ud) {
 void luaF_newtbcupval (lua_State *L, StkId level) {
   int status = luaD_rawrunprotected(L, trynewtbcupval, level);
   if (unlikely(status != LUA_OK)) {  /* memory error creating upvalue? */
-    StkId func = level + 1;
     lua_assert(status == LUA_ERRMEM);
-    setobjs2s(L, func, level);  /* open space for error message */
-    luaD_seterrorobj(L, status, level);  /* save error message */
-    if (prepclosingmethod(L, func))
-      callclose(L, func);  /* call closing method */
+    luaD_seterrorobj(L, LUA_ERRMEM, level + 1);  /* save error message */
+    if (prepclosingmethod(L, s2v(level), s2v(level + 1)))
+      callclose(L, NULL);  /* call closing method */
     luaD_throw(L, LUA_ERRMEM);  /* throw memory error */
   }
 }
@@ -201,7 +203,7 @@ int luaF_close (lua_State *L, StkId level, int status) {
     luaC_barrier(L, uv, slot);
     if (status >= 0 && uv->tt == LUA_TUPVALTBC) {  /* must be closed? */
       ptrdiff_t levelrel = savestack(L, level);
-      status = closeupval(L, uv->v, upl, status);  /* may reallocate the stack */
+      status = closeupval(L, uv->v, upl, status);  /* may realloc. the stack */
       level = restorestack(L, levelrel);
     }
   }
diff --git a/lparser.c b/lparser.c
index 6b14b800..c0c40eae 100644
--- a/lparser.c
+++ b/lparser.c
@@ -561,7 +561,7 @@ static void close_func (LexState *ls) {
   lua_State *L = ls->L;
   FuncState *fs = ls->fs;
   Proto *f = fs->f;
-  luaK_ret(fs, 0, 0);  /* final return */
+  luaK_ret(fs, fs->nactvar, 0);  /* final return */
   leaveblock(fs);
   lua_assert(fs->bl == NULL);
   luaK_finish(fs);
@@ -1602,9 +1602,10 @@ static void retstat (LexState *ls) {
   /* stat -> RETURN [explist] [';'] */
   FuncState *fs = ls->fs;
   expdesc e;
-  int first, nret;  /* registers with returned values */
+  int nret;  /* number of values being returned */
+  int first = fs->nactvar;  /* first slot to be returned */
   if (block_follow(ls, 1) || ls->t.token == ';')
-    first = nret = 0;  /* return no values */
+    nret = 0;  /* return no values */
   else {
     nret = explist(ls, &e);  /* optional return values */
     if (hasmultret(e.k)) {
@@ -1613,15 +1614,13 @@ static void retstat (LexState *ls) {
         SET_OPCODE(getinstruction(fs,&e), OP_TAILCALL);
         lua_assert(GETARG_A(getinstruction(fs,&e)) == fs->nactvar);
       }
-      first = fs->nactvar;
       nret = LUA_MULTRET;  /* return all values */
     }
     else {
       if (nret == 1)  /* only one single value? */
-        first = luaK_exp2anyreg(fs, &e);
-      else {
-        luaK_exp2nextreg(fs, &e);  /* values must go to the stack */
-        first = fs->nactvar;  /* return all active values */
+        first = luaK_exp2anyreg(fs, &e);  /* can use original slot */
+      else {  /* values must go to the top of the stack */
+        luaK_exp2nextreg(fs, &e);
         lua_assert(nret == fs->freereg - first);
       }
     }
diff --git a/lvm.c b/lvm.c
index 0d82756b..2a1ee175 100644
--- a/lvm.c
+++ b/lvm.c
@@ -1452,7 +1452,8 @@ void luaV_execute (lua_State *L, CallInfo *ci) {
         vmbreak;
       }
       vmcase(OP_CLOSE) {
-        luaF_close(L, ra, LUA_OK);
+        L->top = ra + 1;  /* everything is free after this slot */
+        ProtectNT(luaF_close(L, ra, LUA_OK));
         vmbreak;
       }
       vmcase(OP_TBC) {
@@ -1619,13 +1620,14 @@ void luaV_execute (lua_State *L, CallInfo *ci) {
           n = cast_int(L->top - ra);  /* get what is available */
         else
           L->top = ra + n;  /* set call for 'luaD_poscall' */
+        savepc(ci);
         if (TESTARG_k(i)) {
           int nparams1 = GETARG_C(i);
           if (nparams1)  /* vararg function? */
             ci->func -= ci->u.l.nextraargs + nparams1;
           luaF_close(L, base, LUA_OK);  /* there may be open upvalues */
         }
-        halfProtect(luaD_poscall(L, ci, n));
+        luaD_poscall(L, ci, n);
         return;
       }
       vmcase(OP_RETURN0) {
diff --git a/testes/locals.lua b/testes/locals.lua
index f21fa2ec..65b145db 100644
--- a/testes/locals.lua
+++ b/testes/locals.lua
@@ -175,6 +175,9 @@ assert(x==20)
 
 print"testing to-be-closed variables"
 
+local function stack(n) n = ((n == 0) or stack(n - 1)) end
+
+
 do
   local a = {}
   do
@@ -187,6 +190,57 @@ do
   assert(a[1] == "in" and a[2] == "y" and a[3] == "x" and a[4] == "out")
 end
 
+do
+  local X = false
+
+  local function closescope () stack(10); X = true end
+
+  -- closing functions do not corrupt returning values
+  local function foo (x)
+    local scoped _ = closescope
+    return x, X, 23
+  end
+
+  local a, b, c = foo(1.5)
+  assert(a == 1.5 and b == false and c == 23 and X == true)
+
+  X = false
+  foo = function (x)
+    local scoped _ = closescope
+    local y = 15
+    return y
+  end
+
+  assert(foo() == 15 and X == true)
+
+  X = false
+  foo = function ()
+    local scoped x = closescope
+    return x
+  end
+
+  assert(foo() == closescope and X == true)
+
+end
+
+
+do
+  -- to-be-closed variables must be closed in tail calls
+  local X, Y
+  local function foo ()
+    local scoped _ = function () Y = 10 end
+    assert(X == 20 and Y == nil)
+    return 1,2,3
+  end
+
+  local function bar ()
+    local scoped _ = function () X = 20 end
+    return foo()
+  end
+
+  local a, b, c, d = bar()
+  assert(a == 1 and b == 2 and c == 3 and X == 20 and Y == 10 and d == nil)
+end
 
 do   -- errors in __close
   local log = {}
@@ -211,7 +265,6 @@ do   -- errors in __close
 end
 
 if rawget(_G, "T") then
-  local function stack(n) n = (n == 0) or stack(n - 1); end;
   -- memory error inside closing function
   local function foo ()
     local scoped y = function () T.alloccount() end