1 files changed, 127 insertions, 210 deletions

diff --git a/src/protocol.c b/src/protocol.c
index 4b832ab..ca1e086 100644
--- a/src/protocol.c
+++ b/src/protocol.c
@@ -1,15 +1,13 @@
 #include <errno.h>
 #include <json.h>
 #include <libwebsockets.h>
-#include <signal.h>
 #include <stdbool.h>
 #include <stdio.h>
 #include <stdlib.h>
 #include <string.h>
-#include <sys/wait.h>
 
+#include "pty.h"
 #include "server.h"
-#include "terminal.h"
 #include "utils.h"
 
 // initial message list
@@ -37,34 +35,22 @@ static int send_initial_message(struct lws *wsi, int index) {
   return lws_write(wsi, p, (size_t)n, LWS_WRITE_BINARY);
 }
 
-static bool parse_window_size(struct pss_tty *pss, int *cols, int *rows) {
-  char json[pss->len];
-  strncpy(json, pss->buffer + 1, pss->len - 1);
-  json[pss->len - 1] = '\0';
-
-  json_object *obj = json_tokener_parse(json);
+static json_object *parse_window_size(const char *buf, size_t len, uint16_t *cols, uint16_t *rows) {
+  json_tokener *tok = json_tokener_new();
+  json_object *obj = json_tokener_parse_ex(tok, buf, len);
   struct json_object *o = NULL;
 
-  if (!json_object_object_get_ex(obj, "columns", &o)) {
-    lwsl_err("columns field not exists, json: %s\n", json);
-    return false;
-  }
-  *cols = json_object_get_int(o);
-  if (!json_object_object_get_ex(obj, "rows", &o)) {
-    lwsl_err("rows field not exists, json: %s\n", json);
-    return false;
-  }
-  *rows = json_object_get_int(o);
-  json_object_put(obj);
+  if (json_object_object_get_ex(obj, "columns", &o)) *cols = (uint16_t)json_object_get_int(o);
+  if (json_object_object_get_ex(obj, "rows", &o)) *rows = (uint16_t)json_object_get_int(o);
 
-  return true;
+  json_tokener_free(tok);
+  return obj;
 }
 
 static bool check_host_origin(struct lws *wsi) {
-  int origin_length = lws_hdr_total_length(wsi, WSI_TOKEN_ORIGIN);
-  char buf[origin_length + 1];
+  char buf[256];
   memset(buf, 0, sizeof(buf));
-  int len = lws_hdr_copy(wsi, buf, sizeof(buf), WSI_TOKEN_ORIGIN);
+  int len = lws_hdr_copy(wsi, buf, (int)sizeof(buf), WSI_TOKEN_ORIGIN);
   if (len <= 0) {
     return false;
   }
@@ -78,154 +64,123 @@ static bool check_host_origin(struct lws *wsi) {
     sprintf(buf, "%s:%d", address, port);
   }
 
-  int host_length = lws_hdr_total_length(wsi, WSI_TOKEN_HOST);
-  if (host_length != strlen(buf)) return false;
-  char host_buf[host_length + 1];
+  char host_buf[256];
   memset(host_buf, 0, sizeof(host_buf));
-  len = lws_hdr_copy(wsi, host_buf, sizeof(host_buf), WSI_TOKEN_HOST);
+  len = lws_hdr_copy(wsi, host_buf, (int)sizeof(host_buf), WSI_TOKEN_HOST);
 
   return len > 0 && strcasecmp(buf, host_buf) == 0;
 }
 
-static void close_cb(uv_handle_t *handle) {
-  struct pty_proc *proc = container_of((uv_pipe_t *)handle, struct pty_proc, pipe);
-  free(proc);
-}
-
-static void pty_proc_free(struct pty_proc *proc) {
-  uv_read_stop((uv_stream_t *)&proc->pipe);
-  close(proc->pty);
-  if (proc->pty_buffer != NULL) {
-    free(proc->pty_buffer);
-    proc->pty_buffer = NULL;
-  }
-  for (int i = 0; i < proc->argc; i++) {
-    free(proc->args[i]);
-  }
-  uv_close((uv_handle_t *)&proc->pipe, close_cb);
-}
-
-static void alloc_cb(uv_handle_t *handle, size_t suggested_size, uv_buf_t *buf) {
-  buf->base = xmalloc(suggested_size);
-  buf->len = suggested_size;
-}
-
-static void read_cb(uv_stream_t *stream, ssize_t nread, const uv_buf_t *buf) {
-  struct pss_tty *pss = (struct pss_tty *)stream->data;
-  struct pty_proc *proc = pss->proc;
-  proc->pty_len = nread;
-
-  uv_read_stop(stream);
-
-  if (nread == UV_ENOBUFS || nread == 0) return;
-  if (nread > 0) {
-    proc->pty_buffer = xmalloc(LWS_PRE + 1 + (size_t)nread);
-    memcpy(proc->pty_buffer + LWS_PRE + 1, buf->base, (size_t)nread);
-  } 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));
-    }
-  }
-  free(buf->base);
+static void process_read_cb(void *ctx, pty_buf_t *buf, bool eof) {
+  struct pss_tty *pss = (struct pss_tty *)ctx;
+  if (eof && !process_running(pss->process))
+    pss->lws_close_status = pss->process->exit_code == 0 ? 1000 : 1006;
+  else
+    pss->pty_buf = buf;
 
   lws_callback_on_writable(pss->wsi);
 }
 
-static void child_cb(uv_signal_t *handle, int signum) {
-  pid_t pid;
-  int stat;
-
-  struct pty_proc *proc;
-  LIST_HEAD(proc, pty_proc) *procs = handle->data;
-  LIST_FOREACH(proc, procs, entry) {
-    do
-      pid = waitpid(proc->pid, &stat, WNOHANG);
-    while (pid == -1 && errno == EINTR);
-
-    if (pid <= 0) continue;
-
-    if (WIFEXITED(stat)) {
-      proc->status = WEXITSTATUS(stat);
-      lwsl_notice("process exited with code %d, pid: %d\n", proc->status, proc->pid);
-    } else if (WIFSIGNALED(stat)) {
-      int sig = WTERMSIG(stat);
-      char sig_name[20];
-
-      proc->status = 128 + sig;
-      get_sig_name(sig, sig_name, sizeof(sig_name));
-      lwsl_notice("process killed with signal %d (%s), pid: %d\n", sig, sig_name, proc->pid);
-    }
-
-    LIST_REMOVE(proc, entry);
-    if (proc->state == STATE_KILL) {
-      pty_proc_free(proc);
-    } else {
-      proc->state = STATE_EXIT;
-    }
+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->lws_close_status = process->exit_code == 0 ? 1000 : 1006;
+    lws_callback_on_writable(pss->wsi);
   }
 }
 
-static int spawn_process(struct pss_tty *pss) {
-  struct pty_proc *proc = pss->proc;
-  // append url args to arguments
-  char *argv[server->argc + proc->argc + 1];
+static char **build_args(struct pss_tty *pss) {
   int i, n = 0;
+  char **argv = xmalloc((server->argc + pss->argc + 1) * sizeof(char *));
+
   for (i = 0; i < server->argc; i++) {
     argv[n++] = server->argv[i];
   }
-  for (i = 0; i < proc->argc; i++) {
-    argv[n++] = proc->args[i];
+
+  for (i = 0; i < pss->argc; i++) {
+    argv[n++] = pss->args[i];
   }
-  argv[n] = NULL;
 
-  uv_signal_start(&server->watcher, child_cb, SIGCHLD);
+  argv[n] = NULL;
 
-  // ensure the lws socket fd close-on-exec
-  fd_set_cloexec(lws_get_socket_fd(pss->wsi));
+  return argv;
+}
 
-  // create process with pseudo-tty
-  proc->pid = pty_fork(&proc->pty, argv[0], argv, server->terminal_type);
-  if (proc->pid < 0) {
-    lwsl_err("pty_fork: %d (%s)\n", errno, strerror(errno));
-    return 1;
+static char **build_env(struct pss_tty *pss) {
+  int i = 0, n = 2;
+  char **envp = xmalloc(n * sizeof(char *));
+
+  // TERM
+  envp[i] = xmalloc(36);
+  snprintf(envp[i], 36, "TERM=%s", server->terminal_type);
+  i++;
+
+  // TTYD_USER
+  if (strlen(pss->user) > 0) {
+    envp = xrealloc(envp, (++n) * sizeof(char *));
+    envp[i] = xmalloc(40);
+    snprintf(envp[i], 40, "TTYD_USER=%s", pss->user);
+    i++;
   }
 
-  lwsl_notice("started process, pid: %d\n", proc->pid);
+  envp[i] = NULL;
 
-  proc->pipe.data = pss;
-  uv_pipe_open(&proc->pipe, proc->pty);
+  return envp;
+}
 
+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));
+  if (server->cwd != NULL) process->cwd = strdup(server->cwd);
+  if (columns > 0) process->columns = columns;
+  if (rows > 0) process->rows = rows;
+  if (pty_spawn(process, process_read_cb, process_exit_cb) != 0) {
+    lwsl_err("pty_spawn: %d (%s)\n", errno, strerror(errno));
+    process_free(process);
+    return false;
+  }
+  lwsl_notice("started process, pid: %d\n", process->pid);
+  pss->process = process;
   lws_callback_on_writable(pss->wsi);
 
-  return 0;
+  return true;
 }
 
-static void kill_process(struct pty_proc *proc) {
-  if (proc->pid <= 0) return;
+static void wsi_output(struct lws *wsi, pty_buf_t *buf) {
+  if (buf == NULL) return;
+  char *message = xmalloc(LWS_PRE + 1 + buf->len);
+  char *ptr = message + LWS_PRE;
 
-  pid_t pid = proc->pid;
-  int sig = server->sig_code;
-  char *sig_name = server->sig_name;
+  *ptr = OUTPUT;
+  memcpy(ptr + 1, buf->base, buf->len);
+  size_t n = buf->len + 1;
 
-  lwsl_notice("killing process %d with signal: %d (%s)\n", pid, sig, sig_name);
-  int pgid = getpgid(pid);
-  if (uv_kill(pgid > 0 ? -pgid : pid, sig) != 0) {
-    lwsl_err("kill: %d, errno: %d (%s)\n", pid, errno, strerror(errno));
+  if (lws_write(wsi, (unsigned char *)ptr, n, LWS_WRITE_BINARY) < n) {
+    lwsl_err("write OUTPUT to WS\n");
   }
+
+  free(message);
 }
 
-static void write_cb(uv_write_t *req, int status) {
-  if (status != 0) lwsl_warn("uv_write callback returned status: %d\n", status);
-  free(req->data);
-  free(req);
+static bool check_auth(struct lws *wsi, struct pss_tty *pss) {
+  if (server->auth_header != NULL) {
+    return lws_hdr_custom_copy(wsi, pss->user, sizeof(pss->user), server->auth_header, strlen(server->auth_header)) > 0;
+  }
+
+  if (server->credential != NULL) {
+    char buf[256];
+    size_t n = lws_hdr_copy(wsi, buf, sizeof(buf), WSI_TOKEN_HTTP_AUTHORIZATION);
+    return n >= 7 && strstr(buf, "Basic ") && !strcmp(buf + 6, server->credential);
+  }
+
+  return true;
 }
 
-int callback_tty(struct lws *wsi, enum lws_callback_reasons reason, void *user, void *in,
-                 size_t len) {
+int callback_tty(struct lws *wsi, enum lws_callback_reasons reason, void *user, void *in, size_t len) {
   struct pss_tty *pss = (struct pss_tty *)user;
-  struct pty_proc *proc;
   char buf[256];
   size_t n = 0;
 
@@ -239,13 +194,14 @@ int callback_tty(struct lws *wsi, enum lws_callback_reasons reason, void *user,
         lwsl_warn("refuse to serve WS client due to the --max-clients option.\n");
         return 1;
       }
+      if (!check_auth(wsi, pss)) return 1;
 
       n = lws_hdr_copy(wsi, pss->path, sizeof(pss->path), WSI_TOKEN_GET_URI);
 #if defined(LWS_ROLE_H2)
       if (n <= 0) n = lws_hdr_copy(wsi, pss->path, sizeof(pss->path), WSI_TOKEN_HTTP_COLON_PATH);
 #endif
       if (strncmp(pss->path, endpoints.ws, n) != 0) {
-        lwsl_warn("refuse to serve WS client for illegal ws path: %s\n", buf);
+        lwsl_warn("refuse to serve WS client for illegal ws path: %s\n", pss->path);
         return 1;
       }
 
@@ -259,46 +215,31 @@ int callback_tty(struct lws *wsi, enum lws_callback_reasons reason, void *user,
 
     case LWS_CALLBACK_ESTABLISHED:
       pss->initialized = false;
-      pss->initial_cmd_index = 0;
       pss->authenticated = false;
       pss->wsi = wsi;
-      pss->buffer = NULL;
-
-      pss->proc = proc = xmalloc(sizeof(struct pty_proc));
-      memset(proc, 0, sizeof(struct pty_proc));
-      proc->status = -1;
-      proc->state = STATE_INIT;
-      uv_pipe_init(server->loop, &proc->pipe, 0);
+      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) {
           if (strncmp(buf, "arg=", 4) == 0) {
-            proc->args = xrealloc(proc->args, (proc->argc + 1) * sizeof(char *));
-            proc->args[proc->argc] = strdup(&buf[4]);
-            proc->argc++;
+            pss->args = xrealloc(pss->args, (pss->argc + 1) * sizeof(char *));
+            pss->args[pss->argc] = strdup(&buf[4]);
+            pss->argc++;
           }
         }
       }
 
-      LIST_INSERT_HEAD(&server->procs, proc, entry);
       server->client_count++;
 
-#if LWS_LIBRARY_VERSION_NUMBER >= 2004000
       lws_get_peer_simple(lws_get_network_wsi(wsi), pss->address, sizeof(pss->address));
-#else
-      char name[100];
-      lws_get_peer_addresses(wsi, lws_get_socket_fd(wsi), name, sizeof(name), pss->address,
-                             sizeof(pss->address));
-#endif
       lwsl_notice("WS   %s - %s, clients: %d\n", pss->path, pss->address, server->client_count);
       break;
 
     case LWS_CALLBACK_SERVER_WRITEABLE:
-      proc = pss->proc;
       if (!pss->initialized) {
         if (pss->initial_cmd_index == sizeof(initial_cmds)) {
           pss->initialized = true;
-          uv_read_start((uv_stream_t *)&proc->pipe, alloc_cb, read_cb);
+          pty_resume(pss->process);
           break;
         }
         if (send_initial_message(wsi, pss->initial_cmd_index) < 0) {
@@ -311,26 +252,17 @@ int callback_tty(struct lws *wsi, enum lws_callback_reasons reason, void *user,
         break;
       }
 
-      // read error or client exited, close connection
-      if (proc->status == 0 || proc->pty_len == UV_EOF) {
-        lws_close_reason(wsi, 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;
-      } 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) {
-        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;
+      if (pss->pty_buf != NULL) {
+        wsi_output(wsi, pss->pty_buf);
+        pty_buf_free(pss->pty_buf);
+        pss->pty_buf = NULL;
+        pty_resume(pss->process);
       }
-
-      uv_read_start((uv_stream_t *)&proc->pipe, alloc_cb, read_cb);
       break;
 
     case LWS_CALLBACK_RECEIVE:
@@ -357,49 +289,32 @@ int callback_tty(struct lws *wsi, enum lws_callback_reasons reason, void *user,
         return 0;
       }
 
-      proc = pss->proc;
       switch (command) {
         case INPUT:
-          if (proc->pty == 0) break;
           if (server->readonly) break;
-
-          char *data = xmalloc(pss->len - 1);
-          memcpy(data, pss->buffer + 1, pss->len - 1);
-
-          uv_buf_t b = {data, pss->len - 1};
-          uv_write_t *req = xmalloc(sizeof(uv_write_t));
-          req->data = data;
-
-          int err = uv_write(req, (uv_stream_t *)&proc->pipe, &b, 1, write_cb);
+          int err = pty_write(pss->process, pty_buf_init(pss->buffer + 1, pss->len - 1));
           if (err) {
             lwsl_err("uv_write: %s (%s)\n", uv_err_name(err), uv_strerror(err));
             return -1;
           }
           break;
-        case RESIZE_TERMINAL: {
-          int cols, rows;
-          if (parse_window_size(pss, &cols, &rows)) {
-            if (pty_resize(proc->pty, cols, rows) < 0) {
-              lwsl_err("pty_resize: %d (%s)\n", errno, strerror(errno));
-            }
-          }
-        } break;
+        case RESIZE_TERMINAL:
+          json_object_put(
+              parse_window_size(pss->buffer + 1, pss->len - 1, &pss->process->columns, &pss->process->rows));
+          pty_resize(pss->process);
+          break;
         case PAUSE:
-          if (proc->state == STATE_INIT) {
-            uv_read_stop((uv_stream_t *)&proc->pipe);
-            proc->state = STATE_PAUSE;
-          }
+          pty_pause(pss->process);
           break;
         case RESUME:
-          if (proc->state == STATE_PAUSE) {
-            uv_read_start((uv_stream_t *)&proc->pipe, alloc_cb, read_cb);
-            proc->state = STATE_INIT;
-          }
+          pty_resume(pss->process);
           break;
         case JSON_DATA:
-          if (proc->pid > 0) break;
+          if (pss->process != NULL) break;
+          uint16_t columns = 0;
+          uint16_t rows = 0;
+          json_object *obj = parse_window_size(pss->buffer, pss->len, &columns, &rows);
           if (server->credential != NULL) {
-            json_object *obj = json_tokener_parse(pss->buffer);
             struct json_object *o = NULL;
             if (json_object_object_get_ex(obj, "AuthToken", &o)) {
               const char *token = json_object_get_string(o);
@@ -409,11 +324,13 @@ int callback_tty(struct lws *wsi, enum lws_callback_reasons reason, void *user,
                 lwsl_warn("WS authentication failed with token: %s\n", token);
             }
             if (!pss->authenticated) {
+              json_object_put(obj);
               lws_close_reason(wsi, LWS_CLOSE_STATUS_POLICY_VIOLATION, NULL, 0);
               return -1;
             }
           }
-          if (spawn_process(pss) != 0) return 1;
+          json_object_put(obj);
+          if (!spawn_process(pss, columns, rows)) return 1;
           break;
         default:
           lwsl_warn("ignored unknown message type: %c\n", command);
@@ -434,14 +351,14 @@ int callback_tty(struct lws *wsi, enum lws_callback_reasons reason, void *user,
       if (pss->buffer != NULL) {
         free(pss->buffer);
       }
+      for (int i = 0; i < pss->argc; i++) {
+        free(pss->args[i]);
+      }
 
-      proc = pss->proc;
-      if (proc->state == STATE_EXIT) {
-        pty_proc_free(proc);
-      } else {
-        proc->state = STATE_KILL;
-        uv_read_stop((uv_stream_t *)&proc->pipe);
-        kill_process(proc);
+      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) {