SimpleRemoter/common/PTYHandler.h

/**
 * PTYHandler.h - Unix Pseudo Terminal Handler
 *
 * This file provides pseudo terminal (PTY) functionality for remote shell access.
 *
 * PLATFORM SUPPORT:
 *   - Linux:  Supported
 *   - macOS:  Supported
 *   - Windows: NOT SUPPORTED (Windows uses different terminal APIs)
 *
 * USAGE:
 *   #include "common/PTYHandler.h"
 *
 *   PTYHandler* handler = new PTYHandler(clientObject);
 *   clientObject->setManagerCallBack(handler, ...);
 */

#pragma once

#if defined(_WIN32) || defined(_WIN64)
#error "PTYHandler.h is not supported on Windows. Use Windows ConPTY or other APIs instead."
#endif

// Platform-specific includes
#ifdef __APPLE__
#include <util.h>       // macOS: openpty()
#else
#include <pty.h>        // Linux: openpty()
#endif

// Common Unix includes
#include <unistd.h>
#include <fcntl.h>
#include <termios.h>
#include <sys/ioctl.h>
#include <sys/wait.h>
#include <signal.h>
#include <errno.h>
#include <string.h>
#include <thread>
#include <atomic>
#include <stdexcept>

#include "commands.h"
#include "../client/IOCPClient.h"

/**
 * PTYHandler - Pseudo Terminal Handler
 *
 * Manages a pseudo terminal for remote shell access.
 * Inherits from IOCPManager to integrate with the IOCP client framework.
 */
class PTYHandler : public IOCPManager
{
public:
    // Non-copyable, non-movable (owns system resources)
    PTYHandler(const PTYHandler&) = delete;
    PTYHandler& operator=(const PTYHandler&) = delete;
    PTYHandler(PTYHandler&&) = delete;
    PTYHandler& operator=(PTYHandler&&) = delete;

    PTYHandler(IOCPClient* client) : m_client(client), m_running(false), m_master_fd(-1), m_slave_fd(-1), m_child_pid(-1)
    {
        if (!client) {
            throw std::invalid_argument("IOCPClient pointer cannot be null");
        }

        // Create pseudo terminal pair
        if (openpty(&m_master_fd, &m_slave_fd, nullptr, nullptr, nullptr) == -1) {
            throw std::runtime_error("Failed to create pseudo terminal");
        }

        // Set master fd to non-blocking mode
        int flags = fcntl(m_master_fd, F_GETFL, 0);
        fcntl(m_master_fd, F_SETFL, flags | O_NONBLOCK);

        // Start shell process
        startShell();
    }

    ~PTYHandler()
    {
        m_running = false;
        if (m_readThread.joinable()) {
            m_readThread.join();
        }
        if (m_master_fd >= 0) {
            close(m_master_fd);
        }
        if (m_slave_fd >= 0) {
            close(m_slave_fd);
        }
        if (m_child_pid > 0) {
            // Check if child is still running before killing
            int status;
            pid_t result = waitpid(m_child_pid, &status, WNOHANG);
            if (result == 0) {
                // Child still running, terminate it
                kill(m_child_pid, SIGTERM);
                waitpid(m_child_pid, nullptr, 0);
            }
            // If result == m_child_pid, child already exited and was reaped
            // If result == -1, child was already reaped elsewhere
        }
    }

    // Start the PTY read thread
    void Start()
    {
        bool expected = false;
        if (!m_running.compare_exchange_strong(expected, true)) return;
        m_readThread = std::thread(&PTYHandler::readFromPTY, this);
    }

    // Handle incoming data from server
    virtual VOID OnReceive(PBYTE data, ULONG size) override
    {
        if (size && data[0] == COMMAND_NEXT) {
            Start();
            return;
        }

        // Handle terminal resize command
        if (size >= 5 && data[0] == CMD_TERMINAL_RESIZE) {
            short cols, rows;
            memcpy(&cols, data + 1, sizeof(short));
            memcpy(&rows, data + 3, sizeof(short));
            SetWindowSize(cols, rows);
            return;
        }

        // Write data to PTY
        if (size > 0) {
            ssize_t total = 0;
            while (total < (ssize_t)size) {
                ssize_t written = write(m_master_fd, (char*)data + total, size - total);
                if (written == -1) {
                    if (errno == EAGAIN || errno == EINTR) continue;
                    break;
                }
                total += written;
            }
        }
    }

    // Set terminal window size
    void SetWindowSize(int cols, int rows)
    {
        struct winsize ws;
        ws.ws_col = cols;
        ws.ws_row = rows;
        ws.ws_xpixel = 0;
        ws.ws_ypixel = 0;

        if (ioctl(m_master_fd, TIOCSWINSZ, &ws) == -1) {
            return;
        }

        // Send SIGWINCH to child process to notify window size change
        if (m_child_pid > 0) {
            kill(m_child_pid, SIGWINCH);
        }
    }

private:
    int m_master_fd;
    int m_slave_fd;
    IOCPClient* m_client;
    std::thread m_readThread;
    std::atomic<bool> m_running;
    pid_t m_child_pid;

    void startShell()
    {
        m_child_pid = fork();
        if (m_child_pid == -1) {
            close(m_master_fd);
            close(m_slave_fd);
            throw std::runtime_error("Failed to fork shell process");
        }

        if (m_child_pid == 0) {
            // Child process
            setsid();  // Create new session, become session leader

            // Set slave PTY as controlling terminal (required for Ctrl+C to work)
            // This must be done after setsid() and before dup2()
            ioctl(m_slave_fd, TIOCSCTTY, 0);

            // Redirect stdin/stdout/stderr to slave PTY
            dup2(m_slave_fd, STDIN_FILENO);
            dup2(m_slave_fd, STDOUT_FILENO);
            dup2(m_slave_fd, STDERR_FILENO);
            close(m_master_fd);
            close(m_slave_fd);

            // Set terminal environment for xterm.js compatibility
            setenv("TERM", "xterm-256color", 1);
            setenv("COLORTERM", "truecolor", 1);

#ifdef __APPLE__
            // macOS locale settings
            setenv("LANG", "en_US.UTF-8", 1);
            setenv("LC_ALL", "en_US.UTF-8", 1);
            // Disable zsh session save/restore (causes errors in PTY)
            setenv("SHELL_SESSIONS_DISABLE", "1", 1);

            // Try zsh first (macOS default), fallback to bash
            if (access("/bin/zsh", X_OK) == 0) {
                execl("/bin/zsh", "zsh", "-i", nullptr);
            }
            execl("/bin/bash", "bash", "-i", nullptr);
#else
            // Linux locale settings (C.UTF-8 is most portable)
            setenv("LANG", "C.UTF-8", 1);
            setenv("LC_ALL", "C.UTF-8", 1);

            // Start interactive bash
            execl("/bin/bash", "bash", "-i", nullptr);
#endif
            _exit(1);
        }
    }

    void readFromPTY()
    {
        char buffer[4096];
        while (m_running) {
            // Check if child process has exited
            int status;
            pid_t result = waitpid(m_child_pid, &status, WNOHANG);
            if (result == m_child_pid) {
                // Shell exited, send close notification
                if (m_client) {
                    BYTE closeToken = TOKEN_TERMINAL_CLOSE;
                    m_client->Send2Server((char*)&closeToken, 1);
                }
                m_running = false;
                break;
            }

            ssize_t bytes_read = read(m_master_fd, buffer, sizeof(buffer) - 1);
            if (bytes_read > 0) {
                if (m_client) {
                    m_client->Send2Server(buffer, bytes_read);
                }
            } else if (bytes_read == 0) {
                // EOF - PTY closed
                if (m_client) {
                    BYTE closeToken = TOKEN_TERMINAL_CLOSE;
                    m_client->Send2Server((char*)&closeToken, 1);
                }
                m_running = false;
                break;
            } else if (bytes_read == -1) {
                if (errno == EAGAIN || errno == EWOULDBLOCK) {
                    usleep(10000);  // 10ms
                } else if (errno == EIO) {
                    // EIO typically means PTY slave closed (shell exited)
                    if (m_client) {
                        BYTE closeToken = TOKEN_TERMINAL_CLOSE;
                        m_client->Send2Server((char*)&closeToken, 1);
                    }
                    m_running = false;
                    break;
                } else {
                    break;
                }
            }
        }
    }
};