quic_socket.h

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.
 
#pragma once
 
#include <asio.hpp>
#include <functional>
#include <memory>
#include <vector>
#include <array>
#include <system_error>
#include <mutex>
#include <atomic>
#include <map>
#include <deque>
 
#include "internal/utils/common_defs.h"
#include "internal/protocols/quic/crypto.h"
#include "internal/protocols/quic/frame.h"
#include "internal/protocols/quic/packet.h"
#include "kcenon/network/detail/utils/result_types.h"
 
namespace kcenon::network::internal
{
    enum class quic_role : uint8_t
    {
        client = 0,
        server = 1
    };
 
    enum class quic_connection_state : uint8_t
    {
        idle = 0,              
        handshake_start = 1,   
        handshake = 2,         
        connected = 3,         
        closing = 4,           
        draining = 5,          
        closed = 6             
    };
 
    class quic_socket : public std::enable_shared_from_this<quic_socket>
    {
    public:
        using stream_data_callback = std::function<void(
            uint64_t stream_id,
            std::span<const uint8_t> data,
            bool fin)>;
 
        using connected_callback = std::function<void()>;
 
        using error_callback = std::function<void(std::error_code)>;
 
        using close_callback = std::function<void(uint64_t error_code, const std::string& reason)>;
 
        quic_socket(asio::ip::udp::socket socket, quic_role role);
 
        ~quic_socket();
 
        // Non-copyable
        quic_socket(const quic_socket&) = delete;
        quic_socket& operator=(const quic_socket&) = delete;
 
        // Movable
        quic_socket(quic_socket&& other) noexcept;
        quic_socket& operator=(quic_socket&& other) noexcept;
 
        // =====================================================================
        // Callback Registration
        // =====================================================================
 
        auto set_stream_data_callback(stream_data_callback cb) -> void;
 
        auto set_connected_callback(connected_callback cb) -> void;
 
        auto set_error_callback(error_callback cb) -> void;
 
        auto set_close_callback(close_callback cb) -> void;
 
        // =====================================================================
        // Connection Management
        // =====================================================================
 
        [[nodiscard]] auto connect(const asio::ip::udp::endpoint& endpoint,
                                   const std::string& server_name = "") -> VoidResult;
 
        [[nodiscard]] auto accept(const std::string& cert_file,
                                  const std::string& key_file) -> VoidResult;
 
        [[nodiscard]] auto close(uint64_t error_code = 0,
                                 const std::string& reason = "") -> VoidResult;
 
        // =====================================================================
        // I/O Operations
        // =====================================================================
 
        auto start_receive() -> void;
 
        auto stop_receive() -> void;
 
        [[nodiscard]] auto send_stream_data(uint64_t stream_id,
                                            std::vector<uint8_t>&& data,
                                            bool fin = false) -> VoidResult;
 
        // =====================================================================
        // Stream Management
        // =====================================================================
 
        [[nodiscard]] auto create_stream(bool unidirectional = false) -> Result<uint64_t>;
 
        [[nodiscard]] auto close_stream(uint64_t stream_id) -> VoidResult;
 
        // =====================================================================
        // State Queries
        // =====================================================================
 
        [[nodiscard]] auto is_connected() const noexcept -> bool;
 
        [[nodiscard]] auto is_handshake_complete() const noexcept -> bool;
 
        [[nodiscard]] auto state() const noexcept -> quic_connection_state;
 
        [[nodiscard]] auto role() const noexcept -> quic_role;
 
        [[nodiscard]] auto remote_endpoint() const -> asio::ip::udp::endpoint;
 
        [[nodiscard]] auto local_connection_id() const
            -> const protocols::quic::connection_id&;
 
        [[nodiscard]] auto remote_connection_id() const
            -> const protocols::quic::connection_id&;
 
        // =====================================================================
        // Socket Access
        // =====================================================================
 
        auto socket() -> asio::ip::udp::socket& { return udp_socket_; }
 
        auto socket() const -> const asio::ip::udp::socket& { return udp_socket_; }
 
    private:
        // =====================================================================
        // Internal Methods
        // =====================================================================
 
        auto do_receive() -> void;
 
        auto handle_packet(std::span<const uint8_t> data) -> void;
 
        auto process_frame(const protocols::quic::frame& f) -> void;
 
        auto process_crypto_frame(const protocols::quic::crypto_frame& f) -> void;
 
        auto process_stream_frame(const protocols::quic::stream_frame& f) -> void;
 
        auto process_ack_frame(const protocols::quic::ack_frame& f) -> void;
 
        auto process_connection_close_frame(
            const protocols::quic::connection_close_frame& f) -> void;
 
        auto process_handshake_done_frame() -> void;
 
        auto send_pending_packets() -> void;
 
        [[nodiscard]] auto send_packet(
            protocols::quic::encryption_level level,
            std::vector<protocols::quic::frame>&& frames) -> VoidResult;
 
        auto queue_crypto_data(std::vector<uint8_t>&& data) -> void;
 
        [[nodiscard]] auto determine_encryption_level(
            const protocols::quic::packet_header& header) const noexcept
            -> protocols::quic::encryption_level;
 
        [[nodiscard]] auto generate_connection_id()
            -> protocols::quic::connection_id;
 
        auto on_retransmit_timeout() -> void;
 
        auto transition_state(quic_connection_state new_state) -> void;
 
        // =====================================================================
        // Member Variables
        // =====================================================================
 
        asio::ip::udp::socket udp_socket_;
 
        asio::ip::udp::endpoint remote_endpoint_;
 
        std::array<uint8_t, 65536> recv_buffer_;
 
        quic_role role_;
 
        std::atomic<quic_connection_state> state_{quic_connection_state::idle};
 
        protocols::quic::quic_crypto crypto_;
 
        protocols::quic::connection_id local_conn_id_;
 
        protocols::quic::connection_id remote_conn_id_;
 
        std::array<uint64_t, 4> next_packet_number_{0, 0, 0, 0};
 
        std::array<uint64_t, 4> largest_received_pn_{0, 0, 0, 0};
 
        uint64_t next_stream_id_{0};
 
        // =====================================================================
        // Pending Data Queues
        // =====================================================================
 
        std::array<std::deque<std::vector<uint8_t>>, 4> pending_crypto_data_;
 
        std::map<uint64_t, std::deque<std::pair<std::vector<uint8_t>, bool>>> pending_stream_data_;
 
        // =====================================================================
        // Callbacks
        // =====================================================================
 
        mutable std::mutex callback_mutex_;
 
        stream_data_callback stream_data_cb_;
 
        connected_callback connected_cb_;
 
        error_callback error_cb_;
 
        close_callback close_cb_;
 
        // =====================================================================
        // State Flags
        // =====================================================================
 
        std::atomic<bool> is_receiving_{false};
 
        std::atomic<bool> handshake_complete_{false};
 
        // =====================================================================
        // Timers
        // =====================================================================
 
        asio::steady_timer retransmit_timer_;
 
        asio::steady_timer idle_timer_;
 
        mutable std::mutex state_mutex_;
    };
 
} // namespace kcenon::network::internal