From f61b9aec118bb87c8421c5db5de50c5e8c539177 Mon Sep 17 00:00:00 2001
From: Shuanglei Tao <tsl0922@gmail.com>
Date: Sat, 30 Nov 2019 16:44:15 +0800
Subject: [PATCH] server: improve signal handling

---
 src/protocol.c | 52 ++++++++++++++++++++++++++++++++++----------------
 src/server.c   | 10 +---------
 src/server.h   |  3 ++-
 3 files changed, 39 insertions(+), 26 deletions(-)

diff --git a/src/protocol.c b/src/protocol.c
index 504aad6..252fae5 100644
--- a/src/protocol.c
+++ b/src/protocol.c
@@ -106,21 +106,9 @@ check_host_origin(struct lws *wsi) {
 }
 
 void
-tty_client_destroy(struct tty_client *client) {
-    if (client->pid <= 0)
-        goto cleanup;
-
-    // kill process (group) and free resource
-    int pgid = getpgid(client->pid);
-    int pid = pgid > 0 ? -pgid : client->pid;
-    if (kill(pid, server->sig_code) != 0) {
-        if (errno == ESRCH)
-            goto cleanup;
-        lwsl_err("kill: %d, errno: %d (%s)\n", pid, errno, strerror(errno));
-    }
-
-cleanup:
+tty_client_free(struct tty_client *client) {
     uv_read_stop((uv_stream_t *) &client->pipe);
+    uv_signal_stop(&client->watcher);
 
     close(client->pty);
 
@@ -164,6 +152,21 @@ read_cb(uv_stream_t* stream, ssize_t nread, const uv_buf_t* buf) {
     uv_read_stop(stream);
 }
 
+void
+child_cb(uv_signal_t *handle, int signum) {
+    struct tty_client *client;
+    pid_t pid;
+    int status;
+
+    client = (struct tty_client *) handle->data;
+    status = wait_proc(client->pid, &pid);
+    if (pid > 0) {
+        lwsl_notice("process exited with code %d, pid: %d\n", status, pid);
+        client->pid = 0;
+        tty_client_free(client);
+    }
+}
+
 int
 spawn_process(struct tty_client *client) {
     // append url args to arguments
@@ -177,6 +180,8 @@ spawn_process(struct tty_client *client) {
     }
     argv[n] = NULL;
 
+    uv_signal_start(&client->watcher, child_cb, SIGCHLD);
+
     int pty;
     pid_t pid = forkpty(&pty, NULL, NULL, NULL);
     if (pid < 0) { /* error */
@@ -214,6 +219,18 @@ spawn_process(struct tty_client *client) {
     return 0;
 }
 
+void
+kill_process(pid_t pid, int sig) {
+    if (pid <= 0) return;
+
+    // kill process (group) and free resource
+    lwsl_notice("killing process %d with signal %d\n", pid, sig);
+    int pgid = getpgid(pid);
+    if (kill(pgid > 0 ? -pgid : pid, sig) != 0) {
+        lwsl_err("kill: %d, errno: %d (%s)\n", pid, errno, strerror(errno));
+    }
+}
+
 int
 callback_tty(struct lws *wsi, enum lws_callback_reasons reason,
              void *user, void *in, size_t len) {
@@ -250,8 +267,11 @@ callback_tty(struct lws *wsi, enum lws_callback_reasons reason,
             client->buffer = NULL;
             client->pty_len = 0;
             client->argc = 0;
+            client->loop = server->loop;
 
-            uv_pipe_init(server->loop, &client->pipe, 0);
+            uv_pipe_init(client->loop, &client->pipe, 0);
+            uv_signal_init(client->loop, &client->watcher);
+            client->watcher.data = client;
 
             if (server->url_arg) {
                 while (lws_hdr_copy_fragment(wsi, buf, sizeof(buf), WSI_TOKEN_HTTP_URI_ARGS, n++) > 0) {
@@ -394,7 +414,7 @@ callback_tty(struct lws *wsi, enum lws_callback_reasons reason,
         case LWS_CALLBACK_CLOSED:
             server->client_count--;
             lwsl_notice("WS closed from %s, clients: %d\n", client->address, server->client_count);
-            tty_client_destroy(client);
+            kill_process(client->pid, server->sig_code);
             if (server->once && server->client_count == 0) {
                 lwsl_notice("exiting due to the --once option.\n");
                 force_exit = true;
diff --git a/src/server.c b/src/server.c
index fb7d090..3f9de60 100644
--- a/src/server.c
+++ b/src/server.c
@@ -174,8 +174,6 @@ tty_server_free(struct tty_server *ts) {
 void
 signal_cb(uv_signal_t *watcher, int signum) {
     char sig_name[20];
-    pid_t pid;
-    int status;
 
     switch (watcher->signum) {
         case SIGINT:
@@ -183,12 +181,6 @@ signal_cb(uv_signal_t *watcher, int signum) {
             get_sig_name(watcher->signum, sig_name, sizeof(sig_name));
             lwsl_notice("received signal: %s (%d), exiting...\n", sig_name, watcher->signum);
             break;
-        case SIGCHLD:
-            status = wait_proc(-1, &pid);
-            if (pid > 0) {
-                lwsl_notice("process exited with code %d, pid: %d\n", status, pid);
-            }
-            return;
         default:
             signal(SIGABRT, SIG_DFL);
             abort();
@@ -493,7 +485,7 @@ main(int argc, char **argv) {
     }
 
 #if LWS_LIBRARY_VERSION_MAJOR >= 3
-    int sig_nums[] = {SIGINT, SIGTERM, SIGCHLD};
+    int sig_nums[] = {SIGINT, SIGTERM};
     int ns = sizeof(sig_nums)/sizeof(sig_nums[0]);
     uv_signal_t signals[ns];
     for (int i = 0; i < ns; i++) {
diff --git a/src/server.h b/src/server.h
index 8a4fde0..1d91cee 100644
--- a/src/server.h
+++ b/src/server.h
@@ -34,11 +34,12 @@ struct tty_client {
 
     int pid;
     int pty;
-    int exit_status;
     char *pty_buffer;
     ssize_t pty_len;
 
+    uv_loop_t *loop;
     uv_pipe_t pipe;
+    uv_signal_t watcher;
 };
 
 struct pss_http {
-- 
2.43.4