arborist/tests/test_mesh_chain.py

"""Mesh wire — per-peer audit-chain tracking.

Receiver-side enforcement that gossip envelopes from a given peer
extend a single linear chain of claims. Catches:
- Tampered ``event_hash`` (recomputation mismatch -> 400).
- Forks / replays where ``prev_event_hash`` doesn't match the last
  envelope we accepted from that peer (-> 409).
- Cross-talk between peers' chains (alice's broadcasts must not
  affect bob's chain state and vice versa).

Implements the spec from ``docs/mesh.md``:
  > Audit chains merge by prev_event_hash. A gossip insert that does
  > not extend a chain consistently is rejected — mesh never silently
  > forks history.
"""

from __future__ import annotations

import time
from pathlib import Path

import pytest

from arborist.mesh import init_identity
from arborist.mesh.crypto import sign as _sign
from arborist.mesh.members import add_member
from arborist.mesh.wire import (
    MeshWireServer,
    TYPE_ANNOUNCE_ROOT,
    WireEnvelope,
)
from arborist.store import connect


def _bootstrap_three_peer_db(tmp_path: Path):
    """Alice (admin) + bob + carol on alice's DB, both enrolled at the
    same current_epoch (each ``add_member`` bumps the epoch by one).

    Returns (alice_db, current_epoch, alice_id, bob_id, carol_id). Alice's
    roster at ``current_epoch`` contains bob and carol so signature
    verification finds their sign_pubs when receiving from either.
    """
    from arborist.mesh.state import current_epoch as _current_epoch

    alice_db = tmp_path / "alice.db"
    bob_db = tmp_path / "bob.db"
    carol_db = tmp_path / "carol.db"

    a_conn = connect(alice_db)
    try:
        alice = init_identity(a_conn, group_name="t", member_id="alice")
    finally:
        a_conn.close()

    b_conn = connect(bob_db)
    try:
        bob = init_identity(b_conn, group_name="t", member_id="bob")
    finally:
        b_conn.close()

    c_conn = connect(carol_db)
    try:
        carol = init_identity(c_conn, group_name="t", member_id="carol")
    finally:
        c_conn.close()

    a_conn = connect(alice_db)
    try:
        add_member(a_conn, member_id="bob", sign_pub=bob.sign_pub, dh_pub=bob.dh_pub)
        add_member(a_conn, member_id="carol", sign_pub=carol.sign_pub, dh_pub=carol.dh_pub)
        epoch = _current_epoch(a_conn)
    finally:
        a_conn.close()

    return alice_db, epoch, alice, bob, carol


def _signed_envelope(
    *,
    sign_priv: bytes,
    sender_id: str,
    body: dict,
    prev_event_hash: str | None,
    epoch_id: int,
    ts: int | None = None,
) -> tuple[WireEnvelope, dict]:
    """Build a chain-linked, signed envelope. Returns (env, signed_dict)."""
    env = WireEnvelope(
        type=TYPE_ANNOUNCE_ROOT,
        sender_id=sender_id,
        epoch_id=epoch_id,
        ts=int(time.time()) if ts is None else ts,
        body=body,
    ).with_chain(prev_event_hash=prev_event_hash)
    return env, env.sign_with(sign_priv)


def _root_for(seed: str) -> str:
    """Deterministic 64-char hex doc-root marker for tests (not a real hash)."""
    return (seed * 64)[:64]


# ---------------------------------------------------------------------------
# Sender-side determinism
# ---------------------------------------------------------------------------


def test_with_chain_produces_deterministic_event_hash():
    """Same envelope + same prev_event_hash -> same event_hash, every time."""
    body = {
        "document_root": _root_for("a"),
        "source_uri": "https://x",
        "chunking_version": "tok-512-v1",
        "canonicalization_version": "norm-v1",
        "schema_version": "v9.8.0",
    }
    env_base = WireEnvelope(
        type=TYPE_ANNOUNCE_ROOT,
        sender_id="alice",
        epoch_id=1,
        ts=1700_000_000,
        body=body,
    )
    a = env_base.with_chain(prev_event_hash=None)
    b = env_base.with_chain(prev_event_hash=None)
    assert a.event_hash == b.event_hash
    assert a.event_hash is not None
    assert len(a.event_hash) == 64

    # Different prev -> different hash.
    c = env_base.with_chain(prev_event_hash="aa" * 32)
    assert c.event_hash != a.event_hash


def test_event_hash_excluded_from_canonical_bytes():
    """canonical_bytes() must NOT contain event_hash (avoids circularity).

    The literal hex of event_hash also must not appear, since the JSON
    serialization is the only path it could leak through. (The substring
    ``"event_hash"`` does appear because ``prev_event_hash`` is included;
    that's expected — only the *value* of event_hash is excluded.)
    """
    env = WireEnvelope(
        type=TYPE_ANNOUNCE_ROOT,
        sender_id="alice",
        epoch_id=1,
        ts=0,
        body={"document_root": _root_for("a")},
    ).with_chain(prev_event_hash=None)
    cb = env.canonical_bytes()
    assert env.event_hash is not None
    assert env.event_hash.encode() not in cb
    # The "event_hash" key itself is NOT serialized — only "prev_event_hash" is.
    assert b'"event_hash"' not in cb
    assert b'"prev_event_hash"' in cb


# ---------------------------------------------------------------------------
# Receiver-side: first envelope, chain extension, fork detection
# ---------------------------------------------------------------------------


def test_first_announce_with_prev_none_accepted(tmp_path):
    """A peer's very first envelope, prev_event_hash=None, is accepted."""
    alice_db, epoch, _alice, bob, _carol = _bootstrap_three_peer_db(tmp_path)
    _env, signed = _signed_envelope(
        sign_priv=bob.sign_priv,
        sender_id="bob",
        body={"document_root": _root_for("b"), "source_uri": "https://x"},
        prev_event_hash=None,
        epoch_id=epoch,
    )
    srv = MeshWireServer(alice_db, host="127.0.0.1", port=0)
    try:
        status, body = srv.handle_announce(signed)
    finally:
        srv.stop()
    assert status == 200, body

    a_conn = connect(alice_db)
    try:
        row = a_conn.execute(
            "SELECT last_event_hash, last_seq FROM mesh_peer_chains "
            "WHERE peer_member_id = 'bob'"
        ).fetchone()
    finally:
        a_conn.close()
    assert row is not None
    assert row["last_event_hash"] == _env.event_hash
    assert row["last_seq"] == 1


def test_chain_extension_accepted(tmp_path):
    """Second envelope with prev=first.event_hash extends the chain."""
    alice_db, epoch, _alice, bob, _carol = _bootstrap_three_peer_db(tmp_path)
    env1, signed1 = _signed_envelope(
        sign_priv=bob.sign_priv,
        sender_id="bob",
        body={"document_root": _root_for("1"), "source_uri": "https://x1"},
        prev_event_hash=None,
        epoch_id=epoch,
    )
    env2, signed2 = _signed_envelope(
        sign_priv=bob.sign_priv,
        sender_id="bob",
        body={"document_root": _root_for("2"), "source_uri": "https://x2"},
        prev_event_hash=env1.event_hash,
        epoch_id=epoch,
    )
    srv = MeshWireServer(alice_db, host="127.0.0.1", port=0)
    try:
        s1, _ = srv.handle_announce(signed1)
        s2, _ = srv.handle_announce(signed2)
    finally:
        srv.stop()
    assert s1 == 200
    assert s2 == 200

    a_conn = connect(alice_db)
    try:
        row = a_conn.execute(
            "SELECT last_event_hash, last_seq FROM mesh_peer_chains "
            "WHERE peer_member_id = 'bob'"
        ).fetchone()
    finally:
        a_conn.close()
    assert row["last_event_hash"] == env2.event_hash
    assert row["last_seq"] == 2


def test_fork_detection_rejects_409(tmp_path):
    """Second envelope with prev=GARBAGE is rejected — chain forked."""
    alice_db, epoch, _alice, bob, _carol = _bootstrap_three_peer_db(tmp_path)
    env1, signed1 = _signed_envelope(
        sign_priv=bob.sign_priv,
        sender_id="bob",
        body={"document_root": _root_for("1"), "source_uri": "https://x"},
        prev_event_hash=None,
        epoch_id=epoch,
    )
    # Bob (or someone using bob's key) sends a second envelope claiming
    # prev = arbitrary garbage instead of env1.event_hash.
    _env2, signed2 = _signed_envelope(
        sign_priv=bob.sign_priv,
        sender_id="bob",
        body={"document_root": _root_for("2"), "source_uri": "https://x"},
        prev_event_hash="ff" * 32,
        epoch_id=epoch,
    )
    srv = MeshWireServer(alice_db, host="127.0.0.1", port=0)
    try:
        s1, _ = srv.handle_announce(signed1)
        s2, body = srv.handle_announce(signed2)
    finally:
        srv.stop()
    assert s1 == 200
    assert s2 == 409
    assert "fork detected" in body["error"]
    assert "bob" in body["error"]

    # The fork rejection must NOT advance the peer's tracked chain.
    a_conn = connect(alice_db)
    try:
        row = a_conn.execute(
            "SELECT last_event_hash, last_seq FROM mesh_peer_chains "
            "WHERE peer_member_id = 'bob'"
        ).fetchone()
    finally:
        a_conn.close()
    assert row["last_event_hash"] == env1.event_hash
    assert row["last_seq"] == 1


def test_tampered_event_hash_rejected_400(tmp_path):
    """An envelope whose event_hash doesn't match recomputation is 400."""
    alice_db, epoch, _alice, bob, _carol = _bootstrap_three_peer_db(tmp_path)
    env, signed = _signed_envelope(
        sign_priv=bob.sign_priv,
        sender_id="bob",
        body={"document_root": _root_for("a"), "source_uri": "https://x"},
        prev_event_hash=None,
        epoch_id=epoch,
    )
    # Tamper with the event_hash field without re-signing. The envelope
    # signature itself remains valid (event_hash is excluded from the
    # signed canonical bytes), so the receiver catches this at the
    # event_hash recomputation step, not the signature step.
    signed["envelope"]["event_hash"] = "deadbeef" * 8
    assert signed["envelope"]["event_hash"] != env.event_hash

    srv = MeshWireServer(alice_db, host="127.0.0.1", port=0)
    try:
        status, body = srv.handle_announce(signed)
    finally:
        srv.stop()
    assert status == 400
    assert "event_hash mismatch" in body["error"]

    # No peer-chain row written for a rejected envelope.
    a_conn = connect(alice_db)
    try:
        row = a_conn.execute(
            "SELECT * FROM mesh_peer_chains WHERE peer_member_id = 'bob'"
        ).fetchone()
    finally:
        a_conn.close()
    assert row is None


def test_chain_tracking_advances_across_multiple_announces(tmp_path):
    """Receiver tracks last_event_hash + last_seq correctly over N announces."""
    alice_db, epoch, _alice, bob, _carol = _bootstrap_three_peer_db(tmp_path)
    srv = MeshWireServer(alice_db, host="127.0.0.1", port=0)
    try:
        prev = None
        last_eh = None
        for i in range(5):
            env, signed = _signed_envelope(
                sign_priv=bob.sign_priv,
                sender_id="bob",
                body={"document_root": _root_for(str(i)), "source_uri": f"https://x/{i}"},
                prev_event_hash=prev,
                epoch_id=epoch,
            )
            status, _ = srv.handle_announce(signed)
            assert status == 200, f"announce {i} failed"
            prev = env.event_hash
            last_eh = env.event_hash
    finally:
        srv.stop()

    a_conn = connect(alice_db)
    try:
        row = a_conn.execute(
            "SELECT last_event_hash, last_seq FROM mesh_peer_chains "
            "WHERE peer_member_id = 'bob'"
        ).fetchone()
    finally:
        a_conn.close()
    assert row["last_event_hash"] == last_eh
    assert row["last_seq"] == 5


def test_two_peers_have_independent_chains(tmp_path):
    """Bob's announces don't affect carol's chain state and vice versa."""
    alice_db, epoch, _alice, bob, carol = _bootstrap_three_peer_db(tmp_path)

    # Bob: 2 envelopes
    bob_env1, bob_signed1 = _signed_envelope(
        sign_priv=bob.sign_priv,
        sender_id="bob",
        body={"document_root": _root_for("b1"), "source_uri": "https://b1"},
        prev_event_hash=None,
        epoch_id=epoch,
    )
    bob_env2, bob_signed2 = _signed_envelope(
        sign_priv=bob.sign_priv,
        sender_id="bob",
        body={"document_root": _root_for("b2"), "source_uri": "https://b2"},
        prev_event_hash=bob_env1.event_hash,
        epoch_id=epoch,
    )
    # Carol: 1 envelope, prev=None (her chain starts fresh — should be
    # accepted even though bob has already posted).
    carol_env1, carol_signed1 = _signed_envelope(
        sign_priv=carol.sign_priv,
        sender_id="carol",
        body={"document_root": _root_for("c1"), "source_uri": "https://c1"},
        prev_event_hash=None,
        epoch_id=epoch,
    )
    srv = MeshWireServer(alice_db, host="127.0.0.1", port=0)
    try:
        assert srv.handle_announce(bob_signed1)[0] == 200
        assert srv.handle_announce(carol_signed1)[0] == 200  # carol fresh start
        assert srv.handle_announce(bob_signed2)[0] == 200    # bob extends his own chain
    finally:
        srv.stop()

    a_conn = connect(alice_db)
    try:
        bob_row = a_conn.execute(
            "SELECT last_event_hash, last_seq FROM mesh_peer_chains "
            "WHERE peer_member_id = 'bob'"
        ).fetchone()
        carol_row = a_conn.execute(
            "SELECT last_event_hash, last_seq FROM mesh_peer_chains "
            "WHERE peer_member_id = 'carol'"
        ).fetchone()
    finally:
        a_conn.close()
    assert bob_row["last_event_hash"] == bob_env2.event_hash
    assert bob_row["last_seq"] == 2
    assert carol_row["last_event_hash"] == carol_env1.event_hash
    assert carol_row["last_seq"] == 1


# ---------------------------------------------------------------------------
# Schema migration: peer-chains table appears on legacy DBs
# ---------------------------------------------------------------------------


def test_mesh_peer_chains_table_exists_after_connect(tmp_path):
    """Fresh connect() creates mesh_peer_chains via SCHEMA_SQL + migration."""
    db = tmp_path / "fresh.db"
    conn = connect(db)
    try:
        row = conn.execute(
            "SELECT name FROM sqlite_master "
            "WHERE type='table' AND name='mesh_peer_chains'"
        ).fetchone()
    finally:
        conn.close()
    assert row is not None


def test_mesh_peer_chains_migration_idempotent(tmp_path):
    """Two connects in a row don't ALTER twice."""
    db = tmp_path / "twice.db"
    connect(db).close()
    conn = connect(db)
    try:
        cnt = conn.execute(
            "SELECT COUNT(*) FROM sqlite_master "
            "WHERE type='table' AND name='mesh_peer_chains'"
        ).fetchone()[0]
    finally:
        conn.close()
    assert cnt == 1


def test_legacy_db_without_peer_chains_migrates(tmp_path):
    """A DB created before the peer-chain rollout gains the table on open."""
    import sqlite3 as _s
    db = tmp_path / "legacy.db"
    raw = _s.connect(db)
    raw.execute("CREATE TABLE meta (key TEXT PRIMARY KEY, value TEXT NOT NULL)")
    raw.commit()
    raw.close()

    # Pre-state: no mesh_peer_chains.
    raw = _s.connect(db)
    pre = raw.execute(
        "SELECT name FROM sqlite_master WHERE name='mesh_peer_chains'"
    ).fetchone()
    raw.close()
    assert pre is None

    # Open via connect() — runs SCHEMA_SQL + migration.
    conn = connect(db)
    try:
        post = conn.execute(
            "SELECT name FROM sqlite_master WHERE name='mesh_peer_chains'"
        ).fetchone()
    finally:
        conn.close()
    assert post is not None


# ---------------------------------------------------------------------------
# Backward compat: legacy envelopes (no chain fields) still accepted
# ---------------------------------------------------------------------------


def test_legacy_envelope_without_chain_fields_accepted(tmp_path):
    """A WireEnvelope with event_hash=None (legacy fixture path) is accepted
    without per-peer chain enforcement. Sig check still runs."""
    alice_db, epoch, _alice, bob, _carol = _bootstrap_three_peer_db(tmp_path)
    # Built directly without with_chain() — both prev and event_hash are None.
    env = WireEnvelope(
        type=TYPE_ANNOUNCE_ROOT,
        sender_id="bob",
        epoch_id=epoch,
        ts=int(time.time()),
        body={"document_root": _root_for("a"), "source_uri": "https://x"},
    )
    assert env.event_hash is None
    signed = env.sign_with(bob.sign_priv)
    srv = MeshWireServer(alice_db, host="127.0.0.1", port=0)
    try:
        status, _ = srv.handle_announce(signed)
    finally:
        srv.stop()
    assert status == 200

    # Legacy path doesn't write a peer-chain row.
    a_conn = connect(alice_db)
    try:
        row = a_conn.execute(
            "SELECT * FROM mesh_peer_chains WHERE peer_member_id='bob'"
        ).fetchone()
    finally:
        a_conn.close()
    assert row is None