feat: gracefully shutdown socket

Author: araujo88

include/socket.hpp

@@ -20,6 +20,7 @@ namespace tpt
         virtual void sendData(SOCKET client_socket, const void *buffer, unsigned int buffer_size, int flags) = 0;
         virtual void closeSocket() = 0;
         virtual void closeSocket(SOCKET client_socket) = 0;
+        virtual void setSocketTimeout(SOCKET sock, int timeoutSec) = 0;
     };
 }
 

include/teapot.hpp

@@ -31,6 +31,7 @@
 #endif
 
 #ifdef _WIN32
+#include <windows.h>
 #include "win_socket.hpp"
 #endif
 

include/unix_socket.hpp

@@ -29,6 +29,7 @@ namespace tpt
         std::string ip_address;
         std::string client_ip;
         std::vector<std::string> ip_blacklist;
+        std::vector<SOCKET> client_sockets;
         ConsoleLogger logger;
 
     public:
@@ -46,6 +47,7 @@ namespace tpt
         virtual void sendData(SOCKET client_socket, const void *buffer, unsigned int buffer_size, int flags) override;
         virtual void closeSocket() override;
         virtual void closeSocket(SOCKET client_socket) override;
+        virtual void setSocketTimeout(SOCKET sock, int timeoutSec) override;
     };
 }
 

include/win_socket.hpp

@@ -49,6 +49,7 @@ namespace tpt
         virtual void sendData(SOCKET client_socket, const void *buffer, unsigned int buffer_size, int flags) override;
         virtual void closeSocket() override;
         virtual void closeSocket(SOCKET client_socket) override;
+        virtual void setSocketTimeout(SOCKET sock, int timeoutSec) override;
     };
 }
 

src/teapot.cpp

@@ -3,6 +3,27 @@
 
 using namespace tpt;
 
+volatile std::sig_atomic_t signal_received = 0;
+
+#ifdef __linux__
+extern "C" void signalHandler(int signal)
+{
+    signal_received = signal;
+}
+#endif
+
+#ifdef _WIN32
+BOOL WINAPI consoleHandler(DWORD signal)
+{
+    if (signal == CTRL_C_EVENT || signal == CTRL_CLOSE_EVENT)
+    {
+        signal_received = 1;
+        return TRUE;
+    }
+    return FALSE;
+}
+#endif
+
 std::optional<Request> Teapot::parseRequest(int client_socket)
 {
     char buffer[BUFFER_SIZE] = {0};
@@ -212,9 +233,21 @@ Teapot::Teapot(std::string ip_address, unsigned int port, unsigned int max_conne
 
 void Teapot::run()
 {
+#ifdef __linux__
+    std::signal(SIGINT, signalHandler);
+    std::signal(SIGTERM, signalHandler);
+#endif
+#ifdef _WIN32
+    if (!SetConsoleCtrlHandler(consoleHandler, TRUE))
+    {
+        std::cerr << "ERROR: Could not set control handler" << std::endl;
+        return;
+    }
+#endif
+
     socket.bindSocket();
 
-    while (true)
+    while (!signal_received)
     {
         socket.listenToConnections();
 
@@ -233,6 +266,10 @@ void Teapot::run()
             this->socket.closeSocket(client_socket);
         }
     }
+
+    LOG_INFO(logger, "Shutting down...");
+    this->socket.closeSocket();
+    sleep(3);
 }
 
 void Teapot::serveFile(std::string url, std::string file_path)

src/unix_socket.cpp

@@ -133,6 +133,7 @@ void UnixSocket::bindSocket()
         std::cout << "Error code: " + errno << std::endl;
         exit(EXIT_FAILURE);
     }
+    this->setSocketTimeout(this->server_socket, 5);
     LOG_INFO(logger, "Binding done! Listening to connections ...");
 }
 
@@ -188,6 +189,8 @@ void UnixSocket::acceptConnection(SOCKET &client_socket, void *client_address)
             throw IPBlackListedException();
         }
     }
+
+    this->client_sockets.push_back(client_socket);
 }
 
 ssize_t UnixSocket::receiveData(SOCKET client_socket, char *buffer, unsigned int buffer_size)
@@ -209,7 +212,13 @@ void UnixSocket::sendData(SOCKET client_socket, const void *buffer, unsigned int
 
 void UnixSocket::closeSocket()
 {
-    std::cout << "Closing socket ..." << std::endl;
+    LOG_INFO(logger, "Cleaning up client sockets ...");
+    for (auto &it : this->client_sockets)
+    {
+        this->closeSocket(it);
+    }
+
+    LOG_INFO(logger, "Closing socket ...");
     if (shutdown(this->server_socket, SHUT_RDWR) == -1)
     {
         perror("An error occurred while shutting down the socket: ");
@@ -218,7 +227,7 @@ void UnixSocket::closeSocket()
     }
     if (close(this->server_socket) == 0)
     {
-        std::cout << "Socket closed!" << std::endl;
+        LOG_INFO(logger, "Socket closed!");
         exit(EXIT_SUCCESS);
     }
     else
@@ -231,14 +240,36 @@ void UnixSocket::closeSocket()
 
 void UnixSocket::closeSocket(SOCKET client_socket)
 {
+    if (shutdown(client_socket, SHUT_RDWR) == -1)
+    {
+        perror("An error occurred while shutting down the socket: ");
+        std::cout << "Error code: " + errno << std::endl;
+        exit(EXIT_FAILURE);
+    }
     close(client_socket);
+
+    auto it = std::find(this->client_sockets.begin(), this->client_sockets.end(), client_socket);
+
+    // Check if element was found before erasing
+    if (it != this->client_sockets.end())
+    {
+        this->client_sockets.erase(it);
+    }
 }
 
 std::string UnixSocket::getClientIp()
 {
     return this->client_ip;
 }
 
+void UnixSocket::setSocketTimeout(SOCKET sock, int timeoutSec)
+{
+    struct timeval tv;
+    tv.tv_sec = timeoutSec; // Seconds
+    tv.tv_usec = 0;         // Microseconds
+    setsockopt(sock, SOL_SOCKET, SO_RCVTIMEO, (const char *)&tv, sizeof tv);
+}
+
 UnixSocket::~UnixSocket() {}
 
 #endif

src/win_socket.cpp

@@ -162,6 +162,7 @@ void WinSocket::bindSocket()
         WSACleanup();
         exit(EXIT_FAILURE);
     }
+    this->setSocketTimeout(this->server_socket, 5);
     std::cout << "Binding done!" << std::endl;
     std::cout << "Listening to connections ..." << std::endl;
 }
@@ -249,6 +250,21 @@ void WinSocket::closeSocket(SOCKET client_socket)
     closesocket(client_socket);
 }
 
+void WinSocket::setSocketTimeout(SOCKET sock, int timeoutSec)
+{
+    // Initialize the timeout value
+    DWORD timeout = timeoutSec * 1000; // Convert to milliseconds, as required by Windows
+
+    // Set the receive timeout option
+    int result = setsockopt(sock, SOL_SOCKET, SO_RCVTIMEO, (const char *)&timeout, sizeof(timeout));
+    if (result == SOCKET_ERROR)
+    {
+        std::printf("Error setting timeout: %d\n", WSAGetLastError());
+        WSACleanup();
+        exit(EXIT_FAILURE);
+    }
+}
+
 WinSocket::~WinSocket() {}
 
 #endif