From 4f91d763b4b9df72e9d52e09e9fb0d22fc8dce31 Mon Sep 17 00:00:00 2001
From: Shuanglei Tao <tsl0922@gmail.com>
Date: Thu, 31 Dec 2020 22:18:22 +0800
Subject: [PATCH] protocol: fix ws code on process exit

---
 src/protocol.c | 20 +++++++++++---------
 src/server.h   |  3 +++
 2 files changed, 14 insertions(+), 9 deletions(-)

diff --git a/src/protocol.c b/src/protocol.c
index 86a98238..4b832ab4 100644
--- a/src/protocol.c
+++ b/src/protocol.c
@@ -124,6 +124,7 @@ static void read_cb(uv_stream_t *stream, ssize_t nread, const uv_buf_t *buf) {
   } else {
     proc->pty_buffer = NULL;
     if (nread != UV_EOF) {
+      proc->err_count++;
       lwsl_err("read_cb: %s (%s)\n", uv_err_name(nread), uv_strerror(nread));
     }
   }
@@ -314,20 +315,21 @@ int callback_tty(struct lws *wsi, enum lws_callback_reasons reason, void *user,
       if (proc->status == 0 || proc->pty_len == UV_EOF) {
         lws_close_reason(wsi, LWS_CLOSE_STATUS_NORMAL, NULL, 0);
         return 1;
-      } else if (proc->status > 0 || proc->pty_len < 0) {
+      } else if (proc->status > 0 || (proc->pty_len < 0 && proc->err_count == MAX_READ_RETRY)) {
         lws_close_reason(wsi, LWS_CLOSE_STATUS_UNEXPECTED_CONDITION, NULL, 0);
         return -1;
       }
 
-      if (proc->pty_buffer == NULL || proc->pty_len == 0) break;
-
-      proc->pty_buffer[LWS_PRE] = OUTPUT;
-      n = (size_t)(proc->pty_len + 1);
-      if (lws_write(wsi, (unsigned char *)proc->pty_buffer + LWS_PRE, n, LWS_WRITE_BINARY) < n) {
-        lwsl_err("write OUTPUT to WS\n");
+      if (proc->pty_buffer != NULL && proc->pty_len > 0) {
+        proc->pty_buffer[LWS_PRE] = OUTPUT;
+        n = (size_t)(proc->pty_len + 1);
+        if (lws_write(wsi, (unsigned char *)proc->pty_buffer + LWS_PRE, n, LWS_WRITE_BINARY) < n) {
+          lwsl_err("write OUTPUT to WS\n");
+        }
+        free(proc->pty_buffer);
+        proc->pty_buffer = NULL;
       }
-      free(proc->pty_buffer);
-      proc->pty_buffer = NULL;
+
       uv_read_start((uv_stream_t *)&proc->pipe, alloc_cb, read_cb);
       break;
 
diff --git a/src/server.h b/src/server.h
index 95696bc1..784d445e 100644
--- a/src/server.h
+++ b/src/server.h
@@ -15,6 +15,8 @@
 #define SET_WINDOW_TITLE '1'
 #define SET_PREFERENCES '2'
 
+#define MAX_READ_RETRY 2
+
 // url paths
 struct endpoints {
   char *ws;
@@ -48,6 +50,7 @@ struct pty_proc {
   int pty;
   char *pty_buffer;
   ssize_t pty_len;
+  int err_count;
 
   uv_pipe_t pipe;