NetUtils.cpp

#include "NetUtils.h"
#include <windows.h>
#include <bcrypt.h>
#include <cstring>
#include <memory>
#include <stdexcept>

bool sendAll(SOCKET socket, const char *data, int bytes)
{
    if (bytes < 0 || (!data && bytes != 0))
        return false;
    int sent = 0;
    while (sent < bytes)
    {
        const int result = send(socket, data + sent, bytes - sent, 0);
        if (result <= 0)
            return false;
        sent += result;
    }
    return true;
}
bool recvAll(SOCKET socket, char *data, int bytes)
{
    if (bytes < 0 || (!data && bytes != 0))
        return false;
    int received = 0;
    while (received < bytes)
    {
        const int result = recv(socket, data + received, bytes - received, 0);
        if (result <= 0)
            return false;
        received += result;
    }
    return true;
}
std::vector<unsigned char> generateRandomIV()
{
    std::vector<unsigned char> iv(16);
    if (BCryptGenRandom(nullptr, iv.data(), static_cast<ULONG>(iv.size()), BCRYPT_USE_SYSTEM_PREFERRED_RNG) <
        0)
        throw std::runtime_error("System random generator failed.");
    return iv;
}
bool sendPacket(SOCKET socket, const unsigned char *plaintext, std::size_t length, AES &aes,
                const unsigned char *key)
{
    if (length > MAX_PLAINTEXT || (!plaintext && length != 0) || !key)
        throw std::invalid_argument("Invalid or oversized plaintext.");
    const auto iv = generateRandomIV();
    unsigned encryptedLength = 0;
    const std::unique_ptr<unsigned char[]> encrypted(
        aes.EncryptCBC_PKCS7(plaintext, static_cast<unsigned>(length), key, iv.data(), encryptedLength));
    const auto bodyLength = 16 + encryptedLength;
    const std::uint32_t header = htonl(bodyLength);
    std::vector<char> buffer(4 + bodyLength);
    std::memcpy(buffer.data(), &header, 4);
    std::memcpy(buffer.data() + 4, iv.data(), 16);
    std::memcpy(buffer.data() + 20, encrypted.get(), encryptedLength);
    return sendAll(socket, buffer.data(), static_cast<int>(buffer.size()));
}
std::optional<std::vector<unsigned char>> receivePacket(SOCKET socket, AES &aes, const unsigned char *key)
{
    std::uint32_t header = 0;
    if (!recvAll(socket, reinterpret_cast<char *>(&header), 4))
        return std::nullopt;
    const auto length = ntohl(header);
    if (length < 32 || length > MAX_PACKET_BODY || (length - 16) % 16 != 0)
        throw std::runtime_error("Malformed encrypted packet length.");
    std::vector<unsigned char> body(length);
    if (!recvAll(socket, reinterpret_cast<char *>(body.data()), static_cast<int>(length)))
        return std::nullopt;
    unsigned plainLength = 0;
    const std::unique_ptr<unsigned char[]> plain(
        aes.DecryptCBC_PKCS7(body.data() + 16, length - 16, key, body.data(), plainLength));
    if (plainLength > MAX_PLAINTEXT)
        throw std::runtime_error("Decoded message exceeds limit.");
    return std::vector<unsigned char>(plain.get(), plain.get() + plainLength);
}
std::vector<unsigned char> recvPacket(SOCKET socket, AES &aes, const unsigned char *key)
{
    auto result = receivePacket(socket, aes, key);
    return result ? std::move(*result) : std::vector<unsigned char>{};
}