utun2/src/chat/merkle_tree.c

#include "merkle_tree.h"
#include "../../lib/debug_config.h"
#include <string.h>

static int tree_error(sqlite3* db, const char* op, const char* ns) {
    DEBUG_ERROR(DEBUG_CATEGORY_MEMBER_SYNC, "merkle_tree: %s ns=%s: %s", op, ns ? ns : "-", db ? sqlite3_errmsg(db) : "no DB");
    return -1;
}

uint64_t merkle_sync_level_prefix(uint64_t key, uint8_t level) {
    if (level == 0) return 0;
    if (level > MT_MAX_LEVEL) return 0;
    return key & (~UINT64_C(0) << (64 - level * 5));
}

uint8_t merkle_sync_prefix_bytes(uint8_t level) {
    return level <= MT_MAX_LEVEL ? (uint8_t)((level * 5 + 7) / 8) : 0;
}

int merkle_tree_init(sqlite3* db) {
    const char* sql = "DROP TABLE IF EXISTS merkle_tree_hash;"
        "CREATE TABLE merkle_tree_hash(namespace TEXT NOT NULL, level INTEGER NOT NULL CHECK(level BETWEEN 0 AND 5),"
        "prefix64 INTEGER NOT NULL, hash BLOB NOT NULL CHECK(length(hash)=32),"
        "PRIMARY KEY(namespace,level,prefix64))";
    if (!db || sqlite3_exec(db, sql, NULL, NULL, NULL) != SQLITE_OK) return tree_error(db, "init", NULL);
    DEBUG_INFO(DEBUG_CATEGORY_MEMBER_SYNC, "merkle_tree: empty derived index initialized");
    return 0;
}

int merkle_tree_get(sqlite3* db, const char* ns, uint8_t level, uint64_t prefix, uint8_t hash[MT_HASH_SIZE]) {
    if (!db || !ns || !hash || level > MT_MAX_LEVEL || prefix != merkle_sync_level_prefix(prefix, level))
        return tree_error(db, "invalid get", ns);
    memset(hash, 0, MT_HASH_SIZE);
    sqlite3_stmt* st = NULL;
    int rc = sqlite3_prepare_v2(db, "SELECT hash FROM merkle_tree_hash WHERE namespace=? AND level=? AND prefix64=?", -1, &st, NULL);
    if (rc == SQLITE_OK) rc = sqlite3_bind_text(st, 1, ns, -1, SQLITE_STATIC);
    if (rc == SQLITE_OK) rc = sqlite3_bind_int(st, 2, level);
    if (rc == SQLITE_OK) rc = sqlite3_bind_int64(st, 3, (sqlite3_int64)prefix);
    if (rc == SQLITE_OK) rc = sqlite3_step(st);
    if (rc == SQLITE_ROW && sqlite3_column_bytes(st, 0) == MT_HASH_SIZE) memcpy(hash, sqlite3_column_blob(st, 0), MT_HASH_SIZE);
    else if (rc != SQLITE_DONE) { sqlite3_finalize(st); return tree_error(db, "read hash", ns); }
    sqlite3_finalize(st);
    return 0;
}

int merkle_tree_children(sqlite3* db, const char* ns, uint8_t level, uint64_t prefix,
        uint8_t hashes[MT_BUCKETS][MT_HASH_SIZE]) {
    if (level >= MT_MAX_LEVEL || prefix != merkle_sync_level_prefix(prefix, level)) return tree_error(db, "invalid children", ns);
    if (!db || !ns) return tree_error(db, "invalid children DB", ns);
    memset(hashes, 0, MT_BUCKETS * MT_HASH_SIZE);
    int shift = 64 - (level + 1) * 5;
    const char* sql = level == 0 ? "SELECT prefix64,hash FROM merkle_tree_hash WHERE namespace=? AND level=?"
        : "SELECT prefix64,hash FROM merkle_tree_hash WHERE namespace=? AND level=? AND prefix64>=? AND prefix64<=?";
    sqlite3_stmt* st = NULL;
    int rc = sqlite3_prepare_v2(db, sql, -1, &st, NULL);
    if (rc == SQLITE_OK) rc = sqlite3_bind_text(st, 1, ns, -1, SQLITE_STATIC);
    if (rc == SQLITE_OK) rc = sqlite3_bind_int(st, 2, level + 1);
    if (level && rc == SQLITE_OK) rc = sqlite3_bind_int64(st, 3, (sqlite3_int64)prefix);
    if (level && rc == SQLITE_OK) rc = sqlite3_bind_int64(st, 4, (sqlite3_int64)(prefix | (UINT64_C(31) << shift)));
    if (rc != SQLITE_OK) { sqlite3_finalize(st); return tree_error(db, "children query", ns); }
    while ((rc = sqlite3_step(st)) == SQLITE_ROW) {
        if (sqlite3_column_bytes(st, 1) != MT_HASH_SIZE) { sqlite3_finalize(st); return tree_error(db, "corrupt child", ns); }
        unsigned i = (unsigned)(((uint64_t)sqlite3_column_int64(st, 0) >> shift) & 31);
        memcpy(hashes[i], sqlite3_column_blob(st, 1), MT_HASH_SIZE);
    }
    sqlite3_finalize(st);
    if (rc != SQLITE_DONE) return tree_error(db, "children read", ns);
    return 0;
}

static int tree_store(sqlite3* db, const char* ns, uint8_t level, uint64_t prefix, const uint8_t hash[MT_HASH_SIZE]) {
    static const uint8_t empty[MT_HASH_SIZE];
    uint8_t old[MT_HASH_SIZE];
    if (merkle_tree_get(db, ns, level, prefix, old) < 0) return -1;
    if (!memcmp(old, hash, MT_HASH_SIZE)) return 0;
    int is_empty = !memcmp(hash, empty, MT_HASH_SIZE);
    const char* sql = is_empty ? "DELETE FROM merkle_tree_hash WHERE namespace=? AND level=? AND prefix64=?"
        : "INSERT OR REPLACE INTO merkle_tree_hash(namespace,level,prefix64,hash) VALUES(?,?,?,?)";
    sqlite3_stmt* st = NULL;
    int rc = sqlite3_prepare_v2(db, sql, -1, &st, NULL);
    if (rc == SQLITE_OK) rc = sqlite3_bind_text(st, 1, ns, -1, SQLITE_STATIC);
    if (rc == SQLITE_OK) rc = sqlite3_bind_int(st, 2, level);
    if (rc == SQLITE_OK) rc = sqlite3_bind_int64(st, 3, (sqlite3_int64)prefix);
    if (rc == SQLITE_OK && !is_empty) rc = sqlite3_bind_blob(st, 4, hash, MT_HASH_SIZE, SQLITE_STATIC);
    if (rc == SQLITE_OK) rc = sqlite3_step(st);
    sqlite3_finalize(st);
    if (rc != SQLITE_DONE) return tree_error(db, "write hash", ns);
    return 1;
}

int merkle_tree_update(sqlite3* db, const char* ns, uint64_t key, merkle_leaf_hash_fn leaf_hash, void* ctx) {
    EVP_MD_CTX* sha = EVP_MD_CTX_new();
    uint8_t hash[MT_HASH_SIZE], children[MT_BUCKETS][MT_HASH_SIZE];
    static const uint8_t empty[sizeof(children)];
    int changed = 0;
    if (!db || !ns || !leaf_hash || !sha) goto fail;
    for (int level = MT_MAX_LEVEL; level >= 0; level--) {
        uint64_t prefix = merkle_sync_level_prefix(key, (uint8_t)level);
        uint8_t domain[2] = { 0x4d, (uint8_t)level };
        if (EVP_DigestInit_ex(sha, EVP_sha256(), NULL) != 1 || EVP_DigestUpdate(sha, domain, sizeof(domain)) != 1) goto fail;
        int count;
        if (level == MT_MAX_LEVEL) {
            count = leaf_hash(ctx, ns, (uint8_t)level, prefix, sha);
            if (count < 0) goto fail;
        } else {
            if (merkle_tree_children(db, ns, (uint8_t)level, prefix, children) < 0) goto fail;
            count = memcmp(children, empty, sizeof(children)) != 0;
            if (EVP_DigestUpdate(sha, children, sizeof(children)) != 1) goto fail;
        }
        if (EVP_DigestFinal_ex(sha, hash, NULL) != 1) goto fail;
        if (!count) memset(hash, 0, sizeof(hash));
        int rc = tree_store(db, ns, (uint8_t)level, prefix, hash);
        if (rc < 0) goto fail;
        if (rc == 0) break;
        changed = 1;
    }
    EVP_MD_CTX_free(sha);
    return changed;
fail:
    EVP_MD_CTX_free(sha);
    return tree_error(db, "recompute", ns);
}