smartbotic-database/tests/test_subdb_identity.cpp

// v2.4.4 — sub-db identity sentinel tests.
//
// Regression cover for the production incident in which an invalidated
// MDB_dbi was reused after LMDB reassigned its slot, so writes aimed at
// `image_hashes` landed in `executions` and succeeded silently. 31 documents
// across smartbotic-automation ended up in a sub-db other than the one they
// declared. See storage/subdb_identity.hpp for the full mechanism.
//
// The core test is `misbound_handle_is_refused`: it reproduces the misbinding
// directly by handing the verifier a handle for a different sub-db, which is
// what a stale cache entry amounts to.

#include <cassert>
#include <atomic>
#include <filesystem>
#include <iostream>
#include <cstdio>
#include <set>
#include <string>
#include <unistd.h>

#include <lmdb.h>
#include <nlohmann/json.hpp>

#include "document.hpp"
#include "storage/document_store_lmdb.hpp"
#include "storage/lmdb_env.hpp"
#include "storage/lmdb_txn.hpp"
#include "storage/subdb_identity.hpp"

namespace fs = std::filesystem;

using smartbotic::database::Document;
using smartbotic::db::storage::is_identity_key;
using smartbotic::db::storage::kSubdbIdentityKey;
using smartbotic::db::storage::LmdbDocumentStore;
using smartbotic::db::storage::LmdbEnv;
using smartbotic::db::storage::LmdbEnvOpts;
using smartbotic::db::storage::read_subdb_identity;
using smartbotic::db::storage::ReadTxn;
using smartbotic::db::storage::verify_subdb_identity;
using smartbotic::db::storage::write_subdb_identity;
using smartbotic::db::storage::WriteTxn;

namespace {

int g_pass = 0;
int g_fail = 0;

void check(bool cond, const char* msg) {
    if (cond) {
        ++g_pass;
    } else {
        ++g_fail;
        std::cerr << "FAIL: " << msg << "\n";
    }
}

std::string make_tmpdir(const char* tag) {
    static std::atomic<int> counter{0};
    std::string path = "/tmp/subdb-identity-test-" + std::to_string(::getpid()) +
                       "-" + std::to_string(counter.fetch_add(1)) + "-" + tag;
    std::error_code ec;
    fs::remove_all(path, ec);
    return path;
}

struct TmpEnv {
    std::string path;
    LmdbEnv env;
    explicit TmpEnv(const char* tag)
        : path(make_tmpdir(tag)),
          env(LmdbEnvOpts{path, 64ULL << 20, 256, 126, false}) {}
    ~TmpEnv() {
        std::error_code ec;
        fs::remove_all(path, ec);
    }
    TmpEnv(const TmpEnv&) = delete;
    TmpEnv& operator=(const TmpEnv&) = delete;
};

// Open (creating) a named sub-db inside a write txn and return its handle.
unsigned int open_subdb(WriteTxn& txn, const char* name) {
    MDB_dbi dbi = 0;
    int rc = mdb_dbi_open(txn.raw(), name, MDB_CREATE, &dbi);
    assert(rc == MDB_SUCCESS);
    (void)rc;
    return dbi;
}

Document make_doc(const std::string& id, const std::string& collection) {
    Document d;
    d.id = id;
    d.collection = collection;
    d.set_data(nlohmann::json{{"seenCount", 1}, {"who", collection}});
    return d;
}

// -------------------------------------------------------------------------

void test_sentinel_roundtrip() {
    TmpEnv t("roundtrip");
    {
        WriteTxn w(t.env);
        unsigned int dbi = open_subdb(w, "image_hashes");
        write_subdb_identity(w, dbi, "image_hashes");
        w.commit();
    }
    {
        WriteTxn w(t.env);
        unsigned int dbi = open_subdb(w, "image_hashes");
        bool threw = false;
        try {
            verify_subdb_identity(w, dbi, "image_hashes");
        } catch (const std::exception&) {
            threw = true;
        }
        check(!threw, "matching sentinel must verify without throwing");
        w.commit();
    }
    {
        ReadTxn r(t.env);
        MDB_dbi dbi = 0;
        mdb_dbi_open(r.raw(), "image_hashes", 0, &dbi);
        check(read_subdb_identity(r, dbi) == "image_hashes",
              "read_subdb_identity returns the stamped name");
    }
}

// THE regression test. A stale cache entry is, in effect, a handle that
// addresses someone else's sub-db. Hand the verifier exactly that.
void test_misbound_handle_is_refused() {
    TmpEnv t("misbound");
    unsigned int executions_dbi = 0;
    {
        WriteTxn w(t.env);
        unsigned int ih = open_subdb(w, "image_hashes");
        write_subdb_identity(w, ih, "image_hashes");
        executions_dbi = open_subdb(w, "executions");
        write_subdb_identity(w, executions_dbi, "executions");
        w.commit();
    }

    WriteTxn w(t.env);
    // Re-open so the handle is valid in this txn, then deliberately verify it
    // under the WRONG name — the production misbinding, reproduced.
    unsigned int exec = open_subdb(w, "executions");
    bool threw = false;
    std::string msg;
    try {
        verify_subdb_identity(w, exec, "image_hashes");
    } catch (const std::exception& e) {
        threw = true;
        msg = e.what();
    }
    check(threw, "handle for 'executions' verified as 'image_hashes' must throw");
    check(msg.find("image_hashes") != std::string::npos &&
              msg.find("executions") != std::string::npos,
          "misbinding error names both the requested and actual sub-db");
    w.abort();
}

// Existing deployments have sub-dbs with no sentinel. Those must keep working.
void test_unstamped_subdb_is_permitted() {
    TmpEnv t("unstamped");
    WriteTxn w(t.env);
    unsigned int dbi = open_subdb(w, "legacy");
    bool threw = false;
    try {
        verify_subdb_identity(w, dbi, "legacy");
    } catch (const std::exception&) {
        threw = true;
    }
    check(!threw, "sub-db without a sentinel must verify (absence is unknown, not wrong)");
    w.commit();
}

void test_identity_key_predicate() {
    check(is_identity_key(kSubdbIdentityKey), "sentinel key recognised");
    check(!is_identity_key("__subdb_identity__"),
          "same text without the leading NUL is NOT the sentinel");
    check(!is_identity_key("e63c1b90"), "a document id is not the sentinel");
    check(kSubdbIdentityKey[0] == '\0',
          "sentinel must start with NUL so it cannot collide with a doc id");
}

// The sentinel is an implementation detail: it must never surface through the
// DocumentStore API as a document, nor inflate a count.
void test_sentinel_invisible_through_store() {
    TmpEnv t("invisible");
    LmdbDocumentStore store(t.env);

    store.put("image_hashes", "aaa", make_doc("aaa", "image_hashes"));
    store.put("image_hashes", "bbb", make_doc("bbb", "image_hashes"));

    check(store.count("image_hashes") == 2,
          "count() must exclude the identity sentinel");

    smartbotic::database::Query q;
    q.limit = 100;
    auto res = store.scan("image_hashes", q);
    check(res.documents.size() == 2, "scan() must exclude the identity sentinel");
    check(res.total_matched == 2, "scan() total_matched must exclude the sentinel");

    for (const auto& d : res.documents) {
        check(d.id == "aaa" || d.id == "bbb",
              "scan() must not surface the sentinel as a document");
    }

    // And it really is on disk.
    ReadTxn r(t.env);
    MDB_dbi dbi = 0;
    int rc = mdb_dbi_open(r.raw(), "image_hashes", 0, &dbi);
    check(rc == MDB_SUCCESS, "sub-db exists");
    check(read_subdb_identity(r, dbi) == "image_hashes",
          "store.put() stamps the sentinel on first write");
}

// A vector sub-db gets stamped with its own (prefixed) name, and scan_vectors
// must skip the sentinel rather than trying to read it as float32 bytes.
void test_vector_subdb_sentinel() {
    TmpEnv t("vectors");
    LmdbDocumentStore store(t.env);
    store.put_vector("emb", "v1", {1.0f, 2.0f, 3.0f});
    store.put_vector("emb", "v2", {4.0f, 5.0f, 6.0f});

    int seen = 0;
    bool bad = false;
    store.scan_vectors("emb", [&](std::string_view id, const float*, size_t n) {
        ++seen;
        if (n != 3) bad = true;
        if (is_identity_key(id)) bad = true;
    });
    check(seen == 2, "scan_vectors must skip the sentinel");
    check(!bad, "scan_vectors must not decode the sentinel as float data");

    ReadTxn r(t.env);
    MDB_dbi dbi = 0;
    mdb_dbi_open(r.raw(), "_vectors_emb", 0, &dbi);
    check(read_subdb_identity(r, dbi) == "_vectors_emb",
          "vector sub-db is stamped with its prefixed name");
}

// v2.4.4 Count is LMDB-first. Unfiltered it uses count(); filtered it uses
// scan() with limit=0 and reads total_matched. Pin that contract: total_matched
// is computed BEFORE pagination, so limit=0 must still report the true total
// while returning no documents.
void test_scan_limit_zero_reports_total() {
    TmpEnv t("counting");
    LmdbDocumentStore store(t.env);
    for (int i = 0; i < 5; ++i) {
        Document d;
        d.id = "id" + std::to_string(i);
        d.collection = "things";
        d.set_data(nlohmann::json{{"kind", i < 3 ? "alpha" : "beta"}});
        store.put("things", d.id, d);
    }

    smartbotic::database::Query q;
    q.limit = 0;
    auto all = store.scan("things", q);
    check(all.total_matched == 5, "limit=0 reports the full total");
    check(all.documents.empty(), "limit=0 returns no documents");
    check(store.count("things") == 5, "unfiltered count matches");

    smartbotic::database::Query fq;
    fq.limit = 0;
    smartbotic::database::Filter f;
    f.field = "kind";
    f.op = smartbotic::database::FilterOp::EQ;
    f.value = "alpha";
    fq.filters.push_back(f);
    auto filtered = store.scan("things", fq);
    check(filtered.total_matched == 3,
          "filtered limit=0 reports the matching total, not the collection size");
}

// v2.7.1 — the unfiltered/unsorted fast path in scan() must agree with the
// general path exactly. It exists because the general path decoded every
// document in the collection to return `limit` of them, so cost tracked total
// bytes rather than page size (382ms to return one 510-byte document from a
// 414 MB collection on a live instance). Any divergence here is a paging bug.
void test_scan_fast_path_matches_general_path() {
    TmpEnv t("fastpath");
    LmdbDocumentStore store(t.env);
    for (int i = 0; i < 25; ++i) {
        Document d;
        char buf[16];
        std::snprintf(buf, sizeof(buf), "id%02d", i);
        d.id = buf;
        d.collection = "things";
        d.set_data(nlohmann::json{{"n", i}, {"kind", i % 2 ? "odd" : "even"}});
        store.put("things", d.id, d);
    }

    // total_matched and has_more must match what a full count says.
    smartbotic::database::Query page;
    page.limit = 10;
    page.offset = 0;
    auto p0 = store.scan("things", page);
    check(p0.documents.size() == 10, "fast path returns exactly `limit` docs");
    check(p0.total_matched == 25, "fast path total_matched excludes the sentinel");
    check(p0.has_more, "has_more true when more remain");

    page.offset = 20;
    auto p2 = store.scan("things", page);
    check(p2.documents.size() == 5, "final page returns the remainder");
    check(p2.total_matched == 25, "total_matched stable across pages");
    check(!p2.has_more, "has_more false on the last page");

    page.offset = 25;
    auto p3 = store.scan("things", page);
    check(p3.documents.empty(), "offset past the end returns nothing");
    check(p3.total_matched == 25, "and still reports the true total");

    // Paging must cover every document exactly once, in a stable order.
    std::set<std::string> seen;
    for (uint32_t off = 0; off < 25; off += 7) {
        smartbotic::database::Query q;
        q.limit = 7;
        q.offset = off;
        for (const auto& d : store.scan("things", q).documents) seen.insert(d.id);
    }
    check(seen.size() == 25, "paging the whole collection yields every document once");

    // A filter forces the general path; it must still be correct.
    smartbotic::database::Query fq;
    fq.limit = 100;
    smartbotic::database::Filter f;
    f.field = "kind";
    f.op = smartbotic::database::FilterOp::EQ;
    f.value = "odd";
    fq.filters.push_back(f);
    auto filtered = store.scan("things", fq);
    check(filtered.total_matched == 12, "filtered path still counts matches, not rows");

    // A sort also forces the general path.
    smartbotic::database::Query sq;
    sq.limit = 3;
    sq.sort = smartbotic::database::Sort{"n", true};
    auto sorted = store.scan("things", sq);
    check(sorted.documents.size() == 3, "sorted path paginates");
    check(sorted.total_matched == 25, "sorted path totals all rows");
    check(sorted.documents[0].data().value("n", -1) == 24,
          "descending sort really sorted (fast path must not swallow sorts)");

    // limit=0 keeps meaning "no documents, but a true total" - the contract
    // Count depends on (see test_scan_limit_zero_reports_total).
    smartbotic::database::Query zq;
    zq.limit = 0;
    auto z = store.scan("things", zq);
    check(z.documents.empty(), "limit=0 returns no documents on the fast path");
    check(z.total_matched == 25, "limit=0 still reports the true total");
}

// v2.8.0 — a WRITE that aborts must not poison the collection.
//
// This is the v2.4.3 EINVAL bug in a third failure mode, observed live in
// production on 2.7.1: `find` on smartbotic-automation:workflows failed with
// "LMDB cursor_open: Invalid argument" on every attempt while every other
// collection was fine, and a restart was the only cure.
//
// Cause: open_for_write() cached the MDB_dbi immediately after mdb_dbi_open,
// BEFORE the caller committed. LMDB keeps a handle private to the opening
// transaction until it commits and CLOSES it if that transaction aborts - so any
// write that threw after the handle was cached (a failed mdb_put, a sentinel
// mismatch, a WriteTxn destructing uncommitted) left a closed handle in the
// cache, and every later operation on that collection failed EINVAL for the rest
// of the process's life.
//
// The abort is induced honestly here, with a key past LMDB's 511-byte limit, so
// the test exercises the same path a real failed write takes.
void test_aborted_write_does_not_poison_the_collection() {
    TmpEnv t("abortpoison");
    LmdbDocumentStore store(t.env);

    // Force a write that opens the sub-db and then fails: an oversized key makes
    // mdb_put return MDB_BAD_VALSIZE, which throws, so the WriteTxn aborts.
    const std::string huge_id(600, 'k');
    bool threw = false;
    try {
        store.put("poisoned", huge_id, make_doc(huge_id, "poisoned"));
    } catch (const std::exception&) {
        threw = true;
    }
    check(threw, "an oversized key really does fail the write");

    // The collection must still be usable. Before the fix, every one of these
    // failed with EINVAL because the cache held a handle LMDB had closed.
    bool ok_put = true;
    try {
        store.put("poisoned", "good", make_doc("good", "poisoned"));
    } catch (const std::exception&) {
        ok_put = false;
    }
    check(ok_put, "a later WRITE to the same collection still works");

    bool ok_read = true;
    try {
        smartbotic::database::Query q;
        q.limit = 10;
        auto res = store.scan("poisoned", q);
        check(res.documents.size() == 1, "and the document written after the abort is there");
    } catch (const std::exception&) {
        ok_read = false;
    }
    check(ok_read, "a later SCAN of the same collection still works (cursor_open)");

    bool ok_count = true;
    try {
        check(store.count("poisoned") == 1, "count is right after the abort");
    } catch (const std::exception&) {
        ok_count = false;
    }
    check(ok_count, "and count() does not throw");

    check(store.get("poisoned", "good").has_value(), "get() works after the abort");
}

}  // namespace

int main() {
    std::cout << "=== test_subdb_identity ===\n";
    test_sentinel_roundtrip();
    test_misbound_handle_is_refused();
    test_unstamped_subdb_is_permitted();
    test_identity_key_predicate();
    test_sentinel_invisible_through_store();
    test_vector_subdb_sentinel();
    test_scan_limit_zero_reports_total();
    test_scan_fast_path_matches_general_path();
    test_aborted_write_does_not_poison_the_collection();

    std::cout << "passed: " << g_pass << ", failed: " << g_fail << "\n";
    return g_fail == 0 ? 0 : 1;
}