/*
 * WSSSH Tunnel (wsssht) - Main Program
 * WebSocket tunnel setup tool for manual connections.
 *
 * Copyright (C) 2024 Stefy Lanza <stefy@nexlab.net> and SexHack.me
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <https://www.gnu.org/licenses/>.
 */

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

#include "libwsssht/wssshlib.h"
#include "libwsssht/websocket.h"
#include "libwsssht/wssh_ssl.h"
#include "libwsssht/tunnel.h"

#include "libwsssht/wsssht.h"
#include "libwsssht/utils.h"
#include "libwsssht/modes.h"
#include "libwsssht/threads.h"
#include "libwsssht/control_messages.h"

volatile sig_atomic_t sigint_count = 0;
volatile sig_atomic_t graceful_shutdown = 0;
volatile sig_atomic_t force_exit = 0;
volatile sig_atomic_t tunnel_close_sent = 0;

void signal_handler(int sig __attribute__((unused))) {
    sigint_count++;
    if (sigint_count == 1) {
        // Set flag for graceful shutdown - main thread will handle logging
        graceful_shutdown = 1;
    } else if (sigint_count >= 2) {
        // Force immediate exit on second signal
        force_exit = 1;
        _exit(1); // Use _exit() instead of exit() for signal safety
    }
}

void cleanup_on_exit(void) {
    // Only attempt cleanup if we haven't been forced to exit and haven't already sent tunnel_close
    if (!force_exit && !tunnel_close_sent) {
        pthread_mutex_lock(&tunnel_mutex);
        if (active_tunnel && active_tunnel->active && active_tunnel->ssl) {
            // Send tunnel_close without debug output for cleaner exit
            send_tunnel_close(active_tunnel->ssl, active_tunnel->request_id, 0);
            tunnel_close_sent = 1;

            // Wait up to 3 seconds for tunnel close to complete, but allow interruption
            time_t shutdown_start = time(NULL);
            while (time(NULL) - shutdown_start < 3 && !force_exit) {
                if (!active_tunnel->active) {
                    break;
                }
                usleep(100000); // Sleep 100ms before checking again
            }
        }
        pthread_mutex_unlock(&tunnel_mutex);
    }
}

int main(int argc, char *argv[]) {
    // Read config from wsssht.conf
    char *config_domain = read_config_value_from_file("wssshd-host", "wsssht");
    char *config_clientid = read_config_value_from_file("clientid", "wsssht");
    char *config_wssshd_port = read_config_value_from_file("wssshd-port", "wsssht");
    char *config_mode = read_config_value_from_file("mode", "wsssht");
    char *config_daemon = read_config_value_from_file("daemon", "wsssht");
    char *config_tunnel = read_config_value_from_file("tunnel", "wsssht");
    char *config_tunnel_control = read_config_value_from_file("tunnel-control", "wsssht");
    char *config_service = read_config_value_from_file("service", "wsssht");
    char *config_tunnel_host = read_config_value_from_file("tunnel-host", "wsssht");
    char *config_interval = read_config_value_from_file("interval", "wsssht");
    char *config_enc = read_config_value_from_file("enc", "wsssht");

    // Parse mode from config
    wsssh_mode_t initial_mode = MODE_INTERACTIVE;
    if (config_mode) {
        if (strcmp(config_mode, "silent") == 0) {
            initial_mode = MODE_SILENT;
        } else if (strcmp(config_mode, "bridge") == 0) {
            initial_mode = MODE_BRIDGE;
        } else if (strcmp(config_mode, "script") == 0) {
            initial_mode = MODE_SCRIPT;
        } else if (strcmp(config_mode, "pipe") == 0) {
            initial_mode = MODE_PIPE;
        }
        free(config_mode);
    }

    // Parse daemon from config
    int initial_daemon = 0;
    if (config_daemon) {
        if (strcmp(config_daemon, "true") == 0 || strcmp(config_daemon, "1") == 0) {
            initial_daemon = 1;
        }
        free(config_daemon);
    }

    // Parse encoding from config
    wsssh_encoding_t initial_encoding = ENCODING_HEX;
    if (config_enc) {
        if (strcmp(config_enc, "hex") == 0) {
            initial_encoding = ENCODING_HEX;
        } else if (strcmp(config_enc, "base64") == 0) {
            initial_encoding = ENCODING_BASE64;
        } else if (strcmp(config_enc, "bin") == 0) {
            initial_encoding = ENCODING_BINARY;
        } else {
            fprintf(stderr, "Warning: Invalid encoding in config: %s, using hex\n", config_enc);
        }
        free(config_enc);
    }

    wsssh_config_t config = {
        .local_port = NULL,
        .tunnel_host = config_tunnel_host ? strdup(config_tunnel_host) : NULL,
        .wssshd_host = config_domain ? strdup(config_domain) : NULL,
        .client_id = config_clientid ? strdup(config_clientid) : NULL,
        .wssshd_port = config_wssshd_port ? atoi(config_wssshd_port) : 9898,
        .debug = 0,
        .interval = 5,
        .dev_tunnel = 0,
        .tunnel = config_tunnel ? strdup(config_tunnel) : NULL,
        .tunnel_control = config_tunnel_control ? strdup(config_tunnel_control) : NULL,
        .service = config_service ? strdup(config_service) : NULL,
        .mode = initial_mode,
        .encoding = initial_encoding,
        .daemon = initial_daemon
    };

    // Set defaults if not provided
    if (!config.tunnel) {
        config.tunnel = strdup("any");
    }
    if (!config.tunnel_control) {
        config.tunnel_control = strdup("any");
    }
    if (!config.service) {
        config.service = strdup("ssh");
    }
    if (!config.tunnel_host) {
        config.tunnel_host = strdup("127.0.0.1");
    }
    if (config_interval) {
        config.interval = atoi(config_interval);
        free(config_interval);
    }

    // Easter egg: --support option (only when it's the only argument)
    if (argc == 2 && strcmp(argv[1], "--support") == 0) {
        print_trans_flag();
        free(config_domain);
        free(config_clientid);
        free(config_wssshd_port);
        free(config_mode);
        free(config_daemon);
        return 0;
    }

    // Handle --help and --free before parsing
    for (int i = 1; i < argc; i++) {
        if (strcmp(argv[i], "--help") == 0 || strcmp(argv[i], "-h") == 0) {
            print_usage(argv[0]);
            free(config_domain);
            free(config_clientid);
            free(config_wssshd_port);
            free(config_mode);
            free(config_daemon);
            return 0;
        } else if (strcmp(argv[i], "--free") == 0) {
            print_palestinian_flag();
            free(config_domain);
            free(config_clientid);
            free(config_wssshd_port);
            free(config_mode);
            free(config_daemon);
            return 0;
        }
    }

    pthread_mutex_init(&tunnel_mutex, NULL);
    pthread_mutex_init(&ssl_mutex, NULL);

    // Register cleanup function to run on exit
    atexit(cleanup_on_exit);

    // Set up signal handlers for SIGINT and SIGTERM
    signal(SIGINT, signal_handler);
    signal(SIGTERM, signal_handler);

    // Parse wsssht arguments
    int remaining_argc;
    char **remaining_argv;
    if (!wsssht_parse_args(argc, argv, &config, &remaining_argc, &remaining_argv)) {
        pthread_mutex_destroy(&tunnel_mutex);
        free(config_domain);
        return 1;
    }

    // No positional arguments expected
    if (remaining_argc > 0) {
        fprintf(stderr, "Error: Unexpected arguments\n");
        print_usage(argv[0]);
        pthread_mutex_destroy(&tunnel_mutex);
        free(config_domain);
        free(config_clientid);
        free(config_wssshd_port);
        return 1;
    }

    // Check if client_id and wssshd_host are provided
    if (!config.client_id) {
        fprintf(stderr, "Error: Client ID not specified. Use --clientid or set in config\n");
        print_usage(argv[0]);
        pthread_mutex_destroy(&tunnel_mutex);
        free(config_domain);
        free(config_clientid);
        free(config_wssshd_port);
        return 1;
    }

    if (!config.wssshd_host) {
        fprintf(stderr, "Error: wssshd host not specified. Use --wssshd-host or set in config\n");
        print_usage(argv[0]);
        pthread_mutex_destroy(&tunnel_mutex);
        free(config_domain);
        free(config_clientid);
        free(config_wssshd_port);
        return 1;
    }

    char *client_id = config.client_id;
    char *wssshd_host = config.wssshd_host;
    int wssshd_port = config.wssshd_port;

    if (config.debug) {
        fprintf(stderr, "[DEBUG - Tunnel] Client ID: %s\n", client_id);
        fprintf(stderr, "[DEBUG - Tunnel] WSSSHD Host: %s\n", wssshd_host);
        fprintf(stderr, "[DEBUG - Tunnel] WSSSHD Port: %d\n", wssshd_port);
    }

    // Handle different operating modes
    if (config.mode == MODE_BRIDGE) {
        return run_bridge_mode(&config, client_id, wssshd_host, wssshd_port);
    } else if (config.mode == MODE_SCRIPT) {
        return run_script_mode(&config, client_id, wssshd_host, wssshd_port);
    } else if (config.mode == MODE_PIPE) {
        return run_pipe_mode(&config, client_id, wssshd_host, wssshd_port);
    }
    // MODE_INTERACTIVE and MODE_SILENT continue with normal flow

    // Handle daemon mode - lazy tunnel establishment
    if (config.daemon) {
        return run_daemon_mode(&config, client_id, wssshd_host, wssshd_port);
    }

    // Find available local port
    int local_port = config.local_port ? atoi(config.local_port) : find_available_port();
    if (local_port == 0) {
        fprintf(stderr, "Error: Could not find available local port\n");
        free(config.local_port);
        free(config.tunnel_host);
        pthread_mutex_destroy(&tunnel_mutex);
        return 1;
    }

    if (config.debug) {
        printf("[DEBUG - Tunnel] Using local port: %d\n", local_port);
        fflush(stdout);
    }

    // Parent process: setup tunnel
    if (config.debug) {
        printf("[DEBUG - Tunnel] Setting up tunnel...\n");
        fflush(stdout);
    }

    // Attempt initial tunnel setup with retry logic using --interval
    int listen_sock = -1;
    int setup_attempts = 0;
    int max_setup_attempts = 3;
    tunnel_setup_result_t setup_result = {0};

    while (setup_attempts < max_setup_attempts && listen_sock < 0) {
        if (config.debug && setup_attempts > 0) {
            printf("[DEBUG - Tunnel] Tunnel setup attempt %d/%d\n", setup_attempts + 1, max_setup_attempts);
            fflush(stdout);
        }

        setup_result = setup_tunnel(wssshd_host, wssshd_port, client_id, local_port, config.debug, 0, config.tunnel_host, config.encoding, 0);
        listen_sock = setup_result.listen_sock;

        if (listen_sock < 0) {
            setup_attempts++;
            if (setup_attempts < max_setup_attempts) {
                if (config.debug) {
                    printf("[DEBUG - Tunnel] Tunnel setup failed, waiting %d seconds...\n", config.interval);
                    fflush(stdout);
                }
                sleep(config.interval);
            }
        }
    }

    if (listen_sock < 0) {
        fprintf(stderr, "Error: Failed to establish tunnel after %d attempts\n", max_setup_attempts);
        free(config.local_port);
        pthread_mutex_destroy(&tunnel_mutex);
        return 1;
    }

    // Print tunnel information and connection instructions (unless silent mode)
    if (config.mode != MODE_SILENT) {
        printf("\n");
        printf("========================================\n");
        printf("        WEBSSH TUNNEL READY\n");
        printf("========================================\n");
        printf("Tunnel established successfully!\n");
        printf("Local port: %d\n", local_port);
        printf("Target: %s@%s\n", client_id, wssshd_host);
        printf("\n");
        printf("Connect manually using one of these commands:\n");
        printf("\n");
        printf("  Telnet:\n");
        printf("    telnet localhost %d\n", local_port);
        printf("\n");
        printf("  Netcat:\n");
        printf("    nc localhost %d\n", local_port);
        printf("\n");
        printf("  SSH (if connecting to SSH server):\n");
        printf("    ssh -p %d user@localhost\n", local_port);
        printf("\n");
        printf("  SCP (if connecting to SSH server):\n");
        printf("    scp -P %d user@localhost:/remote/path ./local/path\n", local_port);
        printf("\n");
        printf("  Any TCP client:\n");
        printf("    Connect to localhost:%d\n", local_port);
        printf("\n");
        printf("Press Ctrl+C to close the tunnel and exit.\n");
        printf("========================================\n");
        printf("\n");
    }

    // Start forwarding threads to handle bidirectional communication
    // Parent process: accept connection and start forwarding
    if (config.debug) {
        printf("[DEBUG - Tunnel] Waiting for connection on localhost:%d...\n", local_port);
        fflush(stdout);
    }

    struct sockaddr_in client_addr;
    socklen_t client_len = sizeof(client_addr);
    int accepted_sock = accept(listen_sock, (struct sockaddr *)&client_addr, &client_len);
    if (accepted_sock < 0) {
        perror("Local accept failed");
        close(listen_sock);
        SSL_free(setup_result.ssl);
        if (setup_result.ssl_ctx) SSL_CTX_free(setup_result.ssl_ctx);
        free(config.local_port);
        pthread_mutex_destroy(&tunnel_mutex);
        return 1;
    }

    close(listen_sock); // No longer needed

    // Now send the tunnel request
    if (config.debug) {
        printf("[DEBUG - Tunnel] Local connection accepted, sending tunnel request...\n");
        fflush(stdout);
    }

    // Send tunnel request
    char *expanded_tunnel = expand_transport_list("any", 0); // Data channel transports
    char *expanded_tunnel_control = expand_transport_list("any", 1); // Control channel transports

    // Select best transport based on weight (lowest weight = highest priority)
    char *best_tunnel = select_best_transport(expanded_tunnel);
    char *best_tunnel_control = select_best_transport(expanded_tunnel_control);

    if (!send_tunnel_request_message_with_enc(setup_result.ssl, client_id, setup_result.request_id, best_tunnel ? best_tunnel : expanded_tunnel, best_tunnel_control ? best_tunnel_control : expanded_tunnel_control, config.service ? config.service : "ssh", config.encoding)) {
        free(expanded_tunnel);
        free(expanded_tunnel_control);
        if (best_tunnel) free(best_tunnel);
        if (best_tunnel_control) free(best_tunnel_control);
        close(accepted_sock);
        SSL_free(setup_result.ssl);
        // SSL_CTX_free(ssl_ctx);
        free(config.local_port);
        pthread_mutex_destroy(&tunnel_mutex);
        return 1;
    }

    free(expanded_tunnel);
    free(expanded_tunnel_control);
    if (best_tunnel) free(best_tunnel);
    if (best_tunnel_control) free(best_tunnel_control);

    if (config.debug) {
        printf("[DEBUG] Tunnel request sent for client: %s, request_id: %s\n", client_id, setup_result.request_id);
    }

    // Read acknowledgment
    char buffer[BUFFER_SIZE];
    int bytes_read = SSL_read(setup_result.ssl, buffer, sizeof(buffer));
    if (bytes_read <= 0) {
        if (config.debug) {
            printf("[DEBUG] No acknowledgment received\n");
        }
        close(accepted_sock);
        SSL_free(setup_result.ssl);
        // SSL_CTX_free(ssl_ctx);
        free(config.local_port);
        pthread_mutex_destroy(&tunnel_mutex);
        return 1;
    }

    // Parse WebSocket frame and check for tunnel_ack
    char *payload;
    int payload_len;
    if (!parse_websocket_frame(buffer, bytes_read, &payload, &payload_len)) {
        fprintf(stderr, "Failed to parse WebSocket frame\n");
        close(accepted_sock);
        SSL_free(setup_result.ssl);
        // SSL_CTX_free(ssl_ctx);
        free(config.local_port);
        pthread_mutex_destroy(&tunnel_mutex);
        return 1;
    }

    payload[payload_len] = '\0';

    if (strstr(payload, "tunnel_ack") == NULL) {
        fprintf(stderr, "Tunnel request denied or failed: %s\n", payload);
        close(accepted_sock);
        SSL_free(setup_result.ssl);
        // SSL_CTX_free(ssl_ctx);
        free(config.local_port);
        pthread_mutex_destroy(&tunnel_mutex);
        return 1;
    }

    if (config.debug) {
        printf("[DEBUG] Tunnel established, creating tunnel structure...\n");
    }

    // Create tunnel structure
    tunnel_t *new_tunnel = malloc(sizeof(tunnel_t));
    if (!new_tunnel) {
        perror("Memory allocation failed");
        close(accepted_sock);
        SSL_free(setup_result.ssl);
        // SSL_CTX_free(ssl_ctx);
        free(config.local_port);
        pthread_mutex_destroy(&tunnel_mutex);
        return 1;
    }
    memset(new_tunnel, 0, sizeof(tunnel_t));

    new_tunnel->outgoing_buffer = NULL;
    new_tunnel->incoming_buffer = frame_buffer_init();
    if (!new_tunnel->incoming_buffer) {
        perror("Failed to initialize incoming buffer");
        free(new_tunnel);
        close(accepted_sock);
        SSL_free(setup_result.ssl);
        if (setup_result.ssl_ctx) SSL_CTX_free(setup_result.ssl_ctx);
        free(config.local_port);
        pthread_mutex_destroy(&tunnel_mutex);
        return 1;
    }

    strcpy(new_tunnel->request_id, setup_result.request_id);
    new_tunnel->sock = -1;
    new_tunnel->local_sock = accepted_sock;
    new_tunnel->active = 1;
    new_tunnel->broken = 0;
    new_tunnel->ssl = setup_result.ssl;
    new_tunnel->server_version_sent = 0;
    new_tunnel->bin = (config.encoding == ENCODING_BINARY);
    new_tunnel->encoding = config.encoding;
    new_tunnel->start_time = time(NULL);

    new_tunnel->retransmission_buffer = retransmission_buffer_init();
    if (!new_tunnel->retransmission_buffer) {
        perror("Failed to initialize retransmission buffer");
        frame_buffer_free(new_tunnel->incoming_buffer);
        free(new_tunnel);
        close(accepted_sock);
        SSL_free(setup_result.ssl);
        if (setup_result.ssl_ctx) SSL_CTX_free(setup_result.ssl_ctx);
        free(config.local_port);
        pthread_mutex_destroy(&tunnel_mutex);
        return 1;
    }

    // Initialize keep-alive statistics
    time_t current_time = time(NULL);
    new_tunnel->last_keepalive_sent = current_time;
    new_tunnel->last_keepalive_received = current_time;
    new_tunnel->total_bytes_sent = 0;
    new_tunnel->total_bytes_received = 0;
    new_tunnel->bytes_last_period = 0;
    new_tunnel->last_stats_reset = current_time;

    // Add the new tunnel to the array
    pthread_mutex_lock(&tunnel_mutex);
    if (!add_tunnel(new_tunnel)) {
        frame_buffer_free(new_tunnel->incoming_buffer);
        free(new_tunnel);
        pthread_mutex_unlock(&tunnel_mutex);
        close(accepted_sock);
        SSL_free(setup_result.ssl);
        if (setup_result.ssl_ctx) SSL_CTX_free(setup_result.ssl_ctx);
        free(config.local_port);
        pthread_mutex_destroy(&tunnel_mutex);
        return 1;
    }

    active_tunnel = new_tunnel;
    pthread_mutex_unlock(&tunnel_mutex);

    if (config.debug) {
        printf("[DEBUG - Tunnel] Local connection accepted! Starting data forwarding...\n");
        fflush(stdout);
    }

    // Get initial SSL connection for thread
    pthread_mutex_lock(&tunnel_mutex);
    SSL *current_ssl = active_tunnel ? active_tunnel->ssl : NULL;
    pthread_mutex_unlock(&tunnel_mutex);

    // Start forwarding thread
    thread_args_t *thread_args = malloc(sizeof(thread_args_t));
    if (!thread_args) {
        perror("Memory allocation failed for thread args");
        pthread_mutex_lock(&tunnel_mutex);
        close(active_tunnel->local_sock);
        free(active_tunnel);
        active_tunnel = NULL;
        pthread_mutex_unlock(&tunnel_mutex);
        free(config.local_port);
        pthread_mutex_destroy(&tunnel_mutex);
        return 1;
    }
    thread_args->ssl = current_ssl;
    thread_args->tunnel = active_tunnel;
    thread_args->debug = config.debug;

    pthread_t thread;
    pthread_create(&thread, NULL, forward_tcp_to_ws, thread_args);
    pthread_detach(thread);

    // Main tunnel loop - handle WebSocket messages
    fd_set readfds;
    struct timeval tv;
    int tunnel_broken = 0;

    // Frame accumulation buffer for handling partial WebSocket frames
    static char frame_buffer[BUFFER_SIZE * 4];
    static int frame_buffer_used = 0;

    // Keep-alive message variables
    time_t last_keepalive_time = time(NULL);
    const int keepalive_interval = 30; // Send keep-alive every 30 seconds

    // Keep-alive timeout check variables
    time_t last_timeout_check = time(NULL);
    const int timeout_check_interval = 10; // Check timeouts every 10 seconds

    while (1) {
        // Check for graceful shutdown
        if (graceful_shutdown) {
            fprintf(stderr, "\nReceived signal, attempting graceful shutdown...\n");
            fflush(stderr);
            fprintf(stderr, "Attempting graceful shutdown of active tunnel...\n");

            // Send tunnel_close message for the active tunnel
            pthread_mutex_lock(&tunnel_mutex);
            if (active_tunnel && active_tunnel->active && active_tunnel->ssl) {
                if (config.debug) {
                    printf("[DEBUG] Sending tunnel_close for tunnel %s\n", active_tunnel->request_id);
                    fflush(stdout);
                }
                send_tunnel_close(active_tunnel->ssl, active_tunnel->request_id, config.debug);
                tunnel_close_sent = 1;

                // Wait up to 3 seconds for tunnel close to complete, but allow interruption by second signal
                fprintf(stderr, "Sent close message for tunnel, waiting up to 3 seconds...\n");
                time_t shutdown_start = time(NULL);
                while (time(NULL) - shutdown_start < 3 && !force_exit) {
                    if (!active_tunnel->active) {
                        fprintf(stderr, "Tunnel closed gracefully\n");
                        break;
                    }
                    usleep(100000); // Sleep 100ms before checking again
                }

                // If we were forced to exit, break immediately
                if (force_exit) {
                    fprintf(stderr, "Forced exit during tunnel close wait\n");
                    pthread_mutex_unlock(&tunnel_mutex);
                    break;
                }
            }
            pthread_mutex_unlock(&tunnel_mutex);
            break;
        }

        // Get SSL fd with mutex protection
        pthread_mutex_lock(&tunnel_mutex);
        if (!active_tunnel || !active_tunnel->active) {
            if (active_tunnel && active_tunnel->broken) {

                pthread_mutex_unlock(&tunnel_mutex);
                goto cleanup_and_exit;
            } else {
                // normal closure
                pthread_mutex_unlock(&tunnel_mutex);
                if (config.debug) {
                    printf("[DEBUG - Tunnel] Tunnel is no longer active, exiting main loop\n");
                    fflush(stdout);
                }
                break;
            }
        }
        int ssl_fd = SSL_get_fd(active_tunnel->ssl);
        current_ssl = active_tunnel->ssl;

        // Check if local socket is still valid
        if (active_tunnel->local_sock < 0) {
            if (config.debug) {
                printf("[DEBUG - Tunnel] Local socket is invalid, tunnel broken\n");
                fflush(stdout);
            }

            active_tunnel->broken = 1;
            // Send tunnel_close notification
            if (config.debug) {
                printf("[DEBUG - Tunnel] Sending tunnel_close notification due to invalid local socket...\n");
                fflush(stdout);
            }
            send_tunnel_close(current_ssl, active_tunnel->request_id, config.debug);
            tunnel_close_sent = 1;
            pthread_mutex_unlock(&tunnel_mutex);
            goto cleanup_and_exit;
        }

        // Check if the local socket connection is broken
        char test_buf[1];
        int result = recv(active_tunnel->local_sock, test_buf, 1, MSG_PEEK | MSG_DONTWAIT);
        if (result == 0 || (result < 0 && (errno == ECONNRESET || errno == EPIPE || errno == EBADF))) {
            if (config.debug) {
                printf("[DEBUG - Tunnel] Local socket connection is broken (errno=%d), sending tunnel_close\n", errno);
                fflush(stdout);
            }

            active_tunnel->broken = 1;
            // Send tunnel_close notification
            send_tunnel_close(current_ssl, active_tunnel->request_id, config.debug);
            tunnel_close_sent = 1;
            pthread_mutex_unlock(&tunnel_mutex);
            goto cleanup_and_exit;
        }

        pthread_mutex_unlock(&tunnel_mutex);

        // Use select to wait for data on SSL socket with timeout
        FD_ZERO(&readfds);
        FD_SET(ssl_fd, &readfds);
        tv.tv_sec = 0;
        tv.tv_usec = 50000;  // 50ms timeout

        int retval = select(ssl_fd + 1, &readfds, NULL, NULL, &tv);
        if (retval == -1) {
            if (config.debug) {
                perror("[DEBUG - WebSockets] select on SSL fd failed");
                fflush(stdout);
            }

            // Send tunnel_close notification
            if (config.debug) {
                printf("[DEBUG - Tunnel] Sending tunnel_close notification due to select failure...\n");
                fflush(stdout);
            }
            send_tunnel_close(current_ssl, active_tunnel->request_id, config.debug);
            tunnel_close_sent = 1;
            goto cleanup_and_exit;
        } else if (retval == 0) {
            // Timeout, check if tunnel became inactive
            pthread_mutex_lock(&tunnel_mutex);
            if (!active_tunnel || !active_tunnel->active) {
                pthread_mutex_unlock(&tunnel_mutex);
                if (config.debug) {
                    printf("[DEBUG - Tunnel] Tunnel became inactive during timeout, exiting\n");
                    fflush(stdout);
                }

                goto cleanup_and_exit;
            }
            pthread_mutex_unlock(&tunnel_mutex);

            // Send keep-alive messages if needed
            time_t current_time = time(NULL);
            if (current_time - last_keepalive_time >= keepalive_interval) {
                pthread_mutex_lock(&tunnel_mutex);
                if (active_tunnel && active_tunnel->active) {
                    send_tunnel_keepalive(current_ssl, active_tunnel, config.debug);
                }
                pthread_mutex_unlock(&tunnel_mutex);
                last_keepalive_time = current_time;
            }

            // Check for keep-alive timeouts
            if (current_time - last_timeout_check >= timeout_check_interval) {
                check_keepalive_timeouts(config.debug);
                last_timeout_check = current_time;
            }

            continue;
        }

        // Read more data if we don't have a complete frame
        if (FD_ISSET(ssl_fd, &readfds)) {
            if ((size_t)frame_buffer_used < sizeof(frame_buffer)) {
                // Validate SSL connection state
                if (SSL_get_shutdown(current_ssl) & SSL_RECEIVED_SHUTDOWN) {
                    if (config.debug) {
                        printf("[DEBUG - WebSockets] SSL connection has received shutdown\n");
                        fflush(stdout);
                    }
                    cleanup_tunnel(config.debug);
                    break;
                }

                // Set up timeout for SSL read
                fd_set readfds_timeout;
                struct timeval tv_timeout;
                int sock_fd = SSL_get_fd(current_ssl);

                FD_ZERO(&readfds_timeout);
                FD_SET(sock_fd, &readfds_timeout);
                tv_timeout.tv_sec = 5;
                tv_timeout.tv_usec = 0;

                int select_result = select(sock_fd + 1, &readfds_timeout, NULL, NULL, &tv_timeout);
                if (select_result == -1) {
                    if (config.debug) {
                        perror("[DEBUG - WebSockets] select failed");
                        fflush(stdout);
                    }
                    cleanup_tunnel(config.debug);
                    break;
                } else if (select_result == 0) {
                    if (config.debug) {
                        printf("[DEBUG - WebSockets] SSL read timeout\n");
                        fflush(stdout);
                    }
                    continue;
                }

                bytes_read = SSL_read(current_ssl, frame_buffer + frame_buffer_used, sizeof(frame_buffer) - frame_buffer_used);
                if (bytes_read <= 0) {
                    if (bytes_read < 0) {
                        int ssl_error = SSL_get_error(current_ssl, bytes_read);
                        if (config.debug) {
                            printf("[DEBUG - WebSockets] SSL read error: %d\n", ssl_error);
                            fflush(stdout);
                        }

                        // Handle transient SSL errors
                        if (ssl_error == SSL_ERROR_WANT_READ || ssl_error == SSL_ERROR_WANT_WRITE) {
                            if (config.debug) {
                                printf("[DEBUG - WebSockets] Transient SSL error, retrying...\n");
                                fflush(stdout);
                            }
                            usleep(10000);
                            continue;
                        }

                        // Print SSL error details
                        char error_buf[256];
                        ERR_error_string_n(ssl_error, error_buf, sizeof(error_buf));
                        if (config.debug) {
                            printf("[DEBUG - WebSockets] SSL error details: %s\n", error_buf);
                            fflush(stdout);
                        }
                        fprintf(stderr, "SSL read error (%d): %s\n", ssl_error, error_buf);
                    } else {
                        if (config.debug) {
                            printf("[DEBUG - WebSockets] Connection closed by server (EOF)\n");
                            fflush(stdout);
                        }
                    }

                    if (config.debug) {
                        printf("[DEBUG - WebSockets] WebSocket connection lost, attempting reconnection...\n");
                        fflush(stdout);
                    }

                    // Attempt reconnection
                    int reconnect_attempts = 0;
                    int max_reconnect_attempts = 3;
                    int reconnected = 0;

                    while (reconnect_attempts < max_reconnect_attempts && !reconnected) {
                        if (config.debug) {
                            printf("[DEBUG - WebSockets] WebSocket reconnection attempt %d/%d\n", reconnect_attempts + 1, max_reconnect_attempts);
                            fflush(stdout);
                        }

                        pthread_mutex_lock(&tunnel_mutex);
                        if (!active_tunnel) {
                            pthread_mutex_unlock(&tunnel_mutex);
                            break;
                        }
                        if (reconnect_websocket(active_tunnel, wssshd_host, wssshd_port, client_id, active_tunnel->request_id, config.debug) == 0) {
                            reconnected = 1;
                            if (config.debug) {
                                printf("[DEBUG - WebSockets] WebSocket reconnection successful, continuing tunnel\n");
                                fflush(stdout);
                            }
                            // Update ssl_fd for select
                            ssl_fd = SSL_get_fd(active_tunnel->ssl);
                            current_ssl = active_tunnel->ssl;
                        }
                        pthread_mutex_unlock(&tunnel_mutex);

                        if (!reconnected) {
                            reconnect_attempts++;
                            if (reconnect_attempts < max_reconnect_attempts) {
                                if (config.debug) {
                                    printf("[DEBUG - WebSockets] WebSocket reconnection failed, waiting 1 second...\n");
                                    fflush(stdout);
                                }
                                sleep(1);
                            }
                        }
                    }

                    if (!reconnected) {
                        if (config.debug) {
                            printf("[DEBUG - WebSockets] All reconnection attempts failed, exiting\n");
                            fflush(stdout);
                        }

                        // Send tunnel_close notification
                        if (config.debug) {
                            printf("[DEBUG - Tunnel] Sending tunnel_close notification due to connection failure...\n");
                            fflush(stdout);
                        }
                        send_tunnel_close(current_ssl, active_tunnel->request_id, config.debug);
                        tunnel_close_sent = 1;
                        goto cleanup_and_exit;
                    }

                    // Skip processing this iteration since we just reconnected
                    continue;
                }

                frame_buffer_used += bytes_read;

                if (config.debug) {
                    printf("[DEBUG - WebSockets] Accumulated %d bytes, frame: 0x%02x 0x%02x 0x%02x 0x%02x\n", frame_buffer_used, frame_buffer[0], frame_buffer[1], frame_buffer[2], frame_buffer[3]);
                    fflush(stdout);
                }
            }

            // Try to parse WebSocket frame
            char *payload;
            int payload_len;
            if (parse_websocket_frame(frame_buffer, frame_buffer_used, &payload, &payload_len)) {
                // Frame is complete, determine frame type
                unsigned char frame_type = frame_buffer[0] & 0x8F;

                if (frame_type == 0x88) { // Close frame
                    if (config.debug) {
                        printf("[DEBUG - WebSockets] Received close frame from server\n");
                        fflush(stdout);
                    }

                    // Send tunnel_close notification
                    if (config.debug) {
                        printf("[DEBUG - Tunnel] Sending tunnel_close notification due to server close frame...\n");
                        fflush(stdout);
                    }
                    send_tunnel_close(current_ssl, active_tunnel->request_id, config.debug);
                    tunnel_close_sent = 1;
                    goto cleanup_and_exit;
                } else if (frame_type == 0x89) { // Ping frame
                    if (config.debug) {
                        printf("[DEBUG - WebSockets] Received ping frame, sending pong\n");
                        fflush(stdout);
                    }
                    // Send pong
                    if (!send_pong_frame(current_ssl, payload, payload_len)) {
                        if (config.debug) {
                            printf("[DEBUG - WebSockets] Failed to send pong frame\n");
                            fflush(stdout);
                        }
                    }
                } else if (frame_type == 0x8A) { // Pong frame
                    if (config.debug) {
                        printf("[DEBUG - WebSockets] Received pong frame\n");
                        fflush(stdout);
                    }
                } else if (frame_type == 0x81 || frame_type == 0x82) { // Text or binary frame
                    // Copy payload to buffer first to check message type
                    if ((size_t)payload_len < sizeof(buffer)) {
                        memcpy(buffer, payload, payload_len);
                        buffer[payload_len] = '\0';
                    } else {
                        fprintf(stderr, "Payload too large for processing buffer\n");
                        frame_buffer_used = 0;
                        continue;
                    }

                    // Check if this is a tunnel_data message to suppress verbose logging
                    int is_tunnel_data = (strstr(buffer, "\"type\":\"tunnel_data\"") != NULL);

                    if (config.debug && !is_tunnel_data) {
                        printf("[DEBUG - WebSockets] Received message: %.*s\n", payload_len, payload);
                        fflush(stdout);
                    }

                    // Handle message
                    if (config.debug && !is_tunnel_data) {
                        printf("[DEBUG - WebSockets] Processing message: %s\n", buffer);
                        fflush(stdout);
                    }

                    // Handle tunnel messages
                    if (strstr(payload, "tunnel_data") || strstr(payload, "tunnel_response")) {
                        if (config.debug) {
                            // Suppress tunnel_data debug messages in debug mode
                            if (!strstr(payload, "tunnel_data")) {
                                printf("[DEBUG - Tunnel] Received tunnel_response message\n");
                                fflush(stdout);
                            }
                        }
                        // Extract request_id and data with bounds checking
                        char *id_start = strstr(payload, "\"request_id\"");
                        char *data_start = strstr(payload, "\"data\"");
                        if (id_start && data_start && id_start < data_start) {  // Ensure proper order
                            char *colon = strchr(id_start, ':');
                            if (colon && colon < data_start) {
                                char *open_quote = strchr(colon, '"');
                                if (open_quote && open_quote < data_start) {
                                    id_start = open_quote + 1;
                                    char *close_quote = strchr(id_start, '"');
                                    if (close_quote && close_quote < data_start) {
                                        *close_quote = '\0';
                                        char *data_colon = strchr(data_start, ':');
                                        if (data_colon && data_colon < payload + payload_len) {
                                            char *data_quote = strchr(data_colon, '"');
                                            if (data_quote && data_quote < payload + payload_len) {
                                                data_start = data_quote + 1;
                                                char *data_end = strchr(data_start, '"');
                                                if (data_end && data_end < payload + payload_len && data_end > data_start) {
                                                    *data_end = '\0';
                                                    size_t hex_len = data_end - data_start;
                                                    // Additional validation: reasonable hex data size
                                                    if (hex_len > 0 && hex_len < 1024 * 1024) {  // Max 1MB hex data
                                                        handle_tunnel_data(current_ssl, id_start, data_start, config.debug);
                                                    } else if (config.debug) {
                                                        printf("[DEBUG] Invalid hex data length: %zu\n", hex_len);
                                                        fflush(stdout);
                                                    }
                                                } else if (config.debug) {
                                                    printf("[DEBUG] Malformed data field in JSON\n");
                                                    fflush(stdout);
                                                }
                                            }
                                        }
                                    }
                                }
                            }
                        } else if (config.debug) {
                            printf("[DEBUG] Malformed tunnel message JSON\n");
                            fflush(stdout);
                        }
                    } else if (strstr(payload, "tunnel_close")) {
                        if (config.debug) {
                            printf("[DEBUG - Tunnel] Received tunnel_close message\n");
                            fflush(stdout);
                        }
                        char *id_start = strstr(payload, "\"request_id\"");
                        if (id_start && id_start < payload + payload_len) {
                            char *colon = strchr(id_start, ':');
                            if (colon && colon < payload + payload_len) {
                                char *open_quote = strchr(colon, '"');
                                if (open_quote && open_quote < payload + payload_len) {
                                    id_start = open_quote + 1;
                                    char *close_quote = strchr(id_start, '"');
                                    if (close_quote && close_quote < payload + payload_len && close_quote > id_start) {
                                        *close_quote = '\0';
                                        handle_tunnel_close(current_ssl, id_start, config.debug);
                                    }
                                }
                            }
                        }
                    } else if (strstr(payload, "tunnel_keepalive")) {
                        // Extract request_id, total_bytes, and rate_bps with bounds checking
                        char *id_start = strstr(payload, "\"request_id\"");
                        char *total_start = strstr(payload, "\"total_bytes\"");
                        char *rate_start = strstr(payload, "\"rate_bps\"");

                        if (id_start && id_start < payload + payload_len) {
                            char *colon = strchr(id_start, ':');
                            if (colon && colon < payload + payload_len) {
                                char *open_quote = strchr(colon, '"');
                                if (open_quote && open_quote < payload + payload_len) {
                                    id_start = open_quote + 1;
                                    char *close_quote = strchr(id_start, '"');
                                    if (close_quote && close_quote < payload + payload_len && close_quote > id_start) {
                                        *close_quote = '\0';

                                        unsigned long long total_bytes = 0;
                                        double rate_bps = 0.0;

                                        // Parse total_bytes
                                        if (total_start && total_start < payload + payload_len) {
                                            char *total_colon = strchr(total_start, ':');
                                            if (total_colon && total_colon < payload + payload_len) {
                                                total_bytes = strtoull(total_colon + 1, NULL, 10);
                                            }
                                        }

                                        // Parse rate_bps
                                        if (rate_start && rate_start < payload + payload_len) {
                                            char *rate_colon = strchr(rate_start, ':');
                                            if (rate_colon && rate_colon < payload + payload_len) {
                                                rate_bps = strtod(rate_colon + 1, NULL);
                                            }
                                        }

                                        handle_tunnel_keepalive(current_ssl, id_start, total_bytes, rate_bps, config.debug);
                                    }
                                }
                            }
                        }
                    } else if (strstr(payload, "tunnel_keepalive_ack")) {
                        // Extract request_id with bounds checking
                        char *id_start = strstr(payload, "\"request_id\"");
                        if (id_start && id_start < payload + payload_len) {
                            char *colon = strchr(id_start, ':');
                            if (colon && colon < payload + payload_len) {
                                char *open_quote = strchr(colon, '"');
                                if (open_quote && open_quote < payload + payload_len) {
                                    id_start = open_quote + 1;
                                    char *close_quote = strchr(id_start, '"');
                                    if (close_quote && close_quote < payload + payload_len && close_quote > id_start) {
                                        *close_quote = '\0';
                                        handle_tunnel_keepalive_ack(current_ssl, id_start, config.debug);
                                    }
                                }
                            }
                        }
                    } else if (strstr(payload, "tunnel_ack")) {
                        // Extract request_id and frame_id with bounds checking
                        char *id_start = strstr(payload, "\"request_id\"");
                        char *frame_start = strstr(payload, "\"frame_id\"");
                        if (id_start && frame_start && id_start < frame_start && frame_start < payload + payload_len) {
                            char *colon = strchr(id_start, ':');
                            if (colon && colon < frame_start) {
                                char *open_quote = strchr(colon, '"');
                                if (open_quote && open_quote < frame_start) {
                                    id_start = open_quote + 1;
                                    char *close_quote = strchr(id_start, '"');
                                    if (close_quote && close_quote < frame_start && close_quote > id_start) {
                                        *close_quote = '\0';
                                        // Extract frame_id
                                        char *frame_colon = strchr(frame_start, ':');
                                        if (frame_colon && frame_colon < payload + payload_len) {
                                            uint32_t frame_id = (uint32_t)atoi(frame_colon + 1);
                                            handle_tunnel_ack(current_ssl, id_start, frame_id, config.debug);
                                        }
                                    }
                                }
                            }
                        }
                    } else if (strstr(payload, "tunnel_ko")) {
                        // Extract request_id and frame_id with bounds checking
                        char *id_start = strstr(payload, "\"request_id\"");
                        char *frame_start = strstr(payload, "\"frame_id\"");
                        if (id_start && frame_start && id_start < frame_start && frame_start < payload + payload_len) {
                            char *colon = strchr(id_start, ':');
                            if (colon && colon < frame_start) {
                                char *open_quote = strchr(colon, '"');
                                if (open_quote && open_quote < frame_start) {
                                    id_start = open_quote + 1;
                                    char *close_quote = strchr(id_start, '"');
                                    if (close_quote && close_quote < frame_start && close_quote > id_start) {
                                        *close_quote = '\0';
                                        // Extract frame_id
                                        char *frame_colon = strchr(frame_start, ':');
                                        if (frame_colon && frame_colon < payload + payload_len) {
                                            uint32_t frame_id = (uint32_t)atoi(frame_colon + 1);
                                            handle_tunnel_ko(current_ssl, id_start, frame_id, config.debug);
                                        }
                                    }
                                }
                            }
                        }
                    } else {
                        if (config.debug) {
                            printf("[DEBUG - WebSockets] Received unknown message type: %.*s\n", payload_len, payload);
                            fflush(stdout);
                        }
                    }
                }

                // Remove processed frame from buffer
                // Calculate the actual frame size consumed from the buffer
                int frame_size = (payload - frame_buffer) + payload_len;
                if (frame_size <= frame_buffer_used) {
                    if (frame_size < frame_buffer_used) {
                        memmove(frame_buffer, frame_buffer + frame_size, frame_buffer_used - frame_size);
                        frame_buffer_used -= frame_size;
                    } else {
                        // Frame consumed entire buffer
                        frame_buffer_used = 0;
                    }
                } else {
                    // Safety check: if calculated frame_size is larger than buffer_used,
                    // something went wrong in parsing, reset buffer to be safe
                    if (config.debug) {
                        printf("[DEBUG] Frame size calculation error: frame_size=%d, buffer_used=%d\n", frame_size, frame_buffer_used);
                        fflush(stdout);
                    }
                    frame_buffer_used = 0;
                }
            } else {
                // Frame not complete yet, continue reading
                continue;
            }
        }
    }

cleanup_and_exit:
    // Cleanup section
    if (config.debug) {
        printf("[DEBUG - Tunnel] Performing cleanup and exiting\n");
        fflush(stdout);
    }

    // Cleanup
    pthread_mutex_lock(&tunnel_mutex);
    if (active_tunnel) {
        if (active_tunnel->local_sock >= 0) {
            close(active_tunnel->local_sock);
        }
        if (active_tunnel->ssl) {
            SSL_free(active_tunnel->ssl);
        }
        free(active_tunnel);
        active_tunnel = NULL;
    }
    pthread_mutex_unlock(&tunnel_mutex);

    free(config.local_port);
    free(config.tunnel);
    free(config.tunnel_control);
    free(config.service);
    free(config_domain);
    free(config_clientid);
    free(config_wssshd_port);
    free(config_tunnel);
    free(config_tunnel_control);
    free(config_service);
    pthread_mutex_destroy(&tunnel_mutex);
    pthread_mutex_destroy(&ssl_mutex);

    if (config.debug) {
        printf("[DEBUG - Tunnel] Cleanup complete, exiting with code %d\n", tunnel_broken ? 1 : 0);
        fflush(stdout);
    }

    return tunnel_broken ? 1 : 0;
}