Harden production workflows and agent admission

server/monlet_server/services/ingestion.py

@@ -2,13 +2,22 @@ from datetime import UTC, datetime
 from uuid import UUID, uuid4
 
 from fastapi import HTTPException
-from sqlalchemy import select, text
+from sqlalchemy import func, select, text
 from sqlalchemy.dialects.postgresql import insert as pg_insert
 from sqlalchemy.ext.asyncio import AsyncSession
 
 from .. import metrics
+from ..agent_auth import AgentCredential
 from ..logging_config import get_logger
-from ..models import Agent, Check, Event, EventIngestDedup, Incident, NotificationOutbox
+from ..models import (
+    Agent,
+    AgentBlacklist,
+    Check,
+    Event,
+    EventIngestDedup,
+    Incident,
+    NotificationOutbox,
+)
 from ..redaction import redact
 from ..schemas import CheckResultEvent, HeartbeatRequest
 from ..settings import get_settings
@@ -37,28 +46,143 @@ def _incident_key(agent_id: str, ev: CheckResultEvent) -> str:
     return key
 
 
-async def upsert_agent_heartbeat(session: AsyncSession, hb: HeartbeatRequest) -> None:
+async def _blacklist_row(session: AsyncSession, cred: AgentCredential) -> AgentBlacklist | None:
+    res = await session.execute(
+        select(AgentBlacklist).where(AgentBlacklist.key_hash == cred.key_hash)
+    )
+    return res.scalar_one_or_none()
+
+
+async def reject_if_blocked(
+    session: AsyncSession,
+    cred: AgentCredential,
+    agent_id: str,
+    hostname: str,
+    observed: datetime | None = None,
+) -> None:
+    row = await _blacklist_row(session, cred)
+    if row is None:
+        return
+    raise HTTPException(status_code=403, detail="agent blocked")
+
+
+async def _ensure_pending_capacity(session: AsyncSession) -> None:
+    max_pending = get_settings().max_pending_agents
+    if max_pending <= 0:
+        return
+    count = (
+        await session.execute(
+            select(func.count()).select_from(Agent).where(Agent.pending_key_hash.is_not(None))
+        )
+    ).scalar_one()
+    if count >= max_pending:
+        raise HTTPException(status_code=429, detail="too many pending agents")
+
+
+async def _reject_hostname_collision(session: AsyncSession, agent_id: str, hostname: str) -> None:
+    row = (
+        await session.execute(
+            select(Agent).where(Agent.hostname == hostname, Agent.agent_id != agent_id)
+        )
+    ).scalar_one_or_none()
+    if row is not None:
+        raise HTTPException(status_code=409, detail="hostname already belongs to another agent")
+
+
+def _set_pending(
+    agent: Agent,
+    hb: HeartbeatRequest,
+    cred: AgentCredential,
+    observed: datetime,
+    features: dict,
+) -> None:
+    agent.pending_key_hash = cred.key_hash
+    agent.pending_key_fingerprint = cred.key_fingerprint
+    agent.pending_hostname = hb.hostname
+    agent.pending_seen_at = observed
+    agent.pending_features = features
+    agent.pending_labels = hb.labels
+
+
+async def upsert_agent_heartbeat(
+    session: AsyncSession, hb: HeartbeatRequest, cred: AgentCredential
+) -> str:
     observed = to_utc(hb.observed_at)
-    stmt = pg_insert(Agent).values(
-        agent_id=hb.agent_id,
-        hostname=hb.hostname,
-        status="alive",
-        last_seen_at=observed,
-        features=hb.features.model_dump(),
-        labels=hb.labels,
+    await reject_if_blocked(session, cred, hb.agent_id, hb.hostname, observed)
+
+    features = hb.features.model_dump()
+    res = await session.execute(select(Agent).where(Agent.accepted_key_hash == cred.key_hash))
+    accepted = res.scalar_one_or_none()
+    if accepted is not None:
+        await _reject_hostname_collision(session, accepted.agent_id, hb.hostname)
+        accepted.hostname = hb.hostname
+        accepted.status = "alive"
+        accepted.last_seen_at = observed
+        accepted.features = features
+        accepted.labels = hb.labels
+        metrics.heartbeats_total.labels(result="ok").inc()
+        return "accepted"
+
+    res = await session.execute(select(Agent).where(Agent.pending_key_hash == cred.key_hash))
+    pending = res.scalar_one_or_none()
+    if pending is not None:
+        await _reject_hostname_collision(session, pending.agent_id, hb.hostname)
+        _set_pending(pending, hb, cred, observed, features)
+        if pending.accepted_key_hash is None:
+            pending.hostname = hb.hostname
+            pending.status = "alive"
+            pending.last_seen_at = observed
+            pending.features = features
+            pending.labels = hb.labels
+        metrics.heartbeats_total.labels(result="pending").inc()
+        return "pending"
+
+    hostname_row = (
+        await session.execute(select(Agent).where(Agent.hostname == hb.hostname))
+    ).scalar_one_or_none()
+    if hostname_row is not None:
+        if hostname_row.pending_key_hash is not None:
+            raise HTTPException(
+                status_code=409,
+                detail="agent already has a different pending key",
+            )
+        await _ensure_pending_capacity(session)
+        _set_pending(hostname_row, hb, cred, observed, features)
+        metrics.heartbeats_total.labels(result="pending").inc()
+        return "pending"
+
+    agent_id_row = (
+        await session.execute(select(Agent).where(Agent.agent_id == hb.agent_id))
+    ).scalar_one_or_none()
+    if agent_id_row is not None:
+        raise HTTPException(status_code=409, detail="agent_id already exists")
+
+    await _ensure_pending_capacity(session)
+    session.add(
+        Agent(
+            agent_id=hb.agent_id,
+            hostname=hb.hostname,
+            status="alive",
+            last_seen_at=observed,
+            features=features,
+            labels=hb.labels,
+            pending_key_hash=cred.key_hash,
+            pending_key_fingerprint=cred.key_fingerprint,
+            pending_hostname=hb.hostname,
+            pending_seen_at=observed,
+            pending_features=features,
+            pending_labels=hb.labels,
+        )
     )
-    stmt = stmt.on_conflict_do_update(
-        index_elements=[Agent.agent_id],
-        set_={
-            "hostname": stmt.excluded.hostname,
-            "status": "alive",
-            "last_seen_at": stmt.excluded.last_seen_at,
-            "features": stmt.excluded.features,
-            "labels": stmt.excluded.labels,
-        },
-    )
-    await session.execute(stmt)
-    metrics.heartbeats_total.labels(result="ok").inc()
+    metrics.heartbeats_total.labels(result="pending").inc()
+    return "pending"
+
+
+async def accepted_agent_for_key(session: AsyncSession, cred: AgentCredential) -> Agent | None:
+    if await _blacklist_row(session, cred) is not None:
+        raise HTTPException(status_code=403, detail="agent blocked")
+    res = await session.execute(select(Agent).where(Agent.accepted_key_hash == cred.key_hash))
+    return res.scalar_one_or_none()
 
 
 async def _apply_incident(