diff options
Diffstat (limited to 'src/protocol.c')
| -rw-r--r-- | src/protocol.c | 395 |
1 files changed, 395 insertions, 0 deletions
diff --git a/src/protocol.c b/src/protocol.c new file mode 100644 index 0000000..3a53f96 --- /dev/null +++ b/src/protocol.c @@ -0,0 +1,395 @@ +#include <errno.h> +#include <json.h> +#include <libwebsockets.h> +#include <stdbool.h> +#include <stdio.h> +#include <stdlib.h> +#include <string.h> + +#include "pty.h" +#include "server.h" +#include "utils.h" + +// initial message list +static char initial_cmds[] = {SET_WINDOW_TITLE, SET_PREFERENCES}; + +static int send_initial_message(struct lws *wsi, int index) { + unsigned char message[LWS_PRE + 1 + 4096]; + unsigned char *p = &message[LWS_PRE]; + char buffer[128]; + int n = 0; + + char cmd = initial_cmds[index]; + switch (cmd) { + case SET_WINDOW_TITLE: + gethostname(buffer, sizeof(buffer) - 1); + n = sprintf((char *)p, "%c%s (%s)", cmd, server->command, buffer); + break; + case SET_PREFERENCES: + n = sprintf((char *)p, "%c%s", cmd, server->prefs_json); + break; + default: + break; + } + + return lws_write(wsi, p, (size_t)n, LWS_WRITE_BINARY); +} + +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)) *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); + + json_tokener_free(tok); + return obj; +} + +static bool check_host_origin(struct lws *wsi) { + char buf[256]; + memset(buf, 0, sizeof(buf)); + int len = lws_hdr_copy(wsi, buf, (int)sizeof(buf), WSI_TOKEN_ORIGIN); + if (len <= 0) return false; + + const char *prot, *address, *path; + int port; + if (lws_parse_uri(buf, &prot, &address, &port, &path)) return false; + if (port == 80 || port == 443) { + sprintf(buf, "%s", address); + } else { + sprintf(buf, "%s:%d", address, port); + } + + char host_buf[256]; + memset(host_buf, 0, sizeof(host_buf)); + len = lws_hdr_copy(wsi, host_buf, (int)sizeof(host_buf), WSI_TOKEN_HOST); + + 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) { + pty_ctx_t *ctx = (pty_ctx_t *)process->ctx; + if (ctx->ws_closed) { + pty_buf_free(buf); + return; + } + + if (eof && !process_running(process)) + ctx->pss->lws_close_status = process->exit_code == 0 ? 1000 : 1006; + else + ctx->pss->pty_buf = buf; + lws_callback_on_writable(ctx->pss->wsi); +} + +static void process_exit_cb(pty_process *process) { + 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); + goto done; + } + + lwsl_notice("process exited with code %d, pid: %d\n", process->exit_code, process->pid); + 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) { + 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; + + 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 *)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; + 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 true; +} + +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; + + *ptr = OUTPUT; + memcpy(ptr + 1, buf->base, buf->len); + size_t n = buf->len + 1; + + if (lws_write(wsi, (unsigned char *)ptr, n, LWS_WRITE_BINARY) < n) { + lwsl_err("write OUTPUT to WS\n"); + } + + free(message); +} + +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) { + struct pss_tty *pss = (struct pss_tty *)user; + char buf[256]; + size_t n = 0; + + switch (reason) { + case LWS_CALLBACK_FILTER_PROTOCOL_CONNECTION: + if (server->once && server->client_count > 0) { + lwsl_warn("refuse to serve WS client due to the --once option.\n"); + return 1; + } + if (server->max_clients > 0 && server->client_count == server->max_clients) { + 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", pss->path); + return 1; + } + + if (server->check_origin && !check_host_origin(wsi)) { + lwsl_warn( + "refuse to serve WS client from different origin due to the " + "--check-origin option.\n"); + return 1; + } + break; + + case LWS_CALLBACK_ESTABLISHED: + pss->initialized = false; + pss->authenticated = false; + pss->wsi = wsi; + 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) { + pss->args = xrealloc(pss->args, (pss->argc + 1) * sizeof(char *)); + pss->args[pss->argc] = strdup(&buf[4]); + pss->argc++; + } + } + } + + server->client_count++; + + lws_get_peer_simple(lws_get_network_wsi(wsi), pss->address, sizeof(pss->address)); + lwsl_notice("WS %s - %s, clients: %d\n", pss->path, pss->address, server->client_count); + break; + + case LWS_CALLBACK_SERVER_WRITEABLE: + if (!pss->initialized) { + if (pss->initial_cmd_index == sizeof(initial_cmds)) { + pss->initialized = true; + pty_resume(pss->process); + break; + } + if (send_initial_message(wsi, pss->initial_cmd_index) < 0) { + lwsl_err("failed to send initial message, index: %d\n", pss->initial_cmd_index); + lws_close_reason(wsi, LWS_CLOSE_STATUS_UNEXPECTED_CONDITION, NULL, 0); + return -1; + } + pss->initial_cmd_index++; + lws_callback_on_writable(wsi); + break; + } + + if (pss->lws_close_status > LWS_CLOSE_STATUS_NOSTATUS) { + lws_close_reason(wsi, pss->lws_close_status, NULL, 0); + return 1; + } + + 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); + } + break; + + case LWS_CALLBACK_RECEIVE: + if (pss->buffer == NULL) { + pss->buffer = xmalloc(len); + pss->len = len; + memcpy(pss->buffer, in, len); + } else { + pss->buffer = xrealloc(pss->buffer, pss->len + len); + memcpy(pss->buffer + pss->len, in, len); + pss->len += len; + } + + const char command = pss->buffer[0]; + + // check auth + if (server->credential != NULL && !pss->authenticated && command != JSON_DATA) { + lwsl_warn("WS client not authenticated\n"); + return 1; + } + + // check if there are more fragmented messages + if (lws_remaining_packet_payload(wsi) > 0 || !lws_is_final_fragment(wsi)) { + return 0; + } + + switch (command) { + case INPUT: + if (!server->writable) break; + 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: + if (pss->process == NULL) break; + 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: + pty_pause(pss->process); + break; + case RESUME: + pty_resume(pss->process); + break; + case JSON_DATA: + 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) { + struct json_object *o = NULL; + if (json_object_object_get_ex(obj, "AuthToken", &o)) { + const char *token = json_object_get_string(o); + if (token != NULL && !strcmp(token, server->credential)) + pss->authenticated = true; + else + 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; + } + } + json_object_put(obj); + if (!spawn_process(pss, columns, rows)) return 1; + break; + default: + lwsl_warn("ignored unknown message type: %c\n", command); + break; + } + + if (pss->buffer != NULL) { + free(pss->buffer); + pss->buffer = NULL; + } + break; + + case LWS_CALLBACK_CLOSED: + if (pss->wsi == NULL) break; + + server->client_count--; + lwsl_notice("WS closed from %s, clients: %d\n", pss->address, server->client_count); + if (pss->buffer != NULL) free(pss->buffer); + if (pss->pty_buf != NULL) pty_buf_free(pss->pty_buf); + for (int i = 0; i < pss->argc; i++) { + free(pss->args[i]); + } + + 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) { + lwsl_notice("exiting due to the --once option.\n"); + force_exit = true; + lws_cancel_service(context); + exit(0); + } + break; + + default: + break; + } + + return 0; +} |