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;
    }
    loadTombstones();
    return true;
}

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

bool ChatDatabase::append(const ChatMessage& message)
{
    if (!key_ready_ || !ensureOpen()) return false;
    if (isTombstoned(message)) return false;  // locally deleted — don't re-persist on a chain re-scan

    std::vector<unsigned char> nonce;
    std::vector<unsigned char> cipher;
    std::string plain;
    bool encrypted = false;
    try {
        // serialize() dumps strictly, so an invalid-UTF-8 (attacker-influenceable) body THROWS here;
        // skip the row instead of letting it escape the write path and crash (appendBatch already does).
        plain = serialize(message);
        encrypted = encrypt(plain, nonce, cipher);
    } catch (const std::exception& e) {
        if (!plain.empty()) sodium_memzero(&plain[0], plain.size());
        DEBUG_LOGF("[chat] append skipped a malformed message: %s\n", e.what());
        return false;
    }
    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;
}

int ChatDatabase::appendBatch(const std::vector<ChatMessage>& messages)
{
    if (!key_ready_ || !ensureOpen()) return 0;
    if (messages.empty()) return 0;

    struct Row { std::string dedup; std::vector<unsigned char> nonce, cipher; };
    std::vector<Row> rows;
    rows.reserve(messages.size());
    for (const auto& m : messages) {
        if (isTombstoned(m)) continue;                 // locally deleted — don't re-persist
        try {
            std::string plain = serialize(m);          // strict json.dump() may throw on non-UTF-8 body
            std::vector<unsigned char> nonce, cipher;
            const bool ok = encrypt(plain, nonce, cipher);
            if (!plain.empty()) sodium_memzero(&plain[0], plain.size());
            if (!ok) continue;
            rows.push_back({ dedupHash(m.txid, m.payload_position), std::move(nonce), std::move(cipher) });
        } catch (const std::exception&) {
            continue;                                  // skip only this message; never open a txn with bad data
        }
    }
    if (rows.empty()) return 0;

    if (!exec("BEGIN")) return 0;                      // match deleteMessages() in-file style
    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) {
        exec("ROLLBACK");
        return 0;
    }
    int inserted = 0;
    bool ok = true;
    for (const auto& r : rows) {
        sqlite3_reset(stmt);
        sqlite3_clear_bindings(stmt);
        sqlite3_bind_text(stmt, 1, wallet_tag_.c_str(), -1, SQLITE_TRANSIENT);
        sqlite3_bind_text(stmt, 2, r.dedup.c_str(), -1, SQLITE_TRANSIENT);
        sqlite3_bind_blob(stmt, 3, r.nonce.data(), static_cast<int>(r.nonce.size()), SQLITE_TRANSIENT);
        sqlite3_bind_blob(stmt, 4, r.cipher.data(), static_cast<int>(r.cipher.size()), SQLITE_TRANSIENT);
        if (sqlite3_step(stmt) != SQLITE_DONE) { ok = false; break; }
        if (sqlite3_changes(db_) > 0) ++inserted;
    }
    sqlite3_finalize(stmt);
    if (!exec(ok ? "COMMIT" : "ROLLBACK")) return 0;
    return ok ? inserted : 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;
    bool encrypted = false;
    try {
        plain = serialize(message);   // strict dump → a malformed body throws; skip rather than crash
        encrypted = encrypt(plain, nonce, cipher);
    } catch (const std::exception& e) {
        if (!plain.empty()) sodium_memzero(&plain[0], plain.size());
        DEBUG_LOGF("[chat] upsert skipped a malformed message: %s\n", e.what());
        return false;
    }
    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;
    for (const char* sql : {"DELETE FROM chat_messages WHERE wallet_tag = ?",
                            "DELETE FROM chat_deleted WHERE wallet_tag = ?"}) {
        sqlite3_stmt* stmt = nullptr;
        if (sqlite3_prepare_v2(db_, sql, -1, &stmt, nullptr) != SQLITE_OK) continue;
        sqlite3_bind_text(stmt, 1, wallet_tag_.c_str(), -1, SQLITE_TRANSIENT);
        sqlite3_step(stmt);
        sqlite3_finalize(stmt);
    }
    tombstones_.clear();
}

bool ChatDatabase::deleteMessages(const std::vector<ChatMessage>& messages, bool tombstone)
{
    if (!key_ready_ || !ensureOpen()) return false;
    if (messages.empty()) return true;

    if (!exec("BEGIN")) return false;
    bool ok = true;
    for (const auto& m : messages) {
        const std::string dedup = dedupHash(m.txid, m.payload_position);

        sqlite3_stmt* del = nullptr;
        if (sqlite3_prepare_v2(db_, "DELETE FROM chat_messages WHERE wallet_tag = ? AND dedup_hash = ?",
                               -1, &del, nullptr) == SQLITE_OK) {
            sqlite3_bind_text(del, 1, wallet_tag_.c_str(), -1, SQLITE_TRANSIENT);
            sqlite3_bind_text(del, 2, dedup.c_str(), -1, SQLITE_TRANSIENT);
            if (sqlite3_step(del) != SQLITE_DONE) ok = false;
            sqlite3_finalize(del);
        } else {
            ok = false;
        }

        if (tombstone) {
            sqlite3_stmt* ins = nullptr;
            if (sqlite3_prepare_v2(db_,
                    "INSERT OR IGNORE INTO chat_deleted (wallet_tag, dedup_hash) VALUES (?, ?)",
                    -1, &ins, nullptr) == SQLITE_OK) {
                sqlite3_bind_text(ins, 1, wallet_tag_.c_str(), -1, SQLITE_TRANSIENT);
                sqlite3_bind_text(ins, 2, dedup.c_str(), -1, SQLITE_TRANSIENT);
                if (sqlite3_step(ins) != SQLITE_DONE) ok = false;
                sqlite3_finalize(ins);
            } else {
                ok = false;
            }
        }
    }
    if (!exec(ok ? "COMMIT" : "ROLLBACK")) ok = false;
    // Only reflect the tombstones in the in-memory cache once they are durably committed.
    if (ok && tombstone)
        for (const auto& m : messages) tombstones_.insert(dedupHash(m.txid, m.payload_position));
    return ok;
}

bool ChatDatabase::isTombstoned(const ChatMessage& message) const
{
    if (!key_ready_ || tombstones_.empty()) return false;
    return tombstones_.count(dedupHash(message.txid, message.payload_position)) > 0;
}

void ChatDatabase::loadTombstones()
{
    tombstones_.clear();
    if (!key_ready_ || !ensureOpen()) return;
    sqlite3_stmt* stmt = nullptr;
    if (sqlite3_prepare_v2(db_, "SELECT dedup_hash FROM chat_deleted WHERE wallet_tag = ?",
                           -1, &stmt, nullptr) != SQLITE_OK) {
        return;
    }
    sqlite3_bind_text(stmt, 1, wallet_tag_.c_str(), -1, SQLITE_TRANSIENT);
    while (sqlite3_step(stmt) == SQLITE_ROW) {
        const auto* h = reinterpret_cast<const char*>(sqlite3_column_text(stmt, 0));
        if (h) tombstones_.insert(h);
    }
    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()
{
    if (!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))"))
        return false;
    // Tombstones for locally-deleted messages (dedup_hash only — the same keyed, non-revealing hash the
    // message rows use). A chain re-scan checks this so a deleted message never re-imports.
    return exec("CREATE TABLE IF NOT EXISTS chat_deleted ("
                "wallet_tag TEXT NOT NULL, "
                "dedup_hash TEXT 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();   // strict: an invalid-UTF-8 body throws and is DROPPED by the callers (contract)
}

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