Add proxy functionality

wssshd2/config.c

@@ -116,7 +116,8 @@ static void load_config_file(wssshd_config_t *config, const char *config_file) {
                 } else if (strcmp(key, "web-https") == 0) {
                     config->web_https = (strcmp(value, "true") == 0 || strcmp(value, "1") == 0);
                 } else if (strcmp(key, "web-proxy") == 0) {
-                    config->web_proxy_enabled = (strcmp(value, "true") == 0 || strcmp(value, "1") == 0);
+                    config->web_proxy_enabled = true;
+                    config->web_proxy_port = atoi(value);
                 } else if (strcmp(key, "web-proxy-port") == 0) {
                     config->web_proxy_port = atoi(value);
                 } else if (strcmp(key, "websocket-tls-only") == 0) {

wssshd2/main.c

@@ -24,6 +24,7 @@
 #include <fcntl.h>
 #include <pthread.h>
 #include <time.h>
+#include <arpa/inet.h>
 #include "config.h"
 #include "websocket.h"
 #include "web.h"
@@ -169,22 +170,25 @@ int main(int argc, char *argv[]) {
         fprintf(stderr, "Warning: Failed to start web interface\n");
     }
 
+    // Start web proxy server if enabled
+    printf("Checking web proxy: enabled=%d, port=%d\n", config->web_proxy_enabled, config->web_proxy_port);
+    if (config->web_proxy_enabled) {
+        printf("Starting web proxy server\n");
+        if (web_proxy_start_server(config, state) != 0) {
+            fprintf(stderr, "Warning: Failed to start web proxy server\n");
+        }
+    }
+
     // Start WebSocket server
     if (websocket_start_server(config, state) != 0) {
         fprintf(stderr, "Failed to start WebSocket server\n");
+        web_proxy_stop_server();
         web_stop_server();
         websocket_free_state(state);
         free_config(config);
         return 1;
     }
 
-    // Start web proxy server if enabled
-    if (config->web_proxy_enabled) {
-        if (web_proxy_start_server(config, state) != 0) {
-            fprintf(stderr, "Warning: Failed to start web proxy server\n");
-        }
-    }
-
     printf("WSSSH Daemon running on %s:%d\n", config->host, config->port);
     if (config->web_proxy_enabled) {
         printf("Web proxy server running on 127.0.0.1:%d\n", config->web_proxy_port);
@@ -208,9 +212,9 @@ int main(int argc, char *argv[]) {
     printf("\nShutting down WSSSH Daemon...\n");
 
     // Stop servers
+    websocket_stop_server();
     web_proxy_stop_server();
     web_stop_server();
-    websocket_stop_server();
 
     // Clean up plugin system
     plugin_system_cleanup();

wssshd2/web.c

@@ -1227,14 +1227,32 @@ static const char *validate_session(const char *session_id) {
 
 static int parse_http_request(int client_fd, http_request_t *req) {
     char buffer[8192];
-    ssize_t bytes_read = recv(client_fd, buffer, sizeof(buffer) - 1, 0);
-    if (bytes_read <= 0) return -1;
+    size_t total_read = 0;
+    int found_headers = 0;
 
-    buffer[bytes_read] = '\0';
+    // Read data in chunks until we find \r\n\r\n or reach max size
+    while (total_read < sizeof(buffer) - 1) {
+        ssize_t bytes_read = recv(client_fd, buffer + total_read, sizeof(buffer) - 1 - total_read, 0);
+        if (bytes_read <= 0) {
+            if (total_read == 0) return -1; // No data read
+            break; // End of data
+        }
+
+        total_read += bytes_read;
+        buffer[total_read] = '\0';
+
+        // Check if we have the complete headers
+        if (strstr(buffer, "\r\n\r\n")) {
+            found_headers = 1;
+            break;
+        }
+    }
+
+    if (!found_headers && total_read == 0) return -1;
 
     if (global_config && global_config->debug_web) {
-        printf("[WEB-DEBUG] Raw buffer (%zd bytes): '", bytes_read);
-        for (ssize_t i = 0; i < bytes_read; i++) {
+        printf("[WEB-DEBUG] Raw buffer (%zu bytes): '", total_read);
+        for (size_t i = 0; i < total_read; i++) {
             if (buffer[i] == '\r') printf("\\r");
             else if (buffer[i] == '\n') printf("\\n");
             else if (buffer[i] >= 32 && buffer[i] <= 126) printf("%c", buffer[i]);
@@ -1276,7 +1294,7 @@ static int parse_http_request(int client_fd, http_request_t *req) {
         printf("[WEB-DEBUG] Content-Length from header: %d\n", content_length);
     }
 
-    // Find body start - look for \r\n\r\n or calculate from Content-Length
+    // Find body start - look for \r\n\r\n
     char *body_separator = strstr(buffer, "\r\n\r\n");
     char *body_start = NULL;
     size_t body_length = 0;
@@ -1287,14 +1305,14 @@ static int parse_http_request(int client_fd, http_request_t *req) {
 
     if (body_separator) {
         body_start = body_separator + 4;
-        body_length = bytes_read - (body_start - buffer);
+        body_length = total_read - (body_start - buffer);
         if (global_config && global_config->debug_web) {
             printf("[WEB-DEBUG] Found \\r\\n\\r\\n separator at offset %ld, body_start at %ld, body_length: %zu\n",
                    body_separator - buffer, body_start - buffer, body_length);
         }
     } else if (content_length > 0) {
         // Calculate body position: total bytes - content_length
-        body_start = buffer + bytes_read - content_length;
+        body_start = buffer + total_read - content_length;
         body_length = content_length;
         if (global_config && global_config->debug_web) {
             printf("[WEB-DEBUG] Calculated body start from Content-Length, body_start at %ld, body_length: %zu\n",
@@ -1463,6 +1481,8 @@ static const char *get_cookie(const char *headers, const char *name) {
             *eq = '\0';
             if (strcmp(token, name) == 0) {
                 strcpy(value, eq + 1);
+                // URL decode the cookie value
+                url_decode(value);
                 return value;
             }
         }
@@ -1486,6 +1506,8 @@ static void send_response(int client_fd, int status_code, const char *status_tex
         }
     }
 
+    printf("[DEBUG] Sending response: %d %s, body_len: %zu\n", status_code, status_text, body_len);
+
     char response[8192];
     char date[64];
     time_t now = time(NULL);
@@ -1656,18 +1678,18 @@ static int generate_users_html(const char *username, char *html, size_t max_len)
 // Handle HTTP requests
 // Returns 1 if the connection should be kept open (WebSocket), 0 if it should be closed
 static int handle_request(int client_fd, const http_request_t *req) {
-    if (global_config && global_config->debug_web && !strstr(req->path, "/xterm/data")) {
-        printf("[WEB-DEBUG] Received %s request for %s\n", req->method, req->path);
-        if (req->query[0]) {
-            printf("[WEB-DEBUG] Query: %s\n", req->query);
-        }
-        printf("[WEB-DEBUG] Headers: %s\n", req->headers);
-        if (req->content_length > 0) {
-            printf("[WEB-DEBUG] Content-Length: %d\n", req->content_length);
-        }
-    }
-
-    const char *session_id = get_cookie(req->headers, "session_id");
+    // if (global_config && global_config->debug_web && !strstr(req->path, "/xterm/data")) {
+    //     printf("[WEB-DEBUG] Received %s request for %s\n", req->method, req->path);
+    //     if (req->query[0]) {
+    //         printf("[WEB-DEBUG] Query: %s\n", req->query);
+    //     }
+    //     printf("[WEB-DEBUG] Headers: %s\n", req->headers);
+    //     if (req->content_length > 0) {
+    //         printf("[WEB-DEBUG] Content-Length: %d\n", req->content_length);
+    //     }
+    // }
+
+    const char *session_id = get_cookie(req->headers, "wssshd_session");
     const char *username = NULL;
     int is_admin = 0;
 
@@ -1690,10 +1712,12 @@ static int handle_request(int client_fd, const http_request_t *req) {
 
     // Check if any plugin handles this request
     int plugin_result = plugin_handle_web_request(client_fd, req);
-    if (plugin_result >= 0) {
+    if (plugin_result > 0) {
         return plugin_result; // Plugin handled the request
     }
 
+    printf("[DEBUG] Handling request: %s %s\n", req->method, req->path);
+
     // Route handling
     // Handle WebSocket terminal, VNC, and RDP connections first
     char *headers_copy = NULL;
@@ -2535,9 +2559,15 @@ static int handle_request(int client_fd, const http_request_t *req) {
             char *html = generate_index_html(username, is_admin, error);
             send_response(client_fd, 200, "OK", "text/html", html, strlen(html), NULL, NULL);
         } else if (strcmp(req->path, "/login") == 0) {
-            send_response(client_fd, 200, "OK", "text/html", login_page_html, strlen(login_page_html), NULL, NULL);
+            printf("[DEBUG] Sending login page\n");
+            if (username) {
+                // If already logged in, redirect to home
+                send_response(client_fd, 302, "Found", "text/html", NULL, 0, NULL, "Location: /");
+            } else {
+                send_response(client_fd, 200, "OK", "text/html", login_page_html, strlen(login_page_html), NULL, NULL);
+            }
         } else if (strcmp(req->path, "/logout") == 0) {
-            send_response(client_fd, 302, "Found", "text/html", NULL, 0, "session_id=; Max-Age=0; Path=/", "Location: /");
+            send_response(client_fd, 302, "Found", "text/html", NULL, 0, "wssshd_session=; Max-Age=0; Path=/", "Location: /");
             return 0;
         } else if (strcmp(req->path, "/users") == 0) {
             if (!username || !is_admin) {
@@ -2566,9 +2596,11 @@ static int handle_request(int client_fd, const http_request_t *req) {
             send_response(client_fd, 200, "OK", "application/json", json, len, NULL, NULL);
         } else {
             // Try to serve embedded asset
+            printf("[DEBUG] Trying to serve embedded asset: %s\n", req->path);
             size_t size;
             const char *asset = get_embedded_asset(req->path, &size);
             if (asset) {
+                printf("[DEBUG] Asset found, size: %zu\n", size);
                 const char *content_type = "text/plain";
                 if (strstr(req->path, ".html")) content_type = "text/html";
                 else if (strstr(req->path, ".css")) content_type = "text/css";
@@ -2579,6 +2611,7 @@ static int handle_request(int client_fd, const http_request_t *req) {
                 else if (strstr(req->path, ".ico")) content_type = "image/x-icon";
                 send_response(client_fd, 200, "OK", content_type, asset, size, NULL, NULL);
             } else {
+                printf("[DEBUG] Asset not found for %s\n", req->path);
                 send_response(client_fd, 404, "Not Found", "text/html", "Not found", 9, NULL, NULL);
             }
         }
@@ -2609,7 +2642,7 @@ static int handle_request(int client_fd, const http_request_t *req) {
                                form_username, new_session);
                     }
                     char cookie[256];
-                    snprintf(cookie, sizeof(cookie), "session_id=%s; Path=/; HttpOnly", new_session);
+                    snprintf(cookie, sizeof(cookie), "wssshd_session=%s; Path=/; HttpOnly", new_session);
                     send_response(client_fd, 302, "Found", "text/html", NULL, 0, cookie, "Location: /");
                     return 0;
                 } else {
@@ -2749,6 +2782,8 @@ static void *http_server_thread(void *arg __attribute__((unused))) {
     struct sockaddr_in server_addr, client_addr;
     socklen_t client_len = sizeof(client_addr);
 
+    printf("HTTP server thread starting\n");
+
     server_socket = socket(AF_INET, SOCK_STREAM, 0);
     if (server_socket < 0) {
         perror("Failed to create server socket");
@@ -2766,10 +2801,13 @@ static void *http_server_thread(void *arg __attribute__((unused))) {
 
     if (bind(server_socket, (struct sockaddr *)&server_addr, sizeof(server_addr)) < 0) {
         perror("Failed to bind server socket");
+        printf("Failed to bind to %s:%d\n", global_config->web_host, global_config->web_port);
         close(server_socket);
         return NULL;
     }
 
+    printf("Web interface bound to %s:%d\n", global_config->web_host, global_config->web_port);
+
     if (listen(server_socket, 10) < 0) {
         perror("Failed to listen on server socket");
         close(server_socket);
@@ -2787,12 +2825,16 @@ static void *http_server_thread(void *arg __attribute__((unused))) {
         }
 
         http_request_t req;
+        memset(&req, 0, sizeof(req));
         if (parse_http_request(client_fd, &req) == 0) {
             int keep_open = handle_request(client_fd, &req);
             if (keep_open) {
                 // WebSocket connection - don't close the socket
                 continue;
             }
+        } else {
+            // Send error response if parsing failed
+            send_response(client_fd, 400, "Bad Request", "text/plain", "Bad Request", 11, NULL, NULL);
         }
 
         close(client_fd);
@@ -2813,6 +2855,8 @@ int web_start_server(const wssshd_config_t *config, wssshd_state_t *state) {
 
     init_users(config);
 
+    printf("Web server initialization complete\n");
+
     // Start HTTP server thread
     pthread_t thread;
     if (pthread_create(&thread, NULL, http_server_thread, NULL) != 0) {

wssshd2/web_proxy.c

@@ -34,9 +34,7 @@
 #include "websocket.h"
 #include "websocket_protocol.h"
 
-// Global state
-static wssshd_state_t *global_state = NULL;
-static const wssshd_config_t *global_config = NULL;
+// Global state for shutdown
 static int server_socket = -1;
 static volatile int server_running = 0;
 
@@ -55,17 +53,30 @@ static proxy_connection_t proxy_connections[MAX_PROXY_CONNECTIONS];
 static int proxy_connections_count = 0;
 static pthread_mutex_t proxy_mutex = PTHREAD_MUTEX_INITIALIZER;
 
+// Thread arguments structure
+typedef struct {
+    const wssshd_config_t *config;
+    wssshd_state_t *state;
+} proxy_thread_args_t;
+
+// Connection handler arguments
+typedef struct {
+    int client_fd;
+    wssshd_state_t *state;
+    const wssshd_config_t *config;
+} connection_handler_args_t;
+
 // Forward declaration for proxy connection handler
 static void *proxy_connection_handler(void *arg);
 
 // Find client by hostname (matching subdomain or client_id)
-static client_t *find_client_by_hostname(const char *hostname) {
-    if (!hostname || !global_state) return NULL;
+static client_t *find_client_by_hostname(const char *hostname, wssshd_state_t *state, const wssshd_config_t *config) {
+    if (!hostname || !state) return NULL;
 
-    pthread_mutex_lock(&global_state->client_mutex);
+    pthread_mutex_lock(&state->client_mutex);
 
-    for (size_t i = 0; i < global_state->clients_count; i++) {
-        client_t *client = &global_state->clients[i];
+    for (size_t i = 0; i < state->clients_count; i++) {
+        client_t *client = &state->clients[i];
         if (!client->active) continue;
 
         // Check if client has web service
@@ -76,23 +87,23 @@ static client_t *find_client_by_hostname(const char *hostname) {
 
         // Try exact match with client_id
         if (strcasecmp(hostname, client->client_id) == 0) {
-            pthread_mutex_unlock(&global_state->client_mutex);
+            pthread_mutex_unlock(&state->client_mutex);
             return client;
         }
 
         // Try subdomain match: hostname.client_id.domain
-        if (global_config && global_config->domain) {
+        if (config && config->domain) {
             char expected_hostname[512];
             snprintf(expected_hostname, sizeof(expected_hostname), "%s.%s",
-                     client->client_id, global_config->domain);
+                     client->client_id, config->domain);
             if (strcasecmp(hostname, expected_hostname) == 0) {
-                pthread_mutex_unlock(&global_state->client_mutex);
+                pthread_mutex_unlock(&state->client_mutex);
                 return client;
             }
         }
     }
 
-    pthread_mutex_unlock(&global_state->client_mutex);
+    pthread_mutex_unlock(&state->client_mutex);
     return NULL;
 }
 
@@ -196,7 +207,7 @@ static void *proxy_data_forward(void *arg) {
 }
 
 // Handle incoming HTTP request and establish tunnel
-static int handle_proxy_request(int client_fd) {
+static int handle_proxy_request(int client_fd, wssshd_state_t *state, const wssshd_config_t *config) {
     char request[8192];
     ssize_t bytes_read = recv(client_fd, request, sizeof(request) - 1, 0);
     if (bytes_read <= 0) {
@@ -214,7 +225,7 @@ static int handle_proxy_request(int client_fd) {
     printf("[WEB-PROXY] Received request for hostname: %s\n", hostname);
 
     // Find client by hostname
-    client_t *client = find_client_by_hostname(hostname);
+    client_t *client = find_client_by_hostname(hostname, state, config);
     if (!client) {
         printf("[WEB-PROXY] No client found for hostname: %s\n", hostname);
         send_http_error(client_fd, 404, "Not Found");
@@ -244,12 +255,15 @@ static int handle_proxy_request(int client_fd) {
 
 // Proxy connection handler thread
 static void *proxy_connection_handler(void *arg) {
-    int client_fd = *(int *)arg;
-    free(arg);
+    connection_handler_args_t *args = (connection_handler_args_t *)arg;
+    int client_fd = args->client_fd;
+    wssshd_state_t *state = args->state;
+    const wssshd_config_t *config = args->config;
+    free(args);
 
     printf("[WEB-PROXY] New connection received\n");
 
-    handle_proxy_request(client_fd);
+    handle_proxy_request(client_fd, state, config);
 
     close(client_fd);
     printf("[WEB-PROXY] Connection closed\n");
@@ -257,16 +271,21 @@ static void *proxy_connection_handler(void *arg) {
     return NULL;
 }
 
-// Main HTTP server thread
-static void *http_proxy_thread(void *arg __attribute__((unused))) {
+static void *http_proxy_thread(void *arg) {
+    proxy_thread_args_t *args = (proxy_thread_args_t *)arg;
+    const wssshd_config_t *config = args->config;
+    wssshd_state_t *state = args->state;
+    free(args);
     struct sockaddr_in server_addr, client_addr;
     socklen_t client_len = sizeof(client_addr);
 
+    printf("[WEB-PROXY] Creating socket\n");
     server_socket = socket(AF_INET, SOCK_STREAM, 0);
     if (server_socket < 0) {
         perror("[WEB-PROXY] Failed to create server socket");
         return NULL;
     }
+    printf("[WEB-PROXY] Socket created\n");
 
     // Set socket options
     int opt = 1;
@@ -274,22 +293,25 @@ static void *http_proxy_thread(void *arg __attribute__((unused))) {
 
     memset(&server_addr, 0, sizeof(server_addr));
     server_addr.sin_family = AF_INET;
-    server_addr.sin_addr.s_addr = inet_addr("127.0.0.1");  // Listen only on localhost
-    server_addr.sin_port = htons(global_config->web_proxy_port);
+    server_addr.sin_addr.s_addr = htonl(INADDR_LOOPBACK);  // Listen on localhost only
+    server_addr.sin_port = htons(config->web_proxy_port);
+    printf("[WEB-PROXY] Binding to 127.0.0.1:%d\n", config->web_proxy_port);
 
     if (bind(server_socket, (struct sockaddr *)&server_addr, sizeof(server_addr)) < 0) {
         perror("[WEB-PROXY] Failed to bind server socket");
         close(server_socket);
         return NULL;
     }
+    printf("[WEB-PROXY] Bound successfully\n");
 
     if (listen(server_socket, 10) < 0) {
         perror("[WEB-PROXY] Failed to listen on server socket");
         close(server_socket);
         return NULL;
     }
+    printf("[WEB-PROXY] Listening successfully\n");
 
-    printf("[WEB-PROXY] Web proxy server starting on 127.0.0.1:%d\n", global_config->web_proxy_port);
+    printf("[WEB-PROXY] Web proxy server starting on 127.0.0.1:%d\n", config->web_proxy_port);
     server_running = 1;
 
     // Ignore SIGPIPE to prevent crashes on broken connections
@@ -303,17 +325,19 @@ static void *http_proxy_thread(void *arg __attribute__((unused))) {
         }
 
         // Create a new thread to handle the connection
-        int *client_fd_ptr = malloc(sizeof(int));
-        if (!client_fd_ptr) {
+        connection_handler_args_t *handler_args = malloc(sizeof(connection_handler_args_t));
+        if (!handler_args) {
             close(client_fd);
             continue;
         }
-        *client_fd_ptr = client_fd;
+        handler_args->client_fd = client_fd;
+        handler_args->state = state;
+        handler_args->config = config;
 
         pthread_t thread;
-        if (pthread_create(&thread, NULL, proxy_connection_handler, client_fd_ptr) != 0) {
+        if (pthread_create(&thread, NULL, proxy_connection_handler, handler_args) != 0) {
             perror("[WEB-PROXY] Failed to create connection handler thread");
-            free(client_fd_ptr);
+            free(handler_args);
             close(client_fd);
             continue;
         }
@@ -332,18 +356,20 @@ int web_proxy_start_server(const wssshd_config_t *config, wssshd_state_t *state)
         return 0; // Web proxy not enabled
     }
 
-    global_config = config;
-    global_state = state;
-
-    // Set default port if not specified
-    if (config->web_proxy_port == 0) {
-        // This won't work directly since config is const, but the CLI parsing sets it
+    // Allocate thread arguments
+    proxy_thread_args_t *args = malloc(sizeof(proxy_thread_args_t));
+    if (!args) {
+        perror("[WEB-PROXY] Failed to allocate thread arguments");
+        return -1;
     }
+    args->config = config;
+    args->state = state;
 
     // Start HTTP proxy server thread
     pthread_t thread;
-    if (pthread_create(&thread, NULL, http_proxy_thread, NULL) != 0) {
+    if (pthread_create(&thread, NULL, http_proxy_thread, args) != 0) {
         perror("[WEB-PROXY] Failed to create HTTP proxy server thread");
+        free(args);
         return -1;
     }
 

wssshd2/websocket.c

@@ -643,6 +643,8 @@ int websocket_handle_message(wssshd_state_t *state, ws_connection_t *conn __attr
         if (client_id && password && strcmp(password, state->server_password) == 0) {
             client_t *client = websocket_add_client(state, client_id, (void *)conn, services);
             if (client) {
+                // Set client IP
+                strcpy(client->ip, inet_ntoa(conn->addr.sin_addr));
                 // Send registration success
                 char response[512];
                 snprintf(response, sizeof(response), REGISTERED_MSG, client_id);
@@ -1780,17 +1782,22 @@ int websocket_start_server(const wssshd_config_t *config, wssshd_state_t *state)
 
     // Load or generate certificates
     char cert_file[256], key_file[256];
-    // Try system location first, then fall back to local directory
-    snprintf(cert_file, sizeof(cert_file), "/etc/wssshd/cert.pem");
-    snprintf(key_file, sizeof(key_file), "/etc/wssshd/key.pem");
-
-    // If system location doesn't exist and we can't write there, use local directory
-    if (access(cert_file, F_OK) != 0) {
+    bool use_system_location = false;
+
+    // Check if /etc/wssshd directory exists
+    if (access("/etc/wssshd", F_OK) == 0) {
+        // Try to use system location
+        snprintf(cert_file, sizeof(cert_file), "/etc/wssshd/cert.pem");
+        snprintf(key_file, sizeof(key_file), "/etc/wssshd/key.pem");
+        use_system_location = true;
+    } else {
+        // Fall back to local directory
         char local_dir[256];
         snprintf(local_dir, sizeof(local_dir), "./certs");
         mkdir(local_dir, 0755); // Create directory if it doesn't exist
         snprintf(cert_file, sizeof(cert_file), "./certs/cert.pem");
         snprintf(key_file, sizeof(key_file), "./certs/key.pem");
+        use_system_location = false;
     }
 
     printf("Using certificates: %s, %s\n", cert_file, key_file);
@@ -1798,12 +1805,32 @@ int websocket_start_server(const wssshd_config_t *config, wssshd_state_t *state)
     if (access(cert_file, F_OK) != 0 || access(key_file, F_OK) != 0) {
         printf("Generating self-signed certificate...\n");
         if (ssl_generate_self_signed_cert(cert_file, key_file) != 0) {
-            fprintf(stderr, "Failed to generate certificate\n");
-            SSL_CTX_free(ssl_ctx);
-            close(server_sock);
-            return -1;
+            if (use_system_location) {
+                // If system location failed, fall back to local directory
+                printf("Failed to generate certificate in system location, falling back to local directory...\n");
+                char local_dir[256];
+                snprintf(local_dir, sizeof(local_dir), "./certs");
+                mkdir(local_dir, 0755); // Create directory if it doesn't exist
+                snprintf(cert_file, sizeof(cert_file), "./certs/cert.pem");
+                snprintf(key_file, sizeof(key_file), "./certs/key.pem");
+                printf("Using certificates: %s, %s\n", cert_file, key_file);
+
+                if (ssl_generate_self_signed_cert(cert_file, key_file) != 0) {
+                    fprintf(stderr, "Failed to generate certificate in fallback location\n");
+                    SSL_CTX_free(ssl_ctx);
+                    close(server_sock);
+                    return -1;
+                }
+                printf("Certificate generated successfully in fallback location\n");
+            } else {
+                fprintf(stderr, "Failed to generate certificate\n");
+                SSL_CTX_free(ssl_ctx);
+                close(server_sock);
+                return -1;
+            }
+        } else {
+            printf("Certificate generated successfully\n");
         }
-        printf("Certificate generated successfully\n");
     }
 
     printf("Loading certificates...\n");
@@ -1862,6 +1889,9 @@ int websocket_start_server(const wssshd_config_t *config, wssshd_state_t *state)
             continue;
         }
 
+        // Set the client address
+        conn->addr = client_addr;
+
         // Start client handler thread
         client_thread_args_t *thread_args = malloc(sizeof(client_thread_args_t));
 

wssshd2/websocket.h

@@ -38,6 +38,7 @@ typedef struct {
     char tunnel_control[16];
     char wssshd_private_ip[64];
     char services[256];  // Comma-separated list of services (e.g., "ssh,rdp,vnc")
+    char ip[INET_ADDRSTRLEN];
 } client_t;
 
 // Terminal session (defined in terminal.h)

wssshd2/websocket_protocol.h

@@ -22,6 +22,7 @@
 
 #include <stdint.h>
 #include <stdbool.h>
+#include <netinet/in.h>
 #include <openssl/ssl.h>
 
 // WebSocket frame opcodes
@@ -60,6 +61,7 @@ typedef struct {
     char *recv_buffer;
     size_t recv_buffer_size;
     size_t recv_buffer_used;
+    struct sockaddr_in addr;
 } ws_connection_t;
 
 // Function declarations