pty: set TTYD_USER env

src/protocol.c

@@ -93,22 +93,49 @@ static void process_exit_cb(void *ctx, pty_process *process) {
   }
 }
 
-static bool spawn_process(struct pss_tty *pss, uint16_t columns, uint16_t rows) {
-  // append url args to arguments
-  char **argv = xmalloc((server->argc + pss->argc + 1) * sizeof(char *));
+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 < pss->argc; i++) {
     argv[n++] = pss->args[i];
   }
+
   argv[n] = NULL;
 
-  pty_process *process = process_init((void *)pss, server->loop, argv);
+  return argv;
+}
+
+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++;
+  }
+
+  envp[i] = NULL;
+
+  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 (columns > 0) process->columns = columns;
   if (rows > 0) process->rows = rows;
-  strncpy(process->term, server->terminal_type, sizeof(process->term));
   if (pty_spawn(process, process_read_cb, process_exit_cb) != 0) {
     lwsl_err("pty_spawn: %d (%s)\n", errno, strerror(errno));
     process_free(process);
@@ -137,9 +164,9 @@ static void wsi_output(struct lws *wsi, pty_buf_t *buf) {
   free(message);
 }
 
-static bool check_auth(struct lws *wsi) {
+static bool check_auth(struct lws *wsi, struct pss_tty *pss) {
   if (server->auth_header != NULL) {
-    return lws_hdr_custom_length(wsi, server->auth_header, strlen(server->auth_header)) > 0;
+    return lws_hdr_custom_copy(wsi, pss->user, sizeof(pss->user), server->auth_header, strlen(server->auth_header)) > 0;
   }
 
   if (server->credential != NULL) {
@@ -147,7 +174,7 @@ static bool check_auth(struct lws *wsi) {
     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;
 }
 
@@ -166,7 +193,7 @@ 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)) 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)

src/pty.c

@@ -9,6 +9,7 @@
 #ifndef _WIN32
 #include <sys/ioctl.h>
 #include <sys/wait.h>
+
 #if defined(__OpenBSD__) || defined(__APPLE__)
 #include <util.h>
 #elif defined(__FreeBSD__)
@@ -16,6 +17,13 @@
 #else
 #include <pty.h>
 #endif
+
+#if defined(__APPLE__)
+#include <crt_externs.h>
+#define environ (*_NSGetEnviron())
+#else
+extern char **environ;
+#endif
 #endif
 
 #include "pty.h"
@@ -86,12 +94,13 @@ static void pty_io_free(pty_io_t *io) {
   free(io);
 }
 
-pty_process *process_init(void *ctx, uv_loop_t *loop, char **argv) {
+pty_process *process_init(void *ctx, uv_loop_t *loop, char *argv[], char *envp[]) {
   pty_process *process = xmalloc(sizeof(pty_process));
   memset(process, 0, sizeof(pty_process));
   process->ctx = ctx;
   process->loop = loop;
   process->argv = argv;
+  process->envp = envp;
   process->columns = 80;
   process->rows = 24;
   process->exit_code = -1;
@@ -116,6 +125,9 @@ void process_free(pty_process *process) {
 #endif
   if (process->io != NULL) pty_io_free(process->io);
   if (process->argv != NULL) free(process->argv);
+  char **p = process->envp;
+  for (; *p; p++) free(*p);
+  free(process->envp);
   free(process);
 }
 
@@ -194,15 +206,16 @@ bool conpty_init() {
 // convert argv to cmdline for CreateProcessW
 static WCHAR *join_args(char **argv) {
   char *args = NULL;
-  for (; *argv; argv++) {
-    char *quoted = (char *) quote_arg(*argv);
+  char **ptr = argv;
+  for (; *ptr; ptr++) {
+    char *quoted = (char *) quote_arg(*ptr);
     size_t arg_len = args == NULL ? 1 : strlen(args) + 1;
     size_t quoted_len = strlen(quoted);
     args = xrealloc(args, arg_len + quoted_len);
     if (arg_len == 1) memset(args, 0, 2);
     if (arg_len != 1) strcat(args, " ");
     strncat(args, quoted, quoted_len);
-    if (quoted != *argv) free(quoted);
+    if (quoted != *ptr) free(quoted);
   }
 
   int len = MultiByteToWideChar(CP_UTF8, 0, args, -1, NULL, 0);
@@ -219,6 +232,30 @@ failed:
   return NULL;
 }
 
+static WCHAR *prep_env(char **envp) {
+  WCHAR *env = NULL;
+  size_t env_len = 0;
+  char **ptr = envp;
+  for (; *ptr; ptr++) {
+    int len = MultiByteToWideChar(CP_UTF8, 0, *ptr, -1, NULL, 0);
+    if (len <= 0) goto failed;
+    env_len += (len + 1);
+
+    env = xrealloc(env, env_len * sizeof(WCHAR));
+    if (len != MultiByteToWideChar(CP_UTF8, 0, *ptr, -1, env+(env_len-len-1), len)) goto failed;
+    env[env_len-1] = L'\0';
+  }
+
+  env_len++;
+  env = xrealloc(env, env_len * sizeof(WCHAR));
+  env[env_len-1] = L'\0';
+  return env;
+
+failed:
+  if (env != NULL) free(env);
+  return NULL;
+}
+
 static bool conpty_setup(HPCON *hnd, COORD size, STARTUPINFOEXW *si_ex, char **in_name, char **out_name) {
   static int count = 0;
   char buf[256];
@@ -305,7 +342,6 @@ static void async_cb(uv_async_t *async) {
 }
 
 int pty_spawn(pty_process *process, pty_read_cb read_cb, pty_exit_cb exit_cb) {
-  WCHAR *cmdline = NULL;
   char *in_name = NULL;
   char *out_name = NULL;
   DWORD flags = EXTENDED_STARTUPINFO_PRESENT | CREATE_UNICODE_ENVIRONMENT;
@@ -324,15 +360,13 @@ int pty_spawn(pty_process *process, pty_read_cb read_cb, pty_exit_cb exit_cb) {
   uv_pipe_connect(out_req, io->out, out_name, connect_cb);
 
   PROCESS_INFORMATION pi = {0};
+  WCHAR *cmdline, *env;
   cmdline = join_args(process->argv);
   if (cmdline == NULL) goto cleanup;
+  env = prep_env(process->envp);
+  if (env == NULL) goto cleanup;
 
-  // TODO: restore env after process creation
-  if (!SetEnvironmentVariable("TERM", process->term)) {
-    print_error("SetEnvironmentVariable");
-  }
-
-  if (!CreateProcessW(NULL, cmdline, NULL, NULL, FALSE, flags, NULL, NULL, &process->si.StartupInfo, &pi)) {
+  if (!CreateProcessW(NULL, cmdline, NULL, NULL, FALSE, flags, env, NULL, &process->si.StartupInfo, &pi)) {
     print_error("CreateProcessW");
     goto cleanup;
   }
@@ -356,7 +390,7 @@ cleanup:
   if (in_name != NULL) free(in_name);
   if (out_name != NULL) free(out_name);
   if (cmdline != NULL) free(cmdline);
-
+  if (env != NULL) free(env);
   return status;
 }
 #else
@@ -406,6 +440,14 @@ static void async_cb(uv_async_t *async) {
   process_free(process);
 }
 
+static int pty_execvpe(const char *file, char **argv, char **envp) {
+  char **old = environ;
+  environ = envp;
+  int ret = execvp(file, argv);
+  environ = old;
+  return ret;
+}
+
 int pty_spawn(pty_process *process, pty_read_cb read_cb, pty_exit_cb exit_cb) {
   int status = 0;
 
@@ -419,8 +461,7 @@ int pty_spawn(pty_process *process, pty_read_cb read_cb, pty_exit_cb exit_cb) {
     return status;
   } else if (pid == 0) {
     setsid();
-    setenv("TERM", process->term, true);
-    int ret = execvp(process->argv[0], process->argv);
+    int ret = pty_execvpe(process->argv[0], process->argv, process->envp);
     if (ret < 0) {
       perror("execvp failed\n");
       _exit(-errno);

src/pty.h

@@ -39,7 +39,6 @@ typedef void (*pty_exit_cb)(void *, pty_process *);
 struct pty_process_ {
   int pid, exit_code, exit_signal;
   uint16_t columns, rows;
-  char term[30];
   bool killed;
 #ifdef _WIN32
   STARTUPINFOEXW si;
@@ -51,6 +50,7 @@ struct pty_process_ {
   uv_thread_t tid;
 #endif
   char **argv;
+  char **envp;
 
   uv_loop_t *loop;
   uv_async_t async;
@@ -61,7 +61,7 @@ struct pty_process_ {
 
 pty_buf_t *pty_buf_init(char *base, size_t len);
 void pty_buf_free(pty_buf_t *buf);
-pty_process *process_init(void *ctx, uv_loop_t *loop, char **argv);
+pty_process *process_init(void *ctx, uv_loop_t *loop, char *argv[], char *envp[]);
 bool process_running(pty_process *process);
 void process_free(pty_process *process);
 int pty_spawn(pty_process *process, pty_read_cb read_cb, pty_exit_cb exit_cb);

src/server.c

@@ -201,10 +201,8 @@ static void server_free(struct server *ts) {
   free(ts->command);
   free(ts->prefs_json);
 
-  int i = 0;
-  do {
-    free(ts->argv[i++]);
-  } while (ts->argv[i] != NULL);
+  char **p = ts->argv;
+  for (; *p; p++) free(*p);
   free(ts->argv);
 
   if (strlen(ts->socket_path) > 0) {

src/server.h

@@ -40,6 +40,7 @@ struct pss_tty {
   bool initialized;
   int initial_cmd_index;
   bool authenticated;
+  char user[30];
   char address[50];
   char path[128];
   char **args;