From 968f1e6e707f5434ff8639f3743ac1a31c47f5be Mon Sep 17 00:00:00 2001 From: Shuanglei Tao Date: Sat, 14 Aug 2021 08:52:53 +0800 Subject: [PATCH] pty: set TTYD_USER env --- src/protocol.c | 45 +++++++++++++++++++++++++------- src/pty.c | 69 ++++++++++++++++++++++++++++++++++++++++---------- src/pty.h | 4 +-- src/server.c | 6 ++--- src/server.h | 1 + 5 files changed, 96 insertions(+), 29 deletions(-) diff --git a/src/protocol.c b/src/protocol.c index b6928a29..0cecd26f 100644 --- a/src/protocol.c +++ b/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) diff --git a/src/pty.c b/src/pty.c index 97acad97..8171ac54 100644 --- a/src/pty.c +++ b/src/pty.c @@ -9,6 +9,7 @@ #ifndef _WIN32 #include #include + #if defined(__OpenBSD__) || defined(__APPLE__) #include #elif defined(__FreeBSD__) @@ -16,6 +17,13 @@ #else #include #endif + +#if defined(__APPLE__) +#include +#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 @@ static WCHAR *join_args(char **argv) { 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 @@ int pty_spawn(pty_process *process, pty_read_cb read_cb, pty_exit_cb exit_cb) { 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); diff --git a/src/pty.h b/src/pty.h index f891315b..c4de47b0 100644 --- a/src/pty.h +++ b/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); diff --git a/src/server.c b/src/server.c index b0a33751..f5b985b6 100644 --- a/src/server.c +++ b/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) { diff --git a/src/server.h b/src/server.h index d12bf6a6..b337536f 100644 --- a/src/server.h +++ b/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;