vr2bridge.cpp

#include "bridge_protocol.hpp"

#include <algorithm>
#include <errno.h>
#include <fcntl.h>
#include <signal.h>
#include <stdarg.h>
#include <stdbool.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/ioctl.h>
#include <sys/select.h>
#include <sys/stat.h>
#include <sys/types.h>
#include <sys/uio.h>
#include <sys/wait.h>
#include <termios.h>
#include <unistd.h>
#include <map>
#include <string>
#include <vector>

#define USB_DATA2_PATH "/dev/usb_data2"

static bool keep_running = true;

void debug_log(const char *fmt, ...) {
  FILE *f = fopen("/tmp/daemon_debug.log", "a");
  if (!f)
    return;
  va_list args;
  va_start(args, fmt);
  vfprintf(f, fmt, args);
  fprintf(f, "\n");
  va_end(args);
  fclose(f);
}

void sigint_handler(int signum) {
  keep_running = false;
}

class VR2BridgeDaemon {
public:
    VR2BridgeDaemon() {}
    ~VR2BridgeDaemon() { cleanup(); }

    void run() {
        debug_log("vr2bridge starting...");
        while (keep_running) {
            usb_fd = open(USB_DATA2_PATH, O_RDWR);
            if (usb_fd < 0) {
                sleep(1);
                continue;
            }

            connection_active = true;
            process_events();

            cleanup_connection();
        }
    }

private:
    int usb_fd = -1;
    bool connection_active = false;
    int file_upload_fd = -1;

    struct Session {
        uint32_t id;
        int pty_master_fd;
        pid_t shell_pid;
        char slave_name[64];
        char cmd[256];
    };

    uint32_t next_session_id = 1;
    std::map<uint32_t, Session> sessions;

    bool setup_pty(Session& sess) {
        const char *pty_chars = "0123456789abcdef";
        for (int i = 0; i < 16; i++) {
            char master_name[32];
            snprintf(master_name, sizeof(master_name), "/dev/ptyp%c", pty_chars[i]);
            sess.pty_master_fd = open(master_name, O_RDWR | O_NOCTTY);
            if (sess.pty_master_fd >= 0) {
                strncpy(sess.slave_name, ptsname(sess.pty_master_fd), sizeof(sess.slave_name) - 1);
                return true;
            }
        }
        return false;
    }

    void process_events() {
        fd_set read_fds;
        uint8_t buffer[WORKING_BUFFER_SIZE];

        while (connection_active && keep_running) {
            FD_ZERO(&read_fds);
            FD_SET(usb_fd, &read_fds);
            int max_fd = usb_fd;

            for (auto& pair : sessions) {
                FD_SET(pair.second.pty_master_fd, &read_fds);
                max_fd = std::max(max_fd, pair.second.pty_master_fd);
            }

            struct timeval tv = {0, 50000};
            int activity = select(max_fd + 1, &read_fds, NULL, NULL, &tv);
            if (activity < 0 && errno != EINTR) break;

            if (FD_ISSET(usb_fd, &read_fds)) {
                PacketHeader hdr;
                ssize_t bytes_read = read(usb_fd, &hdr, sizeof(hdr));

                if (bytes_read == sizeof(hdr) && hdr.magic == PKT_MAGIC) {
                    ssize_t payload_read = 0;
                    if (hdr.length > 0) {
                        payload_read = read(usb_fd, buffer, hdr.length);
                    }
                    dispatch_packet(hdr.type, hdr.session_id, buffer, payload_read);
                } else if (bytes_read <= 0) {
                    connection_active = false;
                }
            }

            std::vector<uint32_t> dead_sessions;
            for (auto& pair : sessions) {
                if (FD_ISSET(pair.second.pty_master_fd, &read_fds)) {
                    ssize_t n = read(pair.second.pty_master_fd, buffer, MAX_PAYLOAD_SIZE);
                    if (n > 0) {
                        send_packet(PKT_TTY_OUT, pair.first, buffer, n);
                    } else {
                        dead_sessions.push_back(pair.first);
                    }
                }
                
                int status;
                if (waitpid(pair.second.shell_pid, &status, WNOHANG) > 0) {
                    if (std::find(dead_sessions.begin(), dead_sessions.end(), pair.first) == dead_sessions.end()) {
                        dead_sessions.push_back(pair.first);
                    }
                }
            }

            for (uint32_t id : dead_sessions) {
                close_session(id);
            }
        }
    }

    void dispatch_packet(uint8_t type, uint32_t session_id, uint8_t* buffer, ssize_t payload_read) {
        switch (type) {
            case PKT_TTY_IN:     handle_tty_in(session_id, buffer, payload_read); break;
            case PKT_TTY_START:  handle_tty_start(buffer, payload_read); break;
            case PKT_TTY_STOP:   handle_tty_stop(session_id); break;
            case PKT_TTY_LIST:   handle_tty_list(); break;
            case PKT_FILE_START: handle_file_start(buffer, payload_read); break;
            case PKT_FILE_DATA:  handle_file_data(buffer, payload_read); break;
            case PKT_FILE_END:   handle_file_end(); break;
            case PKT_DOWNLOAD_REQ: handle_download_req(buffer, payload_read); break;
        }
    }

    void handle_tty_in(uint32_t session_id, uint8_t* buffer, ssize_t payload_read) {
        auto it = sessions.find(session_id);
        if (it != sessions.end() && payload_read > 0) {
            if (write(it->second.pty_master_fd, buffer, payload_read) < 0) {}
        }
    }

    void handle_tty_start(uint8_t* buffer, ssize_t payload_read) {
        Session sess;
        sess.id = next_session_id++;
        
        if (payload_read > 0) {
            buffer[payload_read] = '\0';
            strncpy(sess.cmd, (char*)buffer, sizeof(sess.cmd) - 1);
        } else {
            strcpy(sess.cmd, "sh");
        }

        if (!setup_pty(sess)) {
            return;
        }

        sess.shell_pid = fork();

        if (sess.shell_pid == 0) {
            close(usb_fd);
            for (auto& pair : sessions) close(pair.second.pty_master_fd);
            close(sess.pty_master_fd);
            
            setsid();

            int pty_slave_fd = open(sess.slave_name, O_RDWR | O_NOCTTY);
            ioctl(pty_slave_fd, TIOCSCTTY, 1);
            setpgid(0, 0);
            tcsetpgrp(pty_slave_fd, getpgrp());

            struct termios tios;
            tcgetattr(pty_slave_fd, &tios);
            cfmakeraw(&tios);
            tios.c_lflag |= (ISIG | ECHO | ICANON | IEXTEN);
            tios.c_iflag |= (ICRNL | IXON);
            tios.c_oflag |= (OPOST | ONLCR);
            tios.c_cc[VINTR] = 3;
            tios.c_cc[VQUIT] = 28;
            tios.c_cc[VEOF] = 4;
            tios.c_cc[VERASE] = 127;
            tcsetattr(pty_slave_fd, TCSANOW, &tios);

            dup2(pty_slave_fd, STDIN_FILENO);
            dup2(pty_slave_fd, STDOUT_FILENO);
            dup2(pty_slave_fd, STDERR_FILENO);

            setenv("TERM", "xterm", 1);
            execl("/bin/sh", "sh", "-c", sess.cmd, NULL);
            exit(1);
        }

        sessions[sess.id] = sess;
        send_packet(PKT_TTY_STARTED, sess.id, NULL, 0);
    }

    void handle_tty_stop(uint32_t session_id) {
        close_session(session_id);
    }

    void handle_tty_list() {
        std::string list;
        for (auto& pair : sessions) {
            list += std::to_string(pair.first) + ": " + pair.second.cmd + "\n";
        }
        if (list.empty()) {
            list = "No active sessions.\n";
        }
        send_packet(PKT_TTY_LIST_RESP, 0, (const uint8_t*)list.c_str(), list.length());
    }

    void close_session(uint32_t session_id) {
        auto it = sessions.find(session_id);
        if (it != sessions.end()) {
            kill(it->second.shell_pid, SIGKILL);
            waitpid(it->second.shell_pid, NULL, WNOHANG);
            close(it->second.pty_master_fd);
            sessions.erase(it);
            send_packet(PKT_TTY_STOP, session_id, NULL, 0);
        }
    }

    void handle_file_start(uint8_t* buffer, ssize_t payload_read) {
        buffer[payload_read] = '\0';
        file_upload_fd = open((char *)buffer, O_WRONLY | O_CREAT | O_TRUNC, 0755);
    }

    void handle_file_data(uint8_t* buffer, ssize_t payload_read) {
        if (file_upload_fd != -1 && payload_read > 0) {
            if (write(file_upload_fd, buffer, payload_read) < 0) {}
        }
    }

    void handle_file_end() {
        if (file_upload_fd != -1) {
            close(file_upload_fd);
            file_upload_fd = -1;
        }
    }

    void handle_download_req(uint8_t* buffer, ssize_t payload_read) {
        if (payload_read <= 0) {
            send_packet(PKT_FILE_ERR, 0, (const uint8_t*)"Empty remote path", 17);
            return;
        }
        buffer[payload_read] = '\0';
        std::string remote_path((char*)buffer);
        
        FILE* file = fopen(remote_path.c_str(), "rb");
        if (!file) {
            std::string err_msg = "Failed to open file: " + std::string(strerror(errno));
            send_packet(PKT_FILE_ERR, 0, (const uint8_t*)err_msg.c_str(), err_msg.length());
            return;
        }

        if (!send_packet(PKT_FILE_START, 0, nullptr, 0)) {
            fclose(file);
            return;
        }

        uint8_t file_buf[MAX_PAYLOAD_SIZE];
        while (true) {
            size_t bytes_read = fread(file_buf, 1, sizeof(file_buf), file);
            if (bytes_read > 0) {
                if (!send_packet(PKT_FILE_DATA, 0, file_buf, static_cast<uint16_t>(bytes_read))) {
                    fclose(file);
                    return;
                }
            }
            if (bytes_read < sizeof(file_buf)) {
                if (ferror(file)) {
                    std::string err_msg = "Error reading file";
                    send_packet(PKT_FILE_ERR, 0, (const uint8_t*)err_msg.c_str(), err_msg.length());
                    fclose(file);
                    return;
                }
                break;
            }
        }

        fclose(file);
        send_packet(PKT_FILE_END, 0, nullptr, 0);
    }

    bool send_packet(uint8_t type, uint32_t session_id, const uint8_t *payload, uint16_t len) {
        PacketHeader hdr = {PKT_MAGIC, type, len, session_id};
        struct iovec iov[2];
        iov[0].iov_base = &hdr;
        iov[0].iov_len = sizeof(hdr);

        if (len > 0 && payload != NULL) {
            iov[1].iov_base = (void *)payload;
            iov[1].iov_len = len;
            return writev(usb_fd, iov, 2) >= 0;
        } else {
            return writev(usb_fd, iov, 1) >= 0;
        }
    }

    void cleanup_connection() {
        for (auto& pair : sessions) {
            kill(pair.second.shell_pid, SIGKILL);
            waitpid(pair.second.shell_pid, NULL, WNOHANG);
            close(pair.second.pty_master_fd);
        }
        sessions.clear();
        if (usb_fd >= 0) { close(usb_fd); usb_fd = -1; }
        if (file_upload_fd >= 0) { close(file_upload_fd); file_upload_fd = -1; }
    }

    void cleanup() {
        cleanup_connection();
    }
};
 
void secure_fds() {
  int null_fd = open("/dev/null", O_RDWR);
  if (null_fd < 0)
    return;
  for (int i = 0; i <= 2; ++i) {
    if (null_fd != i)
      dup2(null_fd, i);
  }
  if (null_fd > 2)
    close(null_fd);
}

int main() {
  secure_fds();
  signal(SIGINT, sigint_handler);
  signal(SIGCHLD, SIG_IGN);

  VR2BridgeDaemon daemon;
  daemon.run();

  return 0;
}