1 files changed, 98 insertions, 26 deletions
diff --git a/packages/meshbay-hub/src/meshbay_hub/api/federation.py b/packages/meshbay-hub/src/meshbay_hub/api/federation.py
index 7f1262d..b1e0e30 100644
--- a/packages/meshbay-hub/src/meshbay_hub/api/federation.py
+++ b/packages/meshbay-hub/src/meshbay_hub/api/federation.py
@@ -27,7 +27,7 @@ import uuid
 import jwt
 from fastapi import APIRouter, Depends, HTTPException, Header
 from pydantic import BaseModel
-from sqlalchemy import select
+from sqlalchemy import func, select
 from sqlalchemy.ext.asyncio import AsyncSession
 
 from meshbay_common import MHP_VERSION
@@ -41,6 +41,17 @@ log = logging.getLogger(__name__)
 
 router = APIRouter(prefix="/mhp", tags=["federation"])
 
+# One push may not dump the world, and one peer may not fill the table.
+MAX_FEDERATED_GROUPS_PER_PUSH = 500
+MAX_FEDERATED_GROUPS_PER_PEER = 2000
+
+# Seen `jti` values for the state-changing MHP endpoints, pruned lazily. The
+# sending side that would set an `aud` claim is unbuilt, so audience binding is
+# not available; this stops a captured POST /mhp/directory or /mhp/revoke from
+# being replayed inside the token's short TTL. GET /mhp/directory is idempotent
+# and not covered.
+_seen_mhp_jti: dict[str, float] = {}
+
 
 def _issue_mhp_token(target_hub_id: str) -> str:
     """Issue a short-lived JWT for authenticating to a peer hub."""
@@ -57,9 +68,15 @@ def _issue_mhp_token(target_hub_id: str) -> str:
 
 
 async def _verify_mhp_token(
-    token: str, db: AsyncSession, expected_aud: str | None = None,
+    token: str, db: AsyncSession, *, single_use: bool = False,
 ) -> dict:
-    """Verify a JWT from a peer hub using DB-stored public key."""
+    """
+    Verify a JWT from a peer hub against its DB-stored public key and return the
+    payload.
+
+    `single_use=True` (the state-changing endpoints) additionally rejects a
+    replayed `jti` within the token's lifetime.
+    """
     unverified = jwt.decode(token, options={"verify_signature": False})
     sender_id = unverified.get("iss")
 
@@ -67,15 +84,22 @@ async def _verify_mhp_token(
     if not peer:
         raise PermissionError(f"Unknown hub: {sender_id!r}. Register as peer first.")
 
-    options = {}
-    if expected_aud:
-        options["audience"] = expected_aud
-
     decoded = jwt.decode(
         token, peer.pk_hub_pem.encode(),
         algorithms=["EdDSA"],
-        options=options,
+        options={"require": ["exp", "iss"]},
     )
+
+    if single_use:
+        now = time.time()
+        for j, exp in list(_seen_mhp_jti.items()):
+            if exp < now:
+                _seen_mhp_jti.pop(j, None)
+        jti = decoded.get("jti", "")
+        if not jti or jti in _seen_mhp_jti:
+            raise PermissionError("MHP token replay")
+        _seen_mhp_jti[jti] = float(decoded.get("exp", now + 300))
+
     return decoded
 
 
@@ -142,30 +166,50 @@ async def receive_directory(
     db: AsyncSession = Depends(get_db),
 ):
     try:
-        await _verify_mhp_token(authorization.removeprefix("Bearer "), db)
+        payload = await _verify_mhp_token(
+            authorization.removeprefix("Bearer "), db, single_use=True)
     except Exception as e:
         raise HTTPException(status_code=401, detail=str(e))
 
+    # `source_hub` is the signer of this request, never `body.hub_id` — a peer
+    # does not get to relay or spoof a third hub's groups into our directory.
+    sender = payload["iss"]
+    if len(body.groups) > MAX_FEDERATED_GROUPS_PER_PUSH:
+        raise HTTPException(status_code=413, detail="Too many groups in one push")
+
     from datetime import datetime, timezone
     now = datetime.now(timezone.utc)
+    have = await db.scalar(
+        select(func.count()).select_from(FederatedGroup)
+        .where(FederatedGroup.source_hub == sender)) or 0
+
     count = 0
     for g in body.groups:
-        existing = await db.get(FederatedGroup, g["id"])
-        if existing:
-            existing.name        = g.get("name", existing.name)
-            existing.join_policy = g.get("join_policy", existing.join_policy)
-            existing.updated_at  = now
+        gid = str(g.get("id", ""))[:36]
+        name = str(g.get("name", ""))[:128]
+        jp = g.get("join_policy", "invite")
+        if not gid or jp not in ("invite", "open"):
+            continue
+        # A federated id must never shadow a real local group.
+        if await db.get(Group, gid):
+            log.warning("Federated id %s collides with a local group — skipped", gid[:8])
+            continue
+        row = await db.get(FederatedGroup, gid)
+        if row:
+            if row.source_hub != sender:
+                continue          # only the hub that advertised it may update it
+            row.name        = name or row.name
+            row.join_policy = jp
+            row.updated_at  = now
         else:
+            if have + count >= MAX_FEDERATED_GROUPS_PER_PEER:
+                break
             db.add(FederatedGroup(
-                id=g["id"],
-                name=g.get("name", ""),
-                source_hub=body.hub_id,
-                join_policy=g.get("join_policy", "invite"),
-            ))
+                id=gid, name=name, source_hub=sender, join_policy=jp))
         count += 1
     await db.commit()
-    log.info("Persisted %d groups from hub %s", count, body.hub_id[:16])
-    return {"accepted": count, "from_hub": body.hub_id}
+    log.info("Persisted %d groups from hub %s", count, sender[:16])
+    return {"accepted": count, "from_hub": sender}
 
 
 # ── Revocation propagation ────────────────────────────────────────────────────
@@ -179,15 +223,43 @@ async def receive_revocation(
     authorization: str = Header(...),
     db: AsyncSession = Depends(get_db),
 ):
+    """
+    Act on a revocation from a peer hub.
+
+    This does **not** reach local nodes: nothing here hosts a federated group,
+    and a local node would reject a token signed by another hub's key anyway
+    (that path was a silent no-op). What a peer may legitimately revoke is a
+    group *it advertised to us* — so this prunes our copy of the peer's
+    directory. Users are per-hub; a peer does not get to revoke ours.
+    """
     try:
-        await _verify_mhp_token(authorization.removeprefix("Bearer "), db)
+        payload = await _verify_mhp_token(
+            authorization.removeprefix("Bearer "), db, single_use=True)
     except Exception as e:
         raise HTTPException(status_code=401, detail=str(e))
 
-    from meshbay_hub.api.revocation import broadcast_revocation
-    sent = await broadcast_revocation(body.token)
-    log.info("Propagated revocation to %d local nodes", sent)
-    return {"propagated_to": sent}
+    sender = payload["iss"]
+    peer = await db.get(HubPeer, sender)
+    try:
+        inner = jwt.decode(
+            body.token, peer.pk_hub_pem.encode(), algorithms=["EdDSA"],
+            options={"verify_exp": False})
+    except Exception as e:
+        raise HTTPException(status_code=400, detail=f"Bad revocation token: {e}")
+
+    if inner.get("type") != "revocation" or inner.get("target") != "group":
+        return {"pruned": 0, "note": "federation may only revoke groups it advertised"}
+
+    target_id = inner.get("target_id", "")
+    row = await db.get(FederatedGroup, target_id)
+    pruned = 0
+    if row and row.source_hub == sender:
+        await db.delete(row)
+        await db.commit()
+        pruned = 1
+    log.info("Federated group %s revoked by %s (pruned=%d)",
+             target_id[:8], sender[:16], pruned)
+    return {"pruned": pruned}
 
 
 # ── Peer management (admin) ───────────────────────────────────────────────────