select.c 6.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220
  1. /*=========================================================================*\
  2. * Select implementation
  3. * LuaSocket toolkit
  4. \*=========================================================================*/
  5. #include <string.h>
  6. #include "lua.h"
  7. #include "lauxlib.h"
  8. #include "compat.h"
  9. #include "socket.h"
  10. #include "timeout.h"
  11. #include "select.h"
  12. /*=========================================================================*\
  13. * Internal function prototypes.
  14. \*=========================================================================*/
  15. static t_socket getfd(lua_State *L);
  16. static int dirty(lua_State *L);
  17. static void collect_fd(lua_State *L, int tab, int itab,
  18. fd_set *set, t_socket *max_fd);
  19. static int check_dirty(lua_State *L, int tab, int dtab, fd_set *set);
  20. static void return_fd(lua_State *L, fd_set *set, t_socket max_fd,
  21. int itab, int tab, int start);
  22. static void make_assoc(lua_State *L, int tab);
  23. static int global_select(lua_State *L);
  24. /* functions in library namespace */
  25. static luaL_Reg func[] = {
  26. {"select", global_select},
  27. {NULL, NULL}
  28. };
  29. /*=========================================================================*\
  30. * Exported functions
  31. \*=========================================================================*/
  32. /*-------------------------------------------------------------------------*\
  33. * Initializes module
  34. \*-------------------------------------------------------------------------*/
  35. int select_open(lua_State *L) {
  36. lua_pushstring(L, "_SETSIZE");
  37. lua_pushinteger(L, FD_SETSIZE);
  38. lua_rawset(L, -3);
  39. lua_pushstring(L, "_SOCKETINVALID");
  40. lua_pushinteger(L, SOCKET_INVALID);
  41. lua_rawset(L, -3);
  42. luaL_setfuncs(L, func, 0);
  43. return 0;
  44. }
  45. /*=========================================================================*\
  46. * Global Lua functions
  47. \*=========================================================================*/
  48. /*-------------------------------------------------------------------------*\
  49. * Waits for a set of sockets until a condition is met or timeout.
  50. \*-------------------------------------------------------------------------*/
  51. static int global_select(lua_State *L) {
  52. int rtab, wtab, itab, ret, ndirty;
  53. t_socket max_fd = SOCKET_INVALID;
  54. fd_set rset, wset;
  55. t_timeout tm;
  56. double t = luaL_optnumber(L, 3, -1);
  57. FD_ZERO(&rset); FD_ZERO(&wset);
  58. lua_settop(L, 3);
  59. lua_newtable(L); itab = lua_gettop(L);
  60. lua_newtable(L); rtab = lua_gettop(L);
  61. lua_newtable(L); wtab = lua_gettop(L);
  62. collect_fd(L, 1, itab, &rset, &max_fd);
  63. collect_fd(L, 2, itab, &wset, &max_fd);
  64. ndirty = check_dirty(L, 1, rtab, &rset);
  65. t = ndirty > 0? 0.0: t;
  66. timeout_init(&tm, t, -1);
  67. timeout_markstart(&tm);
  68. ret = socket_select(max_fd+1, &rset, &wset, NULL, &tm);
  69. if (ret > 0 || ndirty > 0) {
  70. return_fd(L, &rset, max_fd+1, itab, rtab, ndirty);
  71. return_fd(L, &wset, max_fd+1, itab, wtab, 0);
  72. make_assoc(L, rtab);
  73. make_assoc(L, wtab);
  74. return 2;
  75. } else if (ret == 0) {
  76. lua_pushstring(L, "timeout");
  77. return 3;
  78. } else {
  79. luaL_error(L, "select failed");
  80. return 3;
  81. }
  82. }
  83. /*=========================================================================*\
  84. * Internal functions
  85. \*=========================================================================*/
  86. static t_socket getfd(lua_State *L) {
  87. t_socket fd = SOCKET_INVALID;
  88. lua_pushstring(L, "getfd");
  89. lua_gettable(L, -2);
  90. if (!lua_isnil(L, -1)) {
  91. lua_pushvalue(L, -2);
  92. lua_call(L, 1, 1);
  93. if (lua_isnumber(L, -1)) {
  94. double numfd = lua_tonumber(L, -1);
  95. fd = (numfd >= 0.0)? (t_socket) numfd: SOCKET_INVALID;
  96. }
  97. }
  98. lua_pop(L, 1);
  99. return fd;
  100. }
  101. static int dirty(lua_State *L) {
  102. int is = 0;
  103. lua_pushstring(L, "dirty");
  104. lua_gettable(L, -2);
  105. if (!lua_isnil(L, -1)) {
  106. lua_pushvalue(L, -2);
  107. lua_call(L, 1, 1);
  108. is = lua_toboolean(L, -1);
  109. }
  110. lua_pop(L, 1);
  111. return is;
  112. }
  113. static void collect_fd(lua_State *L, int tab, int itab,
  114. fd_set *set, t_socket *max_fd) {
  115. int i = 1, n = 0;
  116. /* nil is the same as an empty table */
  117. if (lua_isnil(L, tab)) return;
  118. /* otherwise we need it to be a table */
  119. luaL_checktype(L, tab, LUA_TTABLE);
  120. for ( ;; ) {
  121. t_socket fd;
  122. lua_pushnumber(L, i);
  123. lua_gettable(L, tab);
  124. if (lua_isnil(L, -1)) {
  125. lua_pop(L, 1);
  126. break;
  127. }
  128. /* getfd figures out if this is a socket */
  129. fd = getfd(L);
  130. if (fd != SOCKET_INVALID) {
  131. /* make sure we don't overflow the fd_set */
  132. #ifdef _WIN32
  133. if (n >= FD_SETSIZE)
  134. luaL_argerror(L, tab, "too many sockets");
  135. #else
  136. if (fd >= FD_SETSIZE)
  137. luaL_argerror(L, tab, "descriptor too large for set size");
  138. #endif
  139. FD_SET(fd, set);
  140. n++;
  141. /* keep track of the largest descriptor so far */
  142. if (*max_fd == SOCKET_INVALID || *max_fd < fd)
  143. *max_fd = fd;
  144. /* make sure we can map back from descriptor to the object */
  145. lua_pushnumber(L, (lua_Number) fd);
  146. lua_pushvalue(L, -2);
  147. lua_settable(L, itab);
  148. }
  149. lua_pop(L, 1);
  150. i = i + 1;
  151. }
  152. }
  153. static int check_dirty(lua_State *L, int tab, int dtab, fd_set *set) {
  154. int ndirty = 0, i = 1;
  155. if (lua_isnil(L, tab))
  156. return 0;
  157. for ( ;; ) {
  158. t_socket fd;
  159. lua_pushnumber(L, i);
  160. lua_gettable(L, tab);
  161. if (lua_isnil(L, -1)) {
  162. lua_pop(L, 1);
  163. break;
  164. }
  165. fd = getfd(L);
  166. if (fd != SOCKET_INVALID && dirty(L)) {
  167. lua_pushnumber(L, ++ndirty);
  168. lua_pushvalue(L, -2);
  169. lua_settable(L, dtab);
  170. FD_CLR(fd, set);
  171. }
  172. lua_pop(L, 1);
  173. i = i + 1;
  174. }
  175. return ndirty;
  176. }
  177. static void return_fd(lua_State *L, fd_set *set, t_socket max_fd,
  178. int itab, int tab, int start) {
  179. t_socket fd;
  180. for (fd = 0; fd < max_fd; fd++) {
  181. if (FD_ISSET(fd, set)) {
  182. lua_pushnumber(L, ++start);
  183. lua_pushnumber(L, (lua_Number) fd);
  184. lua_gettable(L, itab);
  185. lua_settable(L, tab);
  186. }
  187. }
  188. }
  189. static void make_assoc(lua_State *L, int tab) {
  190. int i = 1, atab;
  191. lua_newtable(L); atab = lua_gettop(L);
  192. for ( ;; ) {
  193. lua_pushnumber(L, i);
  194. lua_gettable(L, tab);
  195. if (!lua_isnil(L, -1)) {
  196. lua_pushnumber(L, i);
  197. lua_pushvalue(L, -2);
  198. lua_settable(L, atab);
  199. lua_pushnumber(L, i);
  200. lua_settable(L, atab);
  201. } else {
  202. lua_pop(L, 1);
  203. break;
  204. }
  205. i = i+1;
  206. }
  207. }