path: root/packages/meshbay-node/src/meshbay_node/chat/store.py

"""
MeshBay Node — SQLite-backed chat message store.

One database per group. The node is a relay and an archive: it stores what it
was handed, serves it back, and — once a group has chat encryption switched on —
cannot read any of it. Decryption happens in the client, which is the only place
that holds the epoch key (`docs/MESHBAY_DESIGN.md` §4.5).

Three things about the schema are load-bearing rather than incidental:

* **`payload` is bytes, always.** It used to be a UTF-8 string in practice, and
  the history path decoded it with `errors="replace"` — which substitutes
  U+FFFD for every byte that is not valid UTF-8, i.e. for most of a ciphertext.
  That would have corrupted history while live messages worked, which reads as
  an intermittent decryption bug rather than as a wire-format error.
* **`format` says how to read a row**, so messages written before a group turned
  encryption on keep rendering. Nothing is ever rewritten in place by the
  switch; see `chat encrypt-history` for the explicit, backed-up alternative.
* **`(device, nonce)` is unique.** The nonce is 96 random bits chosen per
  message by the sending device, so it is already required to be unique for
  AES-GCM to be safe — making it a key costs nothing and turns a replayed
  message (which is validly signed, being a copy of a real one) into an
  integrity error instead of a duplicate.
"""

import logging
import time
from dataclasses import dataclass
from pathlib import Path

import aiosqlite

log = logging.getLogger(__name__)


class ReplayedMessage(Exception):
    """This device has already sent a message under this nonce."""

# Plaintext, as every message was before chat encryption existed. Rows keep it
# for ever; nothing rewrites them.
FORMAT_PLAIN = 0
# Sealed under a chat epoch key: `payload` is the AES-256-GCM ciphertext,
# `nonce` its 96-bit nonce, `sig` the sender device's Ed25519 signature.
FORMAT_SEALED_V1 = 1

_SCHEMA = """
CREATE TABLE IF NOT EXISTS messages (
    id          INTEGER PRIMARY KEY AUTOINCREMENT,
    sender_id   TEXT NOT NULL,
    iteration   INTEGER NOT NULL,
    payload     BLOB NOT NULL,
    timestamp   REAL NOT NULL,
    thread_id   TEXT DEFAULT NULL,
    sender_name TEXT DEFAULT ''
);
CREATE INDEX IF NOT EXISTS idx_messages_ts ON messages(timestamp);
CREATE INDEX IF NOT EXISTS idx_messages_thread ON messages(thread_id);
"""

# Additive, every one with a default, so an existing chat.db opens unchanged and
# a node that is downgraded still reads its own rows. `CREATE TABLE IF NOT
# EXISTS` adds no column to a table that already exists — the same trap
# `create_all()` is recorded for on the hub — so each of these runs on its own
# and a duplicate-column error is the expected outcome on the second start.
_MIGRATIONS = (
    "ALTER TABLE messages ADD COLUMN sender_name TEXT DEFAULT ''",
    "ALTER TABLE messages ADD COLUMN format INTEGER NOT NULL DEFAULT 0",
    "ALTER TABLE messages ADD COLUMN epoch INTEGER NOT NULL DEFAULT 0",
    "ALTER TABLE messages ADD COLUMN device BLOB DEFAULT NULL",
    "ALTER TABLE messages ADD COLUMN nonce BLOB DEFAULT NULL",
    "ALTER TABLE messages ADD COLUMN sig BLOB DEFAULT NULL",
)

# A replay is a validly signed copy of a real message, so nothing about the
# signature refuses it. The nonce does: it is per message, per device, and a
# repeat is either an attack or a bug. Partial, because plaintext rows carry no
# nonce at all and NULLs are distinct in SQLite anyway.
_REPLAY_INDEX = (
    "CREATE UNIQUE INDEX IF NOT EXISTS idx_messages_replay "
    "ON messages(device, nonce) WHERE nonce IS NOT NULL"
)


@dataclass
class StoredMessage:
    id: int
    sender_id: str
    iteration: int
    payload: bytes
    timestamp: float
    thread_id: str | None
    sender_name: str = ""
    format: int = FORMAT_PLAIN
    epoch: int = 0
    device: bytes | None = None
    nonce: bytes | None = None
    sig: bytes | None = None


_COLUMNS = ("id, sender_id, iteration, payload, timestamp, thread_id, "
            "sender_name, format, epoch, device, nonce, sig")


def _row(r) -> StoredMessage:
    return StoredMessage(
        id=r[0], sender_id=r[1], iteration=r[2], payload=r[3],
        timestamp=r[4], thread_id=r[5], sender_name=r[6] or "",
        format=r[7] or FORMAT_PLAIN, epoch=r[8] or 0,
        device=r[9], nonce=r[10], sig=r[11],
    )


class ChatStore:
    """Async SQLite chat store for one group."""

    def __init__(self, db_path: Path):
        self._db_path = db_path
        self._db: aiosqlite.Connection | None = None

    async def open(self) -> None:
        self._db_path.parent.mkdir(parents=True, exist_ok=True)
        self._db = await aiosqlite.connect(str(self._db_path))
        await self._db.executescript(_SCHEMA)
        for statement in _MIGRATIONS:
            try:
                await self._db.execute(statement)
            except Exception:
                # Already applied. Swallowed per column rather than per batch:
                # one loop with a shared `try` would stop at the first
                # already-present column and silently skip every later one, so a
                # node upgraded twice would be missing the newest fields with
                # nothing to show for it.
                pass
        await self._db.execute(_REPLAY_INDEX)
        await self._db.commit()

    async def close(self) -> None:
        if self._db:
            await self._db.close()
            self._db = None

    async def __aenter__(self):
        await self.open()
        return self

    async def __aexit__(self, *_):
        await self.close()

    async def save_message(
        self,
        sender_id: str,
        iteration: int,
        payload: bytes,
        thread_id: str | None = None,
        sender_name: str = "",
        *,
        format: int = FORMAT_PLAIN,
        epoch: int = 0,
        device: bytes | None = None,
        nonce: bytes | None = None,
        sig: bytes | None = None,
    ) -> int:
        """
        Store a message. Returns the row id.

        Raises `ReplayedMessage` if this device has already used this nonce —
        see `_REPLAY_INDEX`. The caller must not turn that into a generic
        failure the sender retries: it means the message is already stored.
        """
        ts = time.time()
        try:
            cursor = await self._db.execute(
                "INSERT INTO messages (sender_id, iteration, payload, timestamp, "
                " thread_id, sender_name, format, epoch, device, nonce, sig) "
                "VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
                (sender_id, iteration, payload, ts, thread_id, sender_name,
                 format, epoch, device, nonce, sig),
            )
        except aiosqlite.IntegrityError as e:
            await self._db.rollback()
            raise ReplayedMessage(str(e)) from e
        await self._db.commit()
        return cursor.lastrowid

    async def get_messages(
        self,
        since: float = 0,
        limit: int = 100,
    ) -> list[StoredMessage]:
        """Get messages after a timestamp, most recent last."""
        cursor = await self._db.execute(
            f"SELECT {_COLUMNS} "
            "FROM messages WHERE timestamp > ? ORDER BY timestamp ASC LIMIT ?",
            (since, limit),
        )
        rows = await cursor.fetchall()
        return [_row(r) for r in rows]

    async def get_recent(self, limit: int = 100) -> list[StoredMessage]:
        """The newest `limit` messages, oldest first so they render in order.

        `get_messages(since=0, limit=n)` cannot answer this: it pages *forward*
        from the oldest, so a client opening a group with 300 messages was shown
        the first 200 and never the conversation it came for. Reaching the newest
        rows means ordering DESC in SQL and reversing here.
        """
        cursor = await self._db.execute(
            f"SELECT {_COLUMNS} "
            "FROM messages ORDER BY id DESC LIMIT ?",
            (limit,),
        )
        rows = await cursor.fetchall()
        return [_row(r) for r in reversed(rows)]

    async def get_before(self, before_id: int, limit: int = 50) -> list[StoredMessage]:
        """The `limit` messages immediately older than `before_id`, oldest first.

        The cursor is the row id, not the timestamp. Nothing makes `timestamp`
        unique — it is a float from `time.time()` — and a cursor on a value two
        rows can share is one that eventually skips a message or repeats it.
        `id` is AUTOINCREMENT: unique, and ordered by insertion.
        """
        cursor = await self._db.execute(
            f"SELECT {_COLUMNS} "
            "FROM messages WHERE id < ? ORDER BY id DESC LIMIT ?",
            (before_id, limit),
        )
        rows = await cursor.fetchall()
        return [_row(r) for r in reversed(rows)]

    async def has_before(self, message_id: int) -> bool:
        """Whether anything older than `message_id` exists — drives "load older"."""
        cursor = await self._db.execute(
            "SELECT 1 FROM messages WHERE id < ? LIMIT 1", (message_id,))
        return await cursor.fetchone() is not None

    async def get_thread(self, thread_id: str, limit: int = 100) -> list[StoredMessage]:
        """Get messages in a thread."""
        cursor = await self._db.execute(
            f"SELECT {_COLUMNS} "
            "FROM messages WHERE thread_id = ? ORDER BY timestamp ASC LIMIT ?",
            (thread_id, limit),
        )
        rows = await cursor.fetchall()
        return [_row(r) for r in rows]

    @property
    def db_path(self) -> Path:
        """Where this store lives — the re-encryption command backs it up."""
        return self._db_path

    async def count_by_format(self) -> tuple[int, int]:
        """(plaintext, sealed). What the operator is deciding from."""
        cursor = await self._db.execute(
            "SELECT format, COUNT(*) FROM messages GROUP BY format")
        counts = {row[0]: row[1] for row in await cursor.fetchall()}
        return (counts.get(FORMAT_PLAIN, 0), counts.get(FORMAT_SEALED_V1, 0))

    async def all_plaintext(self) -> list[StoredMessage]:
        """Every message still stored in the clear, oldest first."""
        cursor = await self._db.execute(
            f"SELECT {_COLUMNS} FROM messages WHERE format = ? ORDER BY id",
            (FORMAT_PLAIN,))
        return [_row(r) for r in await cursor.fetchall()]

    async def reseal(self, message_id: int, *, epoch: int, device: bytes,
                     nonce: bytes, ct: bytes, sig: bytes) -> None:
        """
        Replace one plaintext row with its sealed form. **Does not commit** —
        the caller commits once, so a re-encryption that fails half way leaves
        the database as it was rather than half readable.

        `sender_name` is cleared because it moves inside the envelope; leaving
        it would keep in the clear the one field the sealing was for.
        """
        await self._db.execute(
            "UPDATE messages SET format = ?, epoch = ?, device = ?, "
            " nonce = ?, payload = ?, sig = ?, sender_name = '' "
            "WHERE id = ?",
            (FORMAT_SEALED_V1, epoch, device, nonce, ct, sig, message_id))

    async def commit(self) -> None:
        await self._db.commit()

    async def delete_older_than(self, cutoff: float) -> int:
        """Retention. Returns how many rows went."""
        cursor = await self._db.execute(
            "DELETE FROM messages WHERE timestamp < ?", (cutoff,))
        await self._db.commit()
        return cursor.rowcount

    async def message_count(self) -> int:
        cursor = await self._db.execute("SELECT COUNT(*) FROM messages")
        row = await cursor.fetchone()
        return row[0]