dtls_socket.cpp

Go to the documentation of this file.

// BSD 3-Clause License
// Copyright (c) 2024, 🍀☀🌕🌥 🌊
// See the LICENSE file in the project root for full license information.
 
#include "internal/tcp/dtls_socket.h"
#include "internal/utils/openssl_compat.h"
 
#include <cstring>
 
namespace kcenon::network::internal
{
 
namespace
{
    // Custom error category for SSL errors
    class ssl_error_category : public std::error_category
    {
    public:
        const char* name() const noexcept override { return "ssl"; }
 
        std::string message(int ev) const override
        {
            char buf[256];
            ERR_error_string_n(static_cast<unsigned long>(ev), buf, sizeof(buf));
            return buf;
        }
    };
 
    const ssl_error_category& get_ssl_category()
    {
        static ssl_error_category instance;
        return instance;
    }
 
    std::error_code make_ssl_error_code(int ssl_error)
    {
        return std::error_code(ssl_error, get_ssl_category());
    }
 
} // anonymous namespace
 
dtls_socket::dtls_socket(asio::ip::udp::socket socket, SSL_CTX* ssl_ctx)
    : socket_(std::move(socket))
    , ssl_ctx_(ssl_ctx)
    , ssl_(nullptr)
    , rbio_(nullptr)
    , wbio_(nullptr)
{
    // Create SSL object
    ssl_ = SSL_new(ssl_ctx_);
    if (!ssl_)
    {
        throw std::runtime_error("Failed to create SSL object");
    }
 
    // Create memory BIOs for non-blocking I/O
    rbio_ = BIO_new(BIO_s_mem());
    wbio_ = BIO_new(BIO_s_mem());
    if (!rbio_ || !wbio_)
    {
        if (rbio_) BIO_free(rbio_);
        if (wbio_) BIO_free(wbio_);
        SSL_free(ssl_);
        throw std::runtime_error("Failed to create BIO objects");
    }
 
    // Set BIOs to non-blocking mode
    BIO_set_nbio(rbio_, 1);
    BIO_set_nbio(wbio_, 1);
 
    // Connect BIOs to SSL object (SSL takes ownership)
    SSL_set_bio(ssl_, rbio_, wbio_);
 
    // Enable DTLS cookie exchange for servers (DoS protection)
    // This is optional and can be configured later
}
 
dtls_socket::~dtls_socket()
{
    stop_receive();
 
    if (ssl_)
    {
        SSL_shutdown(ssl_);
        SSL_free(ssl_);  // Also frees the BIOs
    }
}
 
auto dtls_socket::async_handshake(
    handshake_type type,
    std::function<void(std::error_code)> handler) -> void
{
    handshake_type_ = type;
    handshake_callback_ = std::move(handler);
    handshake_in_progress_.store(true);
 
    {
        std::lock_guard<std::mutex> lock(ssl_mutex_);
 
        if (type == handshake_type::client)
        {
            SSL_set_connect_state(ssl_);
        }
        else
        {
            SSL_set_accept_state(ssl_);
        }
    }
 
    // Start receiving for handshake packets
    start_receive();
 
    // Initiate handshake
    continue_handshake();
}
 
auto dtls_socket::continue_handshake() -> void
{
    if (!handshake_in_progress_.load())
    {
        return;
    }
 
    int result;
    {
        std::lock_guard<std::mutex> lock(ssl_mutex_);
        result = SSL_do_handshake(ssl_);
    }
 
    if (result == 1)
    {
        // Handshake complete
        handshake_in_progress_.store(false);
        handshake_complete_.store(true);
 
        // Flush any remaining output
        flush_bio_output();
 
        std::function<void(std::error_code)> callback;
        {
            std::lock_guard<std::mutex> lock(callback_mutex_);
            callback = std::move(handshake_callback_);
        }
 
        if (callback)
        {
            callback(std::error_code{});
        }
        return;
    }
 
    int ssl_error;
    {
        std::lock_guard<std::mutex> lock(ssl_mutex_);
        ssl_error = SSL_get_error(ssl_, result);
    }
 
    // Flush any output generated by the handshake
    flush_bio_output();
 
    if (ssl_error == SSL_ERROR_WANT_READ || ssl_error == SSL_ERROR_WANT_WRITE)
    {
        // Need more data - continue receiving
        return;
    }
 
    // Handshake failed
    handshake_in_progress_.store(false);
 
    std::function<void(std::error_code)> callback;
    {
        std::lock_guard<std::mutex> lock(callback_mutex_);
        callback = std::move(handshake_callback_);
    }
 
    if (callback)
    {
        callback(make_ssl_error());
    }
}
 
auto dtls_socket::set_receive_callback(
    std::function<void(const std::vector<uint8_t>&,
                       const asio::ip::udp::endpoint&)> callback) -> void
{
    std::lock_guard<std::mutex> lock(callback_mutex_);
    receive_callback_ = std::move(callback);
}
 
auto dtls_socket::set_error_callback(
    std::function<void(std::error_code)> callback) -> void
{
    std::lock_guard<std::mutex> lock(callback_mutex_);
    error_callback_ = std::move(callback);
}
 
auto dtls_socket::start_receive() -> void
{
    bool expected = false;
    if (is_receiving_.compare_exchange_strong(expected, true))
    {
        do_receive();
    }
}
 
auto dtls_socket::stop_receive() -> void
{
    is_receiving_.store(false);
}
 
auto dtls_socket::do_receive() -> void
{
    if (!is_receiving_.load())
    {
        return;
    }
 
    auto self = shared_from_this();
    socket_.async_receive_from(
        asio::buffer(read_buffer_),
        sender_endpoint_,
        [this, self](std::error_code ec, std::size_t length)
        {
            if (!is_receiving_.load())
            {
                return;
            }
 
            if (ec)
            {
                std::function<void(std::error_code)> callback;
                {
                    std::lock_guard<std::mutex> lock(callback_mutex_);
                    callback = error_callback_;
                }
                if (callback)
                {
                    callback(ec);
                }
                return;
            }
 
            if (length > 0)
            {
                std::vector<uint8_t> data(read_buffer_.begin(),
                                          read_buffer_.begin() + length);
                process_received_data(data, sender_endpoint_);
            }
 
            // Continue receiving
            if (is_receiving_.load())
            {
                do_receive();
            }
        });
}
 
auto dtls_socket::process_received_data(const std::vector<uint8_t>& data,
                                         const asio::ip::udp::endpoint& sender) -> void
{
    // Write received data to the read BIO
    {
        std::lock_guard<std::mutex> lock(ssl_mutex_);
        int written = BIO_write(rbio_, data.data(), static_cast<int>(data.size()));
        if (written <= 0)
        {
            // BIO write failed
            return;
        }
    }
 
    // If handshake is in progress, continue it
    if (handshake_in_progress_.load())
    {
        continue_handshake();
        return;
    }
 
    // Handshake complete, try to read decrypted data
    if (handshake_complete_.load())
    {
        std::vector<uint8_t> decrypted;
        decrypted.resize(65536);
 
        int read_len;
        {
            std::lock_guard<std::mutex> lock(ssl_mutex_);
            read_len = SSL_read(ssl_, decrypted.data(), static_cast<int>(decrypted.size()));
        }
 
        if (read_len > 0)
        {
            decrypted.resize(static_cast<std::size_t>(read_len));
 
            std::function<void(const std::vector<uint8_t>&, const asio::ip::udp::endpoint&)> callback;
            {
                std::lock_guard<std::mutex> lock(callback_mutex_);
                callback = receive_callback_;
            }
            if (callback)
            {
                callback(decrypted, sender);
            }
        }
        else
        {
            int ssl_error;
            {
                std::lock_guard<std::mutex> lock(ssl_mutex_);
                ssl_error = SSL_get_error(ssl_, read_len);
            }
 
            if (ssl_error != SSL_ERROR_WANT_READ && ssl_error != SSL_ERROR_WANT_WRITE)
            {
                // Real error
                std::function<void(std::error_code)> callback;
                {
                    std::lock_guard<std::mutex> lock(callback_mutex_);
                    callback = error_callback_;
                }
                if (callback)
                {
                    callback(make_ssl_error());
                }
            }
        }
    }
}
 
auto dtls_socket::flush_bio_output() -> void
{
    // Check if there's data to send in the write BIO
    std::vector<uint8_t> output;
    {
        std::lock_guard<std::mutex> lock(ssl_mutex_);
        int pending = BIO_ctrl_pending(wbio_);
        if (pending <= 0)
        {
            return;
        }
 
        output.resize(static_cast<std::size_t>(pending));
        int read_len = BIO_read(wbio_, output.data(), pending);
        if (read_len <= 0)
        {
            return;
        }
        output.resize(static_cast<std::size_t>(read_len));
    }
 
    // Send the encrypted data
    asio::ip::udp::endpoint target;
    {
        std::lock_guard<std::mutex> lock(endpoint_mutex_);
        target = peer_endpoint_;
    }
 
    if (target.port() != 0)
    {
        auto buffer = std::make_shared<std::vector<uint8_t>>(std::move(output));
        socket_.async_send_to(
            asio::buffer(*buffer),
            target,
            [buffer](std::error_code /*ec*/, std::size_t /*bytes*/)
            {
                // Fire and forget for handshake messages
            });
    }
}
 
auto dtls_socket::async_send(
    std::vector<uint8_t>&& data,
    std::function<void(std::error_code, std::size_t)> handler) -> void
{
    asio::ip::udp::endpoint target;
    {
        std::lock_guard<std::mutex> lock(endpoint_mutex_);
        target = peer_endpoint_;
    }
 
    async_send_to(std::move(data), target, std::move(handler));
}
 
auto dtls_socket::async_send_to(
    std::vector<uint8_t>&& data,
    const asio::ip::udp::endpoint& endpoint,
    std::function<void(std::error_code, std::size_t)> handler) -> void
{
    if (!handshake_complete_.load())
    {
        if (handler)
        {
            handler(std::make_error_code(std::errc::not_connected), 0);
        }
        return;
    }
 
    // Encrypt the data
    int written;
    {
        std::lock_guard<std::mutex> lock(ssl_mutex_);
        written = SSL_write(ssl_, data.data(), static_cast<int>(data.size()));
    }
 
    if (written <= 0)
    {
        if (handler)
        {
            handler(make_ssl_error(), 0);
        }
        return;
    }
 
    // Get encrypted data from write BIO
    std::vector<uint8_t> encrypted;
    {
        std::lock_guard<std::mutex> lock(ssl_mutex_);
        int pending = BIO_ctrl_pending(wbio_);
        if (pending <= 0)
        {
            if (handler)
            {
                handler(std::make_error_code(std::errc::io_error), 0);
            }
            return;
        }
 
        encrypted.resize(static_cast<std::size_t>(pending));
        int read_len = BIO_read(wbio_, encrypted.data(), pending);
        if (read_len <= 0)
        {
            if (handler)
            {
                handler(std::make_error_code(std::errc::io_error), 0);
            }
            return;
        }
        encrypted.resize(static_cast<std::size_t>(read_len));
    }
 
    // Send encrypted data
    auto self = shared_from_this();
    auto buffer = std::make_shared<std::vector<uint8_t>>(std::move(encrypted));
    auto original_size = data.size();
 
    socket_.async_send_to(
        asio::buffer(*buffer),
        endpoint,
        [handler = std::move(handler), buffer, original_size](
            std::error_code ec, std::size_t /*bytes_sent*/)
        {
            if (handler)
            {
                // Report original (plaintext) size to the handler
                handler(ec, ec ? 0 : original_size);
            }
        });
}
 
auto dtls_socket::set_peer_endpoint(const asio::ip::udp::endpoint& endpoint) -> void
{
    std::lock_guard<std::mutex> lock(endpoint_mutex_);
    peer_endpoint_ = endpoint;
}
 
auto dtls_socket::peer_endpoint() const -> asio::ip::udp::endpoint
{
    std::lock_guard<std::mutex> lock(const_cast<std::mutex&>(endpoint_mutex_));
    return peer_endpoint_;
}
 
auto dtls_socket::make_ssl_error() const -> std::error_code
{
    unsigned long err = ERR_get_error();
    if (err == 0)
    {
        return std::make_error_code(std::errc::io_error);
    }
    return make_ssl_error_code(static_cast<int>(err));
}
 
} // namespace kcenon::network::internal