Added basic UDP handling - only sending from client supported.

README.md

@@ -10,9 +10,16 @@ It operates on port **48042** for both TCP and UDP.
 
 ## 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. 
+
+
+The **data / payload** section is:
+- For **TCP**: **encrypted** or **plain-text** depending on the **packet type** (packets with **sensitive data** are **encrypted**)
+- For **UDP**: **encrypted**, as they're transfering the actual data
+
 ### TCP Packets
 
-TCP Packets generally follow the structure **Packet ID + Data**
+TCP Packets generally follow the structure **Packet ID + Data**. They're only used for the **inital handshake** and **commands sent between the client and server**.
 
 #### Packet ID
 
@@ -20,14 +27,30 @@ The **Packet ID** is an **8 bit unsigned integer** that is predefined from eithe
 
 **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.
 
-
 #### Data
 
 The data section is unspecified. It may change depending on the **Packet ID**. It is encoded as a **raw byte array**
 
+#### Final General Structure
+
+| Type | Length | Name | Description |
+|:-----|:-------|:-----|:------------|
+| uint8_t | 1 byte | **Header** - Packet Type | General type of packet |
+| uint8_t/byte array | variable | Data | General packet data - changes for packet to packet |
+
 ### UDP Packets
 
-*WIP, fill in later*
+**UDP Packets** follow the same general structure of **Packet ID + Data**, however, they are **encrypted in full** with the exchanged AES key. This is done to prevent any metadata leakage by either the client or the server.
+
+The **Data** is generally just the **raw underlying packet** forwarded to the server/client.
+
+#### Final General Structure
+
+| Type | Length | Name | Description |
+|:-----|:-------|:-----|:------------|
+| uint8_t | 12 bytes | **Header** - Nonce | Random nonce to obfuscate encrypted contents |
+| uint64_t | 8 bytes | **Header** - Session ID | The unique and random session identifier for the client |
+| uint8_t | variable | Data | General data / payload |
 
 ## Legal
 
@@ -50,5 +73,7 @@ DcruBro is the online pseudonym of Jonas Korene Novak. Both refer to the same in
 
 ### Licensing
 
+*See **ATTRIBUTIONS.md** for details.*
+
 This project includes the [ASIO C++ Library](https://think-async.com/Asio/),
 distributed under the [Boost Software License, Version 1.0](https://www.boost.org/LICENSE_1_0.txt).

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

@@ -22,8 +22,10 @@ namespace ColumnLynx::Net::TCP {
             TCPClient(asio::io_context& ioContext,
                       const std::string& host,
                       const std::string& port,
-                      Utils::LibSodiumWrapper* sodiumWrapper)
-                : mResolver(ioContext), mSocket(ioContext), mHost(host), mPort(port), mLibSodiumWrapper(sodiumWrapper) {}
+                      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) {}
 
             void start() {
                 auto self = shared_from_this();
@@ -88,6 +90,10 @@ namespace ColumnLynx::Net::TCP {
                 }
             }
 
+            bool isHandshakeComplete() const {
+                return mHandshakeComplete;
+            }
+
         private:
             void mHandleMessage(ServerMessageType type, const std::string& data) {
                 switch (type) {
@@ -118,6 +124,9 @@ namespace ColumnLynx::Net::TCP {
 
                                 // Generate AES key and send confirmation
                                 mConnectionAESKey = Utils::LibSodiumWrapper::generateRandom256Bit();
+                                if (mGlobalKeyRef) { // Copy to the global reference
+                                    std::copy(mConnectionAESKey.begin(), mConnectionAESKey.end(), mGlobalKeyRef->begin());
+                                }
                                 AsymNonce nonce{};
                                 randombytes_buf(nonce.data(), nonce.size());
 
@@ -166,6 +175,12 @@ namespace ColumnLynx::Net::TCP {
 
                             std::memcpy(&mConnectionSessionID, decrypted.data(), sizeof(mConnectionSessionID));
                             Utils::log("Connection established with Session ID: " + std::to_string(mConnectionSessionID));
+                        
+                            if (mSessionIDRef) { // Copy to the global reference
+                                *mSessionIDRef = mConnectionSessionID;
+                            }
+
+                            mHandshakeComplete = true;
                         }
                     
                         break;
@@ -182,6 +197,7 @@ namespace ColumnLynx::Net::TCP {
             }
 
             bool mConnected = false;
+            bool mHandshakeComplete = false;
             tcp::resolver mResolver;
             tcp::socket mSocket;
             std::shared_ptr<MessageHandler> mHandler;
@@ -191,5 +207,7 @@ namespace ColumnLynx::Net::TCP {
             Utils::LibSodiumWrapper* mLibSodiumWrapper;
             uint64_t mConnectionSessionID;
             SymmetricKey mConnectionAESKey;
+            std::array<uint8_t, 32>* mGlobalKeyRef; // Reference to global AES key
+            uint64_t* mSessionIDRef; // Reference to global Session ID
     };
 }

include/columnlynx/client/net/udp/udp_client.hpp

@@ -0,0 +1,70 @@
+// udp_client.hpp - UDP Client for ColumnLynx
+// Copyright (C) 2025 DcruBro
+// Distributed under the GPLv3 license. See LICENSE for details.
+
+#pragma once
+
+#include <asio/asio.hpp>
+#include <columnlynx/common/net/udp/udp_message_type.hpp>
+#include <columnlynx/common/utils.hpp>
+#include <columnlynx/common/libsodium_wrapper.hpp>
+#include <array>
+
+namespace ColumnLynx::Net::UDP {
+    class UDPClient {
+        public:
+            UDPClient(asio::io_context& ioContext,
+                      const std::string& host,
+                      const std::string& port,
+                      std::array<uint8_t, 32>* aesKeyRef,
+                      uint64_t* sessionIDRef)
+                : mSocket(ioContext), mResolver(ioContext), mHost(host), mPort(port), mAesKeyRef(aesKeyRef), mSessionIDRef(sessionIDRef) {}
+
+            void start() {
+                auto endpoints = mResolver.resolve(asio::ip::udp::v4(), mHost, mPort);
+                mRemoteEndpoint = *endpoints.begin();
+                mSocket.open(asio::ip::udp::v4());
+                Utils::log("UDP Client ready to send to " + mRemoteEndpoint.address().to_string() + ":" + std::to_string(mRemoteEndpoint.port()));
+            }
+
+            void sendMessage(const std::string& data = "") {
+                UDPPacketHeader hdr{};
+                randombytes_buf(hdr.nonce.data(), hdr.nonce.size());
+
+                if (mAesKeyRef == nullptr || mSessionIDRef == nullptr) {
+                    Utils::error("UDP Client AES key or Session ID reference is null!");
+                    return;
+                }
+
+                auto encryptedPayload = Utils::LibSodiumWrapper::encryptMessage(
+                    reinterpret_cast<const uint8_t*>(data.data()), data.size(),
+                    *mAesKeyRef, hdr.nonce, "udp-data"
+                );
+
+                std::vector<uint8_t> packet;
+                packet.reserve(sizeof(UDPPacketHeader) + sizeof(uint64_t) + encryptedPayload.size());
+                packet.insert(packet.end(), 
+                    reinterpret_cast<uint8_t*>(&hdr),
+                    reinterpret_cast<uint8_t*>(&hdr) + sizeof(UDPPacketHeader)
+                );
+                uint64_t sid = *mSessionIDRef;
+                packet.insert(packet.end(),
+                    reinterpret_cast<uint8_t*>(&sid),
+                    reinterpret_cast<uint8_t*>(&sid) + sizeof(sid)
+                );
+                packet.insert(packet.end(), encryptedPayload.begin(), encryptedPayload.end());
+
+                mSocket.send_to(asio::buffer(packet), mRemoteEndpoint);
+                Utils::log("Sent UDP packet of size " + std::to_string(packet.size()));
+            }
+
+        private:
+            asio::ip::udp::socket mSocket;
+            asio::ip::udp::resolver mResolver;
+            asio::ip::udp::endpoint mRemoteEndpoint;
+            std::string mHost;
+            std::string mPort;
+            std::array<uint8_t, 32>* mAesKeyRef;
+            uint64_t* mSessionIDRef;
+    };
+}

include/columnlynx/common/net/session_registry.hpp

@@ -0,0 +1,88 @@
+// session_registry.hpp - Session Registry for ColumnLynx
+// Copyright (C) 2025 DcruBro
+// Distributed under the GPLv3 license. See LICENSE for details.
+
+#pragma once
+#include <unordered_set>
+#include <shared_mutex>
+#include <memory>
+#include <chrono>
+#include <array>
+#include <sodium.h>
+#include <columnlynx/common/utils.hpp>
+#include <columnlynx/common/libsodium_wrapper.hpp>
+
+namespace ColumnLynx::Net {
+    struct SessionState {
+        SymmetricKey aesKey; // Immutable after creation
+        std::atomic<uint64_t> send_ctr{0}; // Per-direction counters
+        std::atomic<uint64_t> recv_ctr{0};
+        asio::ip::udp::endpoint udpEndpoint;
+        std::atomic<uint64_t> sendCounter{0};
+        std::chrono::steady_clock::time_point created = std::chrono::steady_clock::now();
+        std::chrono::steady_clock::time_point expires{};
+        Nonce base_nonce{};
+
+        ~SessionState() { sodium_memzero(aesKey.data(), aesKey.size()); }
+        SessionState(const SessionState&) = delete;
+        SessionState& operator=(const SessionState&) = delete;
+        SessionState(SessionState&&) = default;
+        SessionState& operator=(SessionState&&) = default;
+
+        explicit SessionState(const SymmetricKey& k, std::chrono::seconds ttl = std::chrono::hours(24)) : aesKey(k) {
+            expires = created + ttl;
+        }
+
+        void setUDPEndpoint(const asio::ip::udp::endpoint& ep) {
+            udpEndpoint = ep;
+        }
+    };
+
+    class SessionRegistry {
+        public:
+            static SessionRegistry& getInstance() { static SessionRegistry instance; return instance; }
+
+            // Insert or replace
+            void put(uint64_t sessionID, std::shared_ptr<SessionState> state) {
+                std::unique_lock lock(mMutex);
+                mSessions[sessionID] = std::move(state);
+            }
+
+            // Lookup
+            std::shared_ptr<const SessionState> get(uint64_t sessionID) const {
+                std::shared_lock lock(mMutex);
+                auto it = mSessions.find(sessionID);
+                return (it == mSessions.end()) ? nullptr : it->second;
+            }
+
+            std::unordered_map<uint64_t, std::shared_ptr<SessionState>> snapshot() const {
+                std::unordered_map<uint64_t, std::shared_ptr<SessionState>> snap;
+                std::shared_lock lock(mMutex);
+                snap = mSessions;
+                return snap;
+            }
+
+            // Remove
+            void erase(uint64_t sessionID) {
+                std::unique_lock lock(mMutex);
+                mSessions.erase(sessionID);
+            }
+
+            // Cleanup expired sessions
+            void cleanupExpired() {
+                std::unique_lock lock(mMutex);
+                auto now = std::chrono::steady_clock::now();
+                for (auto it = mSessions.begin(); it != mSessions.end(); ) {
+                    if (it->second && it->second->expires <= now) {
+                        it = mSessions.erase(it);
+                    } else {
+                        ++it;
+                    }
+                }
+            }
+
+        private:
+            mutable std::shared_mutex mMutex;
+            std::unordered_map<uint64_t, std::shared_ptr<SessionState>> mSessions;
+    };
+}

include/columnlynx/common/net/udp/udp_message_type.hpp

@@ -0,0 +1,33 @@
+// udp_message_type.hpp - UDP Message Types for ColumnLynx
+// Copyright (C) 2025 DcruBro
+// Distributed under the GPLv3 license. See LICENSE for details.
+
+#pragma once
+
+#include <cstdint>
+#include <variant>
+#include <array>
+
+namespace ColumnLynx::Net::UDP {
+    // Shared between server and client
+    enum class MessageType : uint8_t {
+        PING = 0x01,
+        PONG = 0x02,
+        DATA = 0x03
+    };
+
+    struct UDPPacketHeader {
+        std::array<uint8_t, 12> nonce;
+    };
+
+    /*enum class ServerMessageType : uint8_t { // Server to Client
+        
+    };
+
+    enum class ClientMessageType : uint8_t { // Client to Server
+
+    };
+
+    // Make a variant type for either message type
+    using AnyMessageType = std::variant<ServerMessageType, ClientMessageType>;*/
+}

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

@@ -15,6 +15,7 @@
 #include <columnlynx/common/net/tcp/tcp_message_handler.hpp>
 #include <columnlynx/common/utils.hpp>
 #include <columnlynx/common/libsodium_wrapper.hpp>
+#include <columnlynx/common/net/session_registry.hpp>
 
 namespace ColumnLynx::Net::TCP {
     class TCPConnection : public std::enable_shared_from_this<TCPConnection> {
@@ -71,6 +72,14 @@ namespace ColumnLynx::Net::TCP {
                 }
             }
 
+            uint64_t getSessionID() const {
+                return mConnectionSessionID;
+            }
+
+            std::array<uint8_t, 32> getAESKey() const {
+                return mConnectionAESKey;
+            }
+        
         private:
             TCPConnection(asio::ip::tcp::socket socket, Utils::LibSodiumWrapper* sodiumWrapper)
                 : mHandler(std::make_shared<MessageHandler>(std::move(socket))), mLibSodiumWrapper(sodiumWrapper) {}
@@ -140,6 +149,8 @@ namespace ColumnLynx::Net::TCP {
                             // Make a Session ID
                             randombytes_buf(&mConnectionSessionID, sizeof(mConnectionSessionID));
 
+                            // TODO: Make the session ID little-endian for network transmission
+
                             // Encrypt the Session ID with the established AES key (using symmetric encryption, nonce can be all zeros for this purpose)
                             Nonce symNonce{}; // All zeros
                             std::vector<uint8_t> encryptedSessionID = Utils::LibSodiumWrapper::encryptMessage(
@@ -149,6 +160,11 @@ namespace ColumnLynx::Net::TCP {
 
                             mHandler->sendMessage(ServerMessageType::HANDSHAKE_EXCHANGE_KEY_CONFIRM, Utils::uint8ArrayToString(encryptedSessionID.data(), encryptedSessionID.size()));
 
+                            // Add to session registry
+                            Utils::log("Handshake with " + reqAddr + " completed successfully. Session ID assigned.");
+                            auto session = std::make_shared<SessionState>(mConnectionAESKey, std::chrono::hours(12));
+                            SessionRegistry::getInstance().put(mConnectionSessionID, std::move(session));
+
                         } catch (const std::exception& e) {
                             Utils::error("Failed to decrypt HANDSHAKE_EXCHANGE_KEY from " + reqAddr + ": " + e.what());
                             disconnect();

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

@@ -0,0 +1,30 @@
+// udp_server.hpp - UDP Server for ColumnLynx
+// Copyright (C) 2025 DcruBro
+// Distributed under the GPLv3 license. See LICENSE for details.
+
+#pragma once
+
+#include <asio/asio.hpp>
+#include <columnlynx/common/net/udp/udp_message_type.hpp>
+#include <columnlynx/common/utils.hpp>
+#include <array>
+
+namespace ColumnLynx::Net::UDP {
+    class UDPServer {
+        public:
+            UDPServer(asio::io_context& ioContext, uint16_t port)
+                : mSocket(ioContext, asio::ip::udp::endpoint(asio::ip::udp::v4(), port))
+            {
+                Utils::log("Started UDP server on port " + std::to_string(port));
+                mStartReceive();
+            }
+
+        private:
+            void mStartReceive();
+            void mHandlePacket(std::size_t bytes);
+            void mSendData(const uint64_t sessionID, const std::string& data);
+            asio::ip::udp::socket mSocket;
+            asio::ip::udp::endpoint mRemoteEndpoint;
+            std::array<uint8_t, 2048> mRecvBuffer; // Adjust size as needed
+    };
+}

src/client/main.cpp

@@ -8,6 +8,7 @@
 #include <columnlynx/common/utils.hpp>
 #include <columnlynx/common/panic_handler.hpp>
 #include <columnlynx/client/net/tcp/tcp_client.hpp>
+#include <columnlynx/client/net/udp/udp_client.hpp>
 
 using asio::ip::tcp;
 using namespace ColumnLynx::Utils;
@@ -34,10 +35,15 @@ int main(int argc, char** argv) {
     try {
         LibSodiumWrapper sodiumWrapper = LibSodiumWrapper();
 
+        std::array<uint8_t, 32> aesKey = {0}; // Defualt zeroed state until modified by handshake
+        uint64_t sessionID = 0;
+
         asio::io_context io;
-        auto client = std::make_shared<ColumnLynx::Net::TCP::TCPClient>(io, "127.0.0.1", std::to_string(serverPort()), &sodiumWrapper);
+        auto client = std::make_shared<ColumnLynx::Net::TCP::TCPClient>(io, "127.0.0.1", std::to_string(serverPort()), &sodiumWrapper, &aesKey, &sessionID);
+        auto udpClient = std::make_shared<ColumnLynx::Net::UDP::UDPClient>(io, "127.0.0.1", std::to_string(serverPort()), &aesKey, &sessionID);
 
         client->start();
+        udpClient->start();
 
         // Run the IO context in a separate thread
         std::thread ioThread([&io]() {
@@ -50,6 +56,16 @@ int main(int argc, char** argv) {
         // Client is running
         while (!done) {
             std::this_thread::sleep_for(std::chrono::milliseconds(100)); // Temp wait
+
+            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();
+                if (std::chrono::duration_cast<std::chrono::seconds>(now - lastSendTime).count() >= 5) {
+                    udpClient->sendMessage("Hello from UDP client!");
+                    lastSendTime = now;
+                }
+            }
         }
         log("Client shutting down.");
         client->disconnect();

src/server/main.cpp

@@ -7,11 +7,14 @@
 #include <columnlynx/common/utils.hpp>
 #include <columnlynx/common/panic_handler.hpp>
 #include <columnlynx/server/net/tcp/tcp_server.hpp>
+#include <columnlynx/server/net/udp/udp_server.hpp>
 #include <columnlynx/common/libsodium_wrapper.hpp>
+#include <unordered_set>
 
 using asio::ip::tcp;
 using namespace ColumnLynx::Utils;
 using namespace ColumnLynx::Net::TCP;
+using namespace ColumnLynx::Net::UDP;
 
 int main(int argc, char** argv) {
     PanicHandler::init();
@@ -26,6 +29,7 @@ int main(int argc, char** argv) {
 
         asio::io_context io;
         auto server = std::make_shared<TCPServer>(io, serverPort(), &sodiumWrapper);
+        auto udpServer = std::make_shared<UDPServer>(io, serverPort());
 
         // Run the IO context in a separate thread
         std::thread ioThread([&io]() {

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

@@ -0,0 +1,72 @@
+// udp_server.cpp - UDP Server for ColumnLynx
+// Copyright (C) 2025 DcruBro
+// Distributed under the GPLv3 license. See LICENSE for details.
+
+#include <columnlynx/server/net/udp/udp_server.hpp>
+#include <columnlynx/common/libsodium_wrapper.hpp>
+#include <columnlynx/common/net/session_registry.hpp>
+#include <sodium.h>
+#include <memory>
+
+namespace ColumnLynx::Net::UDP {
+    void UDPServer::mStartReceive() {
+        mSocket.async_receive_from(
+            asio::buffer(mRecvBuffer), mRemoteEndpoint,
+            [this](asio::error_code ec, std::size_t bytes) {
+                if (!ec && bytes > 0) {
+                    mHandlePacket(bytes);
+                }
+                mStartReceive(); // Continue receiving
+            }
+        );
+    }
+
+    void UDPServer::mHandlePacket(std::size_t bytes) {
+        if (bytes < sizeof(UDPPacketHeader))
+            return;
+        
+        const auto* hdr = reinterpret_cast<UDPPacketHeader*>(mRecvBuffer.data());
+
+        // Get plaintext session ID (assuming first 8 bytes after nonce (header))
+        uint64_t sessionID = 0;
+        std::memcpy(&sessionID, mRecvBuffer.data() + sizeof(UDPPacketHeader), sizeof(uint64_t));
+
+        auto it = mRecvBuffer.begin() + sizeof(UDPPacketHeader) + sizeof(uint64_t);
+        std::vector<uint8_t> encryptedPayload(it, mRecvBuffer.begin() + bytes);
+
+        // Get associated session state
+        std::shared_ptr<const SessionState> session = SessionRegistry::getInstance().get(sessionID);
+
+        if (!session) {
+            Utils::warn("UDP: Unknown or invalid session from " + mRemoteEndpoint.address().to_string());
+            return;
+        }
+
+        // Decrypt the actual payload
+        try {
+            auto plaintext = Utils::LibSodiumWrapper::decryptMessage(
+                encryptedPayload.data(), encryptedPayload.size(),
+                session->aesKey,
+                hdr->nonce,
+                "udp-data"
+            );
+
+            const_cast<SessionState*>(session.get())->setUDPEndpoint(mRemoteEndpoint); // Update endpoint after confirming decryption
+            // Update recv counter
+            const_cast<SessionState*>(session.get())->recv_ctr.fetch_add(1, std::memory_order_relaxed);
+
+            // TODO: Process the packet payload
+
+            // For now, just log the decrypted payload
+            std::string payloadStr(plaintext.begin(), plaintext.end());
+            Utils::log("UDP: Received packet from " + mRemoteEndpoint.address().to_string() + " - Payload: " + payloadStr);
+        } catch (...) {
+            Utils::warn("UDP: Failed to decrypt payload from " + mRemoteEndpoint.address().to_string());
+            return;
+        }
+    }
+
+    void UDPServer::mSendData(const uint64_t sessionID, const std::string& data) {
+        // TODO: Implement
+    }
+}