Context fill-in and CI tests

This commit adds common units tests and CI sanitasion.
Additional context for commit b64d9c4498:
 - Fixed macOS/Linux non-portable and unsafe shell usage by adding a posix_spawn helper and replacing system() calls in virtual_interface.cpp.
 - Fixed SessionRegistry::erase() to remove mIPSessions and mSessionIPs entries in session_registry.cpp.
 - Prevented message-length truncation in tcp_message_handler.cpp by rejecting payloads > 65535 bytes.
 - Validated handshake message sizes and removed silent truncation in:
  - tcp_connection.cpp
  - tcp_client.cpp
 - Canonicalized and validated config and whitelist paths in utils.cpp using std::filesystem.
 - Hardened environment-provided config path handling in main.cpp.
 - Validated UDP ciphertext lengths and fixed session ID endianness in udp_client.cpp.
 - Scheduled periodic SessionRegistry::cleanupExpired() in main.cpp (every 5 minutes).

.github/workflows/sanitizers.yml

@@ -0,0 +1,39 @@
+name: Sanitizers
+
+on:
+  push:
+    branches: [ main, master ]
+  pull_request:
+    branches: [ main, master ]
+
+jobs:
+  build-and-test:
+    runs-on: ubuntu-latest
+    env:
+      SANITIZERS: "-fsanitize=address,undefined -fno-omit-frame-pointer -g -O1"
+      ASAN_OPTIONS: "detect_leaks=1:abort_on_error=1"
+      UBSAN_OPTIONS: "print_stacktrace=1"
+    steps:
+      - name: Checkout
+        uses: actions/checkout@v4
+
+      - name: Install dependencies
+        run: |
+          sudo apt-get update
+          sudo apt-get install -y cmake build-essential clang
+
+      - name: Configure (CMake)
+        run: |
+          mkdir -p build-sanitizers
+          cd build-sanitizers
+          CC=clang CXX=clang++ cmake -DCMAKE_BUILD_TYPE=Debug -DCMAKE_CXX_FLAGS="$SANITIZERS" -DCMAKE_EXE_LINKER_FLAGS="$SANITIZERS" ..
+
+      - name: Build
+        run: |
+          cd build-sanitizers
+          cmake --build . -- -j
+
+      - name: Run tests
+        run: |
+          cd build-sanitizers
+          ctest --output-on-failure || (echo "ctest failed"; exit 1)

.gitignore

@@ -12,5 +12,6 @@ _deps
 CMakeUserPresets.json
 
 build/
+build*/
 .vscode/
 .DS_Store

CMakeLists.txt

@@ -173,3 +173,28 @@ install(FILES
 	LICENSES/GPL-2.0-only.txt
 	LICENSES/GPL-3.0.txt
         DESTINATION ${CMAKE_INSTALL_DATAROOTDIR}/licenses/${PROJECT_NAME})
+
+# ---------------------------------------------------------
+# Unit tests
+# ---------------------------------------------------------
+option(BUILD_TESTS "Build unit tests" ON)
+
+if(BUILD_TESTS)
+    enable_testing()
+    file(GLOB_RECURSE TEST_SRC CONFIGURE_DEPENDS tests/*.cpp)
+    if(TEST_SRC)
+        foreach(TEST_FILE IN LISTS TEST_SRC)
+            get_filename_component(TEST_NAME ${TEST_FILE} NAME_WE)
+            add_executable(${TEST_NAME} ${TEST_FILE})
+            target_link_libraries(${TEST_NAME} PRIVATE common sodium)
+            target_include_directories(${TEST_NAME} PRIVATE 
+                ${PROJECT_SOURCE_DIR}/include
+                ${asio_SOURCE_DIR}/asio/include
+                ${sodium_SOURCE_DIR}/src/libsodium/include
+                ${sodium_BINARY_DIR}/src/libsodium/include
+            )
+            add_test(NAME ${TEST_NAME} COMMAND ${TEST_NAME})
+            set_target_properties(${TEST_NAME} PROPERTIES OUTPUT_NAME "columnlynx_${TEST_NAME}")
+        endforeach()
+    endif()
+endif()

include/columnlynx/client/client_session.hpp

@@ -87,6 +87,10 @@ namespace ColumnLynx {
 
             void incrementSendCount() {
                 std::unique_lock lock(mMutex);
+                if (mClientState->send_cnt == std::numeric_limits<uint64_t>::max()) {
+                    Utils::error("ClientSession: send counter overflow detected");
+                    return;
+                }
                 mClientState->send_cnt++;
             }
 

include/columnlynx/common/utils.hpp

@@ -15,7 +15,7 @@
 #include <unordered_map>
 #include <unordered_set>
 #include <algorithm>
-#include <bits/stdc++.h>
+
 
 #ifdef _WIN32
     #include <winsock2.h>

src/client/main.cpp

@@ -5,6 +5,7 @@
 #include <asio.hpp>
 #include <csignal>
 #include <iostream>
+#include <filesystem>
 #include <columnlynx/common/utils.hpp>
 #include <columnlynx/common/panic_handler.hpp>
 #include <columnlynx/client/net/tcp/tcp_client.hpp>
@@ -90,7 +91,20 @@ int main(int argc, char** argv) {
         std::string configPath = optionsObj["config-dir"].as<std::string>();
         const char* envConfigPath = std::getenv("COLUMNLYNX_CONFIG_DIR");
         if (envConfigPath != nullptr) {
-            configPath = std::string(envConfigPath);
+            // Validate and canonicalize environment-provided path
+            try {
+                namespace fs = std::filesystem;
+                std::error_code ec;
+                fs::path candidate(envConfigPath);
+                fs::path abs = fs::absolute(candidate, ec);
+                if (!ec) {
+                    configPath = abs.string();
+                } else {
+                    warn(std::string("Invalid COLUMNLYNX_CONFIG_DIR value: ") + envConfigPath + " - using default");
+                }
+            } catch (const std::exception& e) {
+                warn(std::string("Failed to canonicalize COLUMNLYNX_CONFIG_DIR: ") + e.what());
+            }
         }
 
         if (configPath.back() != '/' && configPath.back() != '\\') {

src/client/net/tcp/tcp_client.cpp

@@ -159,7 +159,13 @@ namespace ColumnLynx::Net::TCP {
     void TCPClient::mHandleMessage(ServerMessageType type, const std::string& data) {
         switch (type) {
             case ServerMessageType::HANDSHAKE_IDENTIFY: {
-                    std::memcpy(mServerPublicKey, data.data(), std::min(data.size(), sizeof(mServerPublicKey)));
+                    if (data.size() != sizeof(mServerPublicKey)) {
+                        Utils::warn("HANDSHAKE_IDENTIFY has invalid size: " + std::to_string(data.size()));
+                        disconnect();
+                        return;
+                    }
+
+                    std::memcpy(mServerPublicKey, data.data(), sizeof(mServerPublicKey));
                     std::string hexServerPub = Utils::bytesToHexString(mServerPublicKey, 32);
                     Utils::log("Received server identity. Public Key: " + hexServerPub);
 
@@ -188,7 +194,13 @@ namespace ColumnLynx::Net::TCP {
                 {
                     // Verify the signature
                     Signature sig{};
-                    std::memcpy(sig.data(), data.data(), std::min(data.size(), sig.size()));
+                    if (data.size() != sig.size()) {
+                        Utils::warn("HANDSHAKE_CHALLENGE_RESPONSE has invalid size: " + std::to_string(data.size()));
+                        disconnect();
+                        return;
+                    }
+
+                    std::memcpy(sig.data(), data.data(), sig.size());
                     if (Utils::LibSodiumWrapper::verifyMessage(mSubmittedChallenge.data(), mSubmittedChallenge.size(), sig, mServerPublicKey)) {
                         Utils::log("Challenge response verified successfully.");
                         

src/client/net/udp/udp_client.cpp

@@ -3,6 +3,8 @@
 // Distributed under the terms of the GNU General Public License, either version 2 only or version 3. See LICENSES/ for details.
 
 #include <columnlynx/client/net/udp/udp_client.hpp>
+#include <thread>
+#include <chrono>
 
 namespace ColumnLynx::Net::UDP {
     void UDPClient::start() {
@@ -73,7 +75,7 @@ namespace ColumnLynx::Net::UDP {
             reinterpret_cast<uint8_t*>(&hdr) + sizeof(UDPPacketHeader)
         );
         uint32_t sessionID = static_cast<uint32_t>(ClientSession::getInstance().getSessionID());
-        uint32_t sessionIDNet = sessionID;
+        uint32_t sessionIDNet = htonl(sessionID);
         packet.insert(packet.end(),
             reinterpret_cast<uint8_t*>(&sessionIDNet),
             reinterpret_cast<uint8_t*>(&sessionIDNet) + sizeof(uint32_t)
@@ -102,7 +104,12 @@ namespace ColumnLynx::Net::UDP {
                 if (ec) {
                     if (ec == asio::error::operation_aborted) return; // Socket closed
                     // Other recv error
-                    mStartReceive();
+                    Utils::warn("UDPClient receive error: " + ec.message());
+                    // Back off briefly before restarting receive to avoid busy error loops
+                    asio::post(mSocket.get_executor(), [this]() {
+                        std::this_thread::sleep_for(std::chrono::milliseconds(50));
+                        mStartReceive();
+                    });
                     return;
                 }
 
@@ -136,9 +143,24 @@ namespace ColumnLynx::Net::UDP {
         }
 
         // Decrypt payload
+        // Extract ciphertext safely
+        size_t headerLen = sizeof(UDPPacketHeader) + sizeof(uint32_t);
+        if (bytes < headerLen) {
+            Utils::warn("UDP Client received packet too small after header check.");
+            return;
+        }
+
+        size_t ciphertextLen = bytes - headerLen;
+        // Enforce reasonable maximum (UDP payload practical limit)
+        const size_t MAX_UDP_PAYLOAD = 65507; // 65535 - UDP/IP headers
+        if (ciphertextLen > MAX_UDP_PAYLOAD) {
+            Utils::warn("UDP Client received packet with excessive payload size: " + std::to_string(ciphertextLen));
+            return;
+        }
+
         std::vector<uint8_t> ciphertext(
-            mRecvBuffer.begin() + sizeof(UDPPacketHeader) + sizeof(uint32_t),
-            mRecvBuffer.begin() + bytes
+            mRecvBuffer.begin() + headerLen,
+            mRecvBuffer.begin() + headerLen + ciphertextLen
         );
 
         if (ClientSession::getInstance().getAESKey().empty()) {

src/common/session_registry.cpp

@@ -32,7 +32,26 @@ namespace ColumnLynx::Net {
 
     void SessionRegistry::erase(uint32_t sessionID) {
         std::unique_lock lock(mMutex);
-        mSessions.erase(sessionID);
+        auto it = mSessions.find(sessionID);
+        if (it != mSessions.end()) {
+            // If the session has a client IP mapping, remove it to avoid stale entries
+            if (it->second) {
+                uint32_t ip = it->second->clientTunIP;
+                auto ipIt = mIPSessions.find(ip);
+                if (ipIt != mIPSessions.end()) {
+                    // Only erase if it points to the same session
+                    if (ipIt->second == it->second) {
+                        mIPSessions.erase(ipIt);
+                    }
+                }
+            }
+
+            // Remove any session->ip bookkeeping
+            mSessionIPs.erase(sessionID);
+
+            // Finally erase the session
+            mSessions.erase(it);
+        }
     }
 
     void SessionRegistry::cleanupExpired() {

src/common/tcp_message_handler.cpp

@@ -20,10 +20,16 @@ namespace ColumnLynx::Net::TCP {
 
         auto data = std::make_shared<std::vector<uint8_t>>();
         data->push_back(typeByte);
-        uint16_t length = payload.size();
+        // Ensure payload fits into protocol's 16-bit length field
+        if (payload.size() > static_cast<size_t>(std::numeric_limits<uint16_t>::max())) {
+            Utils::error("sendMessage(): payload too large (>65535 bytes)");
+            return;
+        }
 
-        data->push_back(length >> 8);
-        data->push_back(length & 0xFF);
+        uint16_t length = static_cast<uint16_t>(payload.size());
+
+        data->push_back(static_cast<uint8_t>(length >> 8));
+        data->push_back(static_cast<uint8_t>(length & 0xFF));
 
         data->insert(data->end(), payload.begin(), payload.end());
         auto self = shared_from_this();

src/common/utils.cpp

@@ -3,6 +3,7 @@
 // Distributed under the terms of the GNU General Public License, either version 2 only or version 3. See LICENSES/ for details.
 
 #include <columnlynx/common/utils.hpp>
+#include <filesystem>
 
 namespace ColumnLynx::Utils {
     std::string unixMillisToISO8601(uint64_t unixMillis, bool local) {
@@ -144,19 +145,49 @@ namespace ColumnLynx::Utils {
 
         std::vector<std::string> out;
 
-        std::ifstream file(basePath + "whitelisted_keys");
+        namespace fs = std::filesystem;
+        std::error_code ec;
+
+        fs::path base(basePath);
+        fs::path absBase = fs::absolute(base, ec);
+        if (ec) {
+            warn("getWhitelistedKeys(): failed to resolve base path: " + basePath + " - " + ec.message());
+            return out;
+        }
+
+        fs::path whitelist = absBase / "whitelisted_keys";
+        if (!fs::exists(whitelist, ec) || ec) {
+            warn("getWhitelistedKeys(): whitelist file not found: " + whitelist.string());
+            return out;
+        }
+
+        // Canonicalize to avoid symlink tricks
+        fs::path canon = fs::canonical(whitelist, ec);
+        if (ec) {
+            warn("getWhitelistedKeys(): failed to canonicalize path: " + whitelist.string());
+            return out;
+        }
+
+        std::ifstream file(canon);
         if (!file.is_open()) {
-            warn("Failed to open whitelisted_keys file at path: " + basePath + "whitelisted_keys");
+            warn("getWhitelistedKeys(): failed to open whitelist file: " + canon.string());
             return out;
         }
 
         std::string line;
         while (std::getline(file, line)) {
+            // Trim whitespace
+            while (!line.empty() && isspace(static_cast<unsigned char>(line.back()))) line.pop_back();
+            size_t start = 0;
+            while (start < line.size() && isspace(static_cast<unsigned char>(line[start]))) ++start;
+            if (start >= line.size()) continue;
+            std::string key = line.substr(start);
+
             // Convert to upper case to align with the bytesToHexString() output
-            for (int i = 0; i < line.length(); i++) {
-                line[i] = toupper(line[i]);
+            for (size_t i = 0; i < key.length(); ++i) {
+                key[i] = static_cast<char>(toupper(static_cast<unsigned char>(key[i])));
             }
-            out.push_back(line);
+            out.push_back(key);
         }
 
         return out;
@@ -166,9 +197,26 @@ namespace ColumnLynx::Utils {
         // TODO: Currently re-reads every time.
         std::vector<std::string> readLines;
 
-        std::ifstream file(path);
+        namespace fs = std::filesystem;
+        std::error_code ec;
+        fs::path p(path);
+        fs::path abs = fs::absolute(p, ec);
+        if (ec) {
+            throw std::runtime_error("getConfigMap(): failed to resolve path: " + path + " - " + ec.message());
+        }
+
+        if (!fs::exists(abs, ec) || ec) {
+            throw std::runtime_error("getConfigMap(): config file does not exist: " + abs.string());
+        }
+
+        fs::path canon = fs::canonical(abs, ec);
+        if (ec) {
+            throw std::runtime_error("getConfigMap(): failed to canonicalize config path: " + abs.string());
+        }
+
+        std::ifstream file(canon);
         if (!file.is_open()) {
-            throw std::runtime_error("Failed to open config file at path: " + path);
+            throw std::runtime_error("Failed to open config file at path: " + canon.string());
         }
 
         std::string line;

src/common/virtual_interface.cpp

@@ -4,6 +4,11 @@
 
 #include <columnlynx/common/net/virtual_interface.hpp>
 
+#include <spawn.h>
+#include <sys/wait.h>
+
+extern char **environ;
+
 // This is all fucking voodoo dark magic.
 
 #if defined(_WIN32)
@@ -56,6 +61,33 @@ static void InitializeWintun()
 #endif // _WIN32
 
 namespace ColumnLynx::Net {
+
+// Run a command without invoking a shell. Arguments are passed directly
+// to the underlying process to avoid shell injection vulnerabilities.
+static bool runCommand(const std::vector<std::string>& args) {
+    if (args.empty()) return false;
+
+    std::vector<char*> argv;
+    argv.reserve(args.size() + 1);
+    for (const auto &s : args) {
+        argv.push_back(const_cast<char*>(s.c_str()));
+    }
+    argv.push_back(nullptr);
+
+    pid_t pid;
+    int rc = posix_spawnp(&pid, argv[0], nullptr, nullptr, argv.data(), environ);
+    if (rc != 0) {
+        return false;
+    }
+
+    int status = 0;
+    if (waitpid(pid, &status, 0) == -1) {
+        return false;
+    }
+
+    return WIFEXITED(status) && WEXITSTATUS(status) == 0;
+}
+
     // ------------------------------ Constructor ------------------------------
     VirtualInterface::VirtualInterface(const std::string& ifName)
         : mIfName(ifName), mFd(-1)
@@ -181,8 +213,8 @@ namespace ColumnLynx::Net {
         pfd.fd = mFd;
         pfd.events = POLLIN;
 
-        // timeout in ms; keep it small so shutdown is responsive
-        int ret = poll(&pfd, 1, 200);
+        // timeout in ms; keep it small so shutdown is responsive. Reduced for lower latency.
+        int ret = poll(&pfd, 1, 50);
 
         if (ret == 0) {
             // No data yet
@@ -307,25 +339,10 @@ namespace ColumnLynx::Net {
 
     void VirtualInterface::resetIP() {
     #if defined(__linux__)
-        char cmd[512];
-        snprintf(cmd, sizeof(cmd),
-                 "ip addr flush dev %s",
-                 mIfName.c_str()
-        );
-        system(cmd);
+        runCommand({"ip", "addr", "flush", "dev", mIfName});
     #elif defined(__APPLE__)
-        char cmd[512];
-        snprintf(cmd, sizeof(cmd),
-                 "ifconfig %s inet 0.0.0.0 delete",
-                 mIfName.c_str()
-        );
-        system(cmd);
-
-        snprintf(cmd, sizeof(cmd),
-                 "ifconfig %s inet6 :: delete",
-                 mIfName.c_str()
-        );
-        system(cmd);
+        runCommand({"ifconfig", mIfName, "inet", "0.0.0.0", "delete"});
+        runCommand({"ifconfig", mIfName, "inet6", "::", "delete"});
 
         // Wipe old routes
         //snprintf(cmd, sizeof(cmd),
@@ -357,26 +374,19 @@ namespace ColumnLynx::Net {
     bool VirtualInterface::mApplyLinuxIP(uint32_t clientIP, uint32_t serverIP,
                                         uint8_t prefixLen, uint16_t mtu)
     {
-        char cmd[512];
     
         std::string ipStr = ipv4ToString(clientIP);
         std::string peerStr = ipv4ToString(serverIP);
     
         // Wipe the current config
-        snprintf(cmd, sizeof(cmd),
-                 "ip addr flush dev %s",
-                 mIfName.c_str()
-        );
-        system(cmd);
+        runCommand({"ip", "addr", "flush", "dev", mIfName});
 
-        snprintf(cmd, sizeof(cmd),
-                 "ip addr add %s/%d peer %s dev %s",
-                 ipStr.c_str(), prefixLen, peerStr.c_str(), mIfName.c_str());
-        system(cmd);
-    
-        snprintf(cmd, sizeof(cmd),
-                 "ip link set dev %s up mtu %d", mIfName.c_str(), mtu);
-        system(cmd);
+        // Add address with peer
+        std::string addrArg = ipStr + "/" + std::to_string(prefixLen);
+        runCommand({"ip", "addr", "add", addrArg, "peer", peerStr, "dev", mIfName});
+
+        // Bring link up and set MTU
+        runCommand({"ip", "link", "set", "dev", mIfName, "up", "mtu", std::to_string(mtu)});
     
         return true;
     }
@@ -387,39 +397,23 @@ namespace ColumnLynx::Net {
     bool VirtualInterface::mApplyMacOSIP(uint32_t clientIP, uint32_t serverIP,
                                         uint8_t prefixLen, uint16_t mtu)
     {
-        char cmd[512];
     
         std::string ipStr = ipv4ToString(clientIP);
         std::string peerStr = ipv4ToString(serverIP);
         std::string prefixStr = ipv4ToString(prefixLengthToNetmask(prefixLen), false);
         Utils::debug("Prefix string: " + prefixStr);
     
-        // Reset
-        snprintf(cmd, sizeof(cmd),
-                 "ifconfig %s inet 0.0.0.0 delete",
-                mIfName.c_str()
-        );
-        system(cmd);
+        // Reset IPv4 and IPv6 addresses
+        runCommand({"ifconfig", mIfName, "inet", "0.0.0.0", "delete"});
+        runCommand({"ifconfig", mIfName, "inet6", "::", "delete"});
 
-        snprintf(cmd, sizeof(cmd),
-                 "ifconfig %s inet6 :: delete",
-                mIfName.c_str()
-        );
-        system(cmd);
+        // Set address and netmask
+        std::string netArg = ipStr + " " + peerStr; // ifconfig expects ip peer
+        runCommand({"ifconfig", mIfName, "inet", ipStr, peerStr, "mtu", std::to_string(mtu), "netmask", prefixStr, "up"});
 
-        // Set
-        snprintf(cmd, sizeof(cmd),
-                "ifconfig %s inet %s %s mtu %d netmask %s up",
-                 mIfName.c_str(), ipStr.c_str(), peerStr.c_str(), mtu, prefixStr.c_str());
-        system(cmd);
-
-        // Host bits are auto-normalized by the kernel on macOS, so we don't need to worry about them not being zeroed out.
-        snprintf(cmd, sizeof(cmd),
-                 "route -n add -net %s/%d -interface %s",
-                 ipStr.c_str(), prefixLen, mIfName.c_str());
-        system(cmd);
-
-        Utils::log("Executed command: " + std::string(cmd));
+        // Add route for the network
+        std::string networkArg = ipStr + "/" + std::to_string(prefixLen);
+        runCommand({"route", "-n", "add", "-net", networkArg, "-interface", mIfName});
     
         return true;
     }

src/server/main.cpp

@@ -145,6 +145,22 @@ int main(int argc, char** argv) {
         auto server = std::make_shared<TCPServer>(io, serverPort());
         auto udpServer = std::make_shared<UDPServer>(io, serverPort());
 
+        // Schedule periodic cleanup of expired sessions every 5 minutes
+        auto cleanupTimer = std::make_shared<asio::steady_timer>(io);
+        auto cleanupHandler = std::make_shared<std::function<void(const asio::error_code&)>>();
+        *cleanupHandler = [cleanupTimer, cleanupHandler](const asio::error_code& ec) {
+            if (ec == asio::error::operation_aborted) return; // Timer cancelled
+            try {
+                SessionRegistry::getInstance().cleanupExpired();
+            } catch (const std::exception& e) {
+                Utils::warn(std::string("SessionRegistry::cleanupExpired() threw: ") + e.what());
+            }
+            cleanupTimer->expires_after(std::chrono::minutes(5));
+            cleanupTimer->async_wait(*cleanupHandler);
+        };
+        cleanupTimer->expires_after(std::chrono::minutes(5));
+        cleanupTimer->async_wait(*cleanupHandler);
+
         asio::signal_set signals(io, SIGINT, SIGTERM);
         signals.async_wait([&](const std::error_code&, int) {
             log("Received termination signal. Shutting down server gracefully.");
@@ -153,6 +169,8 @@ int main(int argc, char** argv) {
                 ServerSession::getInstance().setHostRunning(false);
                 server->stop();
                 udpServer->stop();
+                // Cancel cleanup timer
+                cleanupTimer->cancel();
             });
         });
 
@@ -195,8 +213,13 @@ int main(int argc, char** argv) {
             // First, check if destination IP is a registered client (e.g., server responding to client or client-to-client)
             auto dstSession = SessionRegistry::getInstance().getByIP(dstIP);
             if (dstSession) {
-                // Destination is a registered client, forward to that client's session
-                udpServer->sendData(dstSession->sessionID, std::string(packet.begin(), packet.end()));
+                // Destination is a registered client, enforce MTU and forward to that client's session
+                const size_t MTU = 1420; // Enforce configured MTU; TODO: read from server config
+                if (packet.size() > MTU) {
+                    Utils::warn("TUN: Dropping oversized packet (" + std::to_string(packet.size()) + " > MTU " + std::to_string(MTU) + ")");
+                } else {
+                    udpServer->sendData(dstSession->sessionID, std::string(packet.begin(), packet.end()));
+                }
                 continue;
             }
 

src/server/net/tcp/tcp_connection.cpp

@@ -124,8 +124,8 @@ namespace ColumnLynx::Net::TCP {
             case ClientMessageType::HANDSHAKE_INIT: {
                 Utils::log("Received HANDSHAKE_INIT from " + reqAddr);
 
-                if (data.size() < 1 + crypto_box_PUBLICKEYBYTES) {
-                    Utils::warn("HANDSHAKE_INIT from " + reqAddr + " is too short.");
+                if (data.size() != 1 + crypto_sign_PUBLICKEYBYTES) {
+                    Utils::warn("HANDSHAKE_INIT from " + reqAddr + " has invalid size: " + std::to_string(data.size()));
                     disconnect();
                     return;
                 }
@@ -141,7 +141,7 @@ namespace ColumnLynx::Net::TCP {
                 Utils::log("Client protocol version " + std::to_string(clientProtoVer) + " accepted from " + reqAddr + ".");
 
                 PublicKey signPk;
-                std::memcpy(signPk.data(), data.data() + 1, std::min(data.size() - 1, sizeof(signPk)));
+                std::memcpy(signPk.data(), data.data() + 1, sizeof(signPk));
 
                 // We can safely store this without further checking, the client will need to send the encrypted AES key in a way where they must possess the corresponding private key anyways.
                 int r = crypto_sign_ed25519_pk_to_curve25519(mConnectionPublicKey.data(), signPk.data()); // Store the client's public encryption key key (for identification)
@@ -173,9 +173,15 @@ namespace ColumnLynx::Net::TCP {
             case ClientMessageType::HANDSHAKE_CHALLENGE: {
                 Utils::log("Received HANDSHAKE_CHALLENGE from " + reqAddr);
                 
-                // Convert to byte array
+                // Convert to byte array - require exact size
+                if (data.size() != 32) {
+                    Utils::warn("HANDSHAKE_CHALLENGE has invalid size: " + std::to_string(data.size()));
+                    disconnect();
+                    return;
+                }
+
                 uint8_t challengeData[32];
-                std::memcpy(challengeData, data.data(), std::min(data.size(), sizeof(challengeData)));
+                std::memcpy(challengeData, data.data(), sizeof(challengeData));
 
                 // Sign the challenge
                 Signature sig = Utils::LibSodiumWrapper::signMessage(

src/server/net/udp/udp_server.cpp

@@ -95,7 +95,23 @@ namespace ColumnLynx::Net::UDP {
         UDPPacketHeader hdr{};
         uint8_t nonce[12];
         uint32_t prefix = session->noncePrefix;
-        uint64_t sendCount = const_cast<SessionState*>(session.get())->send_ctr.fetch_add(1, std::memory_order_relaxed);
+            // Increment send counter with overflow protection
+            uint64_t sendCount = 0;
+            {
+                auto ptr = const_cast<SessionState*>(session.get());
+                uint64_t old = ptr->send_ctr.load(std::memory_order_relaxed);
+                for (;;) {
+                    if (old == std::numeric_limits<uint64_t>::max()) {
+                        Utils::error("UDP: send counter overflow for session " + std::to_string(sessionID));
+                        return;
+                    }
+                    if (ptr->send_ctr.compare_exchange_weak(old, old + 1, std::memory_order_relaxed)) {
+                        sendCount = old;
+                        break;
+                    }
+                    // old updated by compare_exchange_weak, loop
+                }
+            }
         memcpy(nonce, &prefix, sizeof(uint32_t)); // Prefix nonce
         memcpy(nonce + sizeof(uint32_t), &sendCount, sizeof(uint64_t)); // Use send count as nonce suffix to ensure uniqueness
         std::copy_n(nonce, 12, hdr.nonce.data());

tests/test_libsodium_wrapper.cpp

@@ -0,0 +1,42 @@
+// Tests for LibSodiumWrapper: random, symmetric encrypt/decrypt, sign/verify
+#include <iostream>
+#include <cassert>
+
+#include <columnlynx/common/libsodium_wrapper.hpp>
+
+int main() {
+    using namespace ColumnLynx::Utils;
+
+    // Random bytes uniqueness
+    auto a = LibSodiumWrapper::generateRandom256Bit();
+    auto b = LibSodiumWrapper::generateRandom256Bit();
+    assert(a != b && "generateRandom256Bit() should produce different outputs (very likely)");
+
+    // Symmetric encrypt/decrypt roundtrip
+    ColumnLynx::SymmetricKey key = {};
+    for (size_t i = 0; i < key.size(); ++i) key[i] = static_cast<uint8_t>(i);
+    auto nonce = LibSodiumWrapper::generateNonce();
+
+    std::string plaintext = "The quick brown fox jumps over the lazy dog";
+    auto ct = LibSodiumWrapper::encryptMessage(reinterpret_cast<const uint8_t*>(plaintext.data()), plaintext.size(), key, nonce, "aad");
+    auto pt = LibSodiumWrapper::decryptMessage(ct.data(), ct.size(), key, nonce, "aad");
+    std::string recovered(pt.begin(), pt.end());
+    assert(recovered == plaintext && "decrypt should recover original plaintext");
+
+    // Sign and verify
+    ColumnLynx::PrivateKey sk{}; ColumnLynx::PublicKey pk{};
+    randombytes_buf(sk.data(), sk.size());
+    // naive keypair generation for test purposes: use libsodium functions via wrapper
+    // generate a real keypair using crypto_sign
+    if (crypto_sign_keypair(pk.data(), sk.data()) != 0) {
+        std::cerr << "Failed to generate keypair\n";
+        return 2;
+    }
+
+    auto sig = LibSodiumWrapper::signMessage(plaintext, sk);
+    bool ok = LibSodiumWrapper::verifyMessage(plaintext, sig, pk);
+    assert(ok && "Signature should verify");
+
+    std::cout << "LibSodiumWrapper tests passed\n";
+    return 0;
+}

tests/test_session_registry.cpp

@@ -0,0 +1,49 @@
+// Simple unit tests for SessionRegistry
+#include <cassert>
+#include <iostream>
+#include <chrono>
+
+#include <columnlynx/common/net/session_registry.hpp>
+#include <columnlynx/common/libsodium_wrapper.hpp>
+
+int main() {
+    using namespace ColumnLynx::Net;
+    using namespace ColumnLynx::Utils;
+
+    auto &reg = SessionRegistry::getInstance();
+
+    // Use a unique session id to avoid colliding with any running instance
+    const uint32_t sid = 0xDEADBEEF;
+    // ensure clean state
+    reg.erase(sid);
+
+    auto key = LibSodiumWrapper::generateRandom256Bit();
+    auto state = std::make_shared<SessionState>(key, std::chrono::hours(24), 0xC0A80101 /*192.168.1.1*/, 0, sid);
+
+    reg.put(sid, state);
+
+    assert(reg.exists(sid) && "Session should exist after put()");
+
+    auto got = reg.get(sid);
+    assert(got && got->sessionID == sid && "get() should return stored session");
+
+    auto byip = reg.getByIP(0xC0A80101);
+    assert(byip && byip->sessionID == sid && "getByIP() should find session by client IP");
+
+    // Erase and verify removed
+    reg.erase(sid);
+    assert(!reg.exists(sid) && "Session should not exist after erase()");
+    assert(reg.getByIP(0xC0A80101) == nullptr && "getByIP() should return nullptr after erase");
+
+    // Test cleanupExpired: insert an already-expired session
+    const uint32_t sid2 = 0xFEEDBEEF;
+    reg.erase(sid2);
+    auto expiredState = std::make_shared<SessionState>(key, std::chrono::seconds(0), 0xC0A80102, 0, sid2);
+    reg.put(sid2, expiredState);
+    // Force cleanup
+    reg.cleanupExpired();
+    assert(!reg.exists(sid2) && "Expired session should be removed by cleanupExpired()");
+
+    std::cout << "SessionRegistry tests passed\n";
+    return 0;
+}

tests/test_session_registry_ip.cpp

@@ -0,0 +1,43 @@
+// Tests for SessionRegistry IP allocation and lock/dealloc
+#include <iostream>
+#include <cassert>
+
+#include <columnlynx/common/net/session_registry.hpp>
+#include <columnlynx/common/libsodium_wrapper.hpp>
+
+int main() {
+    using namespace ColumnLynx::Net;
+    using namespace ColumnLynx::Utils;
+
+    auto &reg = SessionRegistry::getInstance();
+
+    const uint32_t sid = 0xABCDEF01;
+    reg.erase(sid);
+
+    auto key = LibSodiumWrapper::generateRandom256Bit();
+    auto state = std::make_shared<SessionState>(key, std::chrono::hours(24), 0, 0, sid);
+    reg.put(sid, state);
+
+    // Lock IP
+    uint32_t ip = 0xC0A80201; // 192.168.2.1
+    reg.lockIP(sid, ip);
+
+    auto byip = reg.getByIP(ip);
+    assert(byip && byip->sessionID == sid && "lockIP should populate mIPSessions");
+
+    // deallocIP
+    reg.deallocIP(sid);
+    assert(reg.getByIP(ip) == nullptr && "deallocIP should remove mapping");
+
+    // getFirstAvailableIP: choose a small /30 range to limit hosts
+    uint32_t base = 0x0A000000; // 10.0.0.0
+    uint8_t mask = 30; // 2 usable hosts
+    uint32_t first = reg.getFirstAvailableIP(base, mask);
+    assert(first != 0 && "Should find available IP in empty registry");
+
+    // cleanup
+    reg.erase(sid);
+
+    std::cout << "SessionRegistry IP tests passed\n";
+    return 0;
+}

tests/test_tcp_message_handler_static.cpp

@@ -0,0 +1,30 @@
+// Tests for TCP MessageHandler static helpers
+#include <cassert>
+#include <iostream>
+
+#include <columnlynx/common/net/tcp/tcp_message_handler.hpp>
+#include <columnlynx/common/net/tcp/tcp_message_type.hpp>
+
+int main() {
+    using namespace ColumnLynx::Net::TCP;
+
+    // server message special codes
+    auto t1 = MessageHandler::decodeMessageType(0xFE);
+    // Expect GRACEFUL_DISCONNECT mapped
+    // Compare by converting back to uint8
+    assert(MessageHandler::toUint8(t1) == 0xFE);
+
+    auto t2 = MessageHandler::decodeMessageType(0xFF);
+    assert(MessageHandler::toUint8(t2) == 0xFF);
+
+    // Client message range (>= 0xA0)
+    auto t3 = MessageHandler::decodeMessageType(0xA5);
+    assert(MessageHandler::toUint8(t3) == 0xA5);
+
+    // Server message range (< 0xA0) and not special
+    auto t4 = MessageHandler::decodeMessageType(0x10);
+    assert(MessageHandler::toUint8(t4) == 0x10);
+
+    std::cout << "TCP MessageHandler static helpers tests passed\n";
+    return 0;
+}