String buffer using to-be-closed variable

The string buffers in the C API now mark their boxes as to-be-closed
variables, to release their buffers in case of errors.

2 files changed, 47 insertions, 11 deletions

diff --git a/lauxlib.c b/lauxlib.c
index 68842a98..6069ed1b 100644
--- a/lauxlib.c
+++ b/lauxlib.c
@@ -470,10 +470,8 @@ static void *resizebox (lua_State *L, int idx, size_t newsize) {
   lua_Alloc allocf = lua_getallocf(L, &ud);
   UBox *box = (UBox *)lua_touserdata(L, idx);
   void *temp = allocf(ud, box->box, box->bsize, newsize);
-  if (temp == NULL && newsize > 0) {  /* allocation error? */
-    resizebox(L, idx, 0);  /* free buffer */
+  if (temp == NULL && newsize > 0)  /* allocation error? */
     luaL_error(L, "not enough memory for buffer allocation");
-  }
   box->box = temp;
   box->bsize = newsize;
   return temp;
@@ -486,16 +484,20 @@ static int boxgc (lua_State *L) {
 }
 
 
-static void *newbox (lua_State *L, size_t newsize) {
+static const luaL_Reg boxmt[] = {  /* box metamethods */
+  {"__gc", boxgc},
+  {"__close", boxgc},
+  {NULL, NULL}
+};
+
+
+static void newbox (lua_State *L) {
   UBox *box = (UBox *)lua_newuserdatauv(L, sizeof(UBox), 0);
   box->box = NULL;
   box->bsize = 0;
-  if (luaL_newmetatable(L, "_UBOX*")) {  /* creating metatable? */
-    lua_pushcfunction(L, boxgc);
-    lua_setfield(L, -2, "__gc");  /* metatable.__gc = boxgc */
-  }
+  if (luaL_newmetatable(L, "_UBOX*"))  /* creating metatable? */
+    luaL_setfuncs(L, boxmt, 0);  /* set its metamethods */
   lua_setmetatable(L, -2);
-  return resizebox(L, -1, newsize);
 }
 
 
@@ -536,9 +538,11 @@ static char *prepbuffsize (luaL_Buffer *B, size_t sz, int boxidx) {
     if (buffonstack(B))  /* buffer already has a box? */
       newbuff = (char *)resizebox(L, boxidx, newsize);  /* resize it */
     else {  /* no box yet */
-      newbuff = (char *)newbox(L, newsize);  /* create a new box */
-      memcpy(newbuff, B->b, B->n * sizeof(char));  /* copy original content */
+      newbox(L);  /* create a new box */
       lua_insert(L, boxidx);  /* move box to its intended position */
+      lua_toclose(L, boxidx);
+      newbuff = (char *)resizebox(L, boxidx, newsize);
+      memcpy(newbuff, B->b, B->n * sizeof(char));  /* copy original content */
     }
     B->b = newbuff;
     B->size = newsize;
diff --git a/testes/locals.lua b/testes/locals.lua
index 28f88e54..8766b3d5 100644
--- a/testes/locals.lua
+++ b/testes/locals.lua
@@ -324,6 +324,38 @@ if rawget(_G, "T") then
   local _, msg = pcall(test)
   assert(msg == 1000)
 
+
+  do    -- testing 'toclose' in C string buffer
+    local s = string.rep("a", 10000)
+    local a = {s, s}
+
+    -- ensure proper initialization (stack space, metatable)
+    table.concat(a)
+    collectgarbage(); collectgarbage()
+
+    local m = T.totalmem()
+
+    -- error in the second buffer allocation
+    T.alloccount(3)
+    assert(not pcall(table.concat, a))
+    T.alloccount()
+    -- first buffer was released by 'toclose'
+    assert(T.totalmem() - m <= 5000)
+
+    -- error in creation of final string
+    T.alloccount(4)
+    assert(not pcall(table.concat, a))
+    T.alloccount()
+    -- second buffer was released by 'toclose'
+    assert(T.totalmem() - m <= 5000)
+
+    -- userdata, upvalue, buffer, buffer, string
+    T.alloccount(5)
+    assert(#table.concat(a) == 20000)
+    T.alloccount()
+
+    print'+'
+  end
 end