wterm/src/main.c

#include "terminal.h"
#include <cJSON.h>
#include <errno.h>
#include <mongoose.h>
#include <signal.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/stat.h>
#include <sys/wait.h>

#define HTTP_PORT "8888"
#define STATIC_DIR "./static/"
#define INPUT_BUFFER_SIZE 256

static const char *s_http_addr = "http://0.0.0.0:8888";
static struct mg_mgr mgr;

struct client_data {
  struct terminal_session term;
  bool terminal_ready;
  bool closing; // Флаг для закрытия соединения
  char input_buffer[INPUT_BUFFER_SIZE];
  size_t input_buffer_len;
};

// Функция для логирования запросов
static void log_request(struct mg_connection *c, struct mg_http_message *hm) {
  printf("HTTP Request: %.*s %.*s\n", (int)hm->method.len, hm->method.buf,
         (int)hm->uri.len, hm->uri.buf);
}

// Проверяем, жив ли дочерний процесс
static bool is_process_alive(pid_t pid) {
  if (pid <= 0)
    return false;

  int status;
  pid_t result = waitpid(pid, &status, WNOHANG);
  if (result == 0) {
    // Процесс все еще работает
    return true;
  } else if (result == -1) {
    // Ошибка
    if (errno == ECHILD) {
      // Процесс не существует
      return false;
    }
    perror("waitpid");
    return false;
  } else {
    // Процесс завершился
    return false;
  }
}

// Перезапускаем терминальную сессию
static void restart_terminal_session(struct mg_connection *c,
                                     struct client_data *cd) {
  if (!cd)
    return;

  printf("Restarting terminal session...\n");

  // Уничтожаем старую сессию
  if (cd->terminal_ready) {
    destroy_terminal_session(&cd->term);
  }

  // Создаем новую сессию
  const char *shell = getenv("SHELL");
  if (!shell)
    shell = "/bin/bash";

  if (create_terminal_session(&cd->term, shell) == 0) {
    cd->terminal_ready = true;
    mg_ws_send(c, "\r\n*** Terminal restarted ***\r\n", 29, WEBSOCKET_OP_TEXT);
  } else {
    mg_ws_send(c, "\r\n*** Failed to restart terminal ***\r\n", 37,
               WEBSOCKET_OP_TEXT);
    cd->terminal_ready = false;
  }
}

// Функция для проверки команды exit в буфере
static bool check_for_exit_command(struct mg_connection *c,
                                   struct client_data *cd) {
  for (size_t i = 0; i < cd->input_buffer_len; i++) {
    if (cd->input_buffer[i] == 0x04) {
      printf("Ctrl+D pressed\n");
      write_to_terminal(&cd->term, &cd->input_buffer[i], 1);
      usleep(100);
      mg_ws_send(c, "CLOSE_TAB", 9, WEBSOCKET_OP_TEXT);
      cd->closing = true;
      cd->input_buffer_len = 0;
      return true;
    }
  }
  // Ищем символ новой строки в буфере
  for (size_t i = 0; i < cd->input_buffer_len; i++) {
    if (cd->input_buffer[i] == '\n' || cd->input_buffer[i] == '\r') {
      // Нашли конец строки, проверяем команду
      cd->input_buffer[i] = '\0'; // Завершаем строку
      printf("Command: %s\n", cd->input_buffer);
      if (cd->input_buffer[i - 1] == 'x') {
        printf("Ctrl+D pressed\n");
        write_to_terminal(&cd->term, &cd->input_buffer[i], 1);
        usleep(100);
        mg_ws_send(c, "CLOSE_TAB", 9, WEBSOCKET_OP_TEXT);
        cd->closing = true;
        cd->input_buffer_len = 0;
        return true;
      }
      // Проверяем, является ли команда "exit"
      if (strcmp(cd->input_buffer, "exit") == 0) {
        printf("Exit command detected, closing terminal session\n");

        // Ждем немного чтобы команда выполнилась
        usleep(100);

        // Отправляем клиенту сообщение о закрытии вкладки
        mg_ws_send(c, "CLOSE_TAB", 9, WEBSOCKET_OP_TEXT);

        // Помечаем соединение для закрытия
        cd->closing = true;

        // Сбрасываем буфер
        cd->input_buffer_len = 0;
        return true;
      }

      // Сдвигаем буфер, удаляя обработанную строку
      size_t remaining = cd->input_buffer_len - i - 1;
      if (remaining > 0) {
        memmove(cd->input_buffer, cd->input_buffer + i + 1, remaining);
      }
      cd->input_buffer_len = remaining;
      return false;
    }
  }
  return false;
}

// Обработчик WebSocket сообщений
static void handle_websocket_message(struct mg_connection *c,
                                     struct mg_ws_message *wm,
                                     struct client_data *cd) {
  if (!cd)
    return;

  if (wm->data.len == 0) {
    return;
  }

  // Парсим JSON для проверки resize сообщения
  cJSON *json = cJSON_ParseWithLength(wm->data.buf, wm->data.len);
  if (json != NULL) {
    cJSON *cols = cJSON_GetObjectItemCaseSensitive(json, "cols");
    cJSON *rows = cJSON_GetObjectItemCaseSensitive(json, "rows");

    if (cJSON_IsNumber(cols) && cJSON_IsNumber(rows)) {
      printf("Resizing terminal to %dx%d\n", rows->valueint, cols->valueint);
      resize_terminal(&cd->term, rows->valueint, cols->valueint);
      cJSON_Delete(json);
      return;
    }
    cJSON_Delete(json);
  }

  // Добавляем данные в буфер ввода
  size_t bytes_to_copy = wm->data.len;
  if (cd->input_buffer_len + bytes_to_copy >= INPUT_BUFFER_SIZE) {
    // Буфер переполнен, сбрасываем его
    cd->input_buffer_len = 0;
  }

  memcpy(cd->input_buffer + cd->input_buffer_len, wm->data.buf, bytes_to_copy);
  cd->input_buffer_len += bytes_to_copy;

  printf("Input received: %s\n", wm->data.buf);
  if (check_for_exit_command(c, cd)) {
    return; // Вкладка будет закрыта, выходим
  }

  // Обычный ввод от пользователя
  write_to_terminal(&cd->term, wm->data.buf, wm->data.len);
}

// Обработчик HTTP запросов
static void fn(struct mg_connection *c, int ev, void *ev_data) {
  if (ev == MG_EV_HTTP_MSG) {
    struct mg_http_message *hm = (struct mg_http_message *)ev_data;
    log_request(c, hm);

    if (mg_match(hm->uri, mg_str("/ws"), nullptr)) {
      // Проверяем, нет ли уже активного соединения
      if (c->fn_data != NULL) {
        printf("WebSocket connection already exists for this connection, "
               "cleaning up...\n");
        struct client_data *old_cd = (struct client_data *)c->fn_data;
        destroy_terminal_session(&old_cd->term);
        free(old_cd);
        c->fn_data = NULL;
      }

      // Upgrade до WebSocket
      struct client_data *cd = calloc(1, sizeof(struct client_data));
      if (!cd) {
        mg_http_reply(c, 500, "Content-Type: text/plain\r\n",
                      "Memory allocation failed");
        return;
      }

      // Создаем терминальную сессию
      const char *shell = getenv("SHELL");
      if (!shell)
        shell = "/bin/bash";

      printf("Creating new terminal session with shell: %s\n", shell);
      if (create_terminal_session(&cd->term, shell) != 0) {
        mg_http_reply(c, 500, "Content-Type: text/plain\r\n",
                      "Failed to create terminal session");
        free(cd);
        return;
      }

      cd->terminal_ready = true;
      cd->closing = false;
      c->fn_data = cd;

      // Upgrade to WebSocket
      struct mg_str *ws_proto =
          mg_http_get_header(hm, "Sec-WebSocket-Protocol");
      if (ws_proto != NULL) {
        mg_ws_upgrade(c, hm, ws_proto->buf);
      } else {
        mg_ws_upgrade(c, hm, NULL);
      }

      printf("WebSocket upgrade completed\n");

    } else if (mg_match(hm->uri, mg_str("/"), nullptr)) {
      // Главная страница
      struct mg_http_serve_opts opts = {.root_dir = STATIC_DIR};
      mg_http_serve_file(c, hm, "./static/index.html", &opts);

    } else if (mg_match(hm->uri, mg_str("/xterm.js"), nullptr)) {
      // xterm.js
      struct mg_http_serve_opts opts = {.root_dir = STATIC_DIR};
      mg_http_serve_file(c, hm, "./static/xterm.js", &opts);

    } else if (mg_match(hm->uri, mg_str("/xterm.css"), nullptr)) {
      // xterm.css
      struct mg_http_serve_opts opts = {.root_dir = STATIC_DIR};
      mg_http_serve_file(c, hm, "./static/xterm.css", &opts);

    } else if (mg_match(hm->uri, mg_str("/xterm-addon-fit.js"), nullptr)) {
      // xterm-addon-fit.js
      struct mg_http_serve_opts opts = {.root_dir = STATIC_DIR};
      mg_http_serve_file(c, hm, "./static/xterm-addon-fit.js", &opts);

    } else {
      mg_http_reply(c, 404, "Content-Type: text/plain\r\n", "Not Found");
    }
  } else if (ev == MG_EV_WS_OPEN) {
    // WebSocket соединение установлено
    printf("WebSocket connection opened\n");
    struct client_data *cd = (struct client_data *)c->fn_data;
    if (cd) {
      mg_ws_send(c, "Terminal ready - type 'exit' to close tab\r\n", 44,
                 WEBSOCKET_OP_TEXT);
    }
  } else if (ev == MG_EV_WS_MSG) {
    // Сообщение от WebSocket клиента
    struct mg_ws_message *wm = (struct mg_ws_message *)ev_data;
    struct client_data *cd = (struct client_data *)c->fn_data;
    handle_websocket_message(c, wm, cd);
  } else if (ev == MG_EV_CLOSE) {
    // Закрытие соединения
    struct client_data *cd = (struct client_data *)c->fn_data;
    if (cd) {
      printf("WebSocket connection closed, cleaning up...\n");
      destroy_terminal_session(&cd->term);
      free(cd);
      c->fn_data = NULL;
    }
  } else if (ev == MG_EV_POLL) {
    // Проверяем данные от терминала и отправляем через WebSocket
    struct mg_connection *conn;
    for (conn = mgr.conns; conn != NULL; conn = conn->next) {
      if (!conn->is_websocket)
        continue;

      struct client_data *cd = (struct client_data *)conn->fn_data;
      if (cd && cd->terminal_ready && !cd->closing) {
        // Проверяем, не завершился ли процесс
        if (!is_process_alive(cd->term.child_pid)) {
          printf("Terminal process died, restarting...\n");
          restart_terminal_session(conn, cd);
          continue;
        }

        // Читаем данные из терминала
        char buffer[1024];
        ssize_t bytes_read =
            read_from_terminal(&cd->term, buffer, sizeof(buffer));
        if (bytes_read > 0) {
          mg_ws_send(conn, buffer, bytes_read, WEBSOCKET_OP_BINARY);
        } else if (bytes_read == 0) {
          // EOF - процесс завершился
          printf("EOF from terminal, restarting...\n");
          restart_terminal_session(conn, cd);
        } else if (bytes_read < 0 && errno != EAGAIN && errno != EWOULDBLOCK) {
          // Ошибка чтения
          printf("Error reading from terminal: %s\n", strerror(errno));
          restart_terminal_session(conn, cd);
        }
      }
    }
  }
}

// Обработчик сигналов для graceful shutdown
static void signal_handler(int sig) {
  printf("\nShutting down server...\n");

  // Очищаем все соединения
  struct mg_connection *conn;
  for (conn = mgr.conns; conn != NULL; conn = conn->next) {
    if (conn->fn_data) {
      struct client_data *cd = (struct client_data *)conn->fn_data;
      destroy_terminal_session(&cd->term);
      free(cd);
      conn->fn_data = NULL;
    }
  }

  mg_mgr_free(&mgr);
  exit(0);
}

// Проверяем существование директории
static bool directory_exists(const char *path) {
  struct stat st;
  return stat(path, &st) == 0 && S_ISDIR(st.st_mode);
}

int main(void) {
  // Устанавливаем обработчики сигналов
  signal(SIGINT, signal_handler);
  signal(SIGTERM, signal_handler);

  // Инициализируем менеджер соединений
  mg_mgr_init(&mgr);
  printf("Mongoose version: %s\n", MG_VERSION);

  // Проверяем существование статической директории
  if (!directory_exists(STATIC_DIR)) {
    fprintf(stderr, "Static directory '%s' not found! Creating it...\n",
            STATIC_DIR);
    if (mkdir(STATIC_DIR, 0755) != 0) {
      fprintf(stderr, "Failed to create static directory\n");
      return 1;
    }
  }

  // Создаем HTTP сервер
  struct mg_connection *c = mg_http_listen(&mgr, s_http_addr, fn, NULL);
  if (c == NULL) {
    fprintf(stderr, "Failed to start server on %s\n", s_http_addr);
    return 1;
  }

  printf("Starting web terminal server on %s\n", s_http_addr);
  printf("Static files directory: %s\n", STATIC_DIR);
  printf("Server is ready. Open http://localhost:8000 in your browser\n");

  // Основной цикл
  while (true) {
    mg_mgr_poll(&mgr, 50); // 50ms timeout
  }

  mg_mgr_free(&mgr);
  return 0;
}