Made a working ssh proxy server

1 files changed, 40 insertions, 21 deletions

diff --git a/src/main.c b/src/main.c
index cd8e74f..0c6b6cb 100644
--- a/src/main.c
+++ b/src/main.c
@@ -1,15 +1,10 @@
 #include <libssh/callbacks.h>
 #include <libssh/server.h>
-//#include <fcntl.h>
-//#include <libutil.h>
-//#include <poll.h>
-//#include <signal.h>
 #include <stdio.h>
 #include <stdlib.h>
-//#include <sys/ioctl.h>
 #include <sys/wait.h>
-//#include <util.h>
 
+#include "../config.h"
 #include "session.h"
 
 /* SIGCHLD handler for cleaning up dead children. */
@@ -18,6 +13,19 @@ static void sigchld_handler(int signo) {
     while (waitpid(-1, NULL, WNOHANG) > 0);
 }
 
+/* SIGINT handler for cleaning up on forced exit. */
+static ssh_bind sshbind;
+static ssh_session session;
+
+__attribute__((noreturn)) void sigint_handler(int signo)
+{
+    (void) signo;
+    ssh_free(session);
+    ssh_bind_free(sshbind);
+    ssh_finalize();
+    exit(0);
+}
+
 int main()
 {
     // Set up SIGCHLD handler
@@ -29,30 +37,40 @@ int main()
         fprintf(stderr, "Failed to register SIGCHLD handler\n");
         return 1;
     }
+    // Set up SIGINT handler
+    struct sigaction sa2;
+    sa2.sa_handler = sigint_handler;
+    sigemptyset(&sa2.sa_mask);
+    sa2.sa_flags = 0;
+    if (sigaction(SIGINT, &sa2, NULL) != 0) {
+        fprintf(stderr, "Failed to register SIGINT handler\n");
+        return 1;
+    }
 
     // Initializing ssh context
-    ssh_threads_set_callbacks(ssh_threads_get_pthread());
     ssh_init();
 
     // Initializing ssh_bind
-    ssh_bind sshbind = ssh_bind_new();
+    sshbind = ssh_bind_new();
     if (sshbind == NULL) {
         fprintf(stderr, "Error initializing ssh_bind\n");
         exit(-1);
     }
     int port = 2222;
     ssh_bind_options_set(sshbind, SSH_BIND_OPTIONS_BINDPORT, &port);
-    ssh_bind_options_set(sshbind, SSH_BIND_OPTIONS_DSAKEY, "ssh_host_dsa_key");
-    ssh_bind_options_set(sshbind, SSH_BIND_OPTIONS_RSAKEY, "ssh_host_rsa_key");
-    ssh_bind_options_set(sshbind, SSH_BIND_OPTIONS_ECDSAKEY, "ssh_host_ecdsa_key");
+    ssh_bind_options_set(sshbind, SSH_BIND_OPTIONS_DSAKEY, DSAKEY_PATH);
+    ssh_bind_options_set(sshbind, SSH_BIND_OPTIONS_RSAKEY, RSAKEY_PATH);
+    ssh_bind_options_set(sshbind, SSH_BIND_OPTIONS_ECDSAKEY, ECDSAKEY_PATH);
 
     if (ssh_bind_listen(sshbind) < 0) {
         printf("Error listening to socket: %s\n", ssh_get_error(sshbind));
+        ssh_bind_free(sshbind);
+        ssh_finalize();
         return 1;
     }
 
     while (1) {
-        ssh_session session = ssh_new();
+        session = ssh_new();
         if (session == NULL) {
             fprintf(stderr, "Error initializing ssh_session\n");
             break;
@@ -67,14 +85,12 @@ int main()
                 /* Remove the SIGCHLD handler inherited from parent. */
                 sa.sa_handler = SIG_DFL;
                 sigaction(SIGCHLD, &sa, NULL);
-                /* Remove socket binding, which allows us to restart the
-                 * parent process, without terminating existing sessions. */
+                /* Remove socket binding, which allows us to restart the parent process, without terminating existing sessions. */
                 ssh_bind_free(sshbind);
 
                 ssh_event event = ssh_event_new();
                 if (event != NULL) {
-                    /* Blocks until the SSH session ends by either
-                     * child process exiting, or client disconnecting. */
+                    /* Blocks until the SSH session ends */
                     handle_session(event, session);
                     ssh_event_free(event);
                 } else {
@@ -82,21 +98,24 @@ int main()
                 }
                 ssh_disconnect(session);
                 ssh_free(session);
+                ssh_finalize();
 
-                exit(0);
+                return 0;
               case -1:
                 fprintf(stderr, "Failed to fork\n");
             }
         } else {
             fprintf(stderr, "Error accepting a connection : %s\n", ssh_get_error(sshbind));
-            exit(1);
+            ssh_disconnect(session);
+            ssh_free(session);
+            ssh_bind_free(sshbind);
+            ssh_finalize();
+            return 1;
         }
-        /* Since the session has been passed to a child fork, do some cleaning
-         * up at the parent process. */
+        /* Since the session has been passed to a child fork, do some cleaning up at the parent process. */
         ssh_disconnect(session);
         ssh_free(session);
     }
-
     ssh_bind_free(sshbind);
     ssh_finalize();
     return 0;