From cbb31ea4aa189061389f5bcff982d318372df37a Mon Sep 17 00:00:00 2001
From: Shuanglei Tao <tsl0922@gmail.com>
Date: Sun, 7 Mar 2021 12:51:31 +0800
Subject: [PATCH] protocol: improve exit code handling

---
 src/protocol.c | 14 ++++++++------
 src/pty.c      |  6 ++++++
 src/server.h   |  3 ++-
 3 files changed, 16 insertions(+), 7 deletions(-)

diff --git a/src/protocol.c b/src/protocol.c
index c2eb016..8523b72 100644
--- a/src/protocol.c
+++ b/src/protocol.c
@@ -79,17 +79,20 @@ static bool check_host_origin(struct lws *wsi) {
 static void process_read_cb(void *ctx, pty_buf_t *buf, bool eof) {
   struct pss_tty *pss = (struct pss_tty *) ctx;
   pss->pty_buf = buf;
-  pss->pty_eof = eof;
+  if (eof && !process_running(pss->process)) {
+    pss->lws_close_status = pss->process->exit_code == 0 ? 1000 : 1006;
+  }
   lws_callback_on_writable(pss->wsi);
 }
 
 static void process_exit_cb(void *ctx, pty_process *process) {
   struct pss_tty *pss = (struct pss_tty *) ctx;
+  pss->process = NULL;
   if (process->killed) {
     lwsl_notice("process killed with signal %d, pid: %d\n", process->exit_signal, process->pid);
   } else {
     lwsl_notice("process exited with code %d, pid: %d\n", process->exit_code, process->pid);
-    pss->pty_eof = true;
+    pss->lws_close_status = process->exit_code == 0 ? 1000 : 1006;
     lws_callback_on_writable(pss->wsi);
   }
 }
@@ -175,6 +178,7 @@ int callback_tty(struct lws *wsi, enum lws_callback_reasons reason, void *user,
       pss->wsi = wsi;
       pss->buffer = NULL;
       pss->pty_buf = NULL;
+      pss->lws_close_status = LWS_CLOSE_STATUS_NOSTATUS;
 
       if (server->url_arg) {
         while (lws_hdr_copy_fragment(wsi, buf, sizeof(buf), WSI_TOKEN_HTTP_URI_ARGS, n++) > 0) {
@@ -215,9 +219,8 @@ int callback_tty(struct lws *wsi, enum lws_callback_reasons reason, void *user,
         break;
       }
 
-      // read error or client exited, close connection
-      if (pss->pty_eof) {
-        lws_close_reason(wsi, pss->process->exit_code ? LWS_CLOSE_STATUS_UNEXPECTED_CONDITION : LWS_CLOSE_STATUS_NORMAL, NULL, 0);
+      if (pss->lws_close_status > LWS_CLOSE_STATUS_NOSTATUS) {
+        lws_close_reason(wsi, pss->lws_close_status, NULL, 0);
         return 1;
       }
 
@@ -269,7 +272,6 @@ int callback_tty(struct lws *wsi, enum lws_callback_reasons reason, void *user,
           }
           break;
         case RESIZE_TERMINAL:
-          if (pss->process == NULL) break;
           json_object_put(parse_window_size(pss->buffer + 1, pss->len - 1,
                           &pss->process->columns, &pss->process->rows));
           pty_resize(pss->process);
diff --git a/src/pty.c b/src/pty.c
index b15c76d..1db5baa 100644
--- a/src/pty.c
+++ b/src/pty.c
@@ -54,6 +54,7 @@ static void read_cb(uv_stream_t *stream, ssize_t n, const uv_buf_t *buf) {
   pty_io_t *io = (pty_io_t *) stream->data;
   if (n <= 0) {
     if (n == UV_ENOBUFS || n == 0) return;
+    if (n != UV_EOF) printf("== uv_read failed with error %ld: %s\n", n, uv_strerror(n));
     io->read_cb(io->ctx, NULL, true);
     goto done;
   }
@@ -120,12 +121,14 @@ void process_free(pty_process *process) {
 }
 
 void pty_pause(pty_process *process) {
+  if (process == NULL) return;
   pty_io_t *io = process->io;
   if (io->paused) return;
   uv_read_stop((uv_stream_t *) io->out);
 }
 
 void pty_resume(pty_process *process) {
+  if (process == NULL) return;
   pty_io_t *io = process->io;
   if (!io->paused) return;
   io->out->data = io;
@@ -133,6 +136,7 @@ void pty_resume(pty_process *process) {
 }
 
 int pty_write(pty_process *process, pty_buf_t *buf) {
+  if (process == NULL) return UV_ESRCH;
   pty_io_t *io = process->io;
   uv_buf_t b = uv_buf_init(buf->base, buf->len);
   uv_write_t *req = xmalloc(sizeof(uv_write_t));
@@ -141,6 +145,7 @@ int pty_write(pty_process *process, pty_buf_t *buf) {
 }
 
 bool pty_resize(pty_process *process) {
+  if (process == NULL) return false;
   if (process->columns <= 0 || process->rows <= 0) return false;
 #ifdef _WIN32
   COORD size = { (int16_t) process->columns, (int16_t) process->rows };
@@ -152,6 +157,7 @@ bool pty_resize(pty_process *process) {
 }
 
 bool pty_kill(pty_process *process, int sig) {
+  if (process == NULL) return false;
   process->killed = true;
 #ifdef _WIN32
   return TerminateProcess(process->handle, 1) != 0;
diff --git a/src/server.h b/src/server.h
index 5d76103..59e49e2 100644
--- a/src/server.h
+++ b/src/server.h
@@ -52,7 +52,8 @@ struct pss_tty {
 
   pty_process *process;
   pty_buf_t *pty_buf;
-  bool pty_eof;
+
+  int lws_close_status;
 };
 
 struct server {
-- 
2.43.4