Relatively graceful shutdowns

README.md

@@ -8,6 +8,14 @@ ColumnLynx makes use of both **TCP** and **UDP**. **TCP** is used for the initia
 
 It operates on port **48042** for both TCP and UDP.
 
+### Handshake Procedure
+
+*wip*
+
+### Packet Exchange
+
+*wip*
+
 ## Packet Structure
 
 These are the general packet structures for both the TCP and UDP sides of the protocol. Generally **headers** are **plain-text (unencrypted)** and do not contain any sensitive data. 

include/columnlynx/client/net/tcp/tcp_client.hpp

@@ -69,9 +69,11 @@ namespace ColumnLynx::Net::TCP {
                 }
             }
 
-            void disconnect() {
+            void disconnect(bool echo = true) {
                 if (mConnected && mHandler) {
-                    mHandler->sendMessage(ClientMessageType::GRACEFUL_DISCONNECT, "Goodbye");
+                    if (echo) {
+                        mHandler->sendMessage(ClientMessageType::GRACEFUL_DISCONNECT, "Goodbye");
+                    }
 
                     asio::error_code ec;
 
@@ -94,6 +96,10 @@ namespace ColumnLynx::Net::TCP {
                 return mHandshakeComplete;
             }
 
+            bool isConnected() const {
+                return mConnected;
+            }
+
         private:
             void mHandleMessage(ServerMessageType type, const std::string& data) {
                 switch (type) {
@@ -187,7 +193,7 @@ namespace ColumnLynx::Net::TCP {
                     case ServerMessageType::GRACEFUL_DISCONNECT:
                         Utils::log("Server is disconnecting: " + data);
                         if (mConnected) { // Prevent Recursion
-                            disconnect();
+                            disconnect(false);
                         }
                         break;
                     default:

include/columnlynx/server/net/tcp/tcp_connection.hpp

@@ -61,6 +61,8 @@ namespace ColumnLynx::Net::TCP {
             void disconnect() {
                 std::string ip = mHandler->socket().remote_endpoint().address().to_string();
 
+                mHandler->sendMessage(ServerMessageType::GRACEFUL_DISCONNECT, "Server initiated disconnect.");
+
                 asio::error_code ec;
                 mHandler->socket().shutdown(asio::ip::tcp::socket::shutdown_both, ec);
                 mHandler->socket().close(ec);

include/columnlynx/server/net/tcp/tcp_server.hpp

@@ -21,19 +21,22 @@ namespace ColumnLynx::Net::TCP {
 
     class TCPServer {
         public:
-            TCPServer(asio::io_context& ioContext, uint16_t port, Utils::LibSodiumWrapper* sodiumWrapper)
+            TCPServer(asio::io_context& ioContext, uint16_t port, Utils::LibSodiumWrapper* sodiumWrapper, bool* hostRunning)
-                : mIoContext(ioContext), mAcceptor(ioContext, asio::ip::tcp::endpoint(asio::ip::tcp::v4(), port)), mSodiumWrapper(sodiumWrapper)
+                : mIoContext(ioContext), mAcceptor(ioContext, asio::ip::tcp::endpoint(asio::ip::tcp::v4(), port)), mSodiumWrapper(sodiumWrapper), mHostRunning(hostRunning)
             {
                 Utils::log("Started TCP server on port " + std::to_string(port));
                 mStartAccept();
             }
 
+            void stop();
+
         private:
             void mStartAccept();
             asio::io_context &mIoContext;
             asio::ip::tcp::acceptor mAcceptor;
             std::unordered_set<TCPConnection::pointer> mClients;
             Utils::LibSodiumWrapper *mSodiumWrapper;
+            bool* mHostRunning;
     };
 
 }

include/columnlynx/server/net/udp/udp_server.hpp

@@ -12,13 +12,15 @@
 namespace ColumnLynx::Net::UDP {
     class UDPServer {
         public:
-            UDPServer(asio::io_context& ioContext, uint16_t port)
+            UDPServer(asio::io_context& ioContext, uint16_t port, bool* hostRunning)
-                : mSocket(ioContext, asio::ip::udp::endpoint(asio::ip::udp::v4(), port))
+                : mSocket(ioContext, asio::ip::udp::endpoint(asio::ip::udp::v4(), port)), mHostRunning(hostRunning)
             {
                 Utils::log("Started UDP server on port " + std::to_string(port));
                 mStartReceive();
             }
 
+            void stop();
+
         private:
             void mStartReceive();
             void mHandlePacket(std::size_t bytes);
@@ -26,5 +28,6 @@ namespace ColumnLynx::Net::UDP {
             asio::ip::udp::socket mSocket;
             asio::ip::udp::endpoint mRemoteEndpoint;
             std::array<uint8_t, 2048> mRecvBuffer; // Adjust size as needed
+            bool* mHostRunning;
     };
 }

src/client/main.cpp

@@ -60,7 +60,7 @@ int main(int argc, char** argv) {
         while (!done) {
             std::this_thread::sleep_for(std::chrono::milliseconds(100)); // Temp wait
 
-            if (client->isHandshakeComplete()) {
+            if (client->isHandshakeComplete() && client->isConnected()) {
                 // Send a test UDP message every 5 seconds after handshake is complete
                 static auto lastSendTime = std::chrono::steady_clock::now();
                 auto now = std::chrono::steady_clock::now();

src/server/main.cpp

@@ -16,11 +16,25 @@ using namespace ColumnLynx::Utils;
 using namespace ColumnLynx::Net::TCP;
 using namespace ColumnLynx::Net::UDP;
 
+volatile sig_atomic_t done = 0;
+
+/*void signalHandler(int signum) {
+    if (signum == SIGINT || signum == SIGTERM) {
+        log("Received termination signal. Shutting down server gracefully.");
+        done = 1;
+    }
+}*/
+
 int main(int argc, char** argv) {
     PanicHandler::init();
 
     try {
-        // TODO: Catch SIGINT and SIGTERM for graceful shutdown
+        // Catch SIGINT and SIGTERM for graceful shutdown
+        /*struct sigaction action;
+        memset(&action, 0, sizeof(struct sigaction));
+        action.sa_handler = signalHandler;
+        sigaction(SIGINT, &action, nullptr);
+        sigaction(SIGTERM, &action, nullptr);*/
 
         log("ColumnLynx Server, Version " + getVersion());
         log("This software is licensed under the GPLv3. See LICENSE for details.");
@@ -30,18 +44,48 @@ int main(int argc, char** argv) {
         log("Server public key: " + bytesToHexString(sodiumWrapper.getPublicKey(), crypto_sign_PUBLICKEYBYTES));
         log("Server private key: " + bytesToHexString(sodiumWrapper.getPrivateKey(), crypto_sign_SECRETKEYBYTES)); // TEMP, remove later
 
+        bool hostRunning = true;
+
         asio::io_context io;
-        auto server = std::make_shared<TCPServer>(io, serverPort(), &sodiumWrapper);
+
-        auto udpServer = std::make_shared<UDPServer>(io, serverPort());
+        auto server = std::make_shared<TCPServer>(io, serverPort(), &sodiumWrapper, &hostRunning);
+        auto udpServer = std::make_shared<UDPServer>(io, serverPort(), &hostRunning);
+
+        asio::signal_set signals(io, SIGINT, SIGTERM);
+        signals.async_wait([&](const std::error_code&, int) {
+            log("Received termination signal. Shutting down server gracefully.");
+            done = 1;
+            asio::post(io, [&]() {
+                hostRunning = false;
+                server->stop();
+                udpServer->stop();
+            });
+        });
 
         // Run the IO context in a separate thread
         std::thread ioThread([&io]() {
             io.run();
         });
 
-        ioThread.join();
+        //ioThread.detach();
 
         log("Server started on port " + std::to_string(serverPort()));
+        
+        while (!done) {
+            std::this_thread::sleep_for(std::chrono::milliseconds(100));
+        }
+
+        log("Shutting down server...");
+        /*hostRunning = false;
+        server->stop();
+        udpServer->stop();*/
+
+        io.stop();
+        if (ioThread.joinable()) {
+            ioThread.join();
+        }
+
+        log("Server stopped.");
     } catch (const std::exception& e) {
         error("Server error: " + std::string(e.what()));
     }

src/server/server/net/tcp/tcp_server.cpp

@@ -20,27 +20,66 @@
 namespace ColumnLynx::Net::TCP {
 
     void TCPServer::mStartAccept() {
+        // A bit of a shotty implementation, might improve later
+        /*std::cout << "Host running pointer: " << *mHostRunning << std::endl;
+
+        if (mHostRunning != nullptr && !(*mHostRunning)) {
+            Utils::log("Server is stopping, not accepting new connections.");
+            return;
+        }*/
+
         mAcceptor.async_accept(
             [this](asio::error_code ec, asio::ip::tcp::socket socket) {
-                if (!NetHelper::isExpectedDisconnect(ec)) {
+                if (ec) {
-                    auto client = TCPConnection::create(std::move(socket),
+                    if (ec == asio::error::operation_aborted) {
-                        mSodiumWrapper,
+                        // Acceptor was cancelled/closed during shutdown
-                        [this](std::shared_ptr<TCPConnection> c) {
+                        return;
-                            mClients.erase(c);
+                    }
-                            Utils::log("Client removed.");
-                        });
-                    
-                    mClients.insert(client);
-                    client->start();
-
-                    Utils::log("Accepted new client connection.");
-                } else {
                     Utils::error("Accept failed: " + ec.message());
+                    // Try again only if still running
+                    if (mHostRunning && *mHostRunning && mAcceptor.is_open())
+                        mStartAccept();
+                    return;
                 }
-
+            
-                TCPServer::mStartAccept(); // Accept next
+                auto client = TCPConnection::create(
+                    std::move(socket),
+                    mSodiumWrapper,
+                    [this](std::shared_ptr<TCPConnection> c) {
+                        mClients.erase(c);
+                        Utils::log("Client removed.");
+                    }
+                );
+                mClients.insert(client);
+                client->start();
+                Utils::log("Accepted new client connection.");
+            
+                if (mHostRunning && *mHostRunning && mAcceptor.is_open())
+                    mStartAccept();
             }
         );
     }
 
+    void TCPServer::stop() {
+        // Stop accepting
+        if (mAcceptor.is_open()) {
+            asio::error_code ec;
+            mAcceptor.cancel(ec);
+            mAcceptor.close(ec);
+            Utils::log("TCP Acceptor closed.");
+        }
+
+        // Snapshot to avoid iterator invalidation while callbacks erase()
+        std::vector<std::shared_ptr<TCPConnection>> snapshot(mClients.begin(), mClients.end());
+        for (auto& client : snapshot) {
+            try {
+                client->disconnect(); // should shutdown+close the socket
+                Utils::log("GRACEFUL_DISCONNECT sent to session: " + std::to_string(client->getSessionID()));
+            } catch (const std::exception& e) {
+                Utils::error(std::string("Error disconnecting client: ") + e.what());
+            }
+        }
+        // Let the erase callback run as sockets close
+        // Do NOT destroy server while io handlers may still reference it
+    }
 }

src/server/server/net/udp/udp_server.cpp

@@ -10,13 +10,23 @@
 
 namespace ColumnLynx::Net::UDP {
     void UDPServer::mStartReceive() {
+        // A bit of a shotty implementation, might improve later
+        /*if (mHostRunning != nullptr && !(*mHostRunning)) {
+            Utils::log("Server is stopping, not receiving new packets.");
+            return;
+        }*/
+
         mSocket.async_receive_from(
             asio::buffer(mRecvBuffer), mRemoteEndpoint,
             [this](asio::error_code ec, std::size_t bytes) {
-                if (!ec && bytes > 0) {
+                if (ec) {
-                    mHandlePacket(bytes);
+                    if (ec == asio::error::operation_aborted) return; // Socket closed
+                    // Other recv error
+                    if (mHostRunning && *mHostRunning) mStartReceive();
+                    return;
                 }
-                mStartReceive(); // Continue receiving
+                if (bytes > 0) mHandlePacket(bytes);
+                if (mHostRunning && *mHostRunning) mStartReceive();
             }
         );
     }
@@ -69,4 +79,13 @@ namespace ColumnLynx::Net::UDP {
     void UDPServer::mSendData(const uint64_t sessionID, const std::string& data) {
         // TODO: Implement
     }
+
+    void UDPServer::stop() {
+        if (mSocket.is_open()) {
+            asio::error_code ec;
+            mSocket.cancel(ec);
+            mSocket.close(ec);
+            Utils::log("UDP Socket closed.");
+        }
+    }
 }