feat(id): P0-A integration — frame-seam charging, HTTP 429+Retry-After, chat limiter consumer, counter visibility (#3882)

pkgs/id/src/abuse.rs

@@ -97,6 +97,13 @@ impl Refusal {
         }
     }
 
+    /// The retry delay as whole seconds, rounded up, for an HTTP
+    /// `Retry-After` header. `None` for permanent blocks.
+    #[must_use]
+    pub fn retry_after_secs_ceil(&self) -> Option<u64> {
+        self.retry_after_ms.map(|ms| ms.div_ceil(1_000))
+    }
+
     /// A permanent operator block: no retry is offered.
     #[must_use]
     pub fn blocked_by_operator() -> Self {
@@ -107,6 +114,29 @@ impl Refusal {
     }
 }
 
+/// A [`Refusal`] as a standard error, so surfaces can propagate it
+/// with `?` / anyhow and callers can downcast back to the structured
+/// refusal (reason + bounded retry value) at their boundary.
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+pub struct RefusalError(pub Refusal);
+
+impl std::fmt::Display for RefusalError {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self.0.retry_after_ms {
+            Some(ms) => write!(f, "refused ({:?}); retry after {ms} ms", self.0.reason),
+            None => write!(f, "refused ({:?}); no retry offered", self.0.reason),
+        }
+    }
+}
+
+impl std::error::Error for RefusalError {}
+
+impl From<Refusal> for RefusalError {
+    fn from(refusal: Refusal) -> Self {
+        Self(refusal)
+    }
+}
+
 /// A token bucket in milli-units (1 token = 1000 milli-tokens), so
 /// refill arithmetic stays in integers: at `rate_per_sec` tokens per
 /// second, each elapsed millisecond adds exactly `rate_per_sec`
@@ -160,6 +190,13 @@ impl TokenBucket {
         }
     }
 
+    /// Refill and check whether `cost` whole tokens are available
+    /// without consuming them.
+    fn can_take(&mut self, cost: u64, now_ms: u64) -> bool {
+        self.refill(now_ms);
+        self.tokens_milli >= cost.saturating_mul(1_000)
+    }
+
     /// Milliseconds until `cost` tokens are available (0 if now).
     fn retry_after_ms(&mut self, cost: u64, now_ms: u64) -> u64 {
         self.refill(now_ms);
@@ -442,7 +479,8 @@ impl AbuseEvaluator {
             &mut state.preauth_frames
         };
         if bucket.try_take(1, now_ms) {
-            self.registry.record(self.local_node, endpoint, 1, 0, now_ms);
+            self.registry
+                .record(self.local_node, endpoint, 1, 0, now_ms);
             Ok(())
         } else {
             let retry = bucket.retry_after_ms(1, now_ms);
@@ -473,14 +511,70 @@ impl AbuseEvaluator {
             &mut state.preauth_bytes
         };
         if bucket.try_take(bytes, now_ms) {
-            self.registry.record(self.local_node, endpoint, 0, bytes, now_ms);
+            self.registry
+                .record(self.local_node, endpoint, 0, bytes, now_ms);
             Ok(())
         } else {
             let retry = bucket.retry_after_ms(bytes, now_ms);
             Err(Refusal::new(RefusalReason::BudgetExhausted, retry))
         }
     }
 
+    /// Atomically admits one frame and its byte cost. Neither bucket
+    /// is charged, and no counter is recorded, unless both budgets
+    /// can admit the frame.
+    ///
+    /// # Errors
+    ///
+    /// Returns the frame-rate refusal first when both budgets are
+    /// exhausted; otherwise returns the byte-budget refusal.
+    pub fn check_frame_and_bytes(
+        &mut self,
+        endpoint: EndpointId,
+        authorized: bool,
+        bytes: u64,
+        now_ms: u64,
+    ) -> Result<(), Refusal> {
+        let (frame_retry, byte_retry) = {
+            let state = self.state(endpoint, now_ms);
+            let (frames, byte_bucket) = if authorized {
+                (&mut state.authorized_frames, &mut state.authorized_bytes)
+            } else {
+                (&mut state.preauth_frames, &mut state.preauth_bytes)
+            };
+            let frame_retry =
+                (!frames.can_take(1, now_ms)).then(|| frames.retry_after_ms(1, now_ms));
+            let byte_retry = (!byte_bucket.can_take(bytes, now_ms))
+                .then(|| byte_bucket.retry_after_ms(bytes, now_ms));
+            (frame_retry, byte_retry)
+        };
+
+        if let Some(retry) = frame_retry {
+            return Err(Refusal::new(RefusalReason::RateLimited, retry));
+        }
+        if let Some(retry) = byte_retry {
+            return Err(Refusal::new(RefusalReason::BudgetExhausted, retry));
+        }
+
+        {
+            let state = self.state(endpoint, now_ms);
+            let (frames, byte_bucket) = if authorized {
+                (&mut state.authorized_frames, &mut state.authorized_bytes)
+            } else {
+                (&mut state.preauth_frames, &mut state.preauth_bytes)
+            };
+            // No other operation can mutate these buckets between the
+            // checks above and these debits: the evaluator is mutably
+            // borrowed for this whole call.
+            let frame_taken = frames.try_take(1, now_ms);
+            let bytes_taken = byte_bucket.try_take(bytes, now_ms);
+            debug_assert!(frame_taken && bytes_taken);
+        }
+        self.registry
+            .record(self.local_node, endpoint, 1, bytes, now_ms);
+        Ok(())
+    }
+
     /// The counter registry this evaluator records into.
     #[must_use]
     pub fn counters(&self) -> &CounterRegistry {

pkgs/id/src/frames.rs

@@ -11,6 +11,8 @@
 
 use anyhow::{Context, Result, bail};
 use iroh::endpoint::{ReadExactError, RecvStream, SendStream};
+
+use crate::abuse::{AbuseEvaluator, EndpointId, RefusalError};
 use tokio::io::{AsyncRead, AsyncReadExt as _, AsyncWrite, AsyncWriteExt as _};
 
 /// Largest accepted or emitted frame, in bytes (1 MiB).
@@ -139,6 +141,116 @@ impl FrameBuffer {
     }
 }
 
+/// Charge one frame of `payload_len` bytes against `endpoint`'s
+/// shared bucket in `abuse` (the #2244 P0-A evaluator): one frame
+/// token plus the byte cost, admitted or refused as a unit. This is
+/// the admission seam every charged transport funnels through — the
+/// read path calls it after the length prefix is known and before
+/// the body is allocated, the write path before the frame is
+/// encoded. A refusal sheds this frame/stream only; nothing here
+/// touches the peer connection.
+///
+/// # Errors
+///
+/// Returns the evaluator's structured [`RefusalError`] (reason +
+/// bounded retry value) when the bucket cannot admit the frame.
+pub fn charge_frame(
+    abuse: &mut AbuseEvaluator,
+    endpoint: EndpointId,
+    authorized: bool,
+    payload_len: usize,
+    now_ms: u64,
+) -> Result<(), RefusalError> {
+    abuse
+        .check_frame_and_bytes(endpoint, authorized, payload_len as u64, now_ms)
+        .map_err(RefusalError::from)
+}
+
+/// Read the length prefix of one frame from an Iroh stream.
+/// `Ok(None)` means the peer finished cleanly between frames. Split
+/// from [`read_frame_body`] so a charged reader can charge the frame
+/// (see [`charge_frame`]) after the prefix and before allocating the
+/// body, without holding any lock across the body read.
+///
+/// # Errors
+///
+/// Fails on a stream error or a length above [`MAX_FRAME_BYTES`].
+pub async fn read_frame_len(recv: &mut RecvStream) -> Result<Option<u32>> {
+    let mut len = [0_u8; 4];
+    match recv.read_exact(&mut len).await {
+        Ok(()) => {}
+        Err(ReadExactError::FinishedEarly(0)) => return Ok(None),
+        Err(e) => return Err(e).context("read frame length"),
+    }
+    let len = u32::from_be_bytes(len);
+    if len as usize > MAX_FRAME_BYTES {
+        bail!("frame of {len} bytes exceeds {MAX_FRAME_BYTES}");
+    }
+    Ok(Some(len))
+}
+
+/// Read the body of one frame whose length prefix ([`read_frame_len`])
+/// has already been consumed.
+///
+/// # Errors
+///
+/// Fails on a truncated frame or a stream error.
+pub async fn read_frame_body(recv: &mut RecvStream, len: u32) -> Result<Vec<u8>> {
+    let mut body = vec![0_u8; len as usize];
+    recv.read_exact(&mut body)
+        .await
+        .context("read frame body")?;
+    Ok(body)
+}
+
+/// Read one frame, charging it against `endpoint`'s shared bucket in
+/// `abuse` after the length prefix and before the body is allocated.
+/// `Ok(None)` means the peer finished cleanly between frames. The
+/// caller owns `abuse` for the duration of the read; transports that
+/// share one evaluator across streams charge via [`read_frame_len`] +
+/// [`charge_frame`] + [`read_frame_body`] instead, so the shared lock
+/// is never held across the body read.
+///
+/// # Errors
+///
+/// Fails with the structured [`RefusalError`] when the bucket refuses
+/// the frame, and on a truncated frame, a stream error, or a length
+/// above [`MAX_FRAME_BYTES`] (the oversized body is never read).
+pub async fn read_frame_charged(
+    recv: &mut RecvStream,
+    abuse: &mut AbuseEvaluator,
+    endpoint: EndpointId,
+    authorized: bool,
+    now_ms: u64,
+) -> Result<Option<Vec<u8>>> {
+    let Some(len) = read_frame_len(recv).await? else {
+        return Ok(None);
+    };
+    charge_frame(abuse, endpoint, authorized, len as usize, now_ms)?;
+    Ok(Some(read_frame_body(recv, len).await?))
+}
+
+/// Write one frame, charging it against `endpoint`'s shared bucket in
+/// `abuse` before the frame is encoded. A refusal sheds this frame
+/// only; the stream stays open.
+///
+/// # Errors
+///
+/// Fails with the structured [`RefusalError`] when the bucket refuses
+/// the frame, if the payload exceeds [`MAX_FRAME_BYTES`], or if the
+/// stream is closed.
+pub async fn write_frame_charged(
+    send: &mut SendStream,
+    abuse: &mut AbuseEvaluator,
+    endpoint: EndpointId,
+    authorized: bool,
+    now_ms: u64,
+    payload: &[u8],
+) -> Result<()> {
+    charge_frame(abuse, endpoint, authorized, payload.len(), now_ms)?;
+    write_frame(send, payload).await
+}
+
 /// Read one frame. `Ok(None)` means the peer finished cleanly between frames.
 ///
 /// This is **not** cancel-safe; run it in its own task, not inside `select!`.
@@ -202,6 +314,77 @@ mod tests {
         assert!(encode_frame(&vec![0; MAX_FRAME_BYTES + 1]).is_err());
     }
 
+    #[tokio::test]
+    async fn charged_frames_round_trip_and_count_both_directions() {
+        use crate::abuse::{AbusePolicy, BucketLimits};
+        let policy = AbusePolicy {
+            authorized: BucketLimits {
+                frames_per_sec: 1,
+                frame_burst: 2,
+                bytes_per_sec: 8,
+                byte_burst: 16,
+            },
+            ..AbusePolicy::default()
+        };
+        let mut abuse = AbuseEvaluator::new(policy, [9; 32]);
+        let endpoint = [7; 32];
+        let (mut client, mut server) = tokio::io::duplex(64);
+        // Write path charges via the encode funnel (write_frame_to on a
+        // byte stream charges the same seam as write_frame_charged).
+        charge_frame(&mut abuse, endpoint, true, 5, 0).unwrap();
+        write_frame_to(&mut client, b"hello").await.unwrap();
+        // Read path charges after the length prefix, before the body.
+        let mut len = [0_u8; 4];
+        server.read_exact(&mut len).await.unwrap();
+        let len = u32::from_be_bytes(len);
+        charge_frame(&mut abuse, endpoint, true, len as usize, 0).unwrap();
+        let mut body = vec![0_u8; len as usize];
+        server.read_exact(&mut body).await.unwrap();
+        assert_eq!(body, b"hello");
+        // Both directions recorded: 2 frames, 10 bytes for the pair.
+        let snapshot = abuse.counters().snapshot();
+        assert_eq!(snapshot.len(), 1);
+        assert_eq!(snapshot[0].1.frames, 2);
+        assert_eq!(snapshot[0].1.bytes, 10);
+        // The burst (2 frames) is spent: the next charge is refused
+        // with a machine-readable retry value.
+        let refusal = charge_frame(&mut abuse, endpoint, true, 1, 0).unwrap_err();
+        assert!(refusal.0.retry_after_ms.is_some());
+    }
+
+    #[test]
+    fn byte_refusal_does_not_spend_a_frame_token_or_increment_counters() {
+        use crate::abuse::{AbusePolicy, BucketLimits};
+        let policy = AbusePolicy {
+            authorized: BucketLimits {
+                frames_per_sec: 1,
+                frame_burst: 2,
+                bytes_per_sec: 1,
+                byte_burst: 4,
+            },
+            ..AbusePolicy::default()
+        };
+        let endpoint = [8; 32];
+        let mut abuse = AbuseEvaluator::new(policy, [9; 32]);
+
+        let refusal = charge_frame(&mut abuse, endpoint, true, 5, 0).unwrap_err();
+        assert_eq!(
+            refusal.0.reason,
+            crate::abuse::RefusalReason::BudgetExhausted
+        );
+        assert!(abuse.counters().get(&[9; 32], &endpoint).is_none());
+
+        // The refused frame did not consume the frame allowance.
+        charge_frame(&mut abuse, endpoint, true, 4, 0).unwrap();
+        assert_eq!(
+            abuse.counters().get(&[9; 32], &endpoint),
+            Some(crate::abuse::Counters {
+                frames: 1,
+                bytes: 4,
+            })
+        );
+    }
+
     #[test]
     fn the_frame_buffer_reassembles_frames_from_any_split() {
         let mut wire = encode_frame(b"hello").unwrap();

pkgs/id/src/web/routes.rs

@@ -1214,9 +1214,16 @@ async fn save_handler(State(state): State<AppState>, Json(req): Json<SaveRequest
             req.name,
             secs
         );
+        // Structured refusal (#2244 P0-A): machine-readable reason and
+        // retry value in the body, `Retry-After` header alongside.
         return (
             StatusCode::TOO_MANY_REQUESTS,
-            format!("Save rate limited. Try again in {secs}s."),
+            [(header::RETRY_AFTER, secs.to_string())],
+            Json(serde_json::json!({
+                "error": format!("Save rate limited. Try again in {secs}s."),
+                "reason": "rate_limited",
+                "retry_after_ms": secs.saturating_mul(1_000),
+            })),
         )
             .into_response();
     }

pkgs/id/src/web/world_ws.rs

@@ -353,6 +353,24 @@ pub(super) fn refusal_status(error: &anyhow::Error) -> StatusCode {
 }
 
 fn refusal_response(error: &anyhow::Error) -> Response {
+    // Structured refusal (#2244 P0-A): a rate-limit refusal carries the
+    // machine-readable reason and a bounded retry value, in the body
+    // and as a `Retry-After` header. The directory limiters behind
+    // this surface compute no bucket value of their own, so the abuse
+    // contract's floor (`RETRY_FLOOR_MS`) is what is emitted.
+    if matches!(error.downcast_ref::<Refusal>(), Some(Refusal::RateLimited(_))) {
+        let retry_ms = crate::abuse::RETRY_FLOOR_MS;
+        return (
+            StatusCode::TOO_MANY_REQUESTS,
+            [(header::RETRY_AFTER, (retry_ms / 1_000).to_string())],
+            Json(serde_json::json!({
+                "error": format!("{error:#}"),
+                "reason": "rate_limited",
+                "retry_after_ms": retry_ms,
+            })),
+        )
+            .into_response();
+    }
     (
         refusal_status(error),
         Json(serde_json::json!({ "error": format!("{error:#}") })),