2 files changed, 32 insertions, 6 deletions

diff --git a/src/lfs.c b/src/lfs.c
index 060034c..13562ee 100644
--- a/src/lfs.c
+++ b/src/lfs.c
@@ -17,7 +17,7 @@
 **   lfs.touch (filepath [, atime [, mtime]])
 **   lfs.unlock (fh)
 **
-** $Id: lfs.c,v 1.59 2009/04/24 22:24:07 mascarenhas Exp $
+** $Id: lfs.c,v 1.60 2009/06/03 20:49:18 mascarenhas Exp $
 */
 
 #ifndef _WIN32
@@ -431,7 +431,7 @@ static int dir_iter (lua_State *L) {
 #else
         struct dirent *entry;
 #endif
-        dir_data *d = (dir_data *)lua_touserdata (L, lua_upvalueindex (1));
+        dir_data *d = (dir_data *)luaL_checkudata (L, 1, DIR_METATABLE);
         luaL_argcheck (L, !d->closed, 1, "closed directory");
 #ifdef _WIN32
         if (d->hFile == 0L) { /* first entry */
@@ -493,7 +493,9 @@ static int dir_close (lua_State *L) {
 */
 static int dir_iter_factory (lua_State *L) {
         const char *path = luaL_checkstring (L, 1);
-        dir_data *d = (dir_data *) lua_newuserdata (L, sizeof(dir_data));
+        dir_data *d;
+        lua_pushcfunction (L, dir_iter);
+        d = (dir_data *) lua_newuserdata (L, sizeof(dir_data));
         d->closed = 0;
 #ifdef _WIN32
         d->hFile = 0L;
@@ -510,8 +512,7 @@ static int dir_iter_factory (lua_State *L) {
         if (d->dir == NULL)
                 luaL_error (L, "cannot open %s: %s", path, strerror (errno));
 #endif
-        lua_pushcclosure (L, dir_iter, 1);
-        return 1;
+        return 2;
 }
 
 
@@ -521,10 +522,18 @@ static int dir_iter_factory (lua_State *L) {
 static int dir_create_meta (lua_State *L) {
         luaL_newmetatable (L, DIR_METATABLE);
         /* set its __gc field */
+        lua_pushstring (L, "__index");
+        lua_newtable(L);
+        lua_pushstring (L, "next");
+        lua_pushcfunction (L, dir_iter);
+        lua_settable(L, -3);
+        lua_pushstring (L, "close");
+        lua_pushcfunction (L, dir_close);
+        lua_settable(L, -3);
+        lua_settable (L, -3);
         lua_pushstring (L, "__gc");
         lua_pushcfunction (L, dir_close);
         lua_settable (L, -3);
-
         return 1;
 }
 
diff --git a/tests/test.lua b/tests/test.lua
index 1fd6157..7111074 100644
--- a/tests/test.lua
+++ b/tests/test.lua
@@ -110,4 +110,21 @@ for i = 1, 4000 do
                 count = count + 1
         end
 end
+
+-- Stressing directory iterator, explicit version
+count = 0
+for i = 1, 4000 do
+  local iter, dir = lfs.dir(tmp)
+  local file = dir:next()
+  while file do
+    count = count + 1
+    file = dir:next()
+  end
+  assert(not pcall(dir.next, dir))
+end
+
+-- directory explicit close
+local iter, dir = lfs.dir(tmp)
+dir:close()
+assert(not pcall(dir.next, dir))
 print"Ok!"