voice-cat/server/src/conn_session.cpp

#include "conn_session.h"

#ifdef VOICECAT_HAS_NET

#include <chrono>
#include <cstdio>

#include "db.h"
#include "session_registry.h"
#include "core/worker_pool.h"
#include "protocol/envelope.h"

namespace voicecat::server {

static voicecat::v1::Envelope make_env(uint64_t req_id = 0) {
    voicecat::v1::Envelope e;
    e.set_request_id(req_id);
    return e;
}

ConnSession::ConnSession(std::shared_ptr<Database>              db,
                         std::shared_ptr<SessionRegistry>       registry,
                         std::shared_ptr<voicecat::WorkerPool>  workers,
                         const std::array<uint8_t, 32>&   server_fp,
                         bool                              allow_guests)
    : db_(std::move(db)),
      registry_(std::move(registry)),
      workers_(std::move(workers)),
      server_fp_(server_fp),
      allow_guests_(allow_guests) {}

void ConnSession::set_io(SendFn send_fn, CloseFn close_fn) {
    send_fn_  = std::move(send_fn);
    close_fn_ = std::move(close_fn);
}

void ConnSession::begin() {
    // Nothing to do at TCP level — wait for ClientHello
}

void ConnSession::on_frame(std::vector<uint8_t> frame) {
    voicecat::v1::Envelope env;
    if (!protocol::decode_envelope(frame, env)) return;

    auto st = state_.load(std::memory_order_acquire);
    switch (env.body_case()) {
        case voicecat::v1::Envelope::kClientHello:
            if (st == State::WaitingHello)
                handle_client_hello(env.request_id(), env.client_hello());
            break;
        case voicecat::v1::Envelope::kAuthRequest:
            if (st == State::WaitingAuth)
                handle_auth_request(env.request_id(), env.auth_request());
            break;
        case voicecat::v1::Envelope::kJoinChannel:
            if (st == State::Authenticated)
                handle_join_channel(env.request_id(), env.join_channel());
            break;
        case voicecat::v1::Envelope::kTextMessage:
            if (st == State::Authenticated)
                handle_text_message(env.text_message());
            break;
        case voicecat::v1::Envelope::kPing:
            handle_ping(env.ping());
            break;
        case voicecat::v1::Envelope::kLeaveChannel:
            if (st == State::Authenticated)
                registry_->set_user_channel(user_id_.load(), 1);
            break;
        default:
            break;
    }
}

void ConnSession::on_disconnect() { close(); }

void ConnSession::send_envelope(const voicecat::v1::Envelope& env) {
    if (!send_fn_ || closed_.load()) return;
    // Serialize to raw protobuf bytes; send_fn_ (→ TcpServerConn::send_frame)
    // adds the [4-byte len] framing, so we must NOT pre-frame here.
    std::string bytes;
    if (!env.SerializeToString(&bytes)) return;
    std::vector<uint8_t> raw(bytes.begin(), bytes.end());
    send_fn_(std::move(raw));
}

void ConnSession::close() {
    if (closed_.exchange(true)) return;
    state_.store(State::Disconnecting, std::memory_order_release);
    uint32_t uid = user_id_.load();
    if (uid) registry_->remove_user(uid);
    if (session_id_) registry_->unregister_session(session_id_);
    if (close_fn_) close_fn_();
}

void ConnSession::handle_client_hello(uint64_t req_id, const voicecat::v1::ClientHello& msg) {
    if (msg.proto_version() != 1) {
        send_disconnect_and_close(1, "unsupported protocol version");
        return;
    }
    auto env = make_env(req_id);
    auto* hello = env.mutable_server_hello();
    hello->set_proto_version(1);
    hello->set_server_name("VoiceCat Server");
    hello->set_server_version("0.1.0");
    if (allow_guests_) hello->add_auth_methods("guest");
    hello->add_auth_methods("password");
    hello->set_server_identity_fingerprint(server_fp_.data(), server_fp_.size());
    send_envelope(env);
    state_.store(State::WaitingAuth, std::memory_order_release);
}

void ConnSession::handle_auth_request(uint64_t req_id, const voicecat::v1::AuthRequest& msg) {
    if (msg.has_guest()) {
        finish_guest_auth(msg.guest(), req_id);
    } else if (msg.has_password()) {
        finish_password_auth(msg.password().username(), msg.password().password(), req_id);
    } else {
        auto env = make_env(req_id);
        env.mutable_auth_result()->set_ok(false);
        env.mutable_auth_result()->set_error("unknown auth method");
        send_envelope(env);
    }
}

void ConnSession::finish_guest_auth(const voicecat::v1::GuestAuth& guest, uint64_t req_id) {
    if (!allow_guests_) {
        auto env = make_env(req_id);
        env.mutable_auth_result()->set_ok(false);
        env.mutable_auth_result()->set_error("guest login not permitted");
        send_envelope(env);
        return;
    }
    voicecat::v1::User user;
    user.set_nickname(guest.nickname().empty() ? "Guest" : guest.nickname());
    user.set_is_guest(true);
    user.set_channel_id(1);

    uint32_t uid = registry_->add_user(session_id_, user);
    user.set_id(uid);
    user_id_.store(uid, std::memory_order_relaxed);
    state_.store(State::Authenticated, std::memory_order_release);

    {
        auto env = make_env(req_id);
        auto* res = env.mutable_auth_result();
        res->set_ok(true);
        res->set_session_id(session_id_);
        *res->mutable_self() = user;
        send_envelope(env);
    }
    broadcast_user_joined(user);
    send_state_snapshot();
}

void ConnSession::finish_password_auth(const std::string& username,
                                        const std::string& password, uint64_t req_id) {
    // Argon2id runs on the worker pool (deliberately slow).
    auto self = shared_from_this();
    workers_->post([self, username, password, req_id] {
        auto acc = self->db_->authenticate(username, password);
        if (!acc) {
            auto env = make_env(req_id);
            env.mutable_auth_result()->set_ok(false);
            env.mutable_auth_result()->set_error("invalid credentials");
            self->send_envelope(env);
            return;
        }
        voicecat::v1::User user;
        user.set_nickname(acc->username);
        user.set_is_guest(false);
        user.set_channel_id(1);

        uint32_t uid = self->registry_->add_user(self->session_id_, user);
        user.set_id(uid);
        self->user_id_.store(uid, std::memory_order_relaxed);
        self->state_.store(State::Authenticated, std::memory_order_release);

        {
            auto env = make_env(req_id);
            auto* res = env.mutable_auth_result();
            res->set_ok(true);
            res->set_session_id(self->session_id_);
            *res->mutable_self() = user;
            auto* perms = res->mutable_permissions();
            perms->set_is_admin(acc->is_admin);
            self->send_envelope(env);
        }
        self->broadcast_user_joined(user);
        self->send_state_snapshot();
    });
}

void ConnSession::send_state_snapshot() {
    auto env = make_env();
    auto* snap = env.mutable_server_state();
    for (auto& ch : registry_->channel_snapshot()) *snap->add_channels() = ch;
    for (auto& u  : registry_->user_snapshot())    *snap->add_users()    = u;
    send_envelope(env);
}

void ConnSession::broadcast_user_joined(const voicecat::v1::User& user) {
    auto bcast = make_env();
    auto* ue = bcast.mutable_user_event();
    ue->set_kind(voicecat::v1::UserEvent::JOINED);
    *ue->mutable_user() = user;
    registry_->broadcast(bcast, session_id_);
}

void ConnSession::handle_join_channel(uint64_t req_id,
                                       const voicecat::v1::JoinChannelRequest& msg) {
    bool ok = registry_->set_user_channel(user_id_.load(), msg.channel_id());
    auto env = make_env(req_id);
    auto* res = env.mutable_join_channel_result();
    res->set_ok(ok);
    if (!ok) res->set_error("channel not found");
    else     res->set_channel_id(msg.channel_id());
    send_envelope(env);
}

void ConnSession::handle_text_message(const voicecat::v1::TextMessage& msg) {
    using namespace std::chrono;
    int64_t now_ms = duration_cast<milliseconds>(
        system_clock::now().time_since_epoch()).count();

    voicecat::v1::TextMessage relay = msg;
    relay.set_sender_id(user_id_.load(std::memory_order_relaxed));
    relay.set_sent_at_unix_ms(now_ms);

    voicecat::v1::Envelope fwd;
    *fwd.mutable_text_message() = relay;

    auto targets = registry_->resolve_text_targets(session_id_, msg.scope(), msg.target_id());
    for (auto& t : targets) t->send_envelope(fwd);

    // Ack
    auto env = make_env();
    auto* ack = env.mutable_text_message_ack();
    ack->set_client_msg_id(msg.client_msg_id());
    ack->set_ok(true);
    send_envelope(env);
}

void ConnSession::handle_ping(const voicecat::v1::Ping& msg) {
    auto env = make_env();
    env.mutable_pong()->set_nonce(msg.nonce());
    send_envelope(env);
}

void ConnSession::send_disconnect_and_close(uint32_t code, const std::string& reason) {
    auto env = make_env();
    auto* d = env.mutable_disconnect();
    d->set_code(code);
    d->set_reason(reason);
    send_envelope(env);
    close();
}

}  // namespace voicecat::server

#endif  // VOICECAT_HAS_NET