voice-cat/core/src/net/transport.h

/*
 * net/transport.h — TCP control channel + UDP media channel.
 *
 * Design: docs/architecture.md (Net thread), docs/protocol.md §1 (framing).
 * Implementation uses standalone Asio for sockets and timers.
 */
#ifndef VOICECAT_NET_TRANSPORT_H
#define VOICECAT_NET_TRANSPORT_H

#include <cstdint>
#include <string>

#define ASIO_STANDALONE 1
#include <asio.hpp>

#include <atomic>
#include <deque>
#include <functional>
#include <memory>
#include <mutex>
#include <thread>
#include <vector>

#include "protocol/protocol.h"

// Forward-declare TlsContext so transport.h does not pull in mbedTLS headers.
namespace voicecat::crypto { class TlsContext; }

namespace voicecat::net {

// Callbacks delivered on the net thread. Callers must not block inside them.
struct TcpChannelCallbacks {
    std::function<void()>                          on_connected;
    std::function<void(std::error_code)>           on_connect_error;
    std::function<void(std::vector<uint8_t>)>      on_frame;   // one decoded frame payload
    std::function<void(std::error_code)>           on_error;
    std::function<void()>                          on_disconnected;
    // Called (on the handshake thread) right after TLS succeeds, before reads begin.
    // Use to export keying material while the handshake context is still fresh.
    std::function<void(voicecat::crypto::TlsContext&)> on_tls_ready;
};

// ── Client-side: owns an io_context + dedicated net thread ──────────────────
class TcpControlChannel {
   public:
    explicit TcpControlChannel(TcpChannelCallbacks cbs);
    ~TcpControlChannel();

    // Async connect; calls on_connected or on_connect_error on the net thread.
    void async_connect(const std::string& host, uint16_t port);

    // Queue a framed send (thread-safe; callable from any thread).
    void send_frame(std::vector<uint8_t> payload);

    // Graceful close; safe to call from any thread. Waits for the net thread to join.
    void close();

    bool connected() const { return connected_.load(std::memory_order_acquire); }

    // Access the io_context so callers can post work back to the net thread.
    asio::io_context& io() { return io_; }

   private:
    void run_loop();
    void start_read();
    void handle_length(std::error_code ec, std::size_t n);
    void handle_body(uint32_t length, std::error_code ec, std::size_t n);
    void do_send();

    asio::io_context                                           io_;
    asio::executor_work_guard<asio::io_context::executor_type> work_guard_;
    asio::ip::tcp::socket                                      socket_;
    asio::strand<asio::io_context::executor_type>              strand_;
    std::thread                                                net_thread_;

    TcpChannelCallbacks cbs_;
    protocol::FrameCodec codec_;

    uint8_t                         len_buf_[4]{};
    std::vector<uint8_t>            body_buf_;
    std::deque<std::vector<uint8_t>> send_queue_;
    bool                            sending_{false};
    std::atomic<bool>               connected_{false};
    std::atomic<bool>               closing_{false};
};

// ── Server-side: one per accepted socket, shares the server's io_context ────
class TcpServerConn : public std::enable_shared_from_this<TcpServerConn> {
   public:
    // Plain TCP constructor (no TLS — for tests or future plaintext paths).
    TcpServerConn(asio::ip::tcp::socket socket, TcpChannelCallbacks cbs);

    // TLS constructor: takes ownership of a TlsContext; start() will run the
    // handshake on a temporary thread then switch to a TLS I/O thread.
    TcpServerConn(asio::ip::tcp::socket socket, TcpChannelCallbacks cbs,
                  std::unique_ptr<voicecat::crypto::TlsContext> tls);

    ~TcpServerConn();

    // Begin reading; must be called once after construction (on the io thread).
    void start();

    // Thread-safe send (safe to call from the server's io thread or another strand).
    void send_frame(std::vector<uint8_t> payload);

    // Close the connection (safe from any thread).
    void close();

    // Block until the TLS I/O thread (if any) has finished. Must be called after close().
    // Safe to call from any thread except the TLS I/O thread itself.
    void wait_closed();

    bool connected() const { return connected_.load(std::memory_order_acquire); }

   private:
    // ── Asio path (no TLS) ───────────────────────────────────────────────────
    void start_read();
    void handle_length(std::error_code ec, std::size_t n);
    void handle_body(uint32_t length, std::error_code ec, std::size_t n);
    void do_send();

    // ── TLS path ─────────────────────────────────────────────────────────────
    void tls_read_loop();
    void tls_drain_sends();

    asio::ip::tcp::socket              socket_;
    asio::strand<asio::any_io_executor> strand_;
    TcpChannelCallbacks                cbs_;
    protocol::FrameCodec               codec_;

    uint8_t                           len_buf_[4]{};
    std::vector<uint8_t>              body_buf_;
    std::deque<std::vector<uint8_t>>  send_queue_;
    bool                              sending_{false};
    std::atomic<bool>                 connected_{false};
    std::atomic<bool>                 closing_{false};

    // TLS members (null in plain-TCP mode)
    std::unique_ptr<voicecat::crypto::TlsContext> tls_;
    std::thread                                    tls_thread_;
    std::mutex                                     tls_send_mutex_;
    std::deque<std::vector<uint8_t>>               tls_send_queue_;
};

// ── Server-side acceptor ─────────────────────────────────────────────────────
// Spawns a TcpServerConn (via factory) for each accepted TCP connection.
class TcpAcceptor {
   public:
    using ConnFactory = std::function<std::shared_ptr<TcpServerConn>(asio::ip::tcp::socket)>;

    TcpAcceptor(asio::io_context& io, uint16_t port, ConnFactory factory);

    // Start accepting. Call once; re-arms itself automatically.
    void start();

    // Stop accepting and close all tracked connections (safe while io_context is alive).
    void stop();

    // Stop accepting, close all connections, and block until every connection's
    // I/O thread has finished. Call BEFORE the io_context is destroyed — the TLS read
    // threads do blocking I/O (not async on the io_context) and will touch the reactor
    // on socket close if the io_context is already gone (manifests as a null-kqueue-reactor
    // segfault on macOS; latent on Windows IOCP / Linux epoll where timing is more forgiving).
    void shutdown();

    // Actual bound port (useful when bind_port=0 lets the OS pick).
    uint16_t local_port() const { return static_cast<uint16_t>(acceptor_.local_endpoint().port()); }

   private:
    void do_accept();

    asio::ip::tcp::acceptor             acceptor_;
    ConnFactory                         factory_;
    bool                                stopped_{false};

    std::mutex                                       conns_mu_;
    std::vector<std::shared_ptr<TcpServerConn>>      conns_;
};

// ── UDP media channel ────────────────────────────────────────────────────────
// Thin async UDP socket. send_to() is thread-safe. Recv callbacks fire on the
// io_context's thread (same thread that runs the io_context::run() loop).
class UdpMediaChannel {
   public:
    using FrameCallback =
        std::function<void(const uint8_t*, size_t, asio::ip::udp::endpoint)>;

    UdpMediaChannel() = default;
    ~UdpMediaChannel() { close(); }

    UdpMediaChannel(const UdpMediaChannel&)            = delete;
    UdpMediaChannel& operator=(const UdpMediaChannel&) = delete;

    // Bind to 0.0.0.0:port (0 = OS-assigned). Must be called before start_recv/send_to.
    bool bind(asio::io_context& io, uint16_t port = 0);

    // Begin the async recv loop. cb is called on the io_context thread.
    void start_recv(FrameCallback cb);

    // Thread-safe fire-and-forget send. Copies data into a heap buffer.
    void send_to(const uint8_t* data, size_t len, asio::ip::udp::endpoint dst);

    // Cancel all async ops and close the socket. Safe to call from any thread.
    void close();

    asio::ip::udp::endpoint local_endpoint() const;
    bool bound() const { return bound_.load(std::memory_order_acquire); }

   private:
    void do_recv();

    // Socket + recv state live here; only accessed from the io_context thread after bind().
    std::unique_ptr<asio::ip::udp::socket>  socket_;
    asio::ip::udp::endpoint                 sender_ep_;
    std::array<uint8_t, 1500>               recv_buf_{};
    FrameCallback                           frame_cb_;
    std::atomic<bool>                       bound_{false};
    std::atomic<bool>                       closed_{false};
};

}  // namespace voicecat::net

#endif  // VOICECAT_NET_TRANSPORT_H