Remove "clean" duplicate

clean

-#!/bin/bash
-
-# WebSocket SSH Tools Clean Script
-# Clean script for removing build artifacts from WebSocket SSH tools
-#
-# 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/>.
-
-# Remove PyInstaller build artifacts
-rm -rf build/
-rm -rf dist/
-rm -f *.spec
-
-# Remove C version build artifacts
-if [ -d "wssshtools" ]; then
-    cd wssshtools
-    make clean 2>/dev/null || true
-    rm -f Makefile
-    rm -f configure.sh.stamp
-    rm -f man/*.1.gz 2>/dev/null || true
-    cd ..
-fi
-
-# Remove Debian packaging artifacts
-rm -f dist/wsssh-tools*.deb
-rm -f dist/wsssh-tools*.dsc
-rm -f dist/wsssh-tools*.tar.gz
-rm -f dist/wsssh-tools*.changes
-rm -f dist/wsssh-tools*.buildinfo
-rm -f wsssh-tools*.deb
-rm -f wsssh-tools*.dsc
-rm -f wsssh-tools*.tar.gz
-rm -f wsssh-tools*.changes
-rm -f wsssh-tools*.buildinfo
-
-# Remove Debian build directory and artifacts
-if [ -d "wssshtools" ]; then
-    rm -rf wssshtools/debian/wsssh-tools/
-    rm -f wssshtools/debian/files
-    rm -f wssshtools/debian/*.debhelper*
-    rm -f wssshtools/debian/*.substvars
-    rm -f wssshtools/debian/debhelper-build-stamp
-fi
-
-# Optionally remove SSL certificates (uncomment if needed)
-# rm -f cert.pem key.pem
-
-echo "Clean complete. Build artifacts removed."
-echo "Note: SSL certificates (cert.pem, key.pem) preserved. Uncomment lines in clean.sh to remove them."
\ No newline at end of file

wssshd.py

@@ -40,7 +40,7 @@ from flask_login import LoginManager, UserMixin, login_user, login_required, log
 from flask_sqlalchemy import SQLAlchemy
 from werkzeug.security import generate_password_hash, check_password_hash
 
-# Client registry: id -> websocket
+# Client registry: id -> {'websocket': ws, 'last_seen': timestamp, 'status': 'active'|'disconnected'}
 clients = {}
 # Active tunnels: request_id -> {'client_ws': ws, 'wsssh_ws': ws, 'client_id': id}
 active_tunnels = {}
@@ -49,6 +49,7 @@ active_terminals = {}
 debug = False
 server_password = None
 args = None
+import time
 
 # Flask app for web interface
 app = Flask(__name__)
@@ -73,6 +74,20 @@ class User(UserMixin, db.Model):
 def load_user(user_id):
     return db.session.get(User, int(user_id))
 
+def cleanup_expired_clients():
+    """Remove clients that have been disconnected for more than 30 seconds"""
+    current_time = time.time()
+    expired_clients = []
+
+    for client_id, client_info in clients.items():
+        if client_info['status'] == 'disconnected':
+            if current_time - client_info['last_seen'] > 30:
+                expired_clients.append(client_id)
+                if debug: print(f"[DEBUG] [WebSocket] Client {client_id} expired and removed")
+
+    for client_id in expired_clients:
+        del clients[client_id]
+
 def openpty_with_fallback():
     """Open a PTY with fallback to different device paths for systems where /dev/pty doesn't exist"""
     # First try the standard pty.openpty()
@@ -152,10 +167,17 @@ with app.app_context():
 @login_required
 def index():
     global args
+    # Get client information with status
+    client_info = {}
+    for client_id, client_data in clients.items():
+        client_info[client_id] = {
+            'status': client_data['status'],
+            'last_seen': client_data['last_seen']
+        }
     return render_template('index.html',
-                         clients=list(clients.keys()),
-                         websocket_port=args.port,
-                         domain=args.domain)
+                          clients=client_info,
+                          websocket_port=args.port,
+                          domain=args.domain)
 
 @app.route('/login', methods=['GET', 'POST'])
 def login():
@@ -241,8 +263,14 @@ def terminal(client_id):
 @app.route('/api/clients')
 @login_required
 def get_clients():
+    client_info = {}
+    for client_id, client_data in clients.items():
+        client_info[client_id] = {
+            'status': client_data['status'],
+            'last_seen': client_data['last_seen']
+        }
     return jsonify({
-        'clients': list(clients.keys()),
+        'clients': client_info,
         'count': len(clients)
     })
 
@@ -377,30 +405,44 @@ def disconnect_terminal(client_id):
 async def handle_websocket(websocket, path=None):
     try:
         async for message in websocket:
-            if debug: print(f"[DEBUG] WebSocket message received: {message[:100]}...")
+            if debug: print(f"[DEBUG] [WebSocket] Message received: {message[:100]}...")
             data = json.loads(message)
             if data.get('type') == 'register':
                 client_id = data['id']
                 client_password = data.get('password', '')
                 if client_password == server_password:
-                    clients[client_id] = websocket
-                    print(f"Client {client_id} registered")
+                    # Check if client was previously disconnected
+                    was_disconnected = False
+                    if client_id in clients and clients[client_id]['status'] == 'disconnected':
+                        was_disconnected = True
+                        print(f"[DEBUG] [WebSocket] Client {client_id} reconnecting (was disconnected)")
+
+                    clients[client_id] = {
+                        'websocket': websocket,
+                        'last_seen': time.time(),
+                        'status': 'active'
+                    }
+
+                    if was_disconnected:
+                        print(f"Client {client_id} reconnected")
+                    else:
+                        print(f"Client {client_id} registered")
                     await websocket.send(json.dumps({"type": "registered", "id": client_id}))
                 else:
-                    print(f"Client {client_id} registration failed: invalid password")
+                    print(f"[DEBUG] [WebSocket] Client {client_id} registration failed: invalid password")
                     await websocket.send(json.dumps({"type": "registration_error", "error": "Invalid password"}))
             elif data.get('type') == 'tunnel_request':
                 client_id = data['client_id']
                 request_id = data['request_id']
-                if client_id in clients:
+                if client_id in clients and clients[client_id]['status'] == 'active':
                     # Store tunnel mapping
                     active_tunnels[request_id] = {
-                        'client_ws': clients[client_id],
+                        'client_ws': clients[client_id]['websocket'],
                         'wsssh_ws': websocket,
                         'client_id': client_id
                     }
                     # Forward tunnel request to client
-                    await clients[client_id].send(json.dumps({
+                    await clients[client_id]['websocket'].send(json.dumps({
                         "type": "tunnel_request",
                         "request_id": request_id
                     }))
@@ -412,7 +454,7 @@ async def handle_websocket(websocket, path=None):
                     await websocket.send(json.dumps({
                         "type": "tunnel_error",
                         "request_id": request_id,
-                        "error": "Client not registered"
+                        "error": "Client not registered or disconnected"
                     }))
             elif data.get('type') == 'tunnel_data':
                 # Forward tunnel data using active tunnel mapping
@@ -420,11 +462,14 @@ async def handle_websocket(websocket, path=None):
                 if request_id in active_tunnels:
                     tunnel = active_tunnels[request_id]
                     # Forward to client
-                    await tunnel['client_ws'].send(json.dumps({
-                        "type": "tunnel_data",
-                        "request_id": request_id,
-                        "data": data['data']
-                    }))
+                    if tunnel['client_id'] in clients and clients[tunnel['client_id']]['status'] == 'active':
+                        await tunnel['client_ws'].send(json.dumps({
+                            "type": "tunnel_data",
+                            "request_id": request_id,
+                            "data": data['data']
+                        }))
+                    else:
+                        if debug: print(f"[DEBUG] [WebSocket] Cannot forward tunnel_data: client {tunnel['client_id']} not active")
             elif data.get('type') == 'tunnel_response':
                 # Forward tunnel response from client to wsssh
                 request_id = data['request_id']
@@ -439,21 +484,24 @@ async def handle_websocket(websocket, path=None):
                 request_id = data['request_id']
                 if request_id in active_tunnels:
                     tunnel = active_tunnels[request_id]
-                    # Forward close to client
-                    await tunnel['client_ws'].send(json.dumps({
-                        "type": "tunnel_close",
-                        "request_id": request_id
-                    }))
+                    # Forward close to client if still active
+                    if tunnel['client_id'] in clients and clients[tunnel['client_id']]['status'] == 'active':
+                        await tunnel['client_ws'].send(json.dumps({
+                            "type": "tunnel_close",
+                            "request_id": request_id
+                        }))
                     # Clean up tunnel
                     del active_tunnels[request_id]
+                    if debug: print(f"[DEBUG] [WebSocket] Tunnel {request_id} closed")
     except websockets.exceptions.ConnectionClosed:
-        # Remove from registry and clean up tunnels
+        # Mark client as disconnected instead of removing immediately
         disconnected_client = None
-        for cid, ws in clients.items():
-            if ws == websocket:
+        for cid, client_info in clients.items():
+            if client_info['websocket'] == websocket:
                 disconnected_client = cid
-                del clients[cid]
-                print(f"Client {cid} disconnected")
+                clients[cid]['status'] = 'disconnected'
+                clients[cid]['last_seen'] = time.time()
+                print(f"[DEBUG] [WebSocket] Client {cid} disconnected (marked for timeout)")
                 break
 
         # Clean up active tunnels for this client
@@ -464,6 +512,13 @@ async def handle_websocket(websocket, path=None):
                     tunnels_to_remove.append(request_id)
             for request_id in tunnels_to_remove:
                 del active_tunnels[request_id]
+                if debug: print(f"[DEBUG] [WebSocket] Tunnel {request_id} cleaned up due to client disconnect")
+
+async def cleanup_task():
+    """Periodic task to clean up expired clients"""
+    while True:
+        await asyncio.sleep(10)  # Run every 10 seconds
+        cleanup_expired_clients()
 
 async def main():
     config = configparser.ConfigParser()
@@ -539,6 +594,9 @@ async def main():
 
     print(f"WebSocket SSH Daemon running on {args.host}:{args.port}")
 
+    # Start cleanup task
+    cleanup_coro = asyncio.create_task(cleanup_task())
+
     # Start web interface if specified
     if args.web_host and args.web_port:
         # Handle HTTPS setup
@@ -565,7 +623,11 @@ async def main():
         flask_thread.start()
         print(f"Web interface available at {protocol}://{args.web_host}:{args.web_port}")
 
+    # Wait for server to close
     await ws_server.wait_closed()
 
+    # Cancel cleanup task
+    cleanup_coro.cancel()
+
 if __name__ == '__main__':
     asyncio.run(main())
\ No newline at end of file

wssshtools/wssshc.c

@@ -30,6 +30,7 @@
 #include <openssl/err.h>
 #include <getopt.h>
 #include <pthread.h>
+#include <errno.h>
 
 #define BUFFER_SIZE 4096
 #define DEFAULT_PORT 9898
@@ -98,6 +99,7 @@ void load_config(wssshc_config_t *config) {
 tunnel_t *active_tunnel = NULL;
 
 void *tunnel_thread(void *arg);
+void cleanup_tunnel();
 
 void handle_tunnel_request(SSL *ssl, const char *request_id, int debug) {
     pthread_mutex_lock(&tunnel_mutex);
@@ -163,7 +165,7 @@ void handle_tunnel_data(SSL *ssl, const char *request_id, const char *data_hex,
         int sock;
 
         if (debug) {
-            printf("[DEBUG] Opening local SSH connection for tunnel %s\n", request_id);
+            printf("[DEBUG] [TCP Tunnel] Opening local SSH connection for tunnel %s\n", request_id);
         }
 
         // Create socket for local SSH connection
@@ -195,12 +197,74 @@ void handle_tunnel_data(SSL *ssl, const char *request_id, const char *data_hex,
             pthread_detach(thread);
 
             if (debug) {
-                printf("[DEBUG] Local SSH connection established and tunnel thread started\n");
+                printf("[DEBUG] [TCP Tunnel] Local SSH connection established and tunnel thread started\n");
             }
         } else {
             close(sock);
         }
         pthread_mutex_unlock(&tunnel_mutex);
+    } else {
+        // SSH connection already exists, verify it's still valid
+        char test_buf[1];
+        int result = recv(sock_fd, test_buf, 1, MSG_PEEK | MSG_DONTWAIT);
+        if (result < 0 && errno != EAGAIN && errno != EWOULDBLOCK) {
+            // SSH connection is broken, close it and create a new one
+            if (debug) {
+                printf("[DEBUG] [TCP Tunnel] Existing SSH connection is broken, creating new one\n");
+            }
+            close(sock_fd);
+
+            struct sockaddr_in addr;
+            int sock;
+
+            // Create new socket for local SSH connection
+            if ((sock = socket(AF_INET, SOCK_STREAM, 0)) < 0) {
+                perror("Local socket creation failed");
+                pthread_mutex_lock(&tunnel_mutex);
+                if (active_tunnel && strcmp(active_tunnel->request_id, request_id) == 0) {
+                    active_tunnel->sock = -1;
+                }
+                pthread_mutex_unlock(&tunnel_mutex);
+                return;
+            }
+
+            memset(&addr, 0, sizeof(addr));
+            addr.sin_family = AF_INET;
+            addr.sin_port = htons(22);
+            addr.sin_addr.s_addr = htonl(INADDR_LOOPBACK);
+
+            // Connect to local SSH
+            if (connect(sock, (struct sockaddr *)&addr, sizeof(addr)) < 0) {
+                perror("Local SSH connection failed");
+                close(sock);
+                pthread_mutex_lock(&tunnel_mutex);
+                if (active_tunnel && strcmp(active_tunnel->request_id, request_id) == 0) {
+                    active_tunnel->sock = -1;
+                }
+                pthread_mutex_unlock(&tunnel_mutex);
+                return;
+            }
+
+            pthread_mutex_lock(&tunnel_mutex);
+            if (active_tunnel && strcmp(active_tunnel->request_id, request_id) == 0) {
+                active_tunnel->sock = sock;
+                sock_fd = sock;
+
+                // Start tunnel thread
+                pthread_t thread;
+                pthread_create(&thread, NULL, tunnel_thread, ssl);
+                pthread_detach(thread);
+
+                if (debug) {
+                    printf("[DEBUG] [TCP Tunnel] New SSH connection established and tunnel thread started\n");
+                }
+            } else {
+                close(sock);
+            }
+            pthread_mutex_unlock(&tunnel_mutex);
+        } else if (debug) {
+            printf("[DEBUG] [TCP Tunnel] Reusing existing SSH connection for tunnel %s\n", request_id);
+        }
     }
 
     if (debug) {
@@ -210,7 +274,10 @@ void handle_tunnel_data(SSL *ssl, const char *request_id, const char *data_hex,
     // Decode hex data
     size_t data_len = strlen(data_hex) / 2;
     char *data = malloc(data_len);
-    if (!data) return;
+    if (!data) {
+        fprintf(stderr, "Memory allocation failed for tunnel data\n");
+        return;
+    }
 
     for (size_t i = 0; i < data_len; i++) {
         sscanf(data_hex + i * 2, "%2hhx", &data[i]);
@@ -229,7 +296,12 @@ void handle_tunnel_data(SSL *ssl, const char *request_id, const char *data_hex,
     }
 
     // Send to local SSH
-    send(sock_fd, data, data_len, 0);
+    ssize_t sent = send(sock_fd, data, data_len, 0);
+    if (sent < 0) {
+        perror("Failed to send data to local SSH");
+    } else if ((size_t)sent != data_len) {
+        fprintf(stderr, "Partial send to local SSH: sent %zd of %zu bytes\n", sent, data_len);
+    }
     free(data);
 }
 
@@ -240,7 +312,7 @@ void handle_tunnel_close(SSL *ssl __attribute__((unused)), const char *request_i
         free(active_tunnel);
         active_tunnel = NULL;
         if (debug) {
-            printf("[DEBUG] Tunnel %s closed\n", request_id);
+            printf("[DEBUG] [TCP Tunnel] Tunnel %s closed\n", request_id);
         }
     }
     pthread_mutex_unlock(&tunnel_mutex);
@@ -264,11 +336,24 @@ void *tunnel_thread(void *arg) {
 
         bytes_read = recv(sock, buffer, sizeof(buffer), 0);
         if (bytes_read <= 0) {
+            if (bytes_read < 0) {
+                perror("Error reading from local SSH");
+                // Check if this is due to tunnel cleanup (ECONNRESET, EPIPE, etc.)
+                if (errno == ECONNRESET || errno == EPIPE || errno == EBADF) {
+                    if (global_debug) {
+                        printf("[DEBUG] [TCP Tunnel] SSH connection closed due to tunnel reset\n");
+                    }
+                }
+            } else {
+                if (global_debug) {
+                    printf("[DEBUG] [TCP Tunnel] Local SSH connection closed normally\n");
+                }
+            }
             break;
         }
 
         if (global_debug) {
-            printf("[DEBUG] Tunnel thread: read %d bytes from local SSH: ", bytes_read);
+            printf("[DEBUG] [TCP Tunnel] Read %d bytes from local SSH: ", bytes_read);
             for (int i = 0; i < bytes_read && i < 50; i++) {
                 if (buffer[i] >= 32 && buffer[i] < 127) {
                     printf("%c", buffer[i]);
@@ -282,7 +367,7 @@ void *tunnel_thread(void *arg) {
         // Send as tunnel_response
         char *hex_data = malloc(bytes_read * 2 + 1);
         if (!hex_data) {
-            perror("Memory allocation failed for hex_data");
+            fprintf(stderr, "Memory allocation failed for hex_data in tunnel_thread\n");
             break;
         }
         for (int i = 0; i < bytes_read; i++) {
@@ -291,7 +376,7 @@ void *tunnel_thread(void *arg) {
         hex_data[bytes_read * 2] = '\0';
 
         if (global_debug) {
-            printf("[DEBUG] Sending tunnel_response, hex len: %zu, hex: %.100s...\n", strlen(hex_data), hex_data);
+            printf("[DEBUG] [WebSocket] Sending tunnel_response, hex len: %zu, hex: %.100s...\n", strlen(hex_data), hex_data);
         }
 
         char response[9000];
@@ -321,7 +406,7 @@ void *tunnel_thread(void *arg) {
 
         char *frame = malloc(header_len + msg_len);
         if (!frame) {
-            perror("Memory allocation failed for frame");
+            fprintf(stderr, "Memory allocation failed for WebSocket frame in tunnel_thread\n");
             free(hex_data);
             break;
         }
@@ -360,7 +445,10 @@ void *tunnel_thread(void *arg) {
 
         int frame_len = header_len + msg_len;
         if (!send_all(ssl, frame, frame_len)) {
-            fprintf(stderr, "Send tunnel_response failed\n");
+            fprintf(stderr, "Send tunnel_response failed - tunnel may hang\n");
+            free(frame);
+            free(hex_data);
+            break;
         }
 
         free(frame);
@@ -577,7 +665,7 @@ int send_json_message(SSL *ssl, const char *type, const char *id, const char *pa
     return 1;
 }
 
-void connect_to_server(const wssshc_config_t *config) {
+int connect_to_server(const wssshc_config_t *config) {
     struct sockaddr_in server_addr;
     struct hostent *he;
     int sock;
@@ -586,16 +674,19 @@ void connect_to_server(const wssshc_config_t *config) {
     char buffer[BUFFER_SIZE];
     int bytes_read;
 
+    // Clean up any leftover tunnel state from previous connection
+    cleanup_tunnel();
+
     // Resolve hostname
     if ((he = gethostbyname(config->server_ip)) == NULL) {
         herror("gethostbyname");
-        return;
+        return 1;
     }
 
     // Create socket
     if ((sock = socket(AF_INET, SOCK_STREAM, 0)) < 0) {
         perror("Socket creation failed");
-        return;
+        return 1;
     }
 
     memset(&server_addr, 0, sizeof(server_addr));
@@ -607,14 +698,14 @@ void connect_to_server(const wssshc_config_t *config) {
     if (connect(sock, (struct sockaddr *)&server_addr, sizeof(server_addr)) < 0) {
         perror("Connection failed");
         close(sock);
-        return;
+        return 1;
     }
 
     // Initialize SSL
     ssl_ctx = init_ssl_ctx();
     if (!ssl_ctx) {
         close(sock);
-        return;
+        return 1;
     }
 
     ssl = SSL_new(ssl_ctx);
@@ -625,7 +716,7 @@ void connect_to_server(const wssshc_config_t *config) {
         SSL_free(ssl);
         SSL_CTX_free(ssl_ctx);
         close(sock);
-        return;
+        return 1;
     }
 
     // Perform WebSocket handshake
@@ -633,7 +724,7 @@ void connect_to_server(const wssshc_config_t *config) {
         SSL_free(ssl);
         SSL_CTX_free(ssl_ctx);
         close(sock);
-        return;
+        return 1;
     }
 
     // Send registration message
@@ -641,7 +732,7 @@ void connect_to_server(const wssshc_config_t *config) {
         SSL_free(ssl);
         SSL_CTX_free(ssl_ctx);
         close(sock);
-        return;
+        return 1;
     }
 
     // Read WebSocket frame with registration response
@@ -651,7 +742,7 @@ void connect_to_server(const wssshc_config_t *config) {
         SSL_free(ssl);
         SSL_CTX_free(ssl_ctx);
         close(sock);
-        return;
+        return 1;
     }
 
     // Parse WebSocket frame with extended length support
@@ -674,7 +765,7 @@ void connect_to_server(const wssshc_config_t *config) {
                 SSL_free(ssl);
                 SSL_CTX_free(ssl_ctx);
                 close(sock);
-                return;
+                return 1;
             }
             payload_len = ((unsigned char)buffer[2] << 8) | (unsigned char)buffer[3];
             header_len = 4;
@@ -684,7 +775,7 @@ void connect_to_server(const wssshc_config_t *config) {
                 SSL_free(ssl);
                 SSL_CTX_free(ssl_ctx);
                 close(sock);
-                return;
+                return 1;
             }
             uint64_t len = 0;
             for (int i = 0; i < 8; i++) {
@@ -715,7 +806,7 @@ void connect_to_server(const wssshc_config_t *config) {
             SSL_free(ssl);
             SSL_CTX_free(ssl_ctx);
             close(sock);
-            return;
+            return 1;
         }
     } else {
         buffer[bytes_read] = '\0';
@@ -727,7 +818,7 @@ void connect_to_server(const wssshc_config_t *config) {
         SSL_free(ssl);
         SSL_CTX_free(ssl_ctx);
         close(sock);
-        return;
+        return 2;
     }
 
     printf("Connected and registered as %s\n", config->client_id);
@@ -738,14 +829,21 @@ void connect_to_server(const wssshc_config_t *config) {
     while (1) {
         bytes_read = SSL_read(ssl, buffer, sizeof(buffer));
         if (bytes_read <= 0) {
-            if (config->debug) {
-                printf("[DEBUG] Connection lost\n");
+            if (bytes_read < 0) {
+                ERR_print_errors_fp(stderr);
+                fprintf(stderr, "SSL read error\n");
+            } else {
+                if (config->debug) {
+                    printf("[DEBUG] Connection closed by server\n");
+                }
             }
+            // Clean up tunnel resources before breaking
+            cleanup_tunnel();
             break;
         }
 
         if (config->debug) {
-            printf("[DEBUG] Read %d bytes, frame: 0x%02x 0x%02x 0x%02x 0x%02x\n", bytes_read, buffer[0], buffer[1], buffer[2], buffer[3]);
+            printf("[DEBUG] [WebSocket] Read %d bytes, frame: 0x%02x 0x%02x 0x%02x 0x%02x\n", bytes_read, buffer[0], buffer[1], buffer[2], buffer[3]);
         }
 
         // Parse WebSocket frame with extended length support
@@ -753,16 +851,26 @@ void connect_to_server(const wssshc_config_t *config) {
             int masked = buffer[1] & 0x80;
             int len_indicator = buffer[1] & 0x7F;
             int header_len = 2;
-            size_t payload_len;
+            size_t payload_len = 0;
 
             if (len_indicator <= 125) {
                 payload_len = len_indicator;
             } else if (len_indicator == 126) {
-                if (bytes_read < 4) continue;
+                if (bytes_read < 4) {
+                    if (config->debug) {
+                        printf("[DEBUG] Incomplete extended length frame\n");
+                    }
+                    continue;
+                }
                 payload_len = ((unsigned char)buffer[2] << 8) | (unsigned char)buffer[3];
                 header_len = 4;
             } else if (len_indicator == 127) {
-                if (bytes_read < 10) continue;
+                if (bytes_read < 10) {
+                    if (config->debug) {
+                        printf("[DEBUG] Incomplete extended length frame\n");
+                    }
+                    continue;
+                }
                 uint64_t len = 0;
                 for (int i = 0; i < 8; i++) {
                     len = (len << 8) | (unsigned char)buffer[2 + i];
@@ -788,7 +896,7 @@ void connect_to_server(const wssshc_config_t *config) {
 
                 // Handle message
                 if (config->debug) {
-                    printf("[DEBUG] Received: %s\n", buffer);
+                    printf("[DEBUG] [WebSocket] Received message: %s\n", buffer);
                 }
 
                 // Parse JSON message
@@ -804,6 +912,9 @@ void connect_to_server(const wssshc_config_t *config) {
                                 char *close_quote = strchr(id_start, '"');
                                 if (close_quote) {
                                     *close_quote = '\0';
+                                    if (config->debug) {
+                                        printf("[DEBUG] [WebSocket] Received tunnel_request for ID: %s\n", id_start);
+                                    }
                                     handle_tunnel_request(ssl, id_start, config->debug);
                                 }
                             }
@@ -811,7 +922,7 @@ void connect_to_server(const wssshc_config_t *config) {
                     }
                 } else if (strstr(buffer, "tunnel_data")) {
                     if (config->debug) {
-                        printf("[DEBUG] Received tunnel_data\n");
+                        printf("[DEBUG] [WebSocket] Received tunnel_data message\n");
                     }
                     // Extract request_id and data
                     char *id_start = strstr(buffer, "\"request_id\"");
@@ -852,6 +963,9 @@ void connect_to_server(const wssshc_config_t *config) {
                                 char *close_quote = strchr(id_start, '"');
                                 if (close_quote) {
                                     *close_quote = '\0';
+                                    if (config->debug) {
+                                        printf("[DEBUG] [WebSocket] Received tunnel_close for ID: %s\n", id_start);
+                                    }
                                     handle_tunnel_close(ssl, id_start, config->debug);
                                 }
                             }
@@ -859,18 +973,20 @@ void connect_to_server(const wssshc_config_t *config) {
                     }
                 }
             } else {
-                if (config->debug) {
-                    printf("[DEBUG] Incomplete frame: need %zu bytes, got %d\n", (size_t)(payload_len + header_len), bytes_read);
-                }
+                fprintf(stderr, "WebSocket frame payload too large or incomplete: payload_len=%zu, buffer_size=%zu, bytes_read=%d\n",
+                        payload_len, sizeof(buffer) - header_len, bytes_read);
+                break;
             }
         } else if (bytes_read >= 2 && (buffer[0] & 0x8F) == 0x88) { // Close frame
             if (config->debug) {
-                printf("[DEBUG] Received close frame\n");
+                printf("[DEBUG] [WebSocket] Received close frame - cleaning up and reconnecting...\n");
             }
-            return;
+            // Clean up tunnel resources before reconnecting
+            cleanup_tunnel();
+            return 0;
         } else if (bytes_read >= 2 && (buffer[0] & 0x8F) == 0x89) { // Ping frame
             if (config->debug) {
-                printf("[DEBUG] Received ping frame\n");
+                printf("[DEBUG] [WebSocket] Received ping frame\n");
             }
             // Parse ping frame and send pong with echoed payload
             int masked = buffer[1] & 0x80;
@@ -881,11 +997,21 @@ void connect_to_server(const wssshc_config_t *config) {
             if (len_indicator <= 125) {
                 payload_len = len_indicator;
             } else if (len_indicator == 126) {
-                if (bytes_read < 4) continue;
+                if (bytes_read < 4) {
+                    if (config->debug) {
+                        printf("[DEBUG] Incomplete ping extended length frame\n");
+                    }
+                    continue;
+                }
                 payload_len = ((unsigned char)buffer[2] << 8) | (unsigned char)buffer[3];
                 header_len = 4;
             } else if (len_indicator == 127) {
-                if (bytes_read < 10) continue;
+                if (bytes_read < 10) {
+                    if (config->debug) {
+                        printf("[DEBUG] Incomplete ping extended length frame\n");
+                    }
+                    continue;
+                }
                 uint64_t len = 0;
                 for (int i = 0; i < 8; i++) {
                     len = (len << 8) | (unsigned char)buffer[2 + i];
@@ -913,11 +1039,15 @@ void connect_to_server(const wssshc_config_t *config) {
 
                 // Send pong frame with same payload, properly masked
                 char pong_frame[14 + 1024]; // Max header + max pong payload
+                if (payload_len > 1024) {
+                    fprintf(stderr, "Pong payload too large: %zu bytes\n", payload_len);
+                    continue;
+                }
                 pong_frame[0] = 0x8A; // Pong opcode
                 int pong_header_len = 2;
                 char pong_mask_key[4];
 
-                size_t pong_payload_len = payload_len > 1024 ? 1024 : payload_len;
+                size_t pong_payload_len = payload_len;
 
                 if (pong_payload_len <= 125) {
                     pong_frame[1] = 0x80 | (int)pong_payload_len; // MASK + length
@@ -953,10 +1083,10 @@ void connect_to_server(const wssshc_config_t *config) {
 
                 int pong_frame_len = pong_header_len + (int)pong_payload_len;
                 if (config->debug) {
-                    printf("[DEBUG] Sending pong frame, len: %d\n", pong_frame_len);
+                    printf("[DEBUG] [WebSocket] Sending pong frame, len: %d\n", pong_frame_len);
                 }
                 if (!send_all(ssl, pong_frame, pong_frame_len)) {
-                    fprintf(stderr, "Send pong failed\n");
+                    fprintf(stderr, "[ERROR] [WebSocket] Send pong failed\n");
                 }
             }
         }
@@ -966,6 +1096,35 @@ void connect_to_server(const wssshc_config_t *config) {
     SSL_free(ssl);
     SSL_CTX_free(ssl_ctx);
     close(sock);
+    return 1;
+}
+
+void cleanup_tunnel() {
+    pthread_mutex_lock(&tunnel_mutex);
+    if (active_tunnel) {
+        if (active_tunnel->sock >= 0) {
+            // Check if SSH connection is still valid before closing
+            char test_buf[1];
+            int result = recv(active_tunnel->sock, test_buf, 1, MSG_PEEK | MSG_DONTWAIT);
+            if (result == 0 || (result < 0 && (errno == ECONNRESET || errno == EPIPE))) {
+                // SSH connection is closed or broken, safe to close
+                close(active_tunnel->sock);
+                if (global_debug) {
+                    printf("[DEBUG] [TCP Tunnel] Closed broken SSH connection\n");
+                }
+            } else {
+                // SSH connection appears valid, don't close it
+                if (global_debug) {
+                    printf("[DEBUG] [TCP Tunnel] Keeping SSH connection alive for potential reuse\n");
+                }
+                // Reset socket to -1 so it will be reconnected if needed
+                active_tunnel->sock = -1;
+            }
+        }
+        free(active_tunnel);
+        active_tunnel = NULL;
+    }
+    pthread_mutex_unlock(&tunnel_mutex);
 }
 
 int main(int argc, char *argv[]) {
@@ -1002,9 +1161,19 @@ int main(int argc, char *argv[]) {
     printf("WebSocket SSH Client starting...\n");
 
     while (1) {
-        connect_to_server(&config);
-        printf("Connection lost, retrying in %d seconds...\n", config.interval);
-        sleep(config.interval);
+        int result = connect_to_server(&config);
+        if (result == 1) {
+            // Error condition - use normal retry interval
+            printf("Connection lost, retrying in %d seconds...\n", config.interval);
+            sleep(config.interval);
+        } else if (result == 0) {
+            // Close frame received - add small delay to prevent rapid reconnection
+            if (config.debug) {
+                printf("[DEBUG] Server initiated disconnect, reconnecting in 2 seconds...\n");
+            }
+            sleep(2);
+        }
+        // result == 2 (registration failed) also uses normal interval
     }
 
     // Cleanup