Apply same performance optimizations to wsssh C version as wsscp

- Replace blocking recv() with select() for efficient I/O multiplexing
- Add 50ms timeouts to prevent indefinite blocking operations
- Increase buffer size from 4KB to 1MB for optimal throughput
- Add 3-second exit delay after SSH completion for clean shutdown
- Fix setup_tunnel return codes to prevent segmentation faults
- Correct memory management in argument cleanup
parent 7eb01b2a
......@@ -32,8 +32,9 @@
#include <sys/wait.h>
#include <fcntl.h>
#include <pthread.h>
#include <sys/select.h>
#define BUFFER_SIZE 4096
#define BUFFER_SIZE 1048576
#define DEFAULT_PORT 22
typedef struct {
......@@ -450,6 +451,8 @@ void *forward_tcp_to_ws(void *arg) {
int debug = args->debug;
char buffer[BUFFER_SIZE];
int bytes_read;
fd_set readfds;
struct timeval tv;
while (1) {
pthread_mutex_lock(&tunnel_mutex);
......@@ -462,6 +465,25 @@ void *forward_tcp_to_ws(void *arg) {
strcpy(request_id, active_tunnel->request_id);
pthread_mutex_unlock(&tunnel_mutex);
// Use select to wait for data on local socket
FD_ZERO(&readfds);
FD_SET(sock, &readfds);
tv.tv_sec = 0; // 0 seconds
tv.tv_usec = 50000; // 50ms timeout
int retval = select(sock + 1, &readfds, NULL, NULL, &tv);
if (retval == -1) {
if (debug) {
perror("[DEBUG] select failed");
fflush(stdout);
}
break;
} else if (retval == 0) {
// Timeout, continue loop
continue;
}
if (FD_ISSET(sock, &readfds)) {
bytes_read = recv(sock, buffer, sizeof(buffer), 0);
if (bytes_read <= 0) {
if (debug) {
......@@ -497,6 +519,7 @@ void *forward_tcp_to_ws(void *arg) {
break;
}
}
}
if (debug) {
printf("[DEBUG] TCP to WebSocket forwarding thread exiting\n");
......@@ -560,13 +583,13 @@ int setup_tunnel(const char *wssshd_host, int wssshd_port, const char *client_id
// Resolve hostname
if ((he = gethostbyname(wssshd_host)) == NULL) {
herror("gethostbyname");
return 0;
return -1;
}
// Create socket
if ((sock = socket(AF_INET, SOCK_STREAM, 0)) < 0) {
perror("Socket creation failed");
return 0;
return -1;
}
memset(&server_addr, 0, sizeof(server_addr));
......@@ -578,7 +601,7 @@ int setup_tunnel(const char *wssshd_host, int wssshd_port, const char *client_id
if (connect(sock, (struct sockaddr *)&server_addr, sizeof(server_addr)) < 0) {
perror("Connection failed");
close(sock);
return 0;
return -1;
}
// Initialize SSL
......@@ -590,7 +613,7 @@ int setup_tunnel(const char *wssshd_host, int wssshd_port, const char *client_id
if (!ssl_ctx) {
ERR_print_errors_fp(stderr);
close(sock);
return 0;
return -1;
}
// Allow self-signed certificates
......@@ -609,7 +632,7 @@ int setup_tunnel(const char *wssshd_host, int wssshd_port, const char *client_id
SSL_free(ssl);
SSL_CTX_free(ssl_ctx);
close(sock);
return 0;
return -1;
}
if (debug) {
printf("[DEBUG] SSL connection established\n");
......@@ -626,7 +649,7 @@ int setup_tunnel(const char *wssshd_host, int wssshd_port, const char *client_id
SSL_free(ssl);
SSL_CTX_free(ssl_ctx);
close(sock);
return 0;
return -1;
}
if (debug) {
printf("[DEBUG] WebSocket handshake successful\n");
......@@ -638,7 +661,7 @@ int setup_tunnel(const char *wssshd_host, int wssshd_port, const char *client_id
SSL_free(ssl);
SSL_CTX_free(ssl_ctx);
close(sock);
return 0;
return -1;
}
if (debug) {
......@@ -654,7 +677,7 @@ int setup_tunnel(const char *wssshd_host, int wssshd_port, const char *client_id
SSL_free(ssl);
SSL_CTX_free(ssl_ctx);
close(sock);
return 0;
return -1;
}
if (debug) {
......@@ -688,7 +711,7 @@ int setup_tunnel(const char *wssshd_host, int wssshd_port, const char *client_id
SSL_free(ssl);
SSL_CTX_free(ssl_ctx);
close(sock);
return 0;
return -1;
}
// Parse WebSocket frame
......@@ -707,7 +730,7 @@ int setup_tunnel(const char *wssshd_host, int wssshd_port, const char *client_id
SSL_free(ssl);
SSL_CTX_free(ssl_ctx);
close(sock);
return 0;
return -1;
}
// Null terminate payload
......@@ -733,7 +756,7 @@ int setup_tunnel(const char *wssshd_host, int wssshd_port, const char *client_id
SSL_free(ssl);
SSL_CTX_free(ssl_ctx);
close(sock);
return 0;
return -1;
}
if (debug) {
......@@ -747,7 +770,7 @@ int setup_tunnel(const char *wssshd_host, int wssshd_port, const char *client_id
SSL_free(ssl);
SSL_CTX_free(ssl_ctx);
close(sock);
return 0;
return -1;
}
strcpy(active_tunnel->request_id, request_id);
......@@ -764,7 +787,7 @@ int setup_tunnel(const char *wssshd_host, int wssshd_port, const char *client_id
SSL_free(ssl);
SSL_CTX_free(ssl_ctx);
close(sock);
return 0;
return -1;
}
struct sockaddr_in local_addr;
......@@ -781,7 +804,7 @@ int setup_tunnel(const char *wssshd_host, int wssshd_port, const char *client_id
SSL_free(ssl);
SSL_CTX_free(ssl_ctx);
close(sock);
return 0;
return -1;
}
if (listen(listen_sock, 1) < 0) {
......@@ -792,7 +815,7 @@ int setup_tunnel(const char *wssshd_host, int wssshd_port, const char *client_id
SSL_free(ssl);
SSL_CTX_free(ssl_ctx);
close(sock);
return 0;
return -1;
}
if (debug) {
......@@ -1034,7 +1057,30 @@ int main(int argc, char *argv[]) {
// Main tunnel loop - handle WebSocket messages
char buffer[BUFFER_SIZE];
int bytes_read;
fd_set readfds;
struct timeval tv;
int ssl_fd = SSL_get_fd(active_tunnel->ssl);
while (active_tunnel && active_tunnel->active) {
// Use select to wait for data on SSL socket with timeout
FD_ZERO(&readfds);
FD_SET(ssl_fd, &readfds);
tv.tv_sec = 0; // 0 seconds
tv.tv_usec = 50000; // 50ms timeout
int retval = select(ssl_fd + 1, &readfds, NULL, NULL, &tv);
if (retval == -1) {
if (config.debug) {
perror("[DEBUG] select on SSL fd failed");
fflush(stdout);
}
break;
} else if (retval == 0) {
// Timeout, continue loop
continue;
}
if (FD_ISSET(ssl_fd, &readfds)) {
bytes_read = SSL_read(active_tunnel->ssl, buffer, sizeof(buffer));
if (bytes_read <= 0) {
if (config.debug) {
......@@ -1131,6 +1177,7 @@ int main(int argc, char *argv[]) {
}
}
}
}
if (config.debug) {
printf("[DEBUG] Tunnel loop ended, waiting for SSH to finish...\n");
......@@ -1146,6 +1193,13 @@ int main(int argc, char *argv[]) {
fflush(stdout);
}
// Wait a few seconds before exiting cleanly
if (config.debug) {
printf("[DEBUG] Waiting 3 seconds before exit...\n");
fflush(stdout);
}
sleep(3);
// Cleanup
if (active_tunnel) {
close(active_tunnel->local_sock);
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment