Skip to content

wsssh
wsssh
Commit 7eb01b2a authored by Stefy Lanza (nextime / spora )'s avatar Stefy Lanza (nextime / spora )
Browse files
Options

Optimize C wsscp performance and add SSH options 

- Replace busy-waiting usleep loops with select() for efficient I/O multiplexing
- Add 50ms timeouts to prevent blocking operations
- Maintain 1MB buffers for optimal throughput
- Add 3-second exit delay after SCP completion for clean shutdown
- Fix critical bug in setup_tunnel return codes to prevent segfaults
- Correct memory management in argument cleanup to avoid corruption
- Add StrictHostKeyChecking=no option to wsscp and wsssh in both Python and C implementations
parent 92294668
Show whitespace changes
Inline Side-by-side
Showing with 286 additions and 228 deletions
wsscp.py
View file @ 7eb01b2a
......@@ -256,8 +256,8 @@ def main():
final_args.append(arg)
i += 1
# Add port argument for local tunnel at the beginning
final_args = ['-P', str(local_port)] + final_args
# Add StrictHostKeyChecking=no and port argument for local tunnel at the beginning
final_args = ['-o', 'StrictHostKeyChecking=no', '-P', str(local_port)] + final_args
if debug: print(f"[DEBUG] Final SCP args: {final_args}")
......
wsssh.py
View file @ 7eb01b2a
......@@ -251,8 +251,8 @@ def main():
else:
final_args.append(arg)
# Add port argument for local tunnel
final_args.extend(['-p', str(local_port)])
# Add StrictHostKeyChecking=no and port argument for local tunnel
final_args.extend(['-o', 'StrictHostKeyChecking=no', '-p', str(local_port)])
if debug: print(f"[DEBUG] Final SSH args: {final_args}")
......
wssshtools/wsscp.c
View file @ 7eb01b2a
......@@ -33,6 +33,7 @@
#include <fcntl.h>
#include <pthread.h>
#include <errno.h>
#include <sys/select.h>
#define BUFFER_SIZE 1048576
#define DEFAULT_PORT 22
......@@ -206,8 +207,8 @@ int parse_scp_args(int argc, char *argv[], char **destination, int *scp_port, in
}
char **modify_scp_args(int argc, char *argv[], const char *original_host, int local_port, int *new_argc) {
// Allocate space for: scp + -P + port + original args + NULL
char **new_args = malloc((argc + 4) * sizeof(char *));
// Allocate space for: scp + -o + StrictHostKeyChecking=no + -P + port + original args + NULL
char **new_args = malloc((argc + 6) * sizeof(char *));
if (!new_args) {
return NULL;
}
......@@ -215,7 +216,11 @@ char **modify_scp_args(int argc, char *argv[], const char *original_host, int lo
int idx = 0;
new_args[idx++] = "scp";
// Add port argument first
// Add StrictHostKeyChecking=no option
new_args[idx++] = "-o";
new_args[idx++] = "StrictHostKeyChecking=no";
// Add port argument
new_args[idx++] = "-P";
char *port_str = malloc(16);
if (!port_str) {
......@@ -594,6 +599,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);
......@@ -626,12 +633,27 @@ void *forward_tcp_to_ws(void *arg) {
pthread_mutex_unlock(&tunnel_mutex);
bytes_read = recv(sock, buffer, sizeof(buffer), MSG_DONTWAIT);
if (bytes_read == -1 && (errno == EAGAIN || errno == EWOULDBLOCK)) {
// No data available, sleep to avoid busy waiting
usleep(50000); // Increased sleep time to reduce CPU usage
// 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;
} else if (bytes_read <= 0) {
}
if (FD_ISSET(sock, &readfds)) {
bytes_read = recv(sock, buffer, sizeof(buffer), 0);
if (bytes_read <= 0) {
if (debug) {
printf("[DEBUG] TCP connection closed or error\n");
fflush(stdout);
......@@ -665,6 +687,7 @@ void *forward_tcp_to_ws(void *arg) {
break;
}
}
}
if (debug) {
printf("[DEBUG] TCP to WebSocket forwarding thread exiting\n");
......@@ -735,13 +758,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));
......@@ -753,7 +776,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
......@@ -784,7 +807,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");
......@@ -801,7 +824,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");
......@@ -813,7 +836,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) {
......@@ -829,7 +852,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) {
......@@ -863,7 +886,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
......@@ -882,7 +905,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
......@@ -908,7 +931,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) {
......@@ -922,7 +945,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;
}
active_tunnel->outgoing_buffer = frame_buffer_init();
......@@ -932,7 +955,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);
......@@ -949,7 +972,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;
......@@ -967,7 +990,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) {
......@@ -979,7 +1002,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) {
......@@ -1234,6 +1257,10 @@ int main(int argc, char *argv[]) {
}
char read_buffer[BUFFER_SIZE];
fd_set readfds;
struct timeval tv;
int ssl_fd = SSL_get_fd(active_tunnel->ssl);
while (active_tunnel && active_tunnel->active) {
// Check if SCP process has finished
int status;
......@@ -1246,6 +1273,25 @@ int main(int argc, char *argv[]) {
break;
}
// 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)) {
int bytes_read = SSL_read(active_tunnel->ssl, read_buffer, sizeof(read_buffer));
if (bytes_read <= 0) {
if (config.debug) {
......@@ -1476,6 +1522,7 @@ int main(int argc, char *argv[]) {
}
}
}
}
frame_buffer_free(frame_buffer);
......@@ -1494,6 +1541,13 @@ cleanup:
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);
......@@ -1509,9 +1563,9 @@ cleanup:
// Free allocated strings in new_scp_args
for (int i = 0; i < new_scp_argc; i++) {
// Free strings that were allocated with malloc/strdup
if (i == 2) { // The port string
if (i == 4) { // The port string
free(new_scp_args[i]);
} else if (i == 3) { // The user@localhost:path string (if allocated)
} else if (i == 5) { // The user@localhost:path string (if allocated)
// Check if this was allocated (contains @localhost)
if (strstr(new_scp_args[i], "@localhost")) {
free(new_scp_args[i]);
......
wssshtools/wsssh.c
View file @ 7eb01b2a
......@@ -194,8 +194,8 @@ int parse_ssh_args(int argc, char *argv[], char **host, int *ssh_port, int debug
}
char **modify_ssh_args(int argc, char *argv[], const char *original_host, int local_port, int *new_argc) {
// Allocate space for: ssh + -p + port + original args + NULL
char **new_args = malloc((argc + 4) * sizeof(char *));
// Allocate space for: ssh + -o + StrictHostKeyChecking=no + -p + port + original args + NULL
char **new_args = malloc((argc + 6) * sizeof(char *));
if (!new_args) {
return NULL;
}
......@@ -203,6 +203,10 @@ char **modify_ssh_args(int argc, char *argv[], const char *original_host, int lo
int idx = 0;
new_args[idx++] = "ssh";
// Add StrictHostKeyChecking=no option
new_args[idx++] = "-o";
new_args[idx++] = "StrictHostKeyChecking=no";
// Add port argument for local tunnel
new_args[idx++] = "-p";
char *port_str = malloc(16);
......@@ -1156,9 +1160,9 @@ int main(int argc, char *argv[]) {
// Free allocated strings in new_ssh_args
for (int i = 0; i < new_ssh_argc; i++) {
// Free strings that were allocated with malloc/strdup
if (i == 2) { // The port string
if (i == 4) { // The port string
free(new_ssh_args[i]);
} else if (i == 3) { // The user@localhost string (if allocated)
} else if (i == 5) { // The user@localhost string (if allocated)
// Check if this was allocated (contains @localhost)
if (strstr(new_ssh_args[i], "@localhost")) {
free(new_ssh_args[i]);
......
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