uloop: add support for user defined signal handlers

Reuse and extend the existing signal waker pipe mechanism to add user
defined signal handling functionality to uloop.

This commit introduces two new api functions `uloop_signal_add()` and
`uloop_signal_remove()` along with a new structure type `uloop_signal`
to allow adding and removing arbitrary signal handlers.

Registered signal handlers are maintained in a linked list and matched
by their signo member value which allows registering multiple handlers
for the same signal numbers.

Upon registering a new signal handler, the existing handler is saved
in the `uloop_signal` structure. When removing the user defined signal
handler, the original behavior is restored.

The Lua binding has been updated as well to support the new signal
handler mechanism.

Signed-off-by: Jo-Philipp Wich <jo@mein.io>

4 files changed, 294 insertions, 16 deletions

diff --git a/examples/uloop-example.lua b/examples/uloop-example.lua
index f3aef60..1b73aed 100755
--- a/examples/uloop-example.lua
+++ b/examples/uloop-example.lua
@@ -63,6 +63,22 @@ uloop.timer(
 	end, 2000
 )
 
+-- SIGINT handler
+uloop.signal(function(signo)
+	print(string.format("Terminating on SIGINT (#%d)!", signo))
+
+	-- end uloop to terminate program
+	uloop.cancel()
+end, uloop.SIGINT)
+
+local sig
+sig = uloop.signal(function(signo)
+	print(string.format("Got SIGUSR2 (#%d)!", signo))
+
+	-- remove signal handler, next SIGUSR2 will terminate program
+	sig:delete()
+end, uloop.SIGUSR2)
+
 -- Keep udp_ev reference, events will be gc'd, even if the callback is still referenced
 -- .delete will manually untrack.
 udp_ev = uloop.fd_add(udp, function(ufd, events)
diff --git a/lua/uloop.c b/lua/uloop.c
index 45c9bc7..7a73f16 100644
--- a/lua/uloop.c
+++ b/lua/uloop.c
@@ -46,6 +46,11 @@ struct lua_uloop_interval {
 	int r;
 };
 
+struct lua_uloop_signal {
+	struct uloop_signal s;
+	int r;
+};
+
 static lua_State *state;
 
 static void *
@@ -493,6 +498,83 @@ static int ul_interval(lua_State *L)
 	return 1;
 }
 
+static void ul_signal_cb(struct uloop_signal *s)
+{
+	struct lua_uloop_signal *sig = container_of(s, struct lua_uloop_signal, s);
+
+	lua_getglobal(state, "__uloop_cb");
+	lua_rawgeti(state, -1, sig->r);
+	lua_remove(state, -2);
+	lua_pushinteger(state, sig->s.signo);
+
+	lua_call(state, 1, 0);
+}
+
+static int ul_signal_signo(lua_State *L)
+{
+	struct lua_uloop_signal *sig = lua_touserdata(L, 1);
+
+	lua_pushinteger(L, sig->s.signo);
+
+	return 1;
+}
+
+static int ul_signal_free(lua_State *L)
+{
+	struct lua_uloop_signal *sig = lua_touserdata(L, 1);
+
+	uloop_signal_delete(&sig->s);
+
+	/* obj.__index.__gc = nil , make sure executing only once*/
+	lua_getfield(L, -1, "__index");
+	lua_pushstring(L, "__gc");
+	lua_pushnil(L);
+	lua_settable(L, -3);
+
+	lua_getglobal(state, "__uloop_cb");
+	luaL_unref(state, -1, sig->r);
+
+	return 1;
+}
+
+static const luaL_Reg signal_m[] = {
+	{ "signo", ul_signal_signo },
+	{ "delete", ul_signal_free },
+	{ NULL, NULL }
+};
+
+static int ul_signal(lua_State *L)
+{
+	struct lua_uloop_signal *sig;
+	int signo = -1;
+	int ref;
+
+	if (lua_isnumber(L, -1)) {
+		signo = lua_tointeger(L, -1);
+		lua_pop(L, 1);
+	}
+
+	if (!lua_isfunction(L, -1) || signo <= 0 || signo > NSIG) {
+		lua_pushstring(L, "invalid arg list");
+		lua_error(L);
+
+		return 0;
+	}
+
+	lua_getglobal(L, "__uloop_cb");
+	lua_pushvalue(L, -2);
+	ref = luaL_ref(L, -2);
+
+	sig = ul_create_userdata(L, sizeof(*sig), signal_m, ul_signal_free);
+	sig->r = ref;
+	sig->s.cb = ul_signal_cb;
+	sig->s.signo = signo;
+
+	uloop_signal_add(&sig->s);
+
+	return 1;
+}
+
 static int ul_init(lua_State *L)
 {
 	uloop_init();
@@ -522,6 +604,7 @@ static luaL_reg uloop_func[] = {
 	{"process", ul_process},
 	{"fd_add", ul_ufd_add},
 	{"interval", ul_interval},
+	{"signal", ul_signal},
 	{"cancel", ul_end},
 	{NULL, NULL},
 };
@@ -561,6 +644,103 @@ int luaopen_uloop(lua_State *L)
 	lua_pushinteger(L, ULOOP_BLOCKING);
 	lua_rawset(L, -3);
 
+#define SIGNAME_CONST(signame) do { \
+	lua_pushstring(L, #signame);    \
+	lua_pushinteger(L, signame);    \
+	lua_rawset(L, -3);	            \
+} while(0)
+
+#if defined(SIGINT)
+	SIGNAME_CONST(SIGINT);
+#endif
+#if defined(SIGILL)
+	SIGNAME_CONST(SIGILL);
+#endif
+#if defined(SIGABRT)
+	SIGNAME_CONST(SIGABRT);
+#endif
+#if defined(SIGFPE)
+	SIGNAME_CONST(SIGFPE);
+#endif
+#if defined(SIGSEGV)
+	SIGNAME_CONST(SIGSEGV);
+#endif
+#if defined(SIGTERM)
+	SIGNAME_CONST(SIGTERM);
+#endif
+#if defined(SIGHUP)
+	SIGNAME_CONST(SIGHUP);
+#endif
+#if defined(SIGQUIT)
+	SIGNAME_CONST(SIGQUIT);
+#endif
+#if defined(SIGTRAP)
+	SIGNAME_CONST(SIGTRAP);
+#endif
+#if defined(SIGKILL)
+	SIGNAME_CONST(SIGKILL);
+#endif
+#if defined(SIGPIPE)
+	SIGNAME_CONST(SIGPIPE);
+#endif
+#if defined(SIGALRM)
+	SIGNAME_CONST(SIGALRM);
+#endif
+#if defined(SIGSTKFLT)
+	SIGNAME_CONST(SIGSTKFLT);
+#endif
+#if defined(SIGPWR)
+	SIGNAME_CONST(SIGPWR);
+#endif
+#if defined(SIGBUS)
+	SIGNAME_CONST(SIGBUS);
+#endif
+#if defined(SIGSYS)
+	SIGNAME_CONST(SIGSYS);
+#endif
+#if defined(SIGURG)
+	SIGNAME_CONST(SIGURG);
+#endif
+#if defined(SIGSTOP)
+	SIGNAME_CONST(SIGSTOP);
+#endif
+#if defined(SIGTSTP)
+	SIGNAME_CONST(SIGTSTP);
+#endif
+#if defined(SIGCONT)
+	SIGNAME_CONST(SIGCONT);
+#endif
+#if defined(SIGCHLD)
+	SIGNAME_CONST(SIGCHLD);
+#endif
+#if defined(SIGTTIN)
+	SIGNAME_CONST(SIGTTIN);
+#endif
+#if defined(SIGTTOU)
+	SIGNAME_CONST(SIGTTOU);
+#endif
+#if defined(SIGPOLL)
+	SIGNAME_CONST(SIGPOLL);
+#endif
+#if defined(SIGXFSZ)
+	SIGNAME_CONST(SIGXFSZ);
+#endif
+#if defined(SIGXCPU)
+	SIGNAME_CONST(SIGXCPU);
+#endif
+#if defined(SIGVTALRM)
+	SIGNAME_CONST(SIGVTALRM);
+#endif
+#if defined(SIGPROF)
+	SIGNAME_CONST(SIGPROF);
+#endif
+#if defined(SIGUSR1)
+	SIGNAME_CONST(SIGUSR1);
+#endif
+#if defined(SIGUSR2)
+	SIGNAME_CONST(SIGUSR2);
+#endif
+
 	return 1;
 }
 
diff --git a/uloop.c b/uloop.c
index a3d3712..89a7029 100644
--- a/uloop.c
+++ b/uloop.c
@@ -57,6 +57,7 @@ static struct uloop_fd_stack *fd_stack = NULL;
 
 static struct list_head timeouts = LIST_HEAD_INIT(timeouts);
 static struct list_head processes = LIST_HEAD_INIT(processes);
+static struct list_head signals = LIST_HEAD_INIT(signals);
 
 static int poll_fd = -1;
 bool uloop_cancelled = false;
@@ -80,18 +81,41 @@ int uloop_fd_add(struct uloop_fd *sock, unsigned int flags);
 #include "uloop-epoll.c"
 #endif
 
-static void waker_consume(struct uloop_fd *fd, unsigned int events)
+static void set_signo(uint64_t *signums, int signo)
 {
-	char buf[4];
+	if (signo >= 1 && signo <= 64)
+		*signums |= (1u << (signo - 1));
+}
+
+static bool get_signo(uint64_t signums, int signo)
+{
+	return (signo >= 1) && (signo <= 64) && (signums & (1u << (signo - 1)));
+}
+
+static void signal_consume(struct uloop_fd *fd, unsigned int events)
+{
+	struct uloop_signal *usig, *usig_next;
+	uint64_t signums = 0;
+	uint8_t buf[32];
+	ssize_t nsigs;
+
+	do {
+		nsigs = read(fd->fd, buf, sizeof(buf));
+
+		for (ssize_t i = 0; i < nsigs; i++)
+			set_signo(&signums, buf[i]);
+	}
+	while (nsigs > 0);
 
-	while (read(fd->fd, buf, 4) > 0)
-		;
+	list_for_each_entry_safe(usig, usig_next, &signals, list)
+		if (get_signo(signums, usig->signo))
+			usig->cb(usig);
 }
 
 static int waker_pipe = -1;
 static struct uloop_fd waker_fd = {
 	.fd = -1,
-	.cb = waker_consume,
+	.cb = signal_consume,
 };
 
 static void waker_init_fd(int fd)
@@ -115,7 +139,7 @@ static int waker_init(void)
 	waker_pipe = fds[1];
 
 	waker_fd.fd = fds[0];
-	waker_fd.cb = waker_consume;
+	waker_fd.cb = signal_consume;
 	uloop_fd_add(&waker_fd, ULOOP_READ);
 
 	return 0;
@@ -438,10 +462,15 @@ int64_t uloop_interval_remaining(struct uloop_interval *timer)
 	return timer_next(timer);
 }
 
-static void uloop_signal_wake(void)
+static void uloop_signal_wake(int signo)
 {
+	uint8_t sigbyte = signo;
+
+	if (signo == ECHILD)
+		do_sigchld = true;
+
 	do {
-		if (write(waker_pipe, "w", 1) < 0) {
+		if (write(waker_pipe, &sigbyte, 1) < 0) {
 			if (errno == EINTR)
 				continue;
 		}
@@ -453,13 +482,7 @@ static void uloop_handle_sigint(int signo)
 {
 	uloop_status = signo;
 	uloop_cancelled = true;
-	uloop_signal_wake();
-}
-
-static void uloop_sigchld(int signo)
-{
-	do_sigchld = true;
-	uloop_signal_wake();
+	uloop_signal_wake(signo);
 }
 
 static void uloop_install_handler(int signum, void (*handler)(int), struct sigaction* old, bool add)
@@ -516,11 +539,55 @@ static void uloop_setup_signals(bool add)
 	uloop_install_handler(SIGTERM, uloop_handle_sigint, &old_sigterm, add);
 
 	if (uloop_handle_sigchld)
-		uloop_install_handler(SIGCHLD, uloop_sigchld, &old_sigchld, add);
+		uloop_install_handler(SIGCHLD, uloop_signal_wake, &old_sigchld, add);
 
 	uloop_ignore_signal(SIGPIPE, add);
 }
 
+int uloop_signal_add(struct uloop_signal *s)
+{
+	struct list_head *h = &signals;
+	struct uloop_signal *tmp;
+	struct sigaction sa;
+
+	if (s->pending)
+		return -1;
+
+	list_for_each_entry(tmp, &signals, list) {
+		if (tmp->signo > s->signo) {
+			h = &tmp->list;
+			break;
+		}
+	}
+
+	list_add_tail(&s->list, h);
+	s->pending = true;
+
+	sigaction(s->signo, NULL, &s->orig);
+
+	if (s->orig.sa_handler != uloop_signal_wake) {
+		sa.sa_handler = uloop_signal_wake;
+		sa.sa_flags = 0;
+		sigaction(s->signo, &sa, NULL);
+	}
+
+	return 0;
+}
+
+int uloop_signal_delete(struct uloop_signal *s)
+{
+	if (!s->pending)
+		return -1;
+
+	list_del(&s->list);
+	s->pending = false;
+
+	if (s->orig.sa_handler != uloop_signal_wake)
+		sigaction(s->signo, &s->orig, NULL);
+
+	return 0;
+}
+
 int uloop_get_next_timeout(void)
 {
 	struct uloop_timeout *timeout;
diff --git a/uloop.h b/uloop.h
index b3f268c..5edeb70 100644
--- a/uloop.h
+++ b/uloop.h
@@ -36,11 +36,13 @@ struct uloop_fd;
 struct uloop_timeout;
 struct uloop_process;
 struct uloop_interval;
+struct uloop_signal;
 
 typedef void (*uloop_fd_handler)(struct uloop_fd *u, unsigned int events);
 typedef void (*uloop_timeout_handler)(struct uloop_timeout *t);
 typedef void (*uloop_process_handler)(struct uloop_process *c, int ret);
 typedef void (*uloop_interval_handler)(struct uloop_interval *t);
+typedef void (*uloop_signal_handler)(struct uloop_signal *s);
 
 #define ULOOP_READ		(1 << 0)
 #define ULOOP_WRITE		(1 << 1)
@@ -99,6 +101,16 @@ struct uloop_interval
 	} private;
 };
 
+struct uloop_signal
+{
+	struct list_head list;
+	struct sigaction orig;
+	bool pending;
+
+	uloop_signal_handler cb;
+	int signo;
+};
+
 extern bool uloop_cancelled;
 extern bool uloop_handle_sigchld;
 extern uloop_fd_handler uloop_fd_set_cb;
@@ -120,6 +132,9 @@ int uloop_interval_set(struct uloop_interval *timer, unsigned int msecs);
 int uloop_interval_cancel(struct uloop_interval *timer);
 int64_t uloop_interval_remaining(struct uloop_interval *timer);
 
+int uloop_signal_add(struct uloop_signal *s);
+int uloop_signal_delete(struct uloop_signal *s);
+
 bool uloop_cancelling(void);
 
 static inline void uloop_end(void)