ObsidianDragon/src/chat/chat_database.cpp

// DragonX Wallet - HushChat persistent message store (implementation).

#include "chat_database.h"

#include "../util/logger.h"
#include "../util/platform.h"

#include <nlohmann/json.hpp>
#include <sodium.h>
#include <sqlite3.h>

#include <cstdint>
#include <filesystem>
#include <utility>

namespace fs = std::filesystem;

namespace dragonx::chat {

namespace {

// Domain-separated KDF contexts (used as the keyed-BLAKE2b key, like chat_identity). Both lengths
// sit inside crypto_generichash's key-length bounds. Bumping a context rotates that derivation.
constexpr char kStorageKeyContext[] = "DragonX-HushChat-Storage-v1";
constexpr char kWalletTagContext[]  = "DragonX-HushChat-WalletId-v1";
constexpr std::size_t kStorageKeyContextLen = sizeof(kStorageKeyContext) - 1;
constexpr std::size_t kWalletTagContextLen  = sizeof(kWalletTagContext) - 1;

std::string toHex(const unsigned char* data, std::size_t len)
{
    static const char* kHex = "0123456789abcdef";
    std::string out;
    out.reserve(len * 2);
    for (std::size_t i = 0; i < len; ++i) {
        out.push_back(kHex[data[i] >> 4]);
        out.push_back(kHex[data[i] & 0x0F]);
    }
    return out;
}

// keyed-BLAKE2b: out = generichash(in=secret, key=context). Deterministic, so the same seed always
// derives the same storage key + wallet tag across sessions.
bool deriveKeyed(const std::string& secret, const char* context, std::size_t contextLen,
                 unsigned char* out, std::size_t outLen)
{
    return crypto_generichash(out, outLen,
                              reinterpret_cast<const unsigned char*>(secret.data()), secret.size(),
                              reinterpret_cast<const unsigned char*>(context), contextLen) == 0;
}

std::string associatedData(const std::string& walletTag)
{
    return std::string("obsidian-dragon-hushchat-v1:") + walletTag;
}

} // namespace

ChatDatabase::ChatDatabase() : database_path_(defaultDatabasePath()) {}
ChatDatabase::ChatDatabase(std::string databasePath) : database_path_(std::move(databasePath)) {}

ChatDatabase::~ChatDatabase()
{
    lock();
    close();
}

std::string ChatDatabase::defaultDatabasePath()
{
    return (fs::path(util::Platform::getConfigDir()) / "chat_messages.sqlite").string();
}

bool ChatDatabase::unlockWithSecret(const std::string& secret)
{
    if (sodium_init() < 0) return false;

    if (!deriveKeyed(secret, kStorageKeyContext, kStorageKeyContextLen, key_.data(), key_.size()))
        return false;

    unsigned char tag[32];
    if (!deriveKeyed(secret, kWalletTagContext, kWalletTagContextLen, tag, sizeof(tag))) {
        sodium_memzero(key_.data(), key_.size());
        return false;
    }
    wallet_tag_ = toHex(tag, sizeof(tag));
    sodium_memzero(tag, sizeof(tag));
    key_ready_ = true;

    if (!ensureOpen()) {
        lock();
        return false;
    }
    return true;
}

void ChatDatabase::lock()
{
    sodium_memzero(key_.data(), key_.size());
    key_ready_ = false;
    wallet_tag_.clear();
}

bool ChatDatabase::append(const ChatMessage& message)
{
    if (!key_ready_ || !ensureOpen()) return false;

    std::vector<unsigned char> nonce;
    std::vector<unsigned char> cipher;
    std::string plain = serialize(message);  // full plaintext (decrypted body + metadata)
    const bool encrypted = encrypt(plain, nonce, cipher);
    if (!plain.empty()) sodium_memzero(&plain[0], plain.size());  // don't leave it on the heap
    if (!encrypted) return false;

    const std::string dedup = dedupHash(message.txid, message.payload_position);

    sqlite3_stmt* stmt = nullptr;
    if (sqlite3_prepare_v2(db_,
            "INSERT OR IGNORE INTO chat_messages (wallet_tag, dedup_hash, nonce, payload) "
            "VALUES (?, ?, ?, ?)",
            -1, &stmt, nullptr) != SQLITE_OK) {
        return false;
    }
    sqlite3_bind_text(stmt, 1, wallet_tag_.c_str(), -1, SQLITE_TRANSIENT);
    sqlite3_bind_text(stmt, 2, dedup.c_str(), -1, SQLITE_TRANSIENT);
    sqlite3_bind_blob(stmt, 3, nonce.data(), static_cast<int>(nonce.size()), SQLITE_TRANSIENT);
    sqlite3_bind_blob(stmt, 4, cipher.data(), static_cast<int>(cipher.size()), SQLITE_TRANSIENT);
    const bool done = sqlite3_step(stmt) == SQLITE_DONE;
    sqlite3_finalize(stmt);
    if (!done) return false;
    return sqlite3_changes(db_) > 0;
}

bool ChatDatabase::upsert(const ChatMessage& message)
{
    if (!key_ready_ || !ensureOpen()) return false;

    std::vector<unsigned char> nonce;
    std::vector<unsigned char> cipher;
    std::string plain = serialize(message);
    const bool encrypted = encrypt(plain, nonce, cipher);
    if (!plain.empty()) sodium_memzero(&plain[0], plain.size());
    if (!encrypted) return false;

    const std::string dedup = dedupHash(message.txid, message.payload_position);

    sqlite3_stmt* stmt = nullptr;
    if (sqlite3_prepare_v2(db_,
            "INSERT INTO chat_messages (wallet_tag, dedup_hash, nonce, payload) VALUES (?, ?, ?, ?) "
            "ON CONFLICT(wallet_tag, dedup_hash) DO UPDATE SET nonce=excluded.nonce, payload=excluded.payload",
            -1, &stmt, nullptr) != SQLITE_OK) {
        return false;
    }
    sqlite3_bind_text(stmt, 1, wallet_tag_.c_str(), -1, SQLITE_TRANSIENT);
    sqlite3_bind_text(stmt, 2, dedup.c_str(), -1, SQLITE_TRANSIENT);
    sqlite3_bind_blob(stmt, 3, nonce.data(), static_cast<int>(nonce.size()), SQLITE_TRANSIENT);
    sqlite3_bind_blob(stmt, 4, cipher.data(), static_cast<int>(cipher.size()), SQLITE_TRANSIENT);
    const bool done = sqlite3_step(stmt) == SQLITE_DONE;
    sqlite3_finalize(stmt);
    return done;
}

std::vector<ChatMessage> ChatDatabase::load()
{
    std::vector<ChatMessage> out;
    if (!key_ready_ || !ensureOpen()) return out;

    sqlite3_stmt* stmt = nullptr;
    if (sqlite3_prepare_v2(db_,
            "SELECT nonce, payload FROM chat_messages WHERE wallet_tag = ? ORDER BY rowid",
            -1, &stmt, nullptr) != SQLITE_OK) {
        return out;
    }
    sqlite3_bind_text(stmt, 1, wallet_tag_.c_str(), -1, SQLITE_TRANSIENT);
    while (sqlite3_step(stmt) == SQLITE_ROW) {
        const auto* noncePtr = static_cast<const unsigned char*>(sqlite3_column_blob(stmt, 0));
        const int nonceLen = sqlite3_column_bytes(stmt, 0);
        const auto* cipherPtr = static_cast<const unsigned char*>(sqlite3_column_blob(stmt, 1));
        const int cipherLen = sqlite3_column_bytes(stmt, 1);
        if (!noncePtr || !cipherPtr) continue;

        std::vector<unsigned char> nonce(noncePtr, noncePtr + nonceLen);
        std::vector<unsigned char> cipher(cipherPtr, cipherPtr + cipherLen);
        std::string plain;
        if (!decrypt(nonce, cipher, plain)) continue;  // wrong wallet / tampered — skip

        ChatMessage message;
        if (deserialize(plain, message)) out.push_back(std::move(message));
        sodium_memzero(&plain[0], plain.size());
    }
    sqlite3_finalize(stmt);
    return out;
}

void ChatDatabase::clearWallet()
{
    if (wallet_tag_.empty() || !ensureOpen()) return;
    sqlite3_stmt* stmt = nullptr;
    if (sqlite3_prepare_v2(db_, "DELETE FROM chat_messages WHERE wallet_tag = ?", -1, &stmt, nullptr)
        != SQLITE_OK) {
        return;
    }
    sqlite3_bind_text(stmt, 1, wallet_tag_.c_str(), -1, SQLITE_TRANSIENT);
    sqlite3_step(stmt);
    sqlite3_finalize(stmt);
}

bool ChatDatabase::ensureOpen()
{
    if (db_) return true;

    try {
        fs::path path(database_path_);
        if (!path.parent_path().empty()) fs::create_directories(path.parent_path());
    } catch (const std::exception& exception) {
        DEBUG_LOGF("Failed to create chat database directory: %s\n", exception.what());
        return false;
    }

    sqlite3* openedDb = nullptr;
    if (sqlite3_open(database_path_.c_str(), &openedDb) != SQLITE_OK) {
        DEBUG_LOGF("Failed to open chat database: %s\n",
                   openedDb ? sqlite3_errmsg(openedDb) : "unknown error");
        if (openedDb) sqlite3_close(openedDb);
        return false;
    }

    db_ = openedDb;
    sqlite3_busy_timeout(db_, 2000);
    exec("PRAGMA journal_mode=WAL");
    exec("PRAGMA synchronous=NORMAL");

    // C3-1: restrict the chat DB and its WAL/SHM sidecars to owner-only. sqlite creates them with
    // umask-derived permissions (often world/group-readable); they hold per-row nonces + AEAD
    // ciphertext of the user's messages. Best-effort (errors swallowed; a no-op-ish on Windows).
    {
        std::error_code perr;
        const auto ownerOnly = std::filesystem::perms::owner_read | std::filesystem::perms::owner_write;
        std::filesystem::permissions(database_path_, ownerOnly, std::filesystem::perm_options::replace, perr);
        std::filesystem::permissions(database_path_ + "-wal", ownerOnly, std::filesystem::perm_options::replace, perr);
        std::filesystem::permissions(database_path_ + "-shm", ownerOnly, std::filesystem::perm_options::replace, perr);
    }

    if (!createSchema()) {
        close();
        return false;
    }
    return true;
}

bool ChatDatabase::exec(const char* sql)
{
    if (!db_) return false;
    char* error = nullptr;
    if (sqlite3_exec(db_, sql, nullptr, nullptr, &error) != SQLITE_OK) {
        DEBUG_LOGF("Chat database SQL error: %s\n", error ? error : sqlite3_errmsg(db_));
        if (error) sqlite3_free(error);
        return false;
    }
    return true;
}

bool ChatDatabase::createSchema()
{
    return exec("CREATE TABLE IF NOT EXISTS chat_messages ("
                "wallet_tag TEXT NOT NULL, "
                "dedup_hash TEXT NOT NULL, "
                "nonce BLOB NOT NULL, "
                "payload BLOB NOT NULL, "
                "PRIMARY KEY (wallet_tag, dedup_hash))");
}

std::string ChatDatabase::dedupHash(const std::string& txid, std::size_t position) const
{
    const std::string input = txid + ":" + std::to_string(position);
    unsigned char hash[32];
    crypto_generichash(hash, sizeof(hash),
                       reinterpret_cast<const unsigned char*>(input.data()), input.size(),
                       key_.data(), key_.size());  // keyed by the storage key → txid stays private
    return toHex(hash, sizeof(hash));
}

std::string ChatDatabase::serialize(const ChatMessage& message) const
{
    nlohmann::json json;
    json["d"] = static_cast<int>(message.direction);
    json["k"] = static_cast<int>(message.kind);
    json["txid"] = message.txid;
    json["cid"] = message.conversation_id;
    json["z"] = message.peer_zaddr;
    json["p"] = message.peer_public_key_hex;
    json["b"] = message.body;
    json["ts"] = message.timestamp;
    json["pos"] = static_cast<std::uint64_t>(message.payload_position);
    json["dl"] = static_cast<int>(message.delivery);
    return json.dump();
}

bool ChatDatabase::deserialize(const std::string& json, ChatMessage& out) const
{
    try {
        const auto parsed = nlohmann::json::parse(json);
        out.direction = static_cast<ChatDirection>(parsed.value("d", 0));
        out.kind = static_cast<ChatMessageKind>(parsed.value("k", 0));
        out.txid = parsed.value("txid", std::string());
        out.conversation_id = parsed.value("cid", std::string());
        out.peer_zaddr = parsed.value("z", std::string());
        out.peer_public_key_hex = parsed.value("p", std::string());
        out.body = parsed.value("b", std::string());
        out.timestamp = parsed.value("ts", static_cast<std::int64_t>(0));
        out.payload_position = static_cast<std::size_t>(parsed.value("pos", static_cast<std::uint64_t>(0)));
        out.delivery = static_cast<ChatDelivery>(parsed.value("dl", 0));  // old rows → Sent (0)
        // A persisted "Sending" means we crashed mid-broadcast; the outcome is unknown. Resolve it
        // optimistically to Sent on load so it can't show a stuck spinner forever.
        if (out.delivery == ChatDelivery::Sending) out.delivery = ChatDelivery::Sent;
        return true;
    } catch (const std::exception&) {
        return false;
    }
}

bool ChatDatabase::encrypt(const std::string& plain,
                           std::vector<unsigned char>& nonce,
                           std::vector<unsigned char>& cipher) const
{
    if (!key_ready_) return false;
    nonce.resize(crypto_aead_xchacha20poly1305_ietf_NPUBBYTES);
    randombytes_buf(nonce.data(), nonce.size());

    const std::string ad = associatedData(wallet_tag_);
    cipher.resize(plain.size() + crypto_aead_xchacha20poly1305_ietf_ABYTES);
    unsigned long long cipherLen = 0;
    if (crypto_aead_xchacha20poly1305_ietf_encrypt(
            cipher.data(), &cipherLen,
            reinterpret_cast<const unsigned char*>(plain.data()), plain.size(),
            reinterpret_cast<const unsigned char*>(ad.data()), ad.size(),
            nullptr, nonce.data(), key_.data()) != 0) {
        return false;
    }
    cipher.resize(static_cast<std::size_t>(cipherLen));
    return true;
}

bool ChatDatabase::decrypt(const std::vector<unsigned char>& nonce,
                           const std::vector<unsigned char>& cipher,
                           std::string& plain) const
{
    if (!key_ready_) return false;
    if (nonce.size() != crypto_aead_xchacha20poly1305_ietf_NPUBBYTES) return false;
    if (cipher.size() < crypto_aead_xchacha20poly1305_ietf_ABYTES) return false;

    const std::string ad = associatedData(wallet_tag_);
    std::vector<unsigned char> out(cipher.size());
    unsigned long long outLen = 0;
    if (crypto_aead_xchacha20poly1305_ietf_decrypt(
            out.data(), &outLen, nullptr,
            cipher.data(), cipher.size(),
            reinterpret_cast<const unsigned char*>(ad.data()), ad.size(),
            nonce.data(), key_.data()) != 0) {
        return false;
    }
    plain.assign(reinterpret_cast<const char*>(out.data()), static_cast<std::size_t>(outLen));
    sodium_memzero(out.data(), out.size());
    return true;
}

void ChatDatabase::close()
{
    if (db_) {
        sqlite3_close(db_);
        db_ = nullptr;
    }
}

} // namespace dragonx::chat