Update README.md, add heartbeat packets to detect dead or hanging connections.

README.md

@@ -10,11 +10,44 @@ It operates on port **48042** for both TCP and UDP.
 
 ### Handshake Procedure
 
-*wip*
+The handshake between the client and server is done over **TCP**. This is to ensure delivery without much hassle.
+
+The procedure will be described packet per packet (line per line) with a **C** or **S** prefix for the client and server, respectfully. After the prefix will be the **Packet ID** and then the data in **<>** tags. Lines without a prefix describe the use / reasoning of the above packet.
+
+```
+C: HANDSHAKE_INIT <Client Identity Public Key>
+S: HANDSHAKE_IDENTIFY <Server Identity Public Key>
+C: HANDSHAKE_CHALLENGE <Random Nonce (32 bytes)>
+S: HANDSHAKE_CHALLENGE_RESPONSE <Signed Nonce (Gotten from Previous Packet)>
+
+The Client now generates a random aesKey (32 bytes long)
+
+C: HANDSHAKE_EXCHANGE_KEY <aesKey Encrypted with Server Public Key>
+
+The Server now assigns a local 8 byte session ID in the Session Registry.
+
+S: HANDSHAKE_EXCHANGE_KEY_CONFIRM <Assigned SessionID>
+```
+
+The **Client** and **Server** have now securely exchanged a symmetric **AES Key** that they'll use to **encrypt all traffic** sent further out.
 
 ### Packet Exchange
 
-*wip*
+Packet exchange and the general data tunneling is done via **Standard UDP** (*see the **UDP Packet** in **Data***).
+
+The **header** of the sent packet always includes a **random 12 byte nonce** used to obscure the **encrypted payload / data** and the **Session ID** assigned by the server to the client (8 bytes). This makes the header **20 bytes long**.
+
+The **payload / data** of the sent packet is **always encrypted** using the exchanged **AES Key** and obscured using the **random nonce**.
+
+*The AES key used is according to the **ChaCha20-Poly1305** algorithm.*
+
+### Connection Termination
+
+The lifetime of a connection is determined based on the lifetime of the **TCP connection** and **Heartbeat packets**.
+
+As soon as the TCP connection terminates, either due to a lost connection, **TCP RST**, **GRACEFUL_DISCONNECT** or **KILL_CONNECTION** packet, etc., the client and server will **stop sending UDP data**. The server will also remove the terminated client from its **Session Registry**.
+
+Additionally, if either party misses **3 of the sent heartbeat packets**, the other party will treat them as dead and remove them.
 
 ## Packet Structure
 
@@ -33,7 +66,7 @@ TCP Packets generally follow the structure **Packet ID + Data**. They're only us
 
 The **Packet ID** is an **8 bit unsigned integer** that is predefined from either the **Client to Server** or **Server to Client** enum set, however they are uniquely numbered as to not collide with each other.
 
-**Server to Client** IDs are always below **0xA0** (exclusive) and **Client to Server** IDs are always above **0xA0** (exclusive). **0xFE** and **OxFF** are shared for **GRACEFUL_DISCONNECT** and **KILL_CONNECTION** respectively.
+**Server to Client** IDs are always below **0xA0** (exclusive) and **Client to Server** IDs are always above **0xA0** (exclusive). **0xF0**, **0xF1**, **0xFE** and **OxFF** are shared for **HEARTBEAT**, **HEARTBEAT_ACK**, **GRACEFUL_DISCONNECT** and **KILL_CONNECTION** respectively.
 
 #### Data
 

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

@@ -25,7 +25,18 @@ namespace ColumnLynx::Net::TCP {
                       Utils::LibSodiumWrapper* sodiumWrapper,
                       std::array<uint8_t, 32>* aesKey,
                       uint64_t* sessionIDRef)
-                : mResolver(ioContext), mSocket(ioContext), mHost(host), mPort(port), mLibSodiumWrapper(sodiumWrapper), mGlobalKeyRef(aesKey), mSessionIDRef(sessionIDRef) {}
+                :
+                mResolver(ioContext),
+                mSocket(ioContext),
+                mHost(host),
+                mPort(port),
+                mLibSodiumWrapper(sodiumWrapper),
+                mGlobalKeyRef(aesKey),
+                mSessionIDRef(sessionIDRef),
+                mHeartbeatTimer(mSocket.get_executor()),
+                mLastHeartbeatReceived(std::chrono::steady_clock::now()),
+                mLastHeartbeatSent(std::chrono::steady_clock::now())
+            {}
 
             void start() {
                 auto self = shared_from_this();
@@ -55,6 +66,8 @@ namespace ColumnLynx::Net::TCP {
                                         );
 
                                         mHandler->sendMessage(ClientMessageType::HANDSHAKE_INIT, Utils::uint8ArrayToString(payload.data(), payload.size()));
+                                    
+                                        mStartHeartbeat();
                                     } else {
                                         Utils::error("Client connect failed: " + ec.message());
                                     }
@@ -85,6 +98,7 @@ namespace ColumnLynx::Net::TCP {
                     }
 
                     asio::error_code ec;
+                    mHeartbeatTimer.cancel();
 
                     mHandler->socket().shutdown(tcp::socket::shutdown_both, ec);
                     if (ec) {
@@ -110,6 +124,42 @@ namespace ColumnLynx::Net::TCP {
             }
 
         private:
+            void mStartHeartbeat() {
+                auto self = shared_from_this();
+                mHeartbeatTimer.expires_after(std::chrono::seconds(5));
+                mHeartbeatTimer.async_wait([this, self](const asio::error_code& ec) {
+                    if (ec == asio::error::operation_aborted) {
+                        return; // Timer was cancelled
+                    }
+
+                    auto now = std::chrono::steady_clock::now();
+                    auto elapsed = std::chrono::duration_cast<std::chrono::seconds>(now - self->mLastHeartbeatReceived).count();
+
+                    if (elapsed >= 15) { // 3 missed heartbeats
+                        Utils::error("Missed 3 heartbeats. I think the other party might have died! Disconnecting.");
+                        
+                        // Close sockets forcefully, server is dead
+                        asio::error_code ec;
+                        mHandler->socket().shutdown(tcp::socket::shutdown_both, ec);
+                        mHandler->socket().close(ec);
+                        mConnected = false;
+
+                        mGlobalKeyRef = nullptr;
+                        if (mSessionIDRef) {
+                            *mSessionIDRef = 0;
+                        }
+
+                        return;
+                    }
+
+                    self->sendMessage(ClientMessageType::HEARTBEAT);
+                    Utils::log("Sent HEARTBEAT to server.");
+                    self->mLastHeartbeatSent = std::chrono::steady_clock::now();
+                    
+                    self->mStartHeartbeat(); // Recursive
+                });
+            }
+
             void mHandleMessage(ServerMessageType type, const std::string& data) {
                 switch (type) {
                     case ServerMessageType::HANDSHAKE_IDENTIFY:
@@ -198,6 +248,15 @@ namespace ColumnLynx::Net::TCP {
                             mHandshakeComplete = true;
                         }
                     
+                        break;
+                    case ServerMessageType::HEARTBEAT:
+                        Utils::log("Received HEARTBEAT from server.");
+                        mHandler->sendMessage(ClientMessageType::HEARTBEAT_ACK, ""); // Send ACK
+                        break;
+                    case ServerMessageType::HEARTBEAT_ACK:
+                        Utils::log("Received HEARTBEAT_ACK from server.");
+                        mLastHeartbeatReceived = std::chrono::steady_clock::now();
+                        mMissedHeartbeats = 0; // Reset missed heartbeat count
                         break;
                     case ServerMessageType::GRACEFUL_DISCONNECT:
                         Utils::log("Server is disconnecting: " + data);
@@ -224,5 +283,9 @@ namespace ColumnLynx::Net::TCP {
             SymmetricKey mConnectionAESKey;
             std::array<uint8_t, 32>* mGlobalKeyRef; // Reference to global AES key
             uint64_t* mSessionIDRef; // Reference to global Session ID
+            asio::steady_timer mHeartbeatTimer;
+            std::chrono::steady_clock::time_point mLastHeartbeatReceived;
+            std::chrono::steady_clock::time_point mLastHeartbeatSent;
+            int mMissedHeartbeats = 0;
     };
 }

include/columnlynx/common/net/tcp/tcp_message_type.hpp

@@ -13,6 +13,9 @@ namespace ColumnLynx::Net::TCP {
         HANDSHAKE_CHALLENGE_RESPONSE = 0x04, // Response to client's challenge
         HANDSHAKE_EXCHANGE_KEY_CONFIRM = 0x06, // If accepted, send encrypted AES key and session ID
 
+        // Shared
+        HEARTBEAT = 0xF0, // Keep-alive message
+        HEARTBEAT_ACK = 0xF1, // Acknowledgement of keep-alive
         GRACEFUL_DISCONNECT = 0xFE, // Notify client of impending disconnection
         KILL_CONNECTION    = 0xFF, // Forecefully terminate the connection (with cleanup if possible), reserved for unrecoverable errors
     };
@@ -22,6 +25,9 @@ namespace ColumnLynx::Net::TCP {
         HANDSHAKE_CHALLENGE = 0xA3, // Challenge ownership of private key
         HANDSHAKE_EXCHANGE_KEY   = 0xA5, // Accept or reject identity, can kill the connection, also sends the AES key
 
+        // Shared
+        HEARTBEAT = 0xF0, // Keep-alive message
+        HEARTBEAT_ACK = 0xF1, // Acknowledgement of keep-alive
         GRACEFUL_DISCONNECT = 0xFE, // Notify server of impending disconnection
         KILL_CONNECTION    = 0xFF, // Forecefully terminate the connection (with cleanup if possible), reserved for unrecoverable errors
     };

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

@@ -43,6 +43,7 @@ namespace ColumnLynx::Net::TCP {
                 });
 
                 mHandler->start();
+                mStartHeartbeat();
 
                 // Placeholder for message handling setup
                 Utils::log("Client connected: " + mHandler->socket().remote_endpoint().address().to_string());
@@ -62,7 +63,7 @@ namespace ColumnLynx::Net::TCP {
                 std::string ip = mHandler->socket().remote_endpoint().address().to_string();
 
                 mHandler->sendMessage(ServerMessageType::GRACEFUL_DISCONNECT, "Server initiated disconnect.");
-
+                mHeartbeatTimer.cancel();
                 asio::error_code ec;
                 mHandler->socket().shutdown(asio::ip::tcp::socket::shutdown_both, ec);
                 mHandler->socket().close(ec);
@@ -84,7 +85,45 @@ namespace ColumnLynx::Net::TCP {
         
         private:
             TCPConnection(asio::ip::tcp::socket socket, Utils::LibSodiumWrapper* sodiumWrapper)
-                : mHandler(std::make_shared<MessageHandler>(std::move(socket))), mLibSodiumWrapper(sodiumWrapper) {}
+                :
+                mHandler(std::make_shared<MessageHandler>(std::move(socket))),
+                mLibSodiumWrapper(sodiumWrapper),
+                mHeartbeatTimer(mHandler->socket().get_executor()),
+                mLastHeartbeatReceived(std::chrono::steady_clock::now()),
+                mLastHeartbeatSent(std::chrono::steady_clock::now())
+            {}
+
+            void mStartHeartbeat() {
+                auto self = shared_from_this();
+                mHeartbeatTimer.expires_after(std::chrono::seconds(5));
+                mHeartbeatTimer.async_wait([this, self](const asio::error_code& ec) {
+                    if (ec == asio::error::operation_aborted) {
+                        return; // Timer was cancelled
+                    }
+
+                    auto now = std::chrono::steady_clock::now();
+                    auto elapsed = std::chrono::duration_cast<std::chrono::seconds>(now - self->mLastHeartbeatReceived).count();
+
+                    if (elapsed >= 15) { // 3 missed heartbeats
+                        Utils::error("Missed 3 heartbeats. I think the other party (client " + std::to_string(self->mConnectionSessionID) + ") might have died! Disconnecting.");
+                        
+                        // Remove socket forcefully, client is dead
+                        asio::error_code ec;
+                        mHandler->socket().shutdown(asio::ip::tcp::socket::shutdown_both, ec);
+                        mHandler->socket().close(ec);
+
+                        SessionRegistry::getInstance().erase(self->mConnectionSessionID);
+
+                        return;
+                    }
+
+                    self->sendMessage(ServerMessageType::HEARTBEAT);
+                    Utils::log("Sent HEARTBEAT to client " + std::to_string(self->mConnectionSessionID));
+                    self->mLastHeartbeatSent = now;
+
+                    self->mStartHeartbeat(); // Recursive
+                });
+            }
 
             void mHandleMessage(ClientMessageType type, const std::string& data) {
                 std::string reqAddr = mHandler->socket().remote_endpoint().address().to_string();
@@ -191,6 +230,17 @@ namespace ColumnLynx::Net::TCP {
 
                         break;
                     }
+                    case ClientMessageType::HEARTBEAT: {
+                        Utils::log("Received HEARTBEAT from " + reqAddr);
+                        mHandler->sendMessage(ServerMessageType::HEARTBEAT_ACK, ""); // Send ACK
+                        break;
+                    }
+                    case ClientMessageType::HEARTBEAT_ACK: {
+                        Utils::log("Received HEARTBEAT_ACK from " + reqAddr);
+                        mLastHeartbeatReceived = std::chrono::steady_clock::now();
+                        mMissedHeartbeats = 0; // Reset missed heartbeat count
+                        break;
+                    }
                     case ClientMessageType::GRACEFUL_DISCONNECT: {
                         Utils::log("Received GRACEFUL_DISCONNECT from " + reqAddr + ": " + data);
                         disconnect();
@@ -208,5 +258,9 @@ namespace ColumnLynx::Net::TCP {
             std::array<uint8_t, 32> mConnectionAESKey;
             uint64_t mConnectionSessionID;
             AsymPublicKey mConnectionPublicKey;
+            asio::steady_timer mHeartbeatTimer;
+            std::chrono::steady_clock::time_point mLastHeartbeatReceived;
+            std::chrono::steady_clock::time_point mLastHeartbeatSent;
+            int mMissedHeartbeats = 0;
     };
 }

src/client/main.cpp

@@ -74,10 +74,10 @@ int main(int argc, char** argv) {
         log("Client connected to " + host + ":" + port);
 
         // Client is running
-        while (!done) {
+        while ((!done && client->isConnected()) || !client->isHandshakeComplete()) {
             std::this_thread::sleep_for(std::chrono::milliseconds(100)); // Temp wait
 
-            if (client->isHandshakeComplete() && client->isConnected()) {
+            if (client->isHandshakeComplete()) {
                 // 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();