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

#include "net/transport.h"

#include <cstring>

#include "crypto/crypto.h"

namespace voicecat::net {

// ── TcpControlChannel ────────────────────────────────────────────────────────

TcpControlChannel::TcpControlChannel(TcpChannelCallbacks cbs)
    : work_guard_(asio::make_work_guard(io_)),
      socket_(io_),
      strand_(io_.get_executor()),
      cbs_(std::move(cbs)) {
    net_thread_ = std::thread([this] { run_loop(); });
}

TcpControlChannel::~TcpControlChannel() { close(); }

void TcpControlChannel::run_loop() { io_.run(); }

void TcpControlChannel::async_connect(const std::string& host, uint16_t port) {
    auto resolver = std::make_shared<asio::ip::tcp::resolver>(io_);
    resolver->async_resolve(
        host, std::to_string(port),
        [this, resolver](std::error_code ec, asio::ip::tcp::resolver::results_type eps) {
            if (ec) {
                if (cbs_.on_connect_error) cbs_.on_connect_error(ec);
                return;
            }
            asio::async_connect(socket_, eps,
                [this](std::error_code ec2, const asio::ip::tcp::endpoint&) {
                    if (ec2) {
                        if (cbs_.on_connect_error) cbs_.on_connect_error(ec2);
                        return;
                    }
                    connected_.store(true, std::memory_order_release);
                    if (cbs_.on_connected) cbs_.on_connected();
                    start_read();
                });
        });
}

void TcpControlChannel::send_frame(std::vector<uint8_t> payload) {
    std::vector<uint8_t> wire;
    protocol::FrameCodec::emit(payload, wire);
    asio::post(strand_, [this, w = std::move(wire)]() mutable {
        send_queue_.push_back(std::move(w));
        if (!sending_) do_send();
    });
}

void TcpControlChannel::do_send() {
    if (send_queue_.empty()) { sending_ = false; return; }
    sending_ = true;
    auto& front = send_queue_.front();
    asio::async_write(socket_,
        asio::buffer(front),
        asio::bind_executor(strand_,
            [this](std::error_code ec, std::size_t) {
                if (ec) {
                    connected_.store(false, std::memory_order_release);
                    if (cbs_.on_error) cbs_.on_error(ec);
                    return;
                }
                send_queue_.pop_front();
                do_send();
            }));
}

void TcpControlChannel::start_read() {
    asio::async_read(socket_, asio::buffer(len_buf_, 4),
        [this](std::error_code ec, std::size_t n) { handle_length(ec, n); });
}

void TcpControlChannel::handle_length(std::error_code ec, std::size_t) {
    if (ec) {
        connected_.store(false, std::memory_order_release);
        if (ec == asio::error::eof || ec == asio::error::connection_reset) {
            if (cbs_.on_disconnected) cbs_.on_disconnected();
        } else {
            if (cbs_.on_error) cbs_.on_error(ec);
        }
        return;
    }
    uint32_t length =
        (static_cast<uint32_t>(len_buf_[0]) << 24) |
        (static_cast<uint32_t>(len_buf_[1]) << 16) |
        (static_cast<uint32_t>(len_buf_[2]) <<  8) |
         static_cast<uint32_t>(len_buf_[3]);

    if (length > protocol::kMaxFrameBytes) {
        if (cbs_.on_error) cbs_.on_error(asio::error::message_size);
        return;
    }
    if (length == 0) {
        if (cbs_.on_frame) cbs_.on_frame({});
        start_read();
        return;
    }
    body_buf_.resize(length);
    asio::async_read(socket_, asio::buffer(body_buf_),
        [this, length](std::error_code ec, std::size_t n) { handle_body(length, ec, n); });
}

void TcpControlChannel::handle_body(uint32_t, std::error_code ec, std::size_t) {
    if (ec) {
        connected_.store(false, std::memory_order_release);
        if (ec == asio::error::eof || ec == asio::error::connection_reset) {
            if (cbs_.on_disconnected) cbs_.on_disconnected();
        } else {
            if (cbs_.on_error) cbs_.on_error(ec);
        }
        return;
    }
    if (cbs_.on_frame) cbs_.on_frame(body_buf_);
    start_read();
}

void TcpControlChannel::close() {
    if (closing_.exchange(true)) return;
    asio::post(io_, [this] {
        std::error_code ignored;
        socket_.shutdown(asio::ip::tcp::socket::shutdown_both, ignored);
        socket_.close(ignored);
    });
    work_guard_.reset();
    if (net_thread_.joinable()) net_thread_.join();
}

// ── TcpServerConn ────────────────────────────────────────────────────────────

TcpServerConn::TcpServerConn(asio::ip::tcp::socket socket, TcpChannelCallbacks cbs)
    : socket_(std::move(socket)),
      strand_(asio::make_strand(socket_.get_executor())),
      cbs_(std::move(cbs)) {}

TcpServerConn::TcpServerConn(asio::ip::tcp::socket socket, TcpChannelCallbacks cbs,
                              std::unique_ptr<crypto::TlsContext> tls)
    : socket_(std::move(socket)),
      strand_(asio::make_strand(socket_.get_executor())),
      cbs_(std::move(cbs)),
      tls_(std::move(tls)) {}

TcpServerConn::~TcpServerConn() {
    close();
    if (tls_thread_.joinable()) {
        if (std::this_thread::get_id() == tls_thread_.get_id()) {
            tls_thread_.detach();  // being destroyed from our own TLS thread — detach safely
        } else {
            tls_thread_.join();
        }
    }
}

void TcpServerConn::start() {
    if (tls_) {
        // Run TLS handshake on a temporary thread so we don't block the io_context.
        auto self = shared_from_this();
        std::thread([self] {
            std::string err;
            int fd = static_cast<int>(self->socket_.native_handle());
            if (!self->tls_->handshake(fd, err)) {
                if (!self->closing_.exchange(true)) {
                    if (self->cbs_.on_error) {
                        asio::post(self->strand_, [self] {
                            self->cbs_.on_error(
                                std::make_error_code(std::errc::connection_reset));
                        });
                    }
                }
                return;
            }
            self->connected_.store(true, std::memory_order_release);
            // Export media keying material before the read loop starts.
            if (self->cbs_.on_tls_ready) self->cbs_.on_tls_ready(*self->tls_);
            // 50 ms timeout so tls_read_loop can drain the send queue between reads.
            self->tls_->set_read_timeout(50);
            self->tls_thread_ = std::thread([self] { self->tls_read_loop(); });
        }).detach();
    } else {
        connected_.store(true, std::memory_order_release);
        start_read();
    }
}

void TcpServerConn::tls_read_loop() {
    std::vector<uint8_t> buf(16384);
    while (!closing_.load(std::memory_order_acquire)) {
        tls_drain_sends();

        int n = tls_->read(buf.data(), buf.size());
        if (crypto::TlsContext::is_timeout_error(n)) continue;
        if (n <= 0) break;

        std::vector<std::vector<uint8_t>> frames;
        if (!codec_.feed(buf.data(), static_cast<size_t>(n), frames)) break;
        for (auto& frame : frames) {
            if (cbs_.on_frame) cbs_.on_frame(std::move(frame));
        }
    }
    connected_.store(false, std::memory_order_release);
    if (cbs_.on_disconnected) cbs_.on_disconnected();
}

void TcpServerConn::tls_drain_sends() {
    while (true) {
        std::vector<uint8_t> frame;
        {
            std::lock_guard lk(tls_send_mutex_);
            if (tls_send_queue_.empty()) return;
            frame = std::move(tls_send_queue_.front());
            tls_send_queue_.pop_front();
        }
        size_t off = 0;
        while (off < frame.size()) {
            int n = tls_->write(frame.data() + off, frame.size() - off);
            if (n <= 0) { closing_.store(true, std::memory_order_release); return; }
            off += static_cast<size_t>(n);
        }
    }
}

void TcpServerConn::start_read() {
    auto self = shared_from_this();
    asio::async_read(socket_, asio::buffer(len_buf_, 4),
        asio::bind_executor(strand_,
            [this, self](std::error_code ec, std::size_t n) { handle_length(ec, n); }));
}

void TcpServerConn::handle_length(std::error_code ec, std::size_t) {
    if (ec) {
        connected_.store(false, std::memory_order_release);
        if (ec == asio::error::eof || ec == asio::error::connection_reset) {
            if (cbs_.on_disconnected) cbs_.on_disconnected();
        } else {
            if (cbs_.on_error) cbs_.on_error(ec);
        }
        return;
    }
    uint32_t length =
        (static_cast<uint32_t>(len_buf_[0]) << 24) |
        (static_cast<uint32_t>(len_buf_[1]) << 16) |
        (static_cast<uint32_t>(len_buf_[2]) <<  8) |
         static_cast<uint32_t>(len_buf_[3]);

    if (length > protocol::kMaxFrameBytes) {
        if (cbs_.on_error) cbs_.on_error(asio::error::message_size);
        return;
    }
    if (length == 0) {
        if (cbs_.on_frame) cbs_.on_frame({});
        start_read();
        return;
    }
    body_buf_.resize(length);
    auto self = shared_from_this();
    asio::async_read(socket_, asio::buffer(body_buf_),
        asio::bind_executor(strand_,
            [this, self, length](std::error_code ec, std::size_t n) {
                handle_body(length, ec, n);
            }));
}

void TcpServerConn::handle_body(uint32_t, std::error_code ec, std::size_t) {
    if (ec) {
        connected_.store(false, std::memory_order_release);
        if (ec == asio::error::eof || ec == asio::error::connection_reset) {
            if (cbs_.on_disconnected) cbs_.on_disconnected();
        } else {
            if (cbs_.on_error) cbs_.on_error(ec);
        }
        return;
    }
    if (cbs_.on_frame) cbs_.on_frame(body_buf_);
    start_read();
}

void TcpServerConn::send_frame(std::vector<uint8_t> payload) {
    std::vector<uint8_t> wire;
    protocol::FrameCodec::emit(payload, wire);
    if (tls_) {
        std::lock_guard lk(tls_send_mutex_);
        tls_send_queue_.push_back(std::move(wire));
    } else {
        auto self = shared_from_this();
        asio::post(strand_, [this, self, w = std::move(wire)]() mutable {
            send_queue_.push_back(std::move(w));
            if (!sending_) do_send();
        });
    }
}

void TcpServerConn::do_send() {
    if (send_queue_.empty()) { sending_ = false; return; }
    sending_ = true;
    auto self = shared_from_this();
    auto& front = send_queue_.front();
    asio::async_write(socket_,
        asio::buffer(front),
        asio::bind_executor(strand_,
            [this, self](std::error_code ec, std::size_t) {
                if (ec) {
                    connected_.store(false, std::memory_order_release);
                    if (cbs_.on_error) cbs_.on_error(ec);
                    return;
                }
                send_queue_.pop_front();
                do_send();
            }));
}

void TcpServerConn::close() {
    if (closing_.exchange(true)) return;
    std::error_code ignored;
    socket_.shutdown(asio::ip::tcp::socket::shutdown_both, ignored);
    socket_.close(ignored);
    connected_.store(false, std::memory_order_release);
    // In non-TLS mode, the Asio async chain will naturally stop when the socket closes.
}

void TcpServerConn::wait_closed() {
    if (tls_thread_.joinable()) {
        if (std::this_thread::get_id() == tls_thread_.get_id()) {
            // Being called from our own TLS thread — detach to avoid self-join deadlock.
            tls_thread_.detach();
        } else {
            tls_thread_.join();
        }
    }
}

// ── TcpAcceptor ─────────────────────────────────────────────────────────────

namespace {
// Try IPv6 dual-stack first (one socket handles both ::1 and 127.0.0.1 — fixes the common
// Windows case where `localhost` resolves to ::1 before 127.0.0.1). Falls back to IPv4-only
// if the OS has IPv6 disabled or the dual-stack bind fails for any reason.
asio::ip::tcp::acceptor make_acceptor(asio::io_context& io, uint16_t port) {
    asio::ip::tcp::acceptor acc(io);
    std::error_code ec;
    acc.open(asio::ip::tcp::v6(), ec);
    if (!ec) {
        acc.set_option(asio::ip::v6_only(false), ec);  // dual-stack
        acc.set_option(asio::ip::tcp::acceptor::reuse_address(true));
        acc.bind(asio::ip::tcp::endpoint(asio::ip::tcp::v6(), port), ec);
        if (!ec) acc.listen(asio::socket_base::max_listen_connections, ec);
    }
    if (ec) {
        if (acc.is_open()) { std::error_code ignored; acc.close(ignored); }
        acc.open(asio::ip::tcp::v4());
        acc.set_option(asio::ip::tcp::acceptor::reuse_address(true));
        acc.bind(asio::ip::tcp::endpoint(asio::ip::tcp::v4(), port));
        acc.listen(asio::socket_base::max_listen_connections);
    }
    return acc;
}
}  // namespace

TcpAcceptor::TcpAcceptor(asio::io_context& io, uint16_t port, ConnFactory factory)
    : acceptor_(make_acceptor(io, port)),
      factory_(std::move(factory)) {}

void TcpAcceptor::start() { do_accept(); }

void TcpAcceptor::stop() {
    stopped_ = true;
    std::error_code ignored;
    acceptor_.close(ignored);
    // Close all tracked connections so their TLS read threads exit. The socket close
    // happens while the io_context (and its reactor) is still alive, preventing the
    // null-reactor use-after-free that manifests on macOS kqueue.
    std::vector<std::shared_ptr<TcpServerConn>> to_close;
    {
        std::lock_guard lk(conns_mu_);
        to_close = conns_;
    }
    for (auto& conn : to_close) conn->close();
}

void TcpAcceptor::shutdown() {
    stop();
    // Wait for every connection's TLS I/O thread to finish. close() (called by stop())
    // set closing_=true and closed the socket, so tls_read_loop is already exiting or has
    // exited; the join is brief. This must complete BEFORE the io_context is destroyed.
    std::vector<std::shared_ptr<TcpServerConn>> to_join;
    {
        std::lock_guard lk(conns_mu_);
        to_join = std::move(conns_);
    }
    for (auto& conn : to_join) {
        conn->wait_closed();
    }
    // to_join drops here — if a thread captured shared_from_this, the TcpServerConn stays
    // alive until that thread releases it; the destructor's close() is a no-op (already
    // closed) and tls_thread_ is already joined, so no reactor access occurs.
}

void TcpAcceptor::do_accept() {
    if (stopped_) return;
    acceptor_.async_accept(
        [this](std::error_code ec, asio::ip::tcp::socket socket) {
            if (ec) {
                if (!stopped_) do_accept();
                return;
            }
            socket.set_option(asio::ip::tcp::no_delay(true));
            auto conn = factory_(std::move(socket));
            if (conn) {
                conn->start();
                // Track so shutdown() can close + join before the io_context is destroyed.
                {
                    std::lock_guard lk(conns_mu_);
                    conns_.push_back(conn);
                }
            }
            do_accept();
        });
}

// ── UdpMediaChannel ──────────────────────────────────────────────────────────

bool UdpMediaChannel::bind(asio::io_context& io, uint16_t port) {
    if (bound_.load()) return false;
    try {
        socket_ = std::make_unique<asio::ip::udp::socket>(io);
        socket_->open(asio::ip::udp::v4());
        socket_->set_option(asio::socket_base::reuse_address(true));
        socket_->bind(asio::ip::udp::endpoint(asio::ip::udp::v4(), port));
        bound_.store(true, std::memory_order_release);
        return true;
    } catch (...) {
        socket_.reset();
        return false;
    }
}

void UdpMediaChannel::start_recv(FrameCallback cb) {
    frame_cb_ = std::move(cb);
    do_recv();
}

void UdpMediaChannel::do_recv() {
    if (!socket_ || closed_.load()) return;
    socket_->async_receive_from(
        asio::buffer(recv_buf_), sender_ep_,
        [this](std::error_code ec, std::size_t n) {
            if (ec || closed_.load()) return;
            if (frame_cb_ && n > 0)
                frame_cb_(recv_buf_.data(), n, sender_ep_);
            do_recv();
        });
}

void UdpMediaChannel::send_to(const uint8_t* data, size_t len,
                               asio::ip::udp::endpoint dst) {
    if (!socket_ || closed_.load() || len == 0) return;
    auto buf = std::make_shared<std::vector<uint8_t>>(data, data + len);
    asio::post(socket_->get_executor(), [this, buf, dst]() mutable {
        if (closed_.load()) return;
        socket_->async_send_to(
            asio::buffer(*buf), dst,
            [buf](std::error_code, std::size_t) {});
    });
}

void UdpMediaChannel::close() {
    if (closed_.exchange(true)) return;
    if (socket_) {
        std::error_code ec;
        socket_->cancel(ec);
        socket_->close(ec);
    }
}

asio::ip::udp::endpoint UdpMediaChannel::local_endpoint() const {
    if (!socket_) return {};
    std::error_code ec;
    return socket_->local_endpoint(ec);
}

}  // namespace voicecat::net