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