Add support to multiples resources and socket select

luawebsocket.c

@@ -161,6 +161,25 @@ websocket_bind(lua_State *L)
 	return 1;
 }
 
+static int
+check_resources(lua_State *L, char *resource)
+{
+	if (lua_isstring(L, 2) && !strcmp(resource, luaL_checkstring(L, 2)))
+		return 1;
+
+	if (lua_istable(L, 2)) {
+		lua_pushnil(L);
+		while (lua_next(L, 2) != 0) {
+			if (!strcmp(resource, luaL_checkstring(L, -1)))
+				return 1;
+
+			lua_pop(L, 1);
+		}
+	}
+
+	return 0;
+}
+
 static int
 websocket_handshake(lua_State *L)
 {
@@ -181,15 +200,15 @@ websocket_handshake(lua_State *L)
 
 	if (wsParseHandshake((unsigned char *)buf, nread, &hs) ==
 	    WS_OPENING_FRAME) {
-		if (!strcmp(hs.resource, luaL_checkstring(L, 2))) {
+		if (check_resources(L, hs.resource)) {
 			wsGetHandshakeAnswer(&hs, (unsigned char *)buf, &nread);
+			lua_pushstring(L, hs.resource);
 			freeHandshake(&hs);
 			if (websock->ssl)
 				SSL_write(websock->ssl, buf, nread);
 			else
 				send(websock->socket, buf, nread, 0);
 			buf[nread] = '\0';
-			lua_pushboolean(L, 1);
 		} else {
 			nread = sprintf(buf, "HTTP/1.1 404 Not Found\r\n\r\n");
 			if (websock->ssl)
@@ -358,6 +377,95 @@ websocket_shutdown(lua_State *L)
 	return 0;
 }
 
+static int
+websocket_select(lua_State *L) {
+	WEBSOCKET *websock;
+	fd_set readfds, writefds;
+	int ret, timeout, rtop, wtop, max_fd;
+	struct timeval tv;
+
+	FD_ZERO(&readfds); FD_ZERO(&writefds);
+
+	websock = luaL_checkudata(L, 1, WEBSOCKET_METATABLE);
+
+	FD_SET(websock->socket, &readfds); FD_SET(websock->socket, &writefds);
+	max_fd = websock->socket;
+
+	if (lua_istable(L, 2)) {
+		lua_pushnil(L);
+		while (lua_next(L, 2) != 0) {
+			if (lua_isnumber(L, -1)) {
+				int fd = (int) lua_tonumber(L, -1);
+				if (fd > 0)
+					FD_SET(fd, &readfds);
+				if (fd > max_fd)
+					max_fd = fd;
+			}
+			lua_pop(L, 1);
+		}
+	}
+
+	if (lua_istable(L, 3)) {
+		lua_pushnil(L);
+		while (lua_next(L, 3) != 0) {
+			if (lua_isnumber(L, -1)) {
+				int fd = (int) lua_tonumber(L, -1);
+				if (fd > 0)
+					FD_SET(fd, &writefds);
+				if (fd > max_fd)
+					max_fd = fd;
+			}
+			lua_pop(L, 1);
+		}
+	}
+
+	if (!lua_isnil(L, 4) && lua_isnumber(L, 4))
+		timeout = (int) lua_tonumber(L, 4);
+	else
+		timeout = -1;
+
+	tv.tv_sec = (int) lua_tonumber(L, 4);
+	tv.tv_usec = 0;
+
+	ret = select(max_fd + 1, &readfds, &writefds, NULL, timeout >= 0 ? &tv : NULL);
+	if (ret > 0) {
+		if (FD_ISSET(websock->socket, &readfds)) {
+			lua_pushinteger(L, websock->socket);
+			lua_pushnil(L);
+			return 2;
+		}
+
+		lua_newtable(L); rtop = lua_gettop(L);
+		lua_newtable(L); wtop = lua_gettop(L);
+
+		int fd;
+		int ir = 0, iw = 0;
+		for (fd = 0; fd < max_fd + 1; fd++) {
+			if (FD_ISSET(fd, &readfds)) {
+				lua_pushnumber(L, (lua_Number) fd);
+				lua_rawseti(L, rtop, ir + 1);
+				ir++;
+			} else if (FD_ISSET(fd, &writefds)){
+				lua_pushnumber(L, (lua_Number) fd);
+				lua_rawseti(L, wtop, iw + 1);
+				iw++;
+			}
+		}
+
+		return 2;
+	} else if (ret == 0) {
+		lua_pushnil(L);
+		lua_pushnil(L);
+		lua_pushstring(L, "timeout");
+		return 3;
+	}
+
+	lua_pushnil(L);
+	lua_pushnil(L);
+	lua_pushstring(L, "select failed");
+	return 3;
+}
+
 int
 luaopen_websocket(lua_State *L)
 {
@@ -373,6 +481,7 @@ luaopen_websocket(lua_State *L)
 		{ "recv", 		websocket_recv},
 		{ "send",		websocket_send },
 		{ "socket",		websocket_socket },
+		{ "select",		websocket_select },
 		{ NULL, NULL }
 	};
 	if (luaL_newmetatable(L, WEBSOCKET_METATABLE)) {