Added basic UDP handling - only sending from client supported.

This commit is contained in:
2025-11-07 21:55:13 +01:00
parent c7c3b1c54c
commit 5127fe1999
10 changed files with 378 additions and 6 deletions

View File

@@ -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
};
}

View File

@@ -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;
};
}

View File

@@ -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;
};
}

View File

@@ -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>;*/
}

View File

@@ -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> {
@@ -70,6 +71,14 @@ namespace ColumnLynx::Net::TCP {
mOnDisconnect(shared_from_this());
}
}
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)
@@ -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();

View File

@@ -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
};
}