Queue Durability for Offline MPC Parties

from __future__ import annotations

import hashlib
import hmac
from dataclasses import dataclass, field
from typing import Iterable, Mapping


class QueueOverflow(RuntimeError):
    """Raised when accepting a share would exceed a party's queue capacity."""


class ShareExpired(RuntimeError):
    """Raised when a share is delivered after its round's deadline."""


@dataclass(frozen=True)
class QueuedShare:
    round_id: str
    party_id: str
    ciphertext: bytes          # encrypted TO the recipient; the queue cannot read it
    enqueued_at: float
    deadline_at: float

    def key(self) -> tuple[str, str]:
        return (self.round_id, self.party_id)


@dataclass
class ShareQueue:
    """A per-party durable queue with bounded retention and idempotent delivery.

    The queue deliberately holds ciphertext it cannot decrypt: it routes and stores,
    and is not a participant in the protocol. That distinction is what keeps the
    threat model to the parties named in it.
    """

    capacity: int
    retention_s: float
    _items: dict[tuple[str, str], QueuedShare] = field(default_factory=dict)
    _delivered: set[tuple[str, str]] = field(default_factory=set)
    dropped: int = 0

    def enqueue(self, share: QueuedShare) -> None:
        if share.key() in self._delivered:
            return                              # idempotent: already applied
        if share.key() in self._items:
            return                              # idempotent: already queued
        if len(self._items) >= self.capacity:
            raise QueueOverflow(
                f"queue for {share.party_id} is full at {self.capacity}; dropping a "
                "share makes its round permanently unreconstructible — shed load "
                "upstream instead"
            )
        self._items[share.key()] = share

    def expire(self, now: float) -> list[QueuedShare]:
        """Remove shares past retention, returning them so the caller can alert."""
        gone = [s for s in self._items.values() if now - s.enqueued_at > self.retention_s]
        for s in gone:
            del self._items[s.key()]
            self.dropped += 1
        return gone

    def deliver(self, party_id: str, now: float) -> list[QueuedShare]:
        """Hand over every live share for a party that has come back online."""
        out: list[QueuedShare] = []
        for key, share in list(self._items.items()):
            if share.party_id != party_id:
                continue
            if now > share.deadline_at:
                del self._items[key]
                self.dropped += 1
                continue
            out.append(share)
            del self._items[key]
            self._delivered.add(key)
        return out

    def depth(self, party_id: str | None = None) -> int:
        if party_id is None:
            return len(self._items)
        return sum(1 for s in self._items.values() if s.party_id == party_id)


def round_is_recoverable(
    delivered: Iterable[QueuedShare], *, threshold: int, round_id: str
) -> bool:
    """A round reconstructs only once `threshold` distinct parties have their shares."""
    parties = {s.party_id for s in delivered if s.round_id == round_id}
    return len(parties) >= threshold

def _validate() -> None:
    q = ShareQueue(capacity=4, retention_s=3_600.0)

    def share(rid: str, pid: str, t: float = 0.0, deadline: float = 7_200.0) -> QueuedShare:
        payload = hashlib.blake2b(f"{rid}:{pid}".encode(), digest_size=32).digest()
        return QueuedShare(rid, pid, payload, t, deadline)

    # 1. Enqueue is idempotent: a redelivered share does not duplicate.
    q.enqueue(share("r1", "p1"))
    q.enqueue(share("r1", "p1"))
    assert q.depth() == 1

    # 2. Capacity is enforced by raising, not by silently dropping.
    for i in range(3):
        q.enqueue(share("r1", f"p{i+2}"))
    try:
        q.enqueue(share("r2", "p9"))
    except QueueOverflow as exc:
        assert "unreconstructible" in str(exc)
    else:
        raise AssertionError("overflow must raise, not drop")

    # 3. Delivery hands over a party's shares and marks them delivered.
    delivered = q.deliver("p1", now=10.0)
    assert len(delivered) == 1 and q.depth("p1") == 0

    # 4. Re-enqueueing a delivered share is a no-op, so an at-least-once broker
    #    cannot corrupt the round by redelivering.
    q.enqueue(share("r1", "p1"))
    assert q.depth("p1") == 0

    # 5. Shares past the round deadline are dropped at delivery, and counted.
    q2 = ShareQueue(capacity=10, retention_s=3_600.0)
    q2.enqueue(share("r3", "p5", t=0.0, deadline=100.0))
    assert q2.deliver("p5", now=500.0) == [] and q2.dropped == 1

    # 6. Retention expiry removes stale shares and reports them for alerting.
    q3 = ShareQueue(capacity=10, retention_s=60.0)
    q3.enqueue(share("r4", "p6", t=0.0))
    expired = q3.expire(now=120.0)
    assert len(expired) == 1 and q3.depth() == 0

    # 7. A round reconstructs only at the threshold, never before.
    shares = [share("r5", f"p{i}") for i in range(3)]
    assert not round_is_recoverable(shares[:2], threshold=3, round_id="r5")
    assert round_is_recoverable(shares, threshold=3, round_id="r5")

    print("share queue durability: all assertions passed")


_validate()