2 files changed, 56 insertions, 4 deletions

diff --git a/packages/meshbay-hub/src/meshbay_hub/api/revocation.py b/packages/meshbay-hub/src/meshbay_hub/api/revocation.py
index c1697a7..bb88283 100644
--- a/packages/meshbay-hub/src/meshbay_hub/api/revocation.py
+++ b/packages/meshbay-hub/src/meshbay_hub/api/revocation.py
@@ -25,6 +25,7 @@ On receipt: immediately refuse JWT tokens matching the revoked user_id,
 and close active connections for that user.
 """
 
+import asyncio
 import base64
 import json
 import logging
@@ -50,6 +51,7 @@ router = APIRouter(tags=["revocation"])
 # ── Connected node registry ───────────────────────────────────────────────────
 
 _connected_nodes: dict[str, WebSocket] = {}   # node_id → websocket
+_punch_events: dict[str, asyncio.Event] = {}  # node_id → signaling event
 
 
 def get_connected_node_count() -> int:
@@ -122,12 +124,16 @@ async def node_websocket(ws: WebSocket):
         log.info("Node WS connected: %s", node_id[:8])
         await ws.send_text(json.dumps({"type": "auth_ok", "node_id": node_id}))
 
-        # Keep-alive loop — wait for ping or disconnect
+        # Message loop — handle ping, punch_ready, etc.
         while True:
             raw = await ws.receive_text()
             msg = json.loads(raw)
             if msg.get("type") == "ping":
                 await ws.send_text(json.dumps({"type": "pong"}))
+            elif msg.get("type") == "punch_ready":
+                event = _punch_events.get(node_id)
+                if event:
+                    event.set()
 
     except WebSocketDisconnect:
         log.info("Node WS disconnected: %s", (node_id or "unknown")[:8])
@@ -140,6 +146,44 @@ async def node_websocket(ws: WebSocket):
 
 # ── Admin revocation endpoint ─────────────────────────────────────────────────
 
+class IncomingRequest(BaseModel):
+    peer_ip:   str
+    peer_port: int
+
+
+@router.post("/v1/nodes/{node_id}/incoming", status_code=200)
+async def notify_incoming(
+    node_id: str,
+    body: IncomingRequest,
+    current_user: User = Depends(get_current_user),
+):
+    """
+    Signal a node that a client wants to connect (NAT punch coordination).
+    Hub forwards the request via WebSocket; node punches NAT and replies punch_ready.
+    """
+    ws = _connected_nodes.get(node_id)
+    if not ws:
+        raise HTTPException(status_code=404, detail="Node not connected")
+
+    event = asyncio.Event()
+    _punch_events[node_id] = event
+
+    await ws.send_text(json.dumps({
+        "type": "client_incoming",
+        "peer_ip": body.peer_ip,
+        "peer_port": body.peer_port,
+    }))
+
+    try:
+        await asyncio.wait_for(event.wait(), timeout=5.0)
+    except asyncio.TimeoutError:
+        raise HTTPException(status_code=504, detail="Node did not respond in time")
+    finally:
+        _punch_events.pop(node_id, None)
+
+    return {"status": "ready", "node_id": node_id}
+
+
 class RevokeRequest(BaseModel):
     target:    str   # "user" or "group"
     target_id: str
diff --git a/packages/meshbay-hub/src/meshbay_hub/api/users.py b/packages/meshbay-hub/src/meshbay_hub/api/users.py
index 5a7a3b4..5a2c7be 100644
--- a/packages/meshbay-hub/src/meshbay_hub/api/users.py
+++ b/packages/meshbay-hub/src/meshbay_hub/api/users.py
@@ -18,7 +18,7 @@ from meshbay_hub.auth import (
 )
 from meshbay_hub.config import HubConfig
 from meshbay_hub.db.engine import get_db
-from meshbay_hub.db.models import IPLog, RefreshToken, User
+from meshbay_hub.db.models import GroupMember, IPLog, RefreshToken, User
 from meshbay_hub.api.deps import get_current_user
 
 router = APIRouter(prefix="/v1/users", tags=["users"])
@@ -136,7 +136,11 @@ async def login(
     if user.status != "active":
         raise HTTPException(status_code=403, detail=f"Account {user.status}")
 
-    access_token = issue_access_token(user.id, user.pk_ed25519, ttl=_ttl())
+    memberships = await db.execute(
+        select(GroupMember.group_id).where(GroupMember.user_id == user.id))
+    group_ids = [gid for (gid,) in memberships.all()]
+    access_token = issue_access_token(
+        user.id, user.pk_ed25519, ttl=_ttl(), groups=group_ids)
     raw_rt, rt_hash = generate_refresh_token()
 
     expires_at = datetime.now(timezone.utc) + timedelta(seconds=_refresh_ttl())
@@ -175,7 +179,11 @@ async def token_refresh(
     if not user or user.status != "active":
         raise HTTPException(status_code=401, detail="User not found or suspended")
 
-    new_token = issue_access_token(user.id, user.pk_ed25519, ttl=_ttl())
+    memberships = await db.execute(
+        select(GroupMember.group_id).where(GroupMember.user_id == user.id))
+    group_ids = [gid for (gid,) in memberships.all()]
+    new_token = issue_access_token(
+        user.id, user.pk_ed25519, ttl=_ttl(), groups=group_ids)
     return {"access_token": new_token, "token_type": "bearer", "expires_in": _ttl()}