Bind select(2)

3 files changed, 277 insertions, 1 deletions

diff --git a/luxio.c b/luxio.c
index d5721c9..102b772 100644
--- a/luxio.c
+++ b/luxio.c
@@ -53,6 +53,7 @@ Not all systems will provide all the functions described here.
 #include <net/if.h>
 #include <fcntl.h>
 #include <poll.h>
+#include <sys/select.h>
 #include <unistd.h>
 #include <dirent.h>
 #include <time.h>
@@ -72,6 +73,25 @@ Not all systems will provide all the functions described here.
 
 #if (LUA_VERSION_NUM == 501)
 # define lua_rawlen(L, idx) lua_objlen((L), (idx))
+
+/* Lua 5.1 doesn't have this :<
+ * so we steal it from lua 5.3, thanks #lua :>
+ */
+void *luaL_testudata (lua_State *L, int i, const char *tname) {
+  void *p = lua_touserdata(L, i);
+  luaL_checkstack(L, 2, "not enough stack slots");
+  if (p == NULL || !lua_getmetatable(L, i))
+    return NULL;
+  else {
+    int res = 0;
+    luaL_getmetatable(L, tname);
+    res = lua_rawequal(L, -1, -2);
+    lua_pop(L, 2);
+    if (!res)
+      p = NULL;
+  }
+  return p;
+}
 #endif
 
 #define INVALID_MODE ((mode_t) -1)
@@ -3459,7 +3479,15 @@ luxio_recvfrom(lua_State *L)
 @section poll
 */
 
+#define LUXIO_TIMEVAL_METATABLE "luxio.timeval"
+
 #define LUXIO_POLLFD_METATABLE "luxio.pollfdarray"
+#define LUXIO_FD_SET_METATABLE "luxio.fd_set"
+
+typedef struct {
+	fd_set *fd_set;
+	int allocated;
+} luxio_fd_set;
 
 typedef struct {
 	struct pollfd *pollfds;
@@ -3467,6 +3495,57 @@ typedef struct {
 } luxio_pollfds;
 
 static int
+luxio__fd_set_gc(lua_State *L)
+{
+	luxio_fd_set *fs = luaL_checkudata(L, 1, LUXIO_FD_SET_METATABLE);
+	free(fs->fd_set);
+	fs->fd_set = NULL;
+	fs->allocated = 0;
+
+	return 0;
+}
+
+static int
+luxio__fd_set_tostring(lua_State *L)
+{
+	luxio_fd_set *fs = luaL_checkudata(L, 1, LUXIO_FD_SET_METATABLE);
+	char out[FD_SETSIZE * 10];
+	char s[100];
+	int i, found = 0, total_set = 0;
+
+	for (i = 0; i < FD_SETSIZE; i++) {
+		if (FD_ISSET(i, fs->fd_set)) {
+			total_set++;
+		}
+	}
+
+	if (total_set == 0) {
+		lua_pushstring(L, "{}");
+		return 1;
+	}
+
+	sprintf(out, "{");
+	for (i = 0; i < FD_SETSIZE; i++) {
+		if (FD_ISSET(i, fs->fd_set)) {
+			sprintf(s, ++found < total_set ? "%d, " : "%d}", i);
+			strcat(out, s);
+		}
+	}
+
+	lua_pushstring(L, out);
+
+	return 1;
+}
+
+static int
+luxio__fd_set_len(lua_State *L)
+{
+	luxio_fd_set *fs = luaL_checkudata(L, 1, LUXIO_FD_SET_METATABLE);
+	lua_pushnumber(L, fs->allocated);
+	return 1;
+}
+
+static int
 luxio__pollfds_gc(lua_State *L)
 {
 	luxio_pollfds *pfds = luaL_checkudata(L, 1, LUXIO_POLLFD_METATABLE);
@@ -3515,6 +3594,94 @@ luxio_pollfds_new(lua_State *L)
 }
 
 static int
+luxio_FD_CLR(lua_State *L)
+{
+	int fd = luaL_checkint(L, 1);
+	luxio_fd_set *fs = luaL_checkudata(L, 2, LUXIO_FD_SET_METATABLE);
+
+	FD_CLR(fd, fs->fd_set);
+
+	return 0;
+}
+
+static int
+luxio_FD_SET(lua_State *L)
+{
+	int fd = luaL_checkint(L, 1);
+	luxio_fd_set *fs = luaL_checkudata(L, 2, LUXIO_FD_SET_METATABLE);
+
+	FD_SET(fd, fs->fd_set);
+
+	return 0;
+}
+
+static int
+luxio_FD_ISSET(lua_State *L)
+{
+	int fd = luaL_checkint(L, 1);
+	luxio_fd_set *fs = luaL_checkudata(L, 2, LUXIO_FD_SET_METATABLE);
+
+	lua_pushboolean(L, FD_ISSET(fd, fs->fd_set));
+
+	return 1;
+}
+
+static int
+luxio_FD_ZERO(lua_State *L)
+{
+	luxio_fd_set *fs = luaL_checkudata(L, 1, LUXIO_FD_SET_METATABLE);
+
+	FD_ZERO(fs->fd_set);
+
+	return 0;
+}
+
+static int
+luxio_fd_set_new(lua_State *L)
+{
+	luxio_fd_set *fs = lua_newuserdata(L, sizeof(*fs));
+	int create_table = luaL_newmetatable(L, LUXIO_FD_SET_METATABLE);
+	fs->fd_set = NULL;
+	fs->allocated = 0;
+
+	if (create_table) {
+		lua_pushcfunction(L, luxio__fd_set_gc);
+		lua_setfield(L, -2, "__gc");
+		lua_pushcfunction(L, luxio__fd_set_tostring);
+		lua_setfield(L, -2, "__tostring");
+		lua_pushcfunction(L, luxio__fd_set_len);
+		lua_setfield(L, -2, "__len");
+	}
+
+	lua_setmetatable(L, -2);
+
+	return 1;
+}
+
+static int
+luxio_fd_set_resize(lua_State *L)
+{
+	luxio_fd_set *fs = luaL_checkudata(L, 1, LUXIO_FD_SET_METATABLE);
+	int desired_size = luaL_checkint(L, 2);
+	int idx;
+	fd_set *new_fs = realloc(fs->fd_set, sizeof(fd_set) * desired_size);
+	if (new_fs != NULL) {
+		if (desired_size > fs->allocated) {
+			for (idx = fs->allocated; idx < desired_size; ++idx) {
+				memset(new_fs + idx, 0, sizeof(*new_fs));
+			}
+		}
+
+		fs->fd_set = new_fs;
+		fs->allocated = desired_size;
+	} else {
+		return luaL_error(L, "Unable to resize fd_set array");
+	}
+	/* Return the fd_set array for neatness */
+	return 1;
+}
+
+static int
 luxio_pollfds_resize(lua_State *L)
 {
 	luxio_pollfds *pfds = luaL_checkudata(L, 1, LUXIO_POLLFD_METATABLE);
@@ -3586,6 +3753,25 @@ luxio_poll(lua_State *L)
 	return 2;
 }
 
+static int
+luxio_select(lua_State *L)
+{
+	int nfds = luaL_checkint(L, 1);
+	luxio_fd_set *readfds = luaL_testudata(L, 2, LUXIO_FD_SET_METATABLE);
+	luxio_fd_set *writefds = luaL_testudata(L, 3, LUXIO_FD_SET_METATABLE);
+	luxio_fd_set *exceptfds = luaL_testudata(L, 4, LUXIO_FD_SET_METATABLE);
+	struct timeval *timeout = luaL_testudata(L, 5, LUXIO_TIMEVAL_METATABLE);
+
+	lua_pushinteger(L, select(nfds,
+		readfds != NULL ? readfds->fd_set : NULL,
+		writefds != NULL ? writefds->fd_set : NULL,
+		exceptfds != NULL ? exceptfds->fd_set : NULL,
+		timeout));
+	lua_pushinteger(L, errno);
+
+	return 2;
+}
+
 /*** Bit and flag operation functions.
 @section bit
 */
@@ -3679,7 +3865,6 @@ as a userdata type, complete with comparison, addition/subtraction, and tostring
 metamethods.  You can set the fields tv_sec, tv_usec, seconds, and useconds.
 @section time
 */
-#define LUXIO_TIMEVAL_METATABLE "luxio.timeval"
 
 static int
 luxio_timeval_lt(lua_State *L)
@@ -4212,6 +4397,15 @@ luxio_functions[] = {
 	{ "pollfds_getslot", luxio_pollfds_get_slot },
 	{ "poll", luxio_poll },
 
+	{ "fd_set_new", luxio_fd_set_new },
+	{ "fd_set_resize", luxio_fd_set_resize },
+	{ "FD_CLR", luxio_FD_CLR },
+	{ "FD_SET", luxio_FD_SET },
+	{ "FD_ISSET", luxio_FD_ISSET },
+	{ "FD_ZERO", luxio_FD_ZERO },
+
+	{ "select", luxio_select },
+
 	{ "zero_timeval", luxio_timeval_zero },
 	{ "gettimeofday", luxio_gettimeofday },
 
diff --git a/luxio_constants.inc.in b/luxio_constants.inc.in
index 0c858df..4c4a3dd 100644
--- a/luxio_constants.inc.in
+++ b/luxio_constants.inc.in
@@ -188,6 +188,9 @@ static const struct {
 	E(POLLNVAL),
 ?	E(POLLMSG),
 
+	/* select-related defines */
+	E(FD_SETSIZE),
+
 	/* useful types */
 	E(SSIZE_MAX),
 ?	E(PATH_MAX),
diff --git a/tests/test-select.lua b/tests/test-select.lua
new file mode 100644
index 0000000..4d810ba
--- /dev/null
+++ b/tests/test-select.lua
@@ -0,0 +1,79 @@
+-- e.g. lua tests/test-select.lua <(while :; do echo hello; sleep 1; done) <(while :; do echo world; sleep 1; done)
+
+local l = require "luxio"
+
+local fds = {}
+
+function eet(path, fd, set)
+   r, errno = l.read(fd, 8192)
+
+   if r == -1 then
+      io.stderr:write(("Couldn't read `%s': %s\n"):format(path, l.strerror(errno)))
+      return -1
+   end
+
+   w, errno = l.write(l.STDOUT_FILENO, r, 0)
+   if w == -1 then
+      io.stderr:write(("Error writing to STDOUT: %s\n"):format(l.strerror(errno)))
+      return -1
+   end
+
+   return w
+end
+
+for _, path in ipairs(arg) do
+   fd, errno = l.open(path, l.O_RDONLY)
+
+   if fd == nil then
+      io.stderr:write(("Couldn't open file `%s': %s\n"):format(path, l.strerror(errno)))
+      os.exit(1)
+   end
+
+   fds[#fds + 1] = fd
+end
+
+local readset = l.fd_set_new()
+l.fd_set_resize(readset, l.FD_SETSIZE)
+
+timeout = l.zero_timeval()
+
+while true do
+   timeout.tv_sec = 1
+   l.FD_ZERO(readset)
+   max_fd = -1
+
+   for _, fd in ipairs(fds) do
+      l.FD_SET(fd, readset)
+
+      max_fd = math.max(fd, max_fd)
+   end
+
+   if max_fd == -1 then
+      break
+   end
+
+   n, errno = l.select(max_fd + 1, readset, nil, nil, timeout)
+
+   if n == -1 then
+      io.stderr:write(("select: %s\n"):format(l.strerror(errno)))
+      os.exit(1)
+   end
+
+   if n > 0 then
+      for i, fd in ipairs(fds) do
+         if fd ~= -1 and l.FD_ISSET(fd, readset) then
+            path = arg[i]
+            n = eet(path, fd)
+
+            if n == 0 then
+               l.close(fd)
+               fds[i] = -1
+            elseif n == -1 then
+               os.exit(1)
+            end
+         end
+      end
+   end
+
+   print(readset)
+end