From: Shuanglei Tao <tsl0922@gmail.com>
Date: Thu, 18 Aug 2022 02:35:23 +0000 (+0800)
Subject: protocol: fix potential pty buf leak
X-Git-Url: http://git.prime8.dev/?a=commitdiff_plain;h=6cfd9b29ae7d47aeb5e7af1ff0e951e7588d30eb;p=ttyd.git

protocol: fix potential pty buf leak
---

diff --git a/src/protocol.c b/src/protocol.c
index 7562d59..731b54f 100644
--- a/src/protocol.c
+++ b/src/protocol.c
@@ -69,29 +69,43 @@ static bool check_host_origin(struct lws *wsi) {
   return len > 0 && strcasecmp(buf, host_buf) == 0;
 }
 
+static pty_ctx_t *pty_ctx_init(struct pss_tty *pss) {
+  pty_ctx_t *ctx = xmalloc(sizeof(pty_ctx_t));
+  ctx->pss = pss;
+  ctx->ws_closed = false;
+  return ctx;
+}
+
+static void pty_ctx_free(pty_ctx_t *ctx) { free(ctx); }
+
 static void process_read_cb(pty_process *process, pty_buf_t *buf, bool eof) {
-  if (process->killed) return ;
+  pty_ctx_t *ctx = (pty_ctx_t *)process->ctx;
+  if (ctx->ws_closed) {
+    pty_buf_free(buf);
+    return;
+  }
 
-  struct pss_tty *pss = (struct pss_tty *)process->ctx;
-  if (eof && !process_running(pss->process))
-    pss->lws_close_status = pss->process->exit_code == 0 ? 1000 : 1006;
+  if (eof && !process_running(process))
+    ctx->pss->lws_close_status = process->exit_code == 0 ? 1000 : 1006;
   else
-    pss->pty_buf = buf;
-
-  lws_callback_on_writable(pss->wsi);
+    ctx->pss->pty_buf = buf;
+  lws_callback_on_writable(ctx->pss->wsi);
 }
 
 static void process_exit_cb(pty_process *process) {
-  if (process->killed) {
+  pty_ctx_t *ctx = (pty_ctx_t *)process->ctx;
+  if (ctx->ws_closed) {
     lwsl_notice("process killed with signal %d, pid: %d\n", process->exit_signal, process->pid);
-    return ;
+    goto done;
   }
 
   lwsl_notice("process exited with code %d, pid: %d\n", process->exit_code, process->pid);
-  struct pss_tty *pss = (struct pss_tty *)process->ctx;
-  pss->process = NULL;
-  pss->lws_close_status = process->exit_code == 0 ? 1000 : 1006;
-  lws_callback_on_writable(pss->wsi);
+  ctx->pss->process = NULL;
+  ctx->pss->lws_close_status = process->exit_code == 0 ? 1000 : 1006;
+  lws_callback_on_writable(ctx->pss->wsi);
+
+done:
+  pty_ctx_free(ctx);
 }
 
 static char **build_args(struct pss_tty *pss) {
@@ -134,7 +148,7 @@ static char **build_env(struct pss_tty *pss) {
 }
 
 static bool spawn_process(struct pss_tty *pss, uint16_t columns, uint16_t rows) {
-  pty_process *process = process_init((void *)pss, server->loop, build_args(pss), build_env(pss));
+  pty_process *process = process_init((void *)pty_ctx_init(pss), server->loop, build_args(pss), build_env(pss));
   if (server->cwd != NULL) process->cwd = strdup(server->cwd);
   if (columns > 0) process->columns = columns;
   if (rows > 0) process->rows = rows;
@@ -356,10 +370,13 @@ int callback_tty(struct lws *wsi, enum lws_callback_reasons reason, void *user,
         free(pss->args[i]);
       }
 
-      if (process_running(pss->process)) {
-        pty_pause(pss->process);
-        lwsl_notice("killing process, pid: %d\n", pss->process->pid);
-        pty_kill(pss->process, server->sig_code);
+      if (pss->process != NULL) {
+        ((pty_ctx_t *)pss->process->ctx)->ws_closed = true;
+        if (process_running(pss->process)) {
+          pty_pause(pss->process);
+          lwsl_notice("killing process, pid: %d\n", pss->process->pid);
+          pty_kill(pss->process, server->sig_code);
+        }
       }
 
       if (server->once && server->client_count == 0) {
diff --git a/src/pty.c b/src/pty.c
index ec06557..2fda996 100644
--- a/src/pty.c
+++ b/src/pty.c
@@ -51,7 +51,8 @@ pty_buf_t *pty_buf_init(char *base, size_t len) {
 }
 
 void pty_buf_free(pty_buf_t *buf) {
-  free(buf->base);
+  if (buf == NULL) return;
+  if (buf->base != NULL) free(buf->base);
   free(buf);
 }
 
@@ -105,6 +106,7 @@ void process_free(pty_process *process) {
 #else
   uv_thread_join(&process->tid);
 #endif
+  close(process->pty);
   if (process->in != NULL) uv_close((uv_handle_t *) process->in, close_cb);
   if (process->out != NULL) uv_close((uv_handle_t *) process->out, close_cb);
   if (process->argv != NULL) free(process->argv);
@@ -153,7 +155,6 @@ 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;
 #else
diff --git a/src/pty.h b/src/pty.h
index 3efa558..c3f761c 100644
--- a/src/pty.h
+++ b/src/pty.h
@@ -29,7 +29,6 @@ typedef void (*pty_exit_cb)(pty_process *);
 struct pty_process_ {
   int pid, exit_code, exit_signal;
   uint16_t columns, rows;
-  bool killed;
 #ifdef _WIN32
   STARTUPINFOEXW si;
   HPCON pty;
diff --git a/src/server.h b/src/server.h
index b122d1c..794965a 100644
--- a/src/server.h
+++ b/src/server.h
@@ -56,6 +56,11 @@ struct pss_tty {
   int lws_close_status;
 };
 
+typedef struct {
+  struct pss_tty *pss;
+  bool ws_closed;
+} pty_ctx_t ;
+
 struct server {
   int client_count;        // client count
   char *prefs_json;        // client preferences