Line data    Source code

       1              : #ifndef THINGER_ASIO_WEBSOCKET_HPP
       2              : #define THINGER_ASIO_WEBSOCKET_HPP
       3              : #pragma once
       4              : 
       5              : #include "socket.hpp"
       6              : #include <queue>
       7              : #include <string_view>
       8              : 
       9              : namespace thinger::asio {
      10              : 
      11              : class websocket : public socket {
      12              : 
      13              : public:
      14              : 
      15              :     static constexpr auto CONNECTION_TIMEOUT_SECONDS = std::chrono::seconds{60};
      16              :     static constexpr int MASK_SIZE_BYTES = 4;
      17              :     static std::atomic<unsigned long> connections;
      18              : 
      19              :     websocket(std::shared_ptr<socket> socket, bool binary = true, bool server = true);
      20              :     virtual ~websocket();
      21              : 
      22              :     // Socket control
      23              :     awaitable<boost::system::error_code> connect(
      24              :         const std::string &host,
      25              :         const std::string &port,
      26              :         std::chrono::seconds timeout) override;
      27              :     void close() override;
      28              :     awaitable<void> close_graceful();
      29              :     void cancel() override;
      30              :     bool requires_handshake() const override;
      31              :     awaitable<boost::system::error_code> handshake(const std::string& host = "") override;
      32              : 
      33              :     // Read operations
      34              :     awaitable<size_t> read_some(uint8_t buffer[], size_t max_size) override;
      35              :     awaitable<size_t> read(uint8_t buffer[], size_t size) override;
      36              :     awaitable<size_t> read(boost::asio::streambuf &buffer, size_t size) override;
      37              :     awaitable<size_t> read_until(boost::asio::streambuf &buffer, std::string_view delim) override;
      38              : 
      39              :     // Write operations
      40              :     awaitable<size_t> write(const uint8_t buffer[], size_t size) override;
      41              :     awaitable<size_t> write(std::string_view str) override;
      42              :     awaitable<size_t> write(const std::vector<boost::asio::const_buffer> &buffers) override;
      43              : 
      44              :     // Wait
      45              :     awaitable<boost::system::error_code> wait(boost::asio::socket_base::wait_type type) override;
      46              : 
      47              :     // WebSocket-specific operations
      48              :     awaitable<void> send_ping(uint8_t buffer[] = nullptr, size_t size = 0);
      49              :     awaitable<void> send_pong(uint8_t buffer[] = nullptr, size_t size = 0);
      50              : 
      51              :     // State getters
      52              :     bool is_open() const override;
      53              :     bool is_secure() const override;
      54              :     size_t available() const override;
      55              :     std::string get_remote_ip() const override;
      56              :     std::string get_local_port() const override;
      57              :     std::string get_remote_port() const override;
      58              :     std::size_t remaining_in_frame() const;
      59              :     bool is_message_complete() const;
      60          106 :     bool is_binary() const { return message_opcode_ == 0x2; }
      61           74 :     void set_binary(bool binary) { binary_ = binary; }
      62              : 
      63              :     // Timeout management
      64              :     void start_timeout();
      65              : 
      66              : private:
      67              :     // Internal helpers
      68              :     void unmask(uint8_t buffer[], size_t size);
      69              :     awaitable<size_t> read_frame(uint8_t buffer[], size_t max_size, boost::system::error_code& ec);
      70              :     awaitable<size_t> send_message(uint8_t opcode, const uint8_t buffer[], size_t size);
      71              :     awaitable<void> send_close(uint8_t buffer[] = nullptr, size_t size = 0);
      72              : 
      73              :     std::shared_ptr<socket> socket_;
      74              :     boost::asio::steady_timer timer_;
      75              :     bool binary_;
      76              :     bool server_role_;
      77              : 
      78              :     // Message state
      79              :     bool new_message_ = true;
      80              :     uint8_t message_opcode_ = 0;
      81              :     std::size_t frame_remaining_ = 0;
      82              : 
      83              :     // Frame state
      84              :     uint8_t opcode_ = 0;
      85              :     bool fin_ = false;
      86              :     bool masked_ = false;
      87              : 
      88              :     // Buffers
      89              :     uint8_t buffer_[140];
      90              :     uint8_t output_[14];
      91              :     uint8_t mask_[MASK_SIZE_BYTES];
      92              : 
      93              :     // Connection state
      94              :     bool close_received_ = false;
      95              :     bool close_sent_ = false;
      96              :     bool data_received_ = true;
      97              :     bool pending_ping_ = false;
      98              : 
      99              :     // Write synchronization (for write ordering)
     100              :     std::mutex write_mutex_;
     101              : };
     102              : 
     103              : }
     104              : 
     105              : #endif