Index: nse_nsock.cc =================================================================== --- nse_nsock.cc (revision 13614) +++ nse_nsock.cc (working copy) @@ -40,6 +40,8 @@ extern NmapOps o; +static int l_nsock_loop(lua_State * L); + static int l_nsock_connect(lua_State * L); static int l_nsock_send(lua_State * L); @@ -182,6 +184,15 @@ return len; } +static void weak_table(lua_State * L, int narr, int nrec, const char *mode) +{ + lua_createtable(L, narr, nrec); + lua_createtable(L, 0, 1); + lua_pushstring(L, mode); + lua_setfield(L, -2, "__mode"); + lua_setmetatable(L, -2); +} + static std::string hexify(const unsigned char *str, size_t len) { size_t num = 0; @@ -237,50 +248,21 @@ /* Some constants used for enforcing a limit on the number of open sockets * in use by threads. The maximum value between MAX_PARALLELISM and * o.maxparallelism is the max # of threads that can have connected sockets - * (open). THREAD_PROXY, SOCKET_PROXY, and CONNECT_WAITING are tables in the - * nsock C functions' environment, at LUA_ENVIRONINDEX, that hold sockets and - * threads used to enforce this. THREAD_PROXY has pairs - * that associate a thread to a proxy userdata. This table has weak keys and - * values so threads and the proxy itself can be collected. SOCKET_PROXY - * has pairs that associate a socket to a proxy userdata. - * SOCKET_PROXY has weak keys (to allow the collection of sockets) and strong - * values, so the proxies are not collected when an associated socket is open. + * (open). * - * All the sockets used by a thread have the same Proxy Userdata. When all - * sockets in use by a thread are closed or collected, the entry in the - * THREAD_PROXY table is cleared, freeing up a slot for another thread - * to make connections. When a slot is freed, proxy_gc is called, via the - * userdata's __gc metamethod, which will add a thread in WAITING to running. + * THREAD_SOCKETS is a weak keyed table of pairs. + * A socket table is a weak keyed table (socket keys with garbage values) of + * sockets the Thread has allocated but not necessarily open). You may + * test for an open socket by checking whether its nsiod field in the + * socket userdata structure is not NULL. + * + * CONNECT_WAITING is a weak keyed table of pairs. + * The table contains threads waiting to make a socket connection. */ #define MAX_PARALLELISM 10 -#define THREAD_PROXY 1 /* */ -#define SOCKET_PROXY 2 /* */ -#define CONNECT_WAITING 3 /* Threads waiting to lock */ -#define PROXY_META 4 /* Proxy userdata's metatable */ +#define THREAD_SOCKETS 1 /* */ +#define CONNECT_WAITING 2 /* Threads waiting to lock */ -static int proxy_gc(lua_State * L) -{ - lua_rawgeti(L, LUA_ENVIRONINDEX, CONNECT_WAITING); - lua_pushnil(L); - if (lua_next(L, -2) != 0) - { - lua_State *thread = lua_tothread(L, -2); - - nse_restore(thread, 0); - lua_pushnil(L); - lua_replace(L, -2); // replace boolean - lua_settable(L, -3); // remove thread from waiting - } - return 0; -} - -static void new_proxy(lua_State * L) -{ - lua_newuserdata(L, 0); - lua_rawgeti(L, LUA_ENVIRONINDEX, PROXY_META); - lua_setmetatable(L, -2); -} - /* int socket_lock (lua_State *L) * * Arguments @@ -292,65 +274,86 @@ static int socket_lock(lua_State * L) { lua_settop(L, 1); - lua_rawgeti(L, LUA_ENVIRONINDEX, THREAD_PROXY); + lua_rawgeti(L, LUA_ENVIRONINDEX, THREAD_SOCKETS); lua_pushthread(L); - lua_gettable(L, -2); - if (!lua_isnil(L, -1)) + lua_rawget(L, -2); + if (lua_istable(L, -1)) { - // Thread already has open sockets. Add the new socket to SOCKET_PROXY - lua_rawgeti(L, LUA_ENVIRONINDEX, SOCKET_PROXY); - lua_pushvalue(L, 1); // socket - lua_pushvalue(L, -3); // proxy userdata - lua_settable(L, -3); - lua_pop(L, 1); // SOCKET_PROXY + /* Thread already has a "lock" with open sockets. Place the new socket + * in its sockets table */ + lua_pushvalue(L, 1); lua_pushboolean(L, true); - } else if (table_length(L, 2) >= MAX(MAX_PARALLELISM, o.max_parallelism)) + lua_rawset(L, -3); + } else if (table_length(L, 2) <= MAX(MAX_PARALLELISM, o.max_parallelism)) { - // Too many threads have sockets open. Add thread to waiting. The caller - // is expected to yield. (see the connect function in luaopen_nsock) - lua_rawgeti(L, LUA_ENVIRONINDEX, CONNECT_WAITING); + /* There is room for this thread to open sockets */ lua_pushthread(L); + weak_table(L, 0, 0, "k"); /* weak socket references */ + lua_pushvalue(L, 1); /* socket */ lua_pushboolean(L, true); - lua_settable(L, -3); - lua_pop(L, 1); // CONNECT_WAITING - return lua_yield(L, 0); + lua_rawset(L, -3); /* add to sockets table */ + lua_rawset(L, 2); /* add new Pair + * to THREAD_SOCKETS */ } else { - // There is room for this thread to open sockets. Make a new proxy userdata - // and add it to the THREAD_PROXY and SOCKET_PROXY tables. - new_proxy(L); - lua_rawgeti(L, LUA_ENVIRONINDEX, THREAD_PROXY); + /* Too many threads have sockets open. Add thread to waiting. The caller + * is expected to yield. (see the connect function in luaopen_nsock) */ + lua_rawgeti(L, LUA_ENVIRONINDEX, CONNECT_WAITING); lua_pushthread(L); - lua_pushvalue(L, -3); // proxy - lua_settable(L, -3); - lua_pop(L, 1); // THREAD_PROXY) - lua_rawgeti(L, LUA_ENVIRONINDEX, SOCKET_PROXY); - lua_pushvalue(L, 1); // Socket - lua_pushvalue(L, -3); // proxy - lua_settable(L, -3); - lua_pop(L, 2); // proxy, SOCKET_PROXY lua_pushboolean(L, true); + lua_settable(L, -3); + return lua_yield(L, 0); } + lua_pushboolean(L, true); return 1; } -/* void socket_unlock (lua_State *L, int index) - * - * index is the location of the userdata on the stack. - * A socket has been closed or collected, remove it from the SOCKET_PROXY - * table. - */ -static void socket_unlock(lua_State * L, int index) +static void socket_unlock(lua_State * L) { - lua_pushvalue(L, index); // socket - lua_rawgeti(L, LUA_ENVIRONINDEX, SOCKET_PROXY); - lua_pushvalue(L, -2); // socket + int top = lua_gettop(L); + + lua_gc(L, LUA_GCSTOP, 0); /* don't collect threads during iteration */ + + lua_rawgeti(L, LUA_ENVIRONINDEX, THREAD_SOCKETS); lua_pushnil(L); - lua_settable(L, -3); - lua_pop(L, 2); // socket, SOCKET_PROXY -} + while (lua_next(L, -2) != 0) + { + unsigned open = 0; + lua_pushnil(L); + while (lua_next(L, -2) != 0) /* for each socket */ + { + lua_pop(L, 1); /* pop garbage boolean */ + if (((struct l_nsock_udata *) lua_touserdata(L, -1))->nsiod != NULL) + open++; + } + if (open == 0) /* thread has no open sockets? */ + { + lua_pushvalue(L, -2); /* thread key */ + lua_pushnil(L); + lua_rawset(L, top+1); /* THREADS_SOCKETS */ + + lua_rawgeti(L, LUA_ENVIRONINDEX, CONNECT_WAITING); + lua_pushnil(L); + if (lua_next(L, -2) != 0) + { + lua_pop(L, 1); /* pop garbage boolean */ + nse_restore(lua_tothread(L, -1), 0); + lua_pushnil(L); + lua_rawset(L, -3); /* remove thread from waiting */ + } + lua_pop(L, 1); /* CONNECT_WAITING */ + } + + lua_pop(L, 1); /* pop sockets table */ + } + + lua_gc(L, LUA_GCRESTART, 0); + + lua_settop(L, top); +} + void l_nsock_clear_buf(lua_State * L, l_nsock_udata * udata); int luaopen_nsock(lua_State * L) @@ -382,36 +385,22 @@ {NULL, NULL} }; - /* Set up an environment for all nsock C functions to share. This is - * especially important to make the THREAD_PROXY, SOCKET_PROXY, and - * CONNECT_WAITING tables available. These values can be accessed at the - * pseudo-index LUA_ENVIRONINDEX. These tables are documented where the - * #defines are above. */ - lua_createtable(L, 5, 0); + /* Set up an environment for all nsock C functions to share. + * This table particularly contains the THREAD_SOCKETS and + * CONNECT_WAITING tables. + * These values are accessed at the Lua pseudo-index LUA_ENVIRONINDEX. + */ + lua_createtable(L, 2, 0); lua_replace(L, LUA_ENVIRONINDEX); - lua_createtable(L, 0, 10); // THREAD_PROXY - lua_createtable(L, 0, 1); // metatable - lua_pushliteral(L, "kv"); // weak keys and values - lua_setfield(L, -2, "__mode"); - lua_setmetatable(L, -2); - lua_rawseti(L, LUA_ENVIRONINDEX, THREAD_PROXY); + weak_table(L, 0, MAX_PARALLELISM, "k"); + lua_rawseti(L, LUA_ENVIRONINDEX, THREAD_SOCKETS); - lua_createtable(L, 0, 193); // SOCKET_PROXY (large amount of room) - lua_createtable(L, 0, 1); // metatable - lua_pushliteral(L, "k"); // weak keys - lua_setfield(L, -2, "__mode"); - lua_setmetatable(L, -2); - lua_rawseti(L, LUA_ENVIRONINDEX, SOCKET_PROXY); - - lua_createtable(L, 0, 499); // CONNECT_WAITING (large amount of - // room) + weak_table(L, 0, 1000, "k"); lua_rawseti(L, LUA_ENVIRONINDEX, CONNECT_WAITING); - lua_createtable(L, 0, 1); // PROXY_META = metatable for proxies - lua_pushcclosure(L, proxy_gc, 0); - lua_setfield(L, -2, "__gc"); - lua_rawseti(L, LUA_ENVIRONINDEX, PROXY_META); + lua_pushcfunction(L, l_nsock_loop); + lua_setfield(L, LUA_REGISTRYINDEX, NSE_NSOCK_LOOP); /* Load the connect function */ if (luaL_loadstring(L, connect) != 0) @@ -476,9 +465,16 @@ return 1; } -int l_nsock_loop(int tout) +static int l_nsock_loop(lua_State * L) { - return nsock_loop(nsp, tout); + int tout = luaL_checkint(L, 1); + + /* clean up old socket locks */ + socket_unlock(L); + + if (nsock_loop(nsp, tout) == NSOCK_LOOP_ERROR) + luaL_error(L, "a fatal error occurred in nsock_loop"); + return 0; } int l_nsock_checkstatus(lua_State * L, nsock_event nse) @@ -513,26 +509,32 @@ static int l_nsock_connect(lua_State * L) { - l_nsock_udata *udata = (l_nsock_udata *) luaL_checkudata(L, 1, "nsock"); + enum type {TCP, UDP, SSL}; + static const char * const op[] = {"tcp", "udp", "ssl", NULL}; + l_nsock_udata *udata = (l_nsock_udata *) luaL_checkudata(L, 1, "nsock"); const char *addr = luaL_checkstring(L, 2); - unsigned short port = (unsigned short) luaL_checkint(L, 3); + int what = luaL_checkoption(L, 4, "tcp", op); - const char *how = luaL_optstring(L, 4, "tcp"); - const char *error; - struct addrinfo *dest; - int error_id; l_nsock_clear_buf(L, udata); +#ifndef HAVE_OPENSSL + if (what == SSL) + { + lua_pushboolean(L, false); + lua_pushstring(L, "Sorry, you don't have OpenSSL\n"); + return 2; + } +#endif + error_id = getaddrinfo(addr, NULL, NULL, &dest); if (error_id) { - socket_unlock(L, 1); error = gai_strerror(error_id); lua_pushboolean(L, false); lua_pushstring(L, error); @@ -540,7 +542,6 @@ } if (dest == NULL) { - socket_unlock(L, 1); lua_pushboolean(L, false); lua_pushstring(L, "getaddrinfo returned success but no addresses"); return 2; @@ -559,48 +560,26 @@ if (o.ipoptionslen) nsi_set_ipoptions(udata->nsiod, o.ipoptions, o.ipoptionslen); - switch (how[0]) + switch (what) { - case 't': - if (strcmp(how, "tcp")) - goto error; + case TCP: nsock_connect_tcp(nsp, udata->nsiod, l_nsock_connect_handler, udata->timeout, &udata->yield, dest->ai_addr, dest->ai_addrlen, port); break; - case 'u': - if (strcmp(how, "udp")) - goto error; + case UDP: nsock_connect_udp(nsp, udata->nsiod, l_nsock_connect_handler, &udata->yield, dest->ai_addr, dest->ai_addrlen, port); break; - case 's': - if (strcmp(how, "ssl")) - goto error; -#ifdef HAVE_OPENSSL + case SSL: nsock_connect_ssl(nsp, udata->nsiod, l_nsock_connect_handler, udata->timeout, &udata->yield, dest->ai_addr, dest->ai_addrlen, port, udata->ssl_session); break; -#else - socket_unlock(L, 1); - lua_pushboolean(L, false); - lua_pushstring(L, "Sorry, you don't have OpenSSL\n"); - return 2; -#endif - default: - goto error; - break; } freeaddrinfo(dest); set_thread(L, 1, udata); return lua_yield(L, 0); - -error: - socket_unlock(L, 1); - freeaddrinfo(dest); - luaL_argerror(L, 4, "invalid connection method"); - return 0; } void l_nsock_connect_handler(nsock_pool nsp, nsock_event nse, void *yield) @@ -900,8 +879,6 @@ { l_nsock_udata *udata = (l_nsock_udata *) luaL_checkudata(L, 1, "nsock"); - socket_unlock(L, 1); // Unlock the socket. - /* Never ever collect nse-pcap connections. */ if (udata->ncap_socket) { Index: nse_main.cc =================================================================== --- nse_main.cc (revision 13614) +++ nse_main.cc (working copy) @@ -154,8 +154,10 @@ static int nsock_loop (lua_State *L) { - if (l_nsock_loop(luaL_checkint(L, 1)) == NSOCK_LOOP_ERROR) - luaL_error(L, "an error occurred in nsock_loop"); + lua_settop(L, 1); + lua_getfield(L, LUA_REGISTRYINDEX, NSE_NSOCK_LOOP); + lua_pushvalue(L, 1); + lua_call(L, 1, 0); return 0; } Index: nse_nsock.h =================================================================== --- nse_nsock.h (revision 13614) +++ nse_nsock.h (working copy) @@ -3,11 +3,12 @@ int luaopen_nsock(lua_State *); int l_nsock_new(lua_State *); -int l_nsock_loop(int tout); int l_nsock_sleep(lua_State *L); int l_dnet_new(lua_State *); int l_dnet_get_interface_link(lua_State *); +#define NSE_NSOCK_LOOP "NSOCK_LOOP" + #endif Index: nse_main.lua =================================================================== --- nse_main.lua (revision 13614) +++ nse_main.lua (working copy) @@ -478,6 +478,11 @@ end end + -- Move pending threads back to running. + for co, thread in pairs(pending) do + pending[co], running[co] = nil, thread; + end + for co, thread in pairs(running) do current, running[co] = thread, nil; cnse.startTimeOutClock(thread.host); @@ -516,12 +521,7 @@ end end - -- Move pending threads back to running. - for co, thread in pairs(pending) do - pending[co], running[co] = nil, thread; - end - - collectgarbage "collect"; -- important for collecting used sockets & proxies + collectgarbage "step"; end progress "endTask";