[NSE] When a script ends for any reason, all mutexes are now unlocked.

Some scripts would fail due to an error (whois.nse) causing other
scripts to become deadlocked on a mutex that would never unlock. This
patch fixes this problem. See [1] for more information.

[1] http://seclists.org/nmap-dev/2009/q2/0533.html

nse_macros.h

@@ -1,8 +1,6 @@
 #ifndef NSE_MACROS
 #define NSE_MACROS
 
-#define NSE_WAITING_TO_RUNNING "WAITING_TO_RUNNING"
-
 #define HOSTRULE	"hostrule"
 #define HOSTTESTS	"hosttests"
 #define PORTRULE	"portrule"

nse_main.cc

@@ -23,6 +23,10 @@
 #define NSE_MAIN "NSE_MAIN" /* the main function */
 #define NSE_TRACEBACK "NSE_TRACEBACK"
 
+/* string keys used in interface with nse_main.lua */
+#define NSE_WAITING_TO_RUNNING "NSE_WAITING_TO_RUNNING"
+#define NSE_DESTRUCTOR "NSE_DESTRUCTOR"
+
 extern NmapOps o;
 
 int current_hosts = LUA_NOREF;
@@ -472,6 +476,35 @@ void nse_restore (lua_State *L, int number)
     fatal("nse_restore: WAITING_TO_RUNNING error!\n%s", lua_tostring(L, -1));
 }
 
+/* This function adds (what = 'a') or removes (what 'r') a destructor
+ * from the Thread owning the running Lua thread (L). We call the nse_main.lua
+ * function _R.NSE_DESTRUCTOR in order to add (or remove) the destructor to
+ * the Thread's close handler table.
+ *
+ * what == 'r', destructor key on stack
+ * what == 'a', destructor key and destructor on stack
+ */
+void nse_destructor (lua_State *L, char what)
+{
+  assert(what == 'a' || what == 'r');
+  lua_getfield(L, LUA_REGISTRYINDEX, NSE_DESTRUCTOR);
+  lua_pushstring(L, what == 'a' ? "add" : "remove");
+  lua_pushthread(L);
+  if (what == 'a')
+  {
+    lua_pushvalue(L, -5); /* destructor key */
+    lua_pushvalue(L, -5); /* destructor */
+  }
+  else
+  {
+    lua_pushvalue(L, -4); /* destructor key */
+    lua_pushnil(L); /* no destructor, we are removing */
+  }
+  if (lua_pcall(L, 4, 0, 0) != 0)
+    fatal("nse_destructor: NSE_DESTRUCTOR error!\n%s", lua_tostring(L, -1));
+  lua_pop(L, what == 'a' ? 2 : 1);
+}
+
 static lua_State *L_NSE = NULL;
 
 int open_nse (void)

nse_main.h

@@ -32,7 +32,9 @@ class Target;
 int script_updatedb();
 void script_scan_free();
 
+/* API */
 void nse_restore (lua_State *, int);
+void nse_destructor (lua_State *, char);
 
 int open_nse (void);
 int script_scan(std::vector<Target *> &targets);

nse_main.lua

@@ -26,6 +26,9 @@
 
 local NAME = "NSE";
 
+local WAITING_TO_RUNNING = "NSE_WAITING_TO_RUNNING";
+local DESTRUCTOR = "NSE_DESTRUCTOR";
+
 local _R = debug.getregistry(); -- The registry
 local _G = _G;
 
@@ -128,6 +131,14 @@ do
     end
   end
 
+  function Thread:close ()
+    local ch = self.close_handlers;
+    for key, destructor_t in pairs(ch) do
+      destructor_t.destructor(destructor_t.thread, key);
+      ch[key] = nil;
+    end
+  end
+
   -- thread = Script:new_thread(rule, ...)
   -- Creates a new thread for the script Script.
   -- Arguments:
@@ -166,6 +177,7 @@ do
         identifier = tostring(co),
         info = format("'%s' (%s)", self.short_basename, tostring(co));
         type = rule == "hostrule" and "host" or "port",
+        close_handlers = {},
       }, {
         __metatable = Thread,
         __index = function (thread, k) return Thread[k] or self[k] end
@@ -396,6 +408,7 @@ local function run (threads)
   -- nse_restore, waiting threads become pending and later are moved all
   -- at once back to running.
   local running, waiting, pending = {}, {}, {};
+  local all = setmetatable({}, {__mode = "kv"}); -- base coroutine to Thread
   -- hosts maps a host to a list of threads for that host.
   local hosts, total = {}, 0;
   local current;
@@ -406,19 +419,34 @@ local function run (threads)
     local thread = remove(threads);
     thread:d("Starting %THREAD against %s%s.", thread.host.ip,
         thread.port and ":"..thread.port.number or "");
-    running[thread.co], total = thread, total + 1;
+    all[thread.co], running[thread.co], total = thread, thread, total+1;
     hosts[thread.host] = hosts[thread.host] or {};
     hosts[thread.host][thread.co] = true;
   end
 
-  -- This WAITING_TO_RUNNING function is called by nse_restore in
-  -- nse_main.cc.
-  _R.WAITING_TO_RUNNING = function (co, ...)
+  -- _R[WAITING_TO_RUNNING] is called by nse_restore in nse_main.cc
+  _R[WAITING_TO_RUNNING] = function (co, ...)
     if waiting[co] then -- ignore a thread not waiting
       pending[co], waiting[co] = waiting[co], nil;
       pending[co].args = {n = select("#", ...), ...};
     end
   end
+  -- _R[DESTRUCTOR] is called by nse_destructor in nse_main.cc
+  _R[DESTRUCTOR] = function (what, co, key, destructor)
+    local thread = all[co] or current;
+    print(thread, what, co, key, destructor)
+    if thread then
+      local ch = thread.close_handlers;
+      if what == "add" then
+        ch[key] = {
+          thread = co,
+          destructor = destructor
+        };
+      elseif what == "remove" then
+        ch[key] = nil;
+      end
+    end
+  end
 
   -- Loop while any thread is running or waiting.
   while next(running) or next(waiting) do
@@ -441,6 +469,7 @@ local function run (threads)
       if cnse.timedOut(thread.host) then
         waiting[co] = nil;
         thread:d("%THREAD target timed out");
+        thread:close();
       end
     end
 
@@ -453,6 +482,7 @@ local function run (threads)
         hosts[thread.host][co] = nil;
         thread:d("%THREAD threw an error!\n%s\n",
             traceback(co, tostring(result)));
+        thread:close();
       elseif status(co) == "suspended" then
         waiting[co] = thread;
       elseif status(co) == "dead" then
@@ -472,6 +502,7 @@ local function run (threads)
         end
         thread:d("Finished %THREAD against %s%s.", thread.host.ip,
             thread.port and ":"..thread.port.number or "");
+        thread:close();
       end
 
       -- Any more threads running for this host?

nse_nmaplib.cc

@@ -216,6 +216,13 @@ static int l_clock_ms (lua_State *L)
   return 1;
 }
 
+/* The actual mutex returned by the nmap.mutex function.
+ * This function has 4 upvalues:
+ * (1) Table (array) of waiting threads.
+ * (2) The running thread or nil.
+ * (3) A unique table key used for destructors.
+ * (4) The destructor function, aux_mutex_done.
+ */
 static int aux_mutex (lua_State *L)
 {
   enum what {LOCK, DONE, TRYLOCK, RUNNING};
@@ -227,6 +234,9 @@ static int aux_mutex (lua_State *L)
       {
         lua_pushthread(L);
         lua_replace(L, lua_upvalueindex(2)); // set running
+        lua_pushvalue(L, lua_upvalueindex(3)); // unique identifier
+        lua_pushvalue(L, lua_upvalueindex(4)); // aux_mutex_done closure
+        nse_destructor(L, 'a');
         return 0;
       }
       lua_pushthread(L);
@@ -236,6 +246,10 @@ static int aux_mutex (lua_State *L)
       lua_pushthread(L);
       if (!lua_equal(L, -1, lua_upvalueindex(2)))
         luaL_error(L, "%s", "do not have a lock on this mutex");
+      /* remove destructor */
+      lua_pushvalue(L, lua_upvalueindex(3));
+      nse_destructor(L, 'r');
+      /* set new thread to lock the mutex */
       lua_getfield(L, LUA_REGISTRYINDEX, "_LOADED");
       lua_getfield(L, -1, "table");
       lua_getfield(L, -1, "remove");
@@ -244,7 +258,14 @@ static int aux_mutex (lua_State *L)
       lua_call(L, 2, 1);
       lua_replace(L, lua_upvalueindex(2));
       if (lua_isthread(L, lua_upvalueindex(2))) // waiting threads had a thread
-        nse_restore(lua_tothread(L, lua_upvalueindex(2)), 0);
+      {
+        lua_State *thread = lua_tothread(L, lua_upvalueindex(2));
+        lua_pushvalue(L, lua_upvalueindex(3)); // destructor key
+        lua_pushvalue(L, lua_upvalueindex(4)); // destructor
+        lua_xmove(L, thread, 2);
+        nse_destructor(thread, 'a');
+        nse_restore(thread, 0);
+      }
       return 0;
     case TRYLOCK:
       if (lua_isnil(L, lua_upvalueindex(2)))
@@ -263,6 +284,20 @@ static int aux_mutex (lua_State *L)
   return 0;
 }
 
+/* This is the mutex destructor called when a thread ends but failed to 
+ * unlock the mutex.
+ * It has 1 upvalue: The nmap.mutex function closure.
+ */
+static int aux_mutex_done (lua_State *L)
+{
+  lua_State *thread = lua_tothread(L, 1);
+  lua_pushvalue(L, lua_upvalueindex(1)); // aux_mutex, actual mutex closure
+  lua_pushliteral(L, "done");
+  lua_xmove(L, thread, 2);
+  if (lua_pcall(thread, 1, 0, 0) != 0) lua_pop(thread, 1); // pop error msg
+  return 0;
+}
+
 static int l_mutex (lua_State *L)
 {
   int t = lua_type(L, 1);
@@ -274,9 +309,14 @@ static int l_mutex (lua_State *L)
   {
     lua_newtable(L); // waiting threads
     lua_pushnil(L); // running thread
-    lua_pushcclosure(L, aux_mutex, 2);
+    lua_newtable(L); // unique object as an identifier
+    lua_pushnil(L); // placeholder for aux_mutex_done
+    lua_pushcclosure(L, aux_mutex, 4);
+    lua_pushvalue(L, -1); // mutex closure
+    lua_pushcclosure(L, aux_mutex_done, 1);
+    lua_setupvalue(L, -2, 4); // replace nil upvalue with aux_mutex_done
     lua_pushvalue(L, 1); // "mutex object"
-    lua_pushvalue(L, -2); // function
+    lua_pushvalue(L, -2); // mutex function
     lua_settable(L, lua_upvalueindex(1)); // Add to mutex table
   }
   return 1; // aux_mutex closure