XCVMKit-OS/post_sbin/getty.c

#define _GNU_SOURCE

#define MIN(X, Y) ((X) < (Y) ? (X) : (Y))
#define MAX(X, Y) ((X) > (Y) ? (X) : (Y))

#if DEBUG
#define DLog(format, ...) log_message(format, __VA_ARGS__);
#else
#define DLog(format, ...)
#endif

#include <stdio.h>
#include <pthread.h>
#include <unistd.h>
#include <stdbool.h>
#include <stdlib.h>
#include <fcntl.h>
#include <sys/stat.h>
#include <sys/types.h>
#include <string.h>
#include <spawn.h>
#include <sys/wait.h>
#include <signal.h>
#include <pty.h>
#include <termios.h>
#include <sys/ioctl.h>
#include <fcntl.h>

#define MAX_BUFFER_SIZE 4096        // 最大缓冲区大小

#define POSIX_SPAWN_SETSID              0x0400
#define DEBUG 0
#define VMP_DEVICE_RX "/etc/.vmpdrvfna_rx"
#define VMP_DEVICE_TX "/etc/.vmpdrvfna_tx"
#define VMP_TID_MAX 64
#define VMP_TX_POOL_SIZE 64
#define VMP_STRBUF_LEN 256
#define MAX_PORTS 16                // 最大端口数量
#define MAX_CONNECTIONS 64          // 最大连接数

/* 消息类型 */
#define MSG_PORT_UPDATE 1           // 端口更新消息
#define MSG_NEW_CONNECTION 2        // 新连接消息
#define MSG_DATA 3                  // 数据传输消息
#define MSG_CLOSE_CONNECTION 4      // 关闭连接消息
#define MSG_HANDSHAKE 5             // 握手消息

typedef struct {
    short width;
    short height;
    short xpixel;
    short ypixel;
} rx_command_win_size;

typedef enum {
    rx_command_type_runproc,
    rx_command_type_runpipe,
    rx_command_type_killproc,
    rx_command_type_lsproc,
    rx_command_type_sstdin,
    rx_command_type_winsize,
    rx_command_type_ping,
    rx_command_type_vmnet
} rx_command_type;

typedef struct {
    rx_command_type type;
    int8_t tid;
    int signal;
    long magic;
    char sstdin[VMP_STRBUF_LEN];
    int sstdin_len;
    rx_command_win_size win_size;
} rx_command;

typedef enum {
    tx_command_type_online,
    tx_command_type_stdout,
    tx_command_type_cb,
    tx_command_type_stoped,
    tx_command_type_vmnet
} tx_command_type;

typedef struct {
    tx_command_type type;
    int8_t tid;
    int8_t error;
    long magic;
    char sstdout[VMP_STRBUF_LEN];
    int sstdout_len;
} tx_command;

/* 消息头结构 */
typedef struct {
    int type;                       // 消息类型
    int connection_id;              // 连接ID
    int port;                       // 端口号
    int data_size;                  // 数据大小
    int target;                     // 接收方标识（0=host, 1=client）
} MessageHeader;

/* 数据传输消息结构 */
typedef struct {
    MessageHeader header;           // 消息头
    char data[MAX_BUFFER_SIZE];     // 数据
} DataMessage;

/* 端口更新消息结构 */
typedef struct {
    MessageHeader header;           // 消息头
    int num_ports;                  // 端口数量
    int ports[MAX_PORTS];           // 端口列表
} PortUpdateMessage;

/* 新连接消息结构 */
typedef struct {
    MessageHeader header;           // 消息头
    char client_ip[16];             // 客户端IP
    int client_port;                // 客户端端口
} NewConnectionMessage;

/* 关闭连接消息结构 */
typedef struct {
    MessageHeader header;           // 消息头
} CloseConnectionMessage;

/* 握手消息结构 */
typedef struct {
    MessageHeader header;           // 消息头
} HandshakeMessage;

typedef struct vmnet_queue_elem {
    char buffer[sizeof(DataMessage)];
    int length;
    struct vmnet_queue_elem *next;
} vmnet_queue_elem;

vmnet_queue_elem *vmnet_queue_front;
vmnet_queue_elem *vmnet_queue_waiting_buffer;
pthread_mutex_t vmnet_queue_lock;

extern int vmnetproxy_main(void);
extern void tx_push(tx_command cmd);
extern void log_message(const char *format, ...);

void vmnet_queue_setup(void) {
    vmnet_queue_front = NULL;
    vmnet_queue_waiting_buffer = NULL;
    pthread_mutex_init(&vmnet_queue_lock, NULL);
}

int vmnet_write_message(const char *buffer, int length) {
    int packetOffset = 0;
    const int maximumSplitPacketSize = VMP_STRBUF_LEN - 1;
    do {
        int currentSplitSize = MIN(length - packetOffset, maximumSplitPacketSize);
        bool hasNext = packetOffset + currentSplitSize < length;
        tx_command tx = {
            .type = tx_command_type_vmnet,
            .tid = hasNext ? 21 : 22,
            .error = 0,
            .magic = 0,
            .sstdout_len = currentSplitSize
        };
        memcpy(&tx.sstdout[0], &buffer[packetOffset], currentSplitSize);
        tx_push(tx);
        packetOffset += currentSplitSize;
    }
    while (packetOffset < length);
}

int vmnet_read_message(char *buffer, int length) {
    vmnet_queue_elem *elem;
    pthread_mutex_lock(&vmnet_queue_lock);
    elem = vmnet_queue_front;
    if (elem) {
        vmnet_queue_front = elem->next;
    }
    pthread_mutex_unlock(&vmnet_queue_lock);

    if (!elem) {
        return 0;
    }

    int write_len = MIN(elem->length, length);
    memcpy(buffer, &elem->buffer[0], write_len);
    free(elem);
    return write_len;
}

typedef int FD;
typedef int pipe_t[2];

typedef struct {
    int8_t tid;
    int8_t inuse;
    char command[VMP_STRBUF_LEN];
    pid_t pid;
    pthread_t io_thread;
    int master_pty;
    int slave_pty;
    FD stdin_fd;
    rx_command_win_size win_size;
    int use_pipe;
    pipe_t pipe_in;
    pipe_t pipe_out;
    FD stdout_fd;
} management_process;

void vmnet_push_data(rx_command command) {
    if (command.type != rx_command_type_vmnet) {
        return;
    }
    bool has_next = command.tid == 21;

    pthread_mutex_lock(&vmnet_queue_lock);

    if (vmnet_queue_waiting_buffer == NULL) {
        vmnet_queue_waiting_buffer = (vmnet_queue_elem *)calloc(1, sizeof(vmnet_queue_elem));
    }

    const int size = command.sstdin_len;
    int approve_size = MIN(size, sizeof(DataMessage) - vmnet_queue_waiting_buffer->length);
    if (approve_size > 0) {
        memcpy(&vmnet_queue_waiting_buffer->buffer[vmnet_queue_waiting_buffer->length], &command.sstdin[0], approve_size);
        vmnet_queue_waiting_buffer->length += approve_size;
    }

    if (!has_next) {
        if (vmnet_queue_front == NULL) {
            vmnet_queue_front = vmnet_queue_waiting_buffer;
        }
        else {
            vmnet_queue_elem *elem, *elem_prev;
            for (elem = vmnet_queue_front; elem != NULL; elem = elem->next) {
                elem_prev = elem;
            }
            elem_prev->next = vmnet_queue_waiting_buffer;
        }
        vmnet_queue_waiting_buffer->next = NULL;
        vmnet_queue_waiting_buffer = NULL;
    }

    pthread_mutex_unlock(&vmnet_queue_lock);
}

void debug_print_rx(rx_command rx) {
#if DEBUG
    printf("RX: (type=%d tid=%d signal=%d magic=%ld sstdin=%s)\n", rx.type, rx.tid, rx.signal, rx.magic, rx.sstdin);
#endif
}

void debug_print_tx(tx_command tx) {
#if DEBUG
    printf("TX: (type=%d tid=%d error=%d magic=%ld sstdout=%s)\n", tx.type, tx.tid, tx.error, tx.magic, tx.sstdout);
#endif
}

static FD RX, TX;
static management_process processes[VMP_TID_MAX];
static pthread_mutex_t processes_lock;
static tx_command tx_cmds[VMP_TX_POOL_SIZE];
static int tx_cmds_ptr = 0;
static pthread_mutex_t tx_lock;

void *tx_thread(void *ptr) {
    while (1) {
        pthread_mutex_lock(&tx_lock);
        for (int i = 0; i < tx_cmds_ptr; i++) {
            write(TX, &tx_cmds[i], sizeof(tx_command));
        }
        tx_cmds_ptr = 0;
        pthread_mutex_unlock(&tx_lock);
        usleep(100);
    }
}

void tx_push(tx_command cmd) {
    debug_print_tx(cmd);
    pthread_mutex_lock(&tx_lock);
    while (tx_cmds_ptr >= VMP_TX_POOL_SIZE) {
        pthread_mutex_unlock(&tx_lock);
        usleep(100);
        pthread_mutex_lock(&tx_lock);
    }
    tx_cmds[tx_cmds_ptr++] = cmd;
    pthread_mutex_unlock(&tx_lock);
}

#define CB_SUCC "succ"
#define CB_NO_TID "no_tid"
#define CB_NO_PIPE "no_pipe"
#define CB_NO_PTY "no_pty"
#define CB_TID_KILLED "tid_killed"
#define CT_LAUNCH_FAIL "launch_fail"
#define CT_PUSH(ptype, pmagic, ptid, perror, reason) \
{ \
    tx_command ct_cmd = { \
        .type = ptype, \
        .tid = ptid, \
        .error = perror, \
        .magic = pmagic, \
        .sstdout_len = strlen(reason) \
    }; \
    bzero(&ct_cmd.sstdout[0], VMP_STRBUF_LEN); \
    memcpy(&ct_cmd.sstdout[0], reason, strlen(reason)); \
    tx_push(ct_cmd); \
}
#define CB_PUSH(pmagic, pterror, reason) CT_PUSH(tx_command_type_cb, pmagic, 0, pterror, reason)

void *mproc_io_thread_pty(management_process *mproc) {
    char slave_name[256];
    if (openpty(&mproc->master_pty, &mproc->slave_pty, 
               slave_name, NULL, &mproc->win_size) == -1) {
        CT_PUSH(tx_command_type_stoped, 0, mproc->tid, 1, CB_NO_PTY);
        mproc->inuse = 0;
        return NULL;
    }

    struct termios tios;
    tcgetattr(mproc->slave_pty, &tios);
    tios.c_lflag &= ~(ECHO | ICANON);
    tcsetattr(mproc->slave_pty, TCSANOW, &tios);

    posix_spawn_file_actions_t actions;
    posix_spawnattr_t attr;
    posix_spawn_file_actions_init(&actions);
    posix_spawnattr_init(&attr);
    posix_spawn_file_actions_adddup2(&actions, mproc->slave_pty, STDIN_FILENO);
    posix_spawn_file_actions_adddup2(&actions, mproc->slave_pty, STDOUT_FILENO);
    posix_spawn_file_actions_adddup2(&actions, mproc->slave_pty, STDERR_FILENO);
    posix_spawn_file_actions_addclose(&actions, mproc->master_pty);
    posix_spawnattr_setflags(&attr, POSIX_SPAWN_SETSID | POSIX_SPAWN_SETPGROUP);
    
    setenv("TERM", "xterm-256color", 1);
    setenv("COLORTERM", "truecolor", 1);

    char *const argv[] = { "/bin/sh", "-c", mproc->command, NULL };
    char *const environment[] = {
        "TERM=xterm-256color", 
        "COLORTERM=truecolor",
        NULL
    };

    int ret = posix_spawnp(&mproc->pid, argv[0], &actions, &attr, argv, environment);
    close(mproc->slave_pty);

    posix_spawn_file_actions_destroy(&actions);
    posix_spawnattr_destroy(&attr);

    pthread_mutex_unlock(&processes_lock);

    if (ret != 0) {
        close(mproc->master_pty);
        CT_PUSH(tx_command_type_stoped, ret, mproc->tid, 1, CT_LAUNCH_FAIL);
        mproc->inuse = 0;
        return NULL;
    }

    fcntl(mproc->master_pty, F_SETFL, O_NONBLOCK);
    mproc->stdin_fd = mproc->master_pty;

    char buf[VMP_STRBUF_LEN];
    ssize_t nread;
    const char *exit_reason = "unknown";
    int8_t exit_error = 0;
    int64_t exit_magic = 0;

    while (1) {
        do {
            bzero(buf, sizeof(buf));
            nread = read(mproc->master_pty, buf, sizeof(buf));
            if (nread > 0) {
                tx_command cmd = {
                    .type = tx_command_type_stdout,
                    .tid = mproc->tid,
                    .error = 0,
                    .magic = 0,
                    .sstdout_len = nread
                };
                bzero(&cmd.sstdout[0], VMP_STRBUF_LEN);
                memcpy(cmd.sstdout, buf, nread);
                tx_push(cmd);
            }
        }
        while (nread == sizeof(buf)); // has next unread

        int pid_status;
        pid_t result = waitpid(mproc->pid, &pid_status, WNOHANG);
        if (result == -1) {
            exit_reason = "err_waitpid";
            break;
        }
        if (result != 0) {
            if (WIFEXITED(pid_status)) {
                exit_reason = "normal";
                exit_magic = WEXITSTATUS(pid_status);
            }
            else if (WIFSIGNALED(pid_status)) {
                exit_reason = "sig";
                exit_error = WIFSIGNALED(pid_status);
            }
            else {
                exit_reason = "exit_unknown";
            }
            break;
        }
        usleep(100);
    }

    close(mproc->master_pty);
    CT_PUSH(tx_command_type_stoped, exit_magic, mproc->tid, exit_error, exit_reason);
    mproc->inuse = 0;
    return NULL;
}

void *mproc_io_thread_pipe(management_process *mproc) {
    posix_spawn_file_actions_t actions;
    posix_spawn_file_actions_init(&actions);
    posix_spawn_file_actions_addclose(&actions, mproc->pipe_in[1]);
    posix_spawn_file_actions_addclose(&actions, mproc->pipe_out[0]);
    posix_spawn_file_actions_adddup2(&actions, mproc->pipe_in[0], STDIN_FILENO);
    posix_spawn_file_actions_adddup2(&actions, mproc->pipe_out[1], STDOUT_FILENO);
    char *const argv[] = {
        "/bin/sh", "-c", mproc->command, NULL
    };
    char *const environment[] = {
        NULL
    };
    int ret = posix_spawnp(&mproc->pid, argv[0], &actions, NULL, argv, environment);
    posix_spawn_file_actions_destroy(&actions);
    pthread_mutex_unlock(&processes_lock);
    
    if (ret != 0) {
        close(mproc->pipe_in[0]);
        close(mproc->pipe_in[1]);
        close(mproc->pipe_out[0]);
        close(mproc->pipe_out[1]);
        CT_PUSH(tx_command_type_stoped, ret, mproc->tid, 1, CT_LAUNCH_FAIL);
        mproc->inuse = 0;
        return NULL;
    }

    close(mproc->pipe_in[0]);
    close(mproc->pipe_out[1]);
    mproc->stdin_fd = mproc->pipe_in[1];
    mproc->stdout_fd = mproc->pipe_out[0];
    int flags = fcntl(mproc->stdout_fd, F_GETFL, 0);
    fcntl(mproc->stdout_fd, F_SETFL, flags | O_NONBLOCK);

    char buf[VMP_STRBUF_LEN];
    ssize_t nread = 0;
    const char *exit_reason = "unknown";
    int8_t exit_error = 0;
    int64_t exit_magic = 0;
    while (1) {
        do {
            bzero(buf, sizeof(buf));
            nread = read(mproc->stdout_fd, buf, sizeof(buf));
            if (nread > 0) {
                tx_command cmd = {
                    .type = tx_command_type_stdout,
                    .tid = mproc->tid,
                    .error = 0,
                    .magic = 0,
                    .sstdout_len = nread
                };
                bzero(&cmd.sstdout[0], VMP_STRBUF_LEN);
                memcpy(cmd.sstdout, buf, nread);
                tx_push(cmd);
            }
        }
        while (nread == sizeof(buf)); // has next unread

        int pid_status;
        pid_t result = waitpid(mproc->pid, &pid_status, WNOHANG);
        if (result == -1) {
            exit_reason = "err_waitpid";
            break;
        }
        if (result != 0) {
            if (WIFEXITED(pid_status)) {
                exit_reason = "normal";
                exit_magic = WEXITSTATUS(pid_status);
            }
            else if (WIFSIGNALED(pid_status)) {
                exit_reason = "sig";
                exit_error = WIFSIGNALED(pid_status);
            }
            else {
                exit_reason = "exit_unknown";
            }
            break;
        }
        usleep(100);
    }

    close(mproc->stdin_fd);
    close(mproc->stdout_fd);
    CT_PUSH(tx_command_type_stoped, exit_magic, mproc->tid, exit_error, exit_reason);
    mproc->inuse = 0;
}

int8_t find_unused_mproc_nolock() {
    for (int8_t tid = 0; tid < VMP_TID_MAX; tid++) {
        if (!processes[tid].inuse) {
            bzero(&processes[tid], sizeof(management_process));
            return tid;
        }
    }
    return -1;
}

void rx_process(rx_command cmd) {
    debug_print_rx(cmd);
    switch (cmd.type) {
        case rx_command_type_runproc: {
            pthread_mutex_lock(&processes_lock);
            int8_t tid = find_unused_mproc_nolock();
            if (tid < 0) {
                pthread_mutex_unlock(&processes_lock);
                CB_PUSH(cmd.magic, 1, CB_NO_TID);
                break;
            }

            management_process mproc;
            bzero(&mproc, sizeof(management_process));
            memcpy(&mproc.command[0], cmd.sstdin, sizeof(mproc.command));

            mproc.tid = tid;
            mproc.inuse = 1;
            processes[tid] = mproc;
            mproc.win_size = cmd.win_size;
            mproc.use_pipe = 0;

            CT_PUSH(tx_command_type_cb, cmd.magic, tid, 0, CB_SUCC);
            pthread_create(&mproc.io_thread, NULL, &mproc_io_thread_pty, &processes[tid]);
            // processes_lock will unlock in mproc_io_thread
            break;
        }
        case rx_command_type_runpipe: {
            pthread_mutex_lock(&processes_lock);
            int8_t tid = find_unused_mproc_nolock();
            if (tid < 0) {
                pthread_mutex_unlock(&processes_lock);
                CB_PUSH(cmd.magic, 1, CB_NO_TID);
                break;
            }
            
            management_process mproc;
            bzero(&mproc, sizeof(management_process));
            memcpy(&mproc.command[0], cmd.sstdin, sizeof(mproc.command));
            if (pipe(mproc.pipe_in) == -1 || pipe(mproc.pipe_out) == -1) {
                pthread_mutex_unlock(&processes_lock);
                close(mproc.pipe_in[0]);
                close(mproc.pipe_in[1]);
                close(mproc.pipe_out[0]);
                close(mproc.pipe_out[1]);
                CB_PUSH(cmd.magic, 1, CB_NO_PIPE);
                break;
            }
            
            mproc.tid = tid;
            mproc.inuse = 1;
            processes[tid] = mproc;
            mproc.use_pipe = 1;
            CT_PUSH(tx_command_type_cb, cmd.magic, tid, 0, CB_SUCC);
            pthread_create(&mproc.io_thread, NULL, &mproc_io_thread_pipe, &processes[tid]);
            // processes_lock will unlock in mproc_io_thread
            break;
        }
        case rx_command_type_killproc: {
            if (cmd.tid < 0 || cmd.tid >= VMP_TID_MAX) {
                CB_PUSH(cmd.magic, 1, CB_TID_KILLED);
                break;
            }
            pthread_mutex_lock(&processes_lock);
            management_process *mproc = &processes[cmd.tid];
            if (!mproc->inuse) {
                pthread_mutex_unlock(&processes_lock);
                CB_PUSH(cmd.magic, 1, CB_TID_KILLED);
                break;
            }
            
            kill(mproc->pid, SIGKILL);
            pthread_mutex_unlock(&processes_lock);
            CB_PUSH(cmd.magic, 0, CB_SUCC);
            break;
        }
        case rx_command_type_lsproc: {
            char bytemap[VMP_TID_MAX];
            bzero(bytemap, sizeof(bytemap));
            pthread_mutex_lock(&processes_lock);
            for (int tid = 0; tid < VMP_TID_MAX; tid++) {
                bytemap[tid] = processes[tid].inuse;
            }
            pthread_mutex_unlock(&processes_lock);
            tx_command tx = {
                .type = tx_command_type_cb,
                .tid = 0,
                .error = 0,
                .magic = cmd.magic,
                .sstdout_len = VMP_TID_MAX
            };
            bzero(&tx.sstdout, VMP_STRBUF_LEN);
            memcpy(&tx.sstdout, bytemap, VMP_TID_MAX);
            tx_push(tx);
            break;
        }
        case rx_command_type_sstdin: {
            if (cmd.tid < 0 || cmd.tid >= VMP_TID_MAX) {
                CB_PUSH(cmd.magic, 1, CB_TID_KILLED);
                break;
            }
            pthread_mutex_lock(&processes_lock);
            management_process *mproc = &processes[cmd.tid];
            if (!mproc->inuse) {
                pthread_mutex_unlock(&processes_lock);
                CB_PUSH(cmd.magic, 1, CB_TID_KILLED);
                break;
            }
            write(mproc->stdin_fd, cmd.sstdin, cmd.sstdin_len);
            pthread_mutex_unlock(&processes_lock);
            CB_PUSH(cmd.magic, 0, CB_SUCC);
            break;
        }
        case rx_command_type_winsize: {
            if (cmd.tid < 0 || cmd.tid >= VMP_TID_MAX) {
                CB_PUSH(cmd.magic, 1, CB_TID_KILLED);
                break;
            }
            pthread_mutex_lock(&processes_lock);
            management_process *mproc = &processes[cmd.tid];
            if (!mproc->inuse) {
                pthread_mutex_unlock(&processes_lock);
                CB_PUSH(cmd.magic, 1, CB_TID_KILLED);
                break;
            }
            if (!mproc->use_pipe) {
                ioctl(mproc->stdin_fd, TIOCGWINSZ, &cmd.win_size);
            }
            pthread_mutex_unlock(&processes_lock);
            CB_PUSH(cmd.magic, 0, CB_SUCC);
            break;
        }
        case rx_command_type_ping: {
            CB_PUSH(cmd.magic, 0, "PONG");
            break;
        }
        case rx_command_type_vmnet: {
            vmnet_push_data(cmd);
            break;
        }
    }
}

void *vmnetproxy_program(void *param) {
    int vret = vmnetproxy_main();
    printf("XCVMKit-OS: vnmetproxy killed with %d", vret);
    return NULL;
}

int main() {
    vmnet_queue_setup();
    printf("Boot Success\n");
    printf("Welcome to XCVMKit-OS!\n");
    RX = open(VMP_DEVICE_RX, 1101824);
    TX = open(VMP_DEVICE_TX, 1101825);
    if (RX < 0) {
        printf("XCVMKit-OS: unable to connect with host.\n");
        return -1;
    }
    if (TX < 0) {
        printf("XCVMKit-OS: unable to connect with host.\n");
        return -1;
    }
    pthread_mutex_init(&tx_lock, NULL);
    pthread_mutex_init(&processes_lock, NULL);
    pthread_t tx_thr;
    pthread_create(&tx_thr, NULL, tx_thread, NULL);
    
    tx_command initial_cmd = {
        .type = tx_command_type_online,
        .tid = 0,
        .magic = 0,
        .sstdout_len = 0
    };
    bzero(&initial_cmd.sstdout, VMP_STRBUF_LEN);
    tx_push(initial_cmd);

    pthread_t vmnetproxy_thr;
    pthread_create(&vmnetproxy_thr, NULL, vmnetproxy_program, NULL);

    while (1) {
        rx_command command;
        size_t read_size = read(RX, &command, sizeof(rx_command));
        if (read_size == sizeof(rx_command)) {
            rx_process(command);
        }
        usleep(100);
    }
    return 0;
}

/// vmnetproxy code
#ifndef COMMON_H
#define COMMON_H

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <fcntl.h>
#include <errno.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <sys/select.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <pthread.h>

/* 函数声明 */
// 文件操作函数
int open_proxy_file(int flags);
void lock_file(int fd);
void unlock_file(int fd);

// 消息处理函数
int write_message(int fd, void *message, int size);
int read_message(int fd, void *buffer, int size, int target);
int clear_message_queue(int fd);
int send_handshake_message(int fd, int target);

// 网络操作函数
int create_server_socket(int port);
int accept_connection(int server_socket);
int connect_to_server(const char *ip, int port);

// 端口扫描函数
int get_listening_ports(int *ports, int max_ports);

// 日志函数
void log_message(const char *format, ...);

#endif /* COMMON_H */

#include <stdarg.h>
#include <time.h>
#include <sys/file.h>
#include <netdb.h>
#include <ifaddrs.h>
#include <sys/ioctl.h>

/* 文件操作函数 */

static pthread_mutex_t file_lock;

__attribute__((constructor))
static void __file_lock_initiator(void) {
    pthread_mutex_init(&file_lock, NULL);
}

/**
 * 打开代理文件
 * @param flags 打开文件的标志
 * @return 文件描述符
 */
int open_proxy_file(int flags) {
    return 0;
}

/**
 * 锁定文件以进行独占访问
 * @param fd 文件描述符
 */
void lock_file(int fd) {
    pthread_mutex_lock(&file_lock);
}

/**
 * 解锁文件
 * @param fd 文件描述符
 */
void unlock_file(int fd) {
    pthread_mutex_unlock(&file_lock);
}

/* 消息处理函数 */

/**
 * 写入消息到文件
 * @param fd 文件描述符
 * @param message 消息指针
 * @param size 消息大小
 * @return 写入的字节数
 */
int write_message(int fd, void *message, int size) {
    lock_file(fd);
    int ret = vmnet_write_message(message, size);
    unlock_file(fd);
    return ret;
}

/**
 * 从文件读取消息
 * @param fd 文件描述符
 * @param buffer 缓冲区指针
 * @param size 缓冲区大小
 * @param target 目标接收方（0=host, 1=client）
 * @return 读取的字节数，如果没有消息则返回0
 */
int read_message(int fd, void *buffer, int size, int target) {
    lock_file(fd);
    int ret = vmnet_read_message(buffer, size);
    unlock_file(fd);
    return ret;
}

/* 网络操作函数 */

/**
 * 创建服务器套接字
 * @param port 端口号
 * @return 套接字描述符
 */
int create_server_socket(int port) {
    int server_socket = socket(AF_INET, SOCK_STREAM, 0);
    if (server_socket < 0) {
        DLog("创建套接字失败: %s", strerror(errno));
        return -1;
    }
    
    // 设置套接字选项，允许地址重用
    int opt = 1;
    if (setsockopt(server_socket, SOL_SOCKET, SO_REUSEADDR, &opt, sizeof(opt)) < 0) {
        DLog("设置套接字选项失败: %s", strerror(errno));
        close(server_socket);
        return -1;
    }
    
    // 绑定地址
    struct sockaddr_in server_addr;
    memset(&server_addr, 0, sizeof(server_addr));
    server_addr.sin_family = AF_INET;
    server_addr.sin_addr.s_addr = INADDR_ANY;
    server_addr.sin_port = htons(port);
    
    if (bind(server_socket, (struct sockaddr *)&server_addr, sizeof(server_addr)) < 0) {
        DLog("绑定套接字失败: %s", strerror(errno));
        close(server_socket);
        return -1;
    }
    
    // 监听连接
    if (listen(server_socket, 5) < 0) {
        DLog("监听套接字失败: %s", strerror(errno));
        close(server_socket);
        return -1;
    }
    
    return server_socket;
}

/**
 * 接受连接
 * @param server_socket 服务器套接字
 * @return 客户端套接字
 */
int accept_connection(int server_socket) {
    struct sockaddr_in client_addr;
    socklen_t client_len = sizeof(client_addr);
    
    int client_socket = accept(server_socket, (struct sockaddr *)&client_addr, &client_len);
    if (client_socket < 0) {
        DLog("接受连接失败: %s", strerror(errno));
        return -1;
    }
    
    DLog("接受来自 %s:%d 的连接", inet_ntoa(client_addr.sin_addr), ntohs(client_addr.sin_port));
    return client_socket;
}

/**
 * 连接到服务器
 * @param ip 服务器IP
 * @param port 服务器端口
 * @return 套接字描述符
 */
int connect_to_server(const char *ip, int port) {
    int client_socket = socket(AF_INET, SOCK_STREAM, 0);
    if (client_socket < 0) {
        DLog("创建套接字失败: %s", strerror(errno));
        return -1;
    }
    
    struct sockaddr_in server_addr;
    memset(&server_addr, 0, sizeof(server_addr));
    server_addr.sin_family = AF_INET;
    server_addr.sin_port = htons(port);
    
    if (inet_pton(AF_INET, ip, &server_addr.sin_addr) <= 0) {
        DLog("无效的IP地址: %s", ip);
        close(client_socket);
        return -1;
    }
    
    if (connect(client_socket, (struct sockaddr *)&server_addr, sizeof(server_addr)) < 0) {
        DLog("连接到服务器失败: %s", strerror(errno));
        close(client_socket);
        return -1;
    }
    
    return client_socket;
}

/* 端口扫描函数 */

/**
 * 获取本地监听的端口列表
 * @param ports 端口数组
 * @param max_ports 最大端口数量
 * @return 找到的端口数量
 */
int get_listening_ports(int *ports, int max_ports) {
    FILE *fp;
    char line[256];
    int count = 0;
    
    // 使用netstat命令获取监听的TCP端口
    fp = popen("netstat -tln | grep LISTEN", "r");
    if (fp == NULL) {
        DLog("执行netstat命令失败: %s", strerror(errno));
        return 0;
    }
    
    while (fgets(line, sizeof(line), fp) != NULL && count < max_ports) {
        // 检查行是否包含 "LISTEN" 关键字，过滤掉错误信息
        if (strstr(line, "LISTEN") == NULL) {
            continue;
        }
        
        char *ptr = strstr(line, ":");
        if (ptr) {
            int port = atoi(ptr + 1);
            if (port > 0) {
                // 检查端口是否已经在列表中
                int i;
                for (i = 0; i < count; i++) {
                    if (ports[i] == port) {
                        break;
                    }
                }
                
                // 如果端口不在列表中，添加它
                if (i == count) {
                    ports[count++] = port;
                }
            }
        }
    }
    
    pclose(fp);
    return count;
}

/* 日志函数 */

/**
 * 记录日志消息
 * @param format 格式字符串
 * @param ... 可变参数
 */
void log_message(const char *format, ...) {
    va_list args;
    char buffer[1024];
    time_t now;
    struct tm *timeinfo;
    
    // 获取当前时间
    time(&now);
    timeinfo = localtime(&now);
    
    // 格式化时间
    char time_str[20];
    strftime(time_str, sizeof(time_str), "%Y-%m-%d %H:%M:%S", timeinfo);
    
    // 格式化日志消息
    va_start(args, format);
    vsnprintf(buffer, sizeof(buffer), format, args);
    va_end(args);
    
    // 输出日志
    printf("[%s] %s\n", time_str, buffer);
    fflush(stdout);
}

/**
 * 清空消息队列
 * @param fd 文件描述符
 * @return 0表示成功，-1表示失败
 */
int clear_message_queue(int fd) {
    lock_file(fd);
    unlock_file(fd);
    return 0;
}

/**
 * 发送握手消息
 * @param fd 文件描述符
 * @param target 目标接收方（0=host, 1=client）
 * @return 写入的字节数
 */
int send_handshake_message(int fd, int target) {
    HandshakeMessage message;
    memset(&message, 0, sizeof(message));
    message.header.type = MSG_HANDSHAKE;
    message.header.target = target;
    
    int result = write_message(fd, &message, sizeof(message));
    if (result > 0) {
        DLog("已发送握手消息到%s", target == 0 ? "host" : "client");
    }
    return result;
}

/* 全局变量 */
int proxy_fd = -1;                  // 代理文件描述符
int ports[MAX_PORTS];               // 监听的端口列表
int num_ports = 0;                  // 端口数量
int connection_map[MAX_CONNECTIONS]; // 连接映射表
int next_connection_id = 1;         // 下一个连接ID

/* 线程函数 */
void *port_monitor_thread(void *arg);
void *connection_handler_thread(void *arg);

/**
 * 初始化连接映射表
 */
void init_connection_map() {
    for (int i = 0; i < MAX_CONNECTIONS; i++) {
        connection_map[i] = -1;
    }
}

/**
 * 分配新的连接ID
 * @return 连接ID
 */
int allocate_connection_id() {
    for (int i = 0; i < MAX_CONNECTIONS; i++) {
        if (connection_map[i] == -1) {
            connection_map[i] = 0; // 标记为已分配但未连接
            return i + 1; // 连接ID从1开始
        }
    }
    return -1; // 没有可用的连接ID
}

/**
 * 释放连接ID
 * @param connection_id 连接ID
 */
void free_connection_id(int connection_id) {
    if (connection_id > 0 && connection_id <= MAX_CONNECTIONS) {
        int socket_fd = connection_map[connection_id - 1];
        if (socket_fd > 0) {
            close(socket_fd);
        }
        connection_map[connection_id - 1] = -1;
    }
}

/**
 * 设置连接套接字
 * @param connection_id 连接ID
 * @param socket_fd 套接字描述符
 */
void set_connection_socket(int connection_id, int socket_fd) {
    if (connection_id > 0 && connection_id <= MAX_CONNECTIONS) {
        connection_map[connection_id - 1] = socket_fd;
    }
}

/**
 * 获取连接套接字
 * @param connection_id 连接ID
 * @return 套接字描述符
 */
int get_connection_socket(int connection_id) {
    if (connection_id > 0 && connection_id <= MAX_CONNECTIONS) {
        return connection_map[connection_id - 1];
    }
    return -1;
}

/**
 * 更新端口列表
 */
void update_ports() {
    int new_ports[MAX_PORTS];
    int new_num_ports = get_listening_ports(new_ports, MAX_PORTS);
    
    // 检查端口列表是否有变化
    int changed = 0;
    if (new_num_ports != num_ports) {
        changed = 1;
    } else {
        for (int i = 0; i < num_ports; i++) {
            int found = 0;
            for (int j = 0; j < new_num_ports; j++) {
                if (ports[i] == new_ports[j]) {
                    found = 1;
                    break;
                }
            }
            if (!found) {
                changed = 1;
                break;
            }
        }
    }
    
    // 如果端口列表有变化，更新并通知host
    if (changed) {
        DLog("端口列表已更新，共 %d 个端口", new_num_ports);
        
        // 更新本地端口列表
        num_ports = new_num_ports;
        for (int i = 0; i < num_ports; i++) {
            ports[i] = new_ports[i];
            DLog("监听端口: %d", ports[i]);
        }
        
        // 发送端口更新消息给host
        PortUpdateMessage message;
        memset(&message, 0, sizeof(message));
        message.header.type = MSG_PORT_UPDATE;
        message.header.target = 0; // 设置接收方为host
        message.num_ports = num_ports;
        for (int i = 0; i < num_ports; i++) {
            message.ports[i] = ports[i];
        }
        
        write_message(proxy_fd, &message, sizeof(message));
    }
}

/**
 * 处理新连接消息
 * @param message 新连接消息
 */
void handle_new_connection(NewConnectionMessage *message) {
    int connection_id = message->header.connection_id;
    int port = message->header.port;
    
    DLog("收到新连接请求: ID=%d, 端口=%d, 客户端=%s:%d",
               connection_id, port, message->client_ip, message->client_port);
    
    // 连接到本地服务
    int socket_fd = connect_to_server("127.0.0.1", port);
    if (socket_fd < 0) {
        DLog("无法连接到本地服务: 端口=%d", port);
        
        // 发送关闭连接消息
        CloseConnectionMessage close_message;
        memset(&close_message, 0, sizeof(close_message));
        close_message.header.type = MSG_CLOSE_CONNECTION;
        close_message.header.connection_id = connection_id;
        close_message.header.target = 0; // 设置接收方为host
        write_message(proxy_fd, &close_message, sizeof(close_message));
        return;
    }
    
    // 设置连接映射
    set_connection_socket(connection_id, socket_fd);
    
    // 创建连接处理线程
    pthread_t thread;
    int *thread_arg = malloc(sizeof(int));
    *thread_arg = connection_id;
    if (pthread_create(&thread, NULL, connection_handler_thread, thread_arg) != 0) {
        DLog("创建连接处理线程失败: %s", strerror(errno));
        free_connection_id(connection_id);
        free(thread_arg);
        return;
    }
    
    pthread_detach(thread);
}

/**
 * 处理数据消息
 * @param message 数据消息
 */
void handle_data_message(DataMessage *message) {
    int connection_id = message->header.connection_id;
    int socket_fd = get_connection_socket(connection_id);
    
    if (socket_fd < 0) {
        DLog("无效的连接ID: %d", connection_id);
        return;
    }
    
    // 将数据发送到本地服务
    int bytes_sent = write(socket_fd, message->data, message->header.data_size);
    if (bytes_sent < 0) {
        DLog("发送数据到本地服务失败: %s", strerror(errno));
        
        // 发送关闭连接消息
        CloseConnectionMessage close_message;
        memset(&close_message, 0, sizeof(close_message));
        close_message.header.type = MSG_CLOSE_CONNECTION;
        close_message.header.connection_id = connection_id;
        close_message.header.target = 0; // 设置接收方为host
        write_message(proxy_fd, &close_message, sizeof(close_message));
        
        free_connection_id(connection_id);
    }
}

/**
 * 处理关闭连接消息
 * @param message 关闭连接消息
 */
void handle_close_connection(CloseConnectionMessage *message) {
    int connection_id = message->header.connection_id;
    DLog("关闭连接: ID=%d", connection_id);
    free_connection_id(connection_id);
}

/**
 * 端口监控线程函数
 * @param arg 线程参数
 * @return NULL
 */
void *port_monitor_thread(void *arg) {
    while (1) {
        update_ports();
        sleep(1);
    }
    return NULL;
}

/**
 * 连接处理线程函数
 * @param arg 线程参数（连接ID）
 * @return NULL
 */
void *connection_handler_thread(void *arg) {
    int connection_id = *((int *)arg);
    free(arg);
    
    int socket_fd = get_connection_socket(connection_id);
    if (socket_fd < 0) {
        return NULL;
    }
    
    char buffer[MAX_BUFFER_SIZE];
    while (1) {
        // 从本地服务读取数据
        int bytes_read = read(socket_fd, buffer, MAX_BUFFER_SIZE);
        if (bytes_read <= 0) {
            if (bytes_read < 0) {
                DLog("从本地服务读取数据失败: %s", strerror(errno));
            }
            break;
        }
        
        // 发送数据消息给host
        DataMessage message;
        memset(&message, 0, sizeof(message));
        message.header.type = MSG_DATA;
        message.header.connection_id = connection_id;
        message.header.data_size = bytes_read;
        message.header.target = 0; // 设置接收方为host
        memcpy(message.data, buffer, bytes_read);
        
        write_message(proxy_fd, &message, sizeof(MessageHeader) + bytes_read);
    }
    
    // 发送关闭连接消息
    CloseConnectionMessage close_message;
    memset(&close_message, 0, sizeof(close_message));
    close_message.header.type = MSG_CLOSE_CONNECTION;
    close_message.header.connection_id = connection_id;
    write_message(proxy_fd, &close_message, sizeof(close_message));
    
    free_connection_id(connection_id);
    return NULL;
}

/**
 * 处理握手消息
 * @param message 握手消息
 */
void handle_handshake_message(HandshakeMessage *message) {
    DLog("收到来自host的握手消息，host已上线");
}

/**
 * 消息处理线程函数
 * @param arg 线程参数
 * @return NULL
 */
void *message_handler_thread(void *arg) {
    char buffer[sizeof(DataMessage)];
    
    while (1) {
        // 读取消息
        int bytes_read = read_message(proxy_fd, buffer, sizeof(buffer), 1); // 1表示client

        if (bytes_read <= 0) {
            usleep(1000 * 1); // 如果没有消息，等待一段时间（1 ms）
            continue;
        }
        
        // 解析消息头
        MessageHeader *header = (MessageHeader *)buffer;
        
        // 根据消息类型处理
        switch (header->type) {
            case MSG_NEW_CONNECTION:
                handle_new_connection((NewConnectionMessage *)buffer);
                break;
                
            case MSG_DATA:
                handle_data_message((DataMessage *)buffer);
                break;
                
            case MSG_CLOSE_CONNECTION:
                handle_close_connection((CloseConnectionMessage *)buffer);
                break;
                
            case MSG_HANDSHAKE:
                handle_handshake_message((HandshakeMessage *)buffer);
                break;
                
            default:
                DLog("未知的消息类型: %d", header->type);
                break;
        }

        usleep(100);
    }
    
    return NULL;
}

/**
 * 主函数
 * @return 退出码
 */
int vmnetproxy_main() {
    DLog("VMNet Proxy Client 启动");
    
    // 初始化连接映射表
    init_connection_map();
    
    // 打开代理文件
    proxy_fd = open_proxy_file(0);
    if (proxy_fd < 0) {
        DLog("无法打开代理文件，退出");
        return 1;
    }
    
    // 清空消息队列
    if (clear_message_queue(proxy_fd) < 0) {
        DLog("清空消息队列失败，继续执行");
    }
    
    // 发送握手消息给host
    if (send_handshake_message(proxy_fd, 0) < 0) {
        DLog("发送握手消息失败，继续执行");
    }
    
    // 创建端口监控线程
    pthread_t port_thread;
    if (pthread_create(&port_thread, NULL, port_monitor_thread, NULL) != 0) {
        DLog("创建端口监控线程失败: %s", strerror(errno));
        return 1;
    }
    
    // 创建消息处理线程
    pthread_t message_thread;
    if (pthread_create(&message_thread, NULL, message_handler_thread, NULL) != 0) {
        DLog("创建消息处理线程失败: %s", strerror(errno));
        return 1;
    }
    
    // 等待线程结束（实际上不会结束）
    pthread_join(port_thread, NULL);
    pthread_join(message_thread, NULL);
    
    return 0;
}