From 8de53eceeb74dbbbf9cdc73da49aaa38904ae435 Mon Sep 17 00:00:00 2001 From: Kyryl R Date: Wed, 28 Oct 2026 10:53:52 +0200 Subject: [PATCH 1/9] Preserve backlog expiry and poll deadlines during drain Keep TTL metadata across reconnect and replacement, discard expired held frames, and drain only available queue slots per read-loop iteration. Use an absolute idle deadline so slow backlog delivery cannot defer token or liveness checks. --- crates/relay/src/connection.rs | 61 +++-- crates/relay/src/server.rs | 81 ++++++ crates/relay/src/session.rs | 435 ++++++++++++++++++++++++++++----- 3 files changed, 498 insertions(+), 79 deletions(-) diff --git a/crates/relay/src/connection.rs b/crates/relay/src/connection.rs index 8ad98f9..291a8f2 100644 --- a/crates/relay/src/connection.rs +++ b/crates/relay/src/connection.rs @@ -319,6 +319,7 @@ async fn write_loop( ping.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay); let mut close_frame = None; + let mut session_ready_sent = false; loop { tokio::select! { received = rx.recv() => match received { @@ -331,9 +332,11 @@ async fn write_loop( if !send_with_timeout(&mut sink, Message::text(text), config.write_timeout).await { break; } - if session_ready { - // The peer has been told it is ready, so the backlog - // held for it may follow. + // Once the peer has been told it is ready, the backlog + // held for it may follow, and every later write makes room + // in the queue for more of it. + session_ready_sent |= session_ready; + if session_ready_sent { backlog_ready.notify_one(); } } @@ -448,7 +451,8 @@ struct ReadContext<'a> { /// Everything a read loop needs to deliver a peer's backlog. struct BacklogDelivery { - /// Armed by the writer after `session_ready` has been written. + /// Armed by the writer after `session_ready` has been written, and after + /// every write from then on. ready: Arc, /// The slot's channel. It is weak because the slot's handle must stay the @@ -480,7 +484,12 @@ async fn read_loop( let mut consecutive_drops = 0u32; let mut redelivery_credit = 0u32; + let mut backlog_pending = true; + let mut idle_deadline = now + config.idle_timeout; + // Every deadline is absolute and polled on each iteration, and no arm + // awaits anything, so neither a slow writer nor a backlog its counterpart + // keeps refilling can hold the connection past them. loop { let message = tokio::select! { biased; @@ -500,29 +509,37 @@ async fn read_loop( ); return Err(ConnectionError::TokenExpired); } - () = backlog.ready.notified() => { - if let Some(tx) = backlog.tx.upgrade() { - let drained = session.drain_backlog(query.role, &tx).await?; - redelivery_credit = redelivery_credit - .saturating_add(u32::try_from(drained).unwrap_or(u32::MAX)); - } + _ = tokio::time::sleep_until(idle_deadline) => { + log::debug!( + "peer on topic {} idle for {:?}; dropping", + query.topic, + config.idle_timeout + ); + return Err(ConnectionError::IdleTimeout); + } + () = backlog.ready.notified(), if backlog_pending => { + // One step: what fits in the queue now. The writer wakes this + // arm again once it has made room. + backlog_pending = match backlog.tx.upgrade() { + Some(tx) => { + let drain = session.drain_backlog(query.role, &tx)?; + redelivery_credit = redelivery_credit + .saturating_add(u32::try_from(drain.delivered).unwrap_or(u32::MAX)); + !drain.finished + } + // The slot let go of the queue: this connection was replaced. + None => false, + }; continue; } - incoming = tokio::time::timeout(config.idle_timeout, source.next()) => match incoming { - Ok(Some(Ok(message))) => message, - Ok(Some(Err(error))) => return Err(error.into()), - Ok(None) => break, - Err(_) => { - log::debug!( - "peer on topic {} idle for {:?}; dropping", - query.topic, - config.idle_timeout - ); - return Err(ConnectionError::IdleTimeout); - } + incoming = source.next() => match incoming { + Some(Ok(message)) => message, + Some(Err(error)) => return Err(error.into()), + None => break, }, }; + idle_deadline = tokio::time::Instant::now() + config.idle_timeout; let _ = ping_deadline.send(None); let frame = match message { diff --git a/crates/relay/src/server.rs b/crates/relay/src/server.rs index 375a760..e023b80 100644 --- a/crates/relay/src/server.rs +++ b/crates/relay/src/server.rs @@ -1819,6 +1819,87 @@ mod tests { Ok(()) } + /// A reader too slow to take its backlog must not hold its connection past + /// the token's expiry: the deadline is observed while the backlog is still + /// being handed over, not once the reader has caught up. + #[tokio::test] + async fn token_expiry_closes_a_connection_stuck_on_its_backlog() -> TestResult { + // Far more than the loopback socket buffers hold. + const BACKLOG: usize = 160; + + let unlimited = RateLimit { + burst: u32::MAX, + interval: Duration::from_secs(1), + }; + // The write deadline is far beyond the test, so only the token can end + // the stuck connection in time. + let addr = serve_with_limits( + NativeRelayTransportConfig::new(AUDIENCE) + .with_outbound_queue_capacity(2) + .with_write_timeout(Duration::from_secs(60)), + RelayLimitsConfig { + inbound_message_rate: unlimited, + inbound_byte_rate: unlimited, + identity_inbound_byte_rate: unlimited, + global_inbound_byte_rate: unlimited, + ..RelayLimitsConfig::default() + }, + ) + .await?; + + let topic = topic()?; + let mut proposer = connect(addr, "proposer", &topic, None).await?; + let mut responder = connect_auth(addr, &token_with(8), "responder", &topic, None).await?; + read_frame(&mut proposer).await?; + read_frame(&mut responder).await?; + + drop(responder); + sleep(Duration::from_millis(100)).await; + let payload = "a".repeat(100 * 1024); + for i in 0..BACKLOG { + send(&mut proposer, &message(&topic, &format!("m{i}"), &payload)).await?; + } + // Let the relay buffer all of it before the responder returns. + sleep(Duration::from_millis(500)).await; + + // It comes back with a token that expires within two seconds, and does + // not read until well after that. + let now = crate::auth::unix_now(); + let short = test_utils::token(&signing_key(8), AUDIENCE, now, now + 2); + let mut stuck = connect_auth(addr, &short, "responder", &topic, None).await?; + assert!( + matches!( + read_frame(&mut proposer).await?.body, + RelayFrameBody::SessionReady + ), + "the responder must reclaim its slot and backlog within the grace window" + ); + sleep(Duration::from_secs(3)).await; + + let mut delivered = 0usize; + let closed = loop { + match timeout(Duration::from_secs(5), stuck.next()).await { + Err(_) => break false, + Ok(None | Some(Err(_)) | Some(Ok(Message::Close(_)))) => break true, + Ok(Some(Ok(Message::Text(text)))) => { + let frame: RelayFrame = serde_json::from_str(&text)?; + if matches!(frame.body, RelayFrameBody::Message { .. }) { + delivered += 1; + } + } + Ok(Some(Ok(_))) => {} + } + }; + assert!(closed, "the relay kept the expired connection open"); + // What the socket buffers held at expiry arrives, but the relay must + // not have gone on handing the backlog to an expired connection. + assert!( + (1..BACKLOG).contains(&delivered), + "expected the relay to stop part-way through the backlog, got {delivered}/{BACKLOG}" + ); + Ok(()) + } + fn tight_inbound_limits() -> RelayLimitsConfig { RelayLimitsConfig { inbound_message_rate: RateLimit { diff --git a/crates/relay/src/session.rs b/crates/relay/src/session.rs index 3f78658..ea555f3 100644 --- a/crates/relay/src/session.rs +++ b/crates/relay/src/session.rs @@ -61,7 +61,9 @@ enum Slot { /// so the per-identity cap cannot refuse a documented takeover (§3). permit: Option, - backlog: VecDeque, + /// Frames not yet handed to the outbound queue. They keep the expiry + /// they were buffered with. + backlog: VecDeque, }, Disconnected { identity: DidKey, @@ -70,12 +72,28 @@ enum Slot { }, } -/// A frame buffered for a disconnected peer, with its expiry time. +impl Slot { + /// Frames held for this slot's peer: its backlog while connected, its + /// buffer while disconnected. + fn held_frames(&mut self) -> &mut VecDeque { + match self { + Slot::Connected { backlog, .. } => backlog, + Slot::Disconnected { buffered, .. } => buffered, + } + } +} + +/// A frame held for a peer, with its expiry time. struct BufferedFrame { frame: RelayFrame, expires_at: Instant, } +/// Drop the frames whose TTL has passed. +fn prune_expired(frames: &mut VecDeque, now: Instant) { + frames.retain(|entry| entry.expires_at > now); +} + /// The two role slots for one topic. #[derive(Default)] struct Slots { @@ -101,6 +119,21 @@ impl Slots { fn is_connected(&self, role: Role) -> bool { matches!(self.get(role), Some(Slot::Connected { .. })) } + + /// Drop the frames held for `role`'s peer: they were sent by a + /// counterpart whose reservation has expired. + fn discard_held_frames(&mut self, role: Role) { + if let Some(slot) = self.get_mut(role) { + let held = slot.held_frames(); + if !held.is_empty() { + log::debug!( + "discarding {} frame(s) from an expired counterpart", + held.len() + ); + held.clear(); + } + } + } } /// Result of expiring a reservation. @@ -109,6 +142,19 @@ pub(crate) struct Expiry { pub empty: bool, } +/// Result of one [`Session::drain_backlog`] step. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) struct Drain { + /// Frames handed to the outbound queue. + pub delivered: usize, + + /// Whether this connection has nothing left to drain: the backlog is + /// empty, or the slot or its queue no longer belongs to the connection. + /// A connection's backlog never grows back once empty, since frames are + /// only queued behind a non-empty one. + pub finished: bool, +} + impl Session { /// Acquire the slot lock fn lock(&self) -> Result, SessionError> { @@ -129,9 +175,10 @@ impl Session { /// peers while the slot lock is held, so it precedes anything /// [`Session::route`] can send afterwards. Frames buffered for this /// identity while it was disconnected, or not yet handed to a replaced - /// connection, become the slot's backlog; hand them over with - /// [`Session::drain_backlog`]. A reservation claimed by another identity - /// keeps no frames. + /// connection, become the slot's backlog with their original expiry; hand + /// them over with [`Session::drain_backlog`]. An expired reservation + /// claimed by another identity keeps no frames, and the frames its previous + /// owner sent to the counterpart are discarded. /// /// `permit` is the connection's per-identity permit. A registration that /// replaces this identity's own live connection inherits that connection's @@ -169,22 +216,20 @@ impl Session { return Err(SessionError::IdentityLimit); } - let (backlog, permit) = match slots.get_mut(role).take() { + let previous = slots.get_mut(role).take(); + let (mut backlog, permit) = match previous { // A reservation reclaimed by the same identity keeps its frames, // whether or not the grace window has passed. Some(Slot::Disconnected { identity, buffered, .. - }) if identity == peer.identity => ( - buffered - .into_iter() - .filter(|entry| entry.expires_at > now) - .map(|entry| entry.frame) - .collect(), - permit, - ), + }) if identity == peer.identity => (buffered, permit), // An expired reservation claimed by another identity must not leak - // the previous owner's frames. - Some(Slot::Disconnected { .. }) => (VecDeque::new(), permit), + // the previous owner's frames, nor deliver what the previous owner + // sent to the counterpart as if it came from the new one. + Some(Slot::Disconnected { .. }) => { + slots.discard_held_frames(role.counterpart()); + (VecDeque::new(), permit) + } Some(Slot::Connected { backlog, permit: inherited, @@ -195,6 +240,7 @@ impl Session { } None => (VecDeque::new(), permit), }; + prune_expired(&mut backlog, now); if !backlog.is_empty() { log::debug!( "redelivering {} buffered frame(s) to {}", @@ -227,31 +273,54 @@ impl Session { }) } - /// Hand the slot's backlog to the peer's outbound queue. - pub(crate) async fn drain_backlog( + /// Hand as much of the slot's backlog to the peer's outbound queue as it + /// has room for, without waiting. + /// + /// Expired frames are dropped, never handed over. What does not fit stays + /// in the backlog, in order, for a later call once the writer has made + /// room. Only the connection owning `tx` may drain the slot, so a replaced + /// connection cannot consume its successor's backlog. + pub(crate) fn drain_backlog( &self, role: Role, tx: &mpsc::Sender, - ) -> Result { - let mut drained = 0usize; - loop { - let Ok(permit) = tx.reserve().await else { - return Ok(drained); - }; + ) -> Result { + let mut slots = self.lock()?; + let Some(Slot::Connected { peer, backlog, .. }) = slots.get_mut(role) else { + return Ok(Drain { + delivered: 0, + finished: true, + }); + }; + if !peer.tx.same_channel(tx) { + return Ok(Drain { + delivered: 0, + finished: true, + }); + } - let mut slots = self.lock()?; - let Some(Slot::Connected { peer, backlog, .. }) = slots.get_mut(role) else { - return Ok(drained); - }; - if !peer.tx.same_channel(tx) { - return Ok(drained); + prune_expired(backlog, Instant::now()); + let mut delivered = 0usize; + let finished = loop { + if backlog.is_empty() { + break true; } - let Some(frame) = backlog.pop_front() else { - return Ok(drained); - }; - permit.send(frame); - drained += 1; - } + match tx.try_reserve() { + Ok(permit) => { + if let Some(entry) = backlog.pop_front() { + permit.send(entry.frame); + delivered += 1; + } + } + Err(mpsc::error::TrySendError::Full(())) => break false, + Err(mpsc::error::TrySendError::Closed(())) => break true, + } + }; + + Ok(Drain { + delivered, + finished, + }) } /// Reserve a peer's slot for `grace`, buffering frames addressed to it. @@ -302,6 +371,11 @@ impl Session { } *slots.get_mut(role) = None; + // Frames still held for the survivor came from the peer that just + // expired: they must not follow `peer_disconnected`, nor a later + // `session_ready` with another counterpart. + slots.discard_held_frames(role.counterpart()); + if let Some(Slot::Connected { peer, .. }) = slots.get(role.counterpart()) && peer.tx.try_send(notification.clone()).is_err() { @@ -318,6 +392,7 @@ impl Session { /// Forward a frame to the counterpart, buffering while it is disconnected. /// /// Buffered frames expire after `buffer_ttl`; a zero TTL drops them immediately. + /// A frame queued behind a connected peer's backlog expires the same way. pub(crate) fn route( &self, from: Role, @@ -327,13 +402,18 @@ impl Session { ) -> Result<(), SessionError> { let mut slots = self.lock()?; let target = from.counterpart(); + let now = Instant::now(); if let Some(Slot::Connected { peer, backlog, .. }) = slots.get_mut(target) { + prune_expired(backlog, now); if !backlog.is_empty() { while backlog.len() >= max_buffered.max(1) { backlog.pop_front(); } - backlog.push_back(frame); + backlog.push_back(BufferedFrame { + frame, + expires_at: now + buffer_ttl, + }); return Ok(()); } return match peer.tx.try_send(frame) { @@ -349,12 +429,13 @@ impl Session { match slots.get_mut(target) { Some(Slot::Disconnected { buffered, .. }) => { if max_buffered > 0 && !buffer_ttl.is_zero() { + prune_expired(buffered, now); if buffered.len() >= max_buffered { buffered.pop_front(); } buffered.push_back(BufferedFrame { frame, - expires_at: Instant::now() + buffer_ttl, + expires_at: now + buffer_ttl, }); } Ok(()) @@ -450,18 +531,25 @@ mod tests { } /// Ids a freshly registered peer receives: queued frames, then the backlog. + /// + /// Drains in steps like a connection does, emptying the queue in between + /// as its writer would, so a backlog larger than the queue gets through. async fn received_ids( session: &Session, role: Role, tx: &mpsc::Sender, rx: &mut mpsc::Receiver, ) -> Result, SessionError> { - session.drain_backlog(role, tx).await?; let mut ids = Vec::new(); - while let Ok(frame) = rx.try_recv() { - ids.push(frame.id); + loop { + let drain = session.drain_backlog(role, tx)?; + while let Ok(frame) = rx.try_recv() { + ids.push(frame.id); + } + if drain.finished { + return Ok(ids); + } } - Ok(ids) } fn register( @@ -827,19 +915,22 @@ mod tests { let tx = responder.tx.clone(); register(&session, Role::Responder, responder)?; - let drain = session.drain_backlog(Role::Responder, &tx); - let collect = async { - let mut ids = Vec::new(); - while ids.len() < total + 1 { - match responder_rx.recv().await { - Some(frame) => ids.push(frame.id), - None => break, - } - } - ids - }; - let (drained, ids) = tokio::join!(drain, collect); - drained?; + // One step fills the queue and leaves the rest for later. + let first = session.drain_backlog(Role::Responder, &tx)?; + assert_eq!( + first, + Drain { + delivered: PEER_QUEUE - 1, + finished: false, + }, + "session_ready holds one slot of the queue" + ); + + let mut ids = Vec::new(); + while let Ok(frame) = responder_rx.try_recv() { + ids.push(frame.id); + } + ids.extend(received_ids(&session, Role::Responder, &tx, &mut responder_rx).await?); let mut expected = vec!["ready".to_owned()]; expected.extend((0..total).map(|i| format!("m{i:02}"))); @@ -1039,12 +1130,12 @@ mod tests { register(&session, Role::Responder, responder)?; assert_eq!( - session.drain_backlog(Role::Responder, &tx).await?, + session.drain_backlog(Role::Responder, &tx)?.delivered, 3, "the count is what the caller owes the peer in inbound budget" ); assert_eq!( - session.drain_backlog(Role::Responder, &tx).await?, + session.drain_backlog(Role::Responder, &tx)?.delivered, 0, "an empty backlog hands over nothing" ); @@ -1165,4 +1256,234 @@ mod tests { ); Ok(()) } + + #[tokio::test(start_paused = true)] + async fn reclaimed_frames_expire_while_waiting_for_the_counterpart() -> TestResult { + let session = Session::default(); + let (proposer, _) = peer(1); + let (responder, _) = peer(2); + register(&session, Role::Proposer, proposer)?; + register(&session, Role::Responder, responder)?; + disconnect_current(&session, Role::Responder, Duration::from_secs(10))?; + + let ttl = Duration::from_secs(5); + session.route(Role::Proposer, frame(&topic()?), MAX_BUFFERED, ttl)?; + disconnect_current(&session, Role::Proposer, Duration::from_secs(60))?; + + // The responder takes its frame back well inside the TTL, but nothing + // may reach it until its counterpart returns. + let (responder, mut responder_rx) = peer(2); + let tx = responder.tx.clone(); + assert!(matches!( + register(&session, Role::Responder, responder)?, + Registration::Waiting + )); + + tokio::time::advance(Duration::from_secs(6)).await; + let (proposer, _proposer_rx) = peer(1); + assert!(matches!( + register(&session, Role::Proposer, proposer)?, + Registration::BothReady + )); + + assert_eq!( + received_ids(&session, Role::Responder, &tx, &mut responder_rx).await?, + ["ready"], + "a frame reclaimed before its TTL must not be delivered after it" + ); + Ok(()) + } + + #[tokio::test(start_paused = true)] + async fn frames_queued_behind_a_backlog_expire_on_their_own_ttl() -> TestResult { + let session = Session::default(); + let (proposer, _proposer_rx) = peer(1); + let (responder, _) = peer(2); + register(&session, Role::Proposer, proposer)?; + register(&session, Role::Responder, responder)?; + disconnect_current(&session, Role::Responder, Duration::from_secs(10))?; + + let topic = topic()?; + let ttl = Duration::from_secs(5); + session.route( + Role::Proposer, + frame_with_id(&topic, "reclaimed"), + MAX_BUFFERED, + ttl, + )?; + + // The responder returns, but its writer is too slow to take anything. + let (responder, mut responder_rx) = peer(2); + let tx = responder.tx.clone(); + register(&session, Role::Responder, responder)?; + + // t=3: queued behind the backlog, expires at t=8. + tokio::time::advance(Duration::from_secs(3)).await; + session.route( + Role::Proposer, + frame_with_id(&topic, "early"), + MAX_BUFFERED, + ttl, + )?; + + // t=6: "reclaimed" has expired; "late" expires at t=11. + tokio::time::advance(Duration::from_secs(3)).await; + session.route( + Role::Proposer, + frame_with_id(&topic, "late"), + MAX_BUFFERED, + ttl, + )?; + + tokio::time::advance(Duration::from_secs(3)).await; + assert_eq!( + received_ids(&session, Role::Responder, &tx, &mut responder_rx).await?, + ["ready", "late"], + "at t=9 only the frame queued at t=6 is still deliverable" + ); + Ok(()) + } + + #[tokio::test(start_paused = true)] + async fn replacement_keeps_the_expiry_of_inherited_frames() -> TestResult { + let session = Session::default(); + let (proposer, _proposer_rx) = peer(1); + let (responder, _) = peer(2); + register(&session, Role::Proposer, proposer)?; + register(&session, Role::Responder, responder)?; + disconnect_current(&session, Role::Responder, Duration::from_secs(10))?; + + let ttl = Duration::from_secs(5); + session.route(Role::Proposer, frame(&topic()?), MAX_BUFFERED, ttl)?; + + // The first reconnect takes the frame back but never drains it. + let (first, _first_rx) = peer(2); + register(&session, Role::Responder, first)?; + + // Its replacement arrives after the TTL has passed. + tokio::time::advance(Duration::from_secs(6)).await; + let (second, mut second_rx) = peer(2); + let second_tx = second.tx.clone(); + register(&session, Role::Responder, second)?; + + assert_eq!( + received_ids(&session, Role::Responder, &second_tx, &mut second_rx).await?, + ["ready"], + "an inherited frame must not outlive its TTL" + ); + Ok(()) + } + + #[tokio::test(start_paused = true)] + async fn backlog_expiring_mid_drain_is_not_handed_over() -> TestResult { + let session = Session::default(); + let (proposer, _proposer_rx) = peer(1); + let (responder, _) = peer(2); + register(&session, Role::Proposer, proposer)?; + register(&session, Role::Responder, responder)?; + disconnect_current(&session, Role::Responder, Duration::from_secs(10))?; + + let topic = topic()?; + let ttl = Duration::from_secs(5); + let total = PEER_QUEUE * 2; + for i in 0..total { + session.route( + Role::Proposer, + frame_with_id(&topic, &format!("m{i:02}")), + total, + ttl, + )?; + } + + let (responder, mut responder_rx) = peer(2); + let tx = responder.tx.clone(); + register(&session, Role::Responder, responder)?; + + // A slow reader: one step fills the queue, then the writer stalls past + // the TTL before it makes room for more. + assert!(!session.drain_backlog(Role::Responder, &tx)?.finished); + tokio::time::advance(Duration::from_secs(6)).await; + + let mut queued = 0; + while responder_rx.try_recv().is_ok() { + queued += 1; + } + assert_eq!(queued, PEER_QUEUE, "session_ready and the first frames"); + assert_eq!( + session.drain_backlog(Role::Responder, &tx)?, + Drain { + delivered: 0, + finished: true, + }, + "the rest of the backlog expired while the writer was stalled" + ); + assert!(responder_rx.try_recv().is_err()); + Ok(()) + } + + #[tokio::test(start_paused = true)] + async fn counterpart_expiry_discards_the_frames_it_sent() -> TestResult { + let session = Session::default(); + let (proposer, _) = peer(1); + let (responder, _) = peer(2); + register(&session, Role::Proposer, proposer)?; + register(&session, Role::Responder, responder)?; + disconnect_current(&session, Role::Responder, Duration::from_secs(10))?; + session.route(Role::Proposer, frame(&topic()?), MAX_BUFFERED, BUFFER_TTL)?; + disconnect_current(&session, Role::Proposer, Duration::from_secs(1))?; + + // The responder takes its frame back while the proposer is away, and + // the proposer never returns. + let (responder, mut responder_rx) = peer(2); + let tx = responder.tx.clone(); + register(&session, Role::Responder, responder)?; + tokio::time::advance(Duration::from_secs(2)).await; + assert!(session.expire(Role::Proposer, ¬ification())?.is_some()); + + let (other, _other_rx) = peer(3); + assert!(matches!( + register(&session, Role::Proposer, other)?, + Registration::BothReady + )); + assert_eq!( + received_ids(&session, Role::Responder, &tx, &mut responder_rx).await?, + ["gone", "ready"], + "frames from an expired counterpart must not follow peer_disconnected" + ); + Ok(()) + } + + #[tokio::test(start_paused = true)] + async fn claiming_a_lapsed_reservation_discards_what_its_owner_sent() -> TestResult { + let session = Session::default(); + let (proposer, _) = peer(1); + let (responder, _) = peer(2); + register(&session, Role::Proposer, proposer)?; + register(&session, Role::Responder, responder)?; + disconnect_current(&session, Role::Responder, Duration::from_secs(10))?; + session.route(Role::Proposer, frame(&topic()?), MAX_BUFFERED, BUFFER_TTL)?; + disconnect_current(&session, Role::Proposer, Duration::from_secs(1))?; + + // The proposer's grace lapses and another identity takes the role + // before the expiry is processed. + tokio::time::advance(Duration::from_secs(2)).await; + let (other, _other_rx) = peer(3); + assert!(matches!( + register(&session, Role::Proposer, other)?, + Registration::Waiting + )); + + let (responder, mut responder_rx) = peer(2); + let tx = responder.tx.clone(); + assert!(matches!( + register(&session, Role::Responder, responder)?, + Registration::BothReady + )); + assert_eq!( + received_ids(&session, Role::Responder, &tx, &mut responder_rx).await?, + ["ready"], + "frames the previous proposer sent must not reach the responder as the new one's" + ); + Ok(()) + } } From c88a38e5131f83f891138fe8abddf09d5cae7d81 Mon Sep 17 00:00:00 2001 From: Kyryl R Date: Wed, 28 Oct 2026 11:29:04 +0200 Subject: [PATCH 2/9] Bound relay outbound frames by bytes and queue age --- crates/relay/README.md | 37 +- crates/relay/configuration/base.yaml | 5 + crates/relay/src/config.rs | 33 ++ crates/relay/src/configuration.rs | 116 +++++ crates/relay/src/connection.rs | 223 +++++++- crates/relay/src/health.rs | 3 +- crates/relay/src/limits.rs | 102 +++- crates/relay/src/registry.rs | 13 +- crates/relay/src/server.rs | 48 +- crates/relay/src/session.rs | 539 ++++++++++++++++++-- deployment/configs/relay/base.yaml.template | 5 + 11 files changed, 1038 insertions(+), 86 deletions(-) diff --git a/crates/relay/README.md b/crates/relay/README.md index 180f5d0..097c154 100644 --- a/crates/relay/README.md +++ b/crates/relay/README.md @@ -33,12 +33,45 @@ See `configuration/base.yaml` for every key. The ones that matter most: | `session.buffer_max_messages` / `session.buffer_ttl_seconds` | frames buffered for it | | `limits.*` | connection caps, per-identity/per-connection rates, and each identity's share of the global inbound byte budget | | `connection.*` | frame size, queue capacity, and the handshake, write, ping, pong and idle timeouts | +| `connection.outbound_queue_max_bytes` | bytes held for one peer and not yet written (see [Outbound memory](#outbound-memory)) | +| `connection.outbound_queue_timeout_seconds` | longest a frame may wait in a peer's outbound queue before that peer is dropped | Values that would make the relay unusable are rejected at startup: zero limits, zero timeouts, a blank audience, a health address equal to the application address, a byte-rate burst below the maximum incoming frame (no frame could -ever fit it), or a maximum incoming frame above `u32::MAX`. Every byte bucket must -hold a whole frame. +ever fit it), an outbound byte budget below the maximum incoming frame, or a +maximum incoming frame above `u32::MAX`. Every byte bucket and outbound budget +must hold a whole frame. + +### Outbound memory + +Frames waiting for a peer are bounded per peer, in frames and in bytes, and +in time: + +- `connection.outbound_queue_max_bytes` (default 1 MiB, four maximal frames) + caps the bytes held for one connection: its outbound queue, the frame being + written, and the backlog of buffered frames waiting to enter the queue. A + reserved slot gets a budget of the same size for what is buffered while its + peer reconnects, and the connection that reclaims it takes those frames over. + Bytes are the frame's JSON text, what the relay writes. A full buffer or + backlog drops its oldest frames; a frame that cannot fit beside the queue is + dropped, like one that finds `outbound_queue_capacity` full. Clients + retransmit what is not acknowledged. `session_ready` and `peer_disconnected` + are not charged, so a full budget never drops them. +- `connection.outbound_queue_timeout_seconds` (default 30) drops a peer whose + queue holds a frame that was not written within that time of being queued, + even if every single write met `write_timeout`. A reader too slow to drain + its queue cannot keep its budget filled indefinitely. + +These are per-peer limits only. The relay has no relay-wide cap on outbound +memory: the worst case is about `limits.max_connections × +outbound_queue_max_bytes` for connected peers (10,000 × 1 MiB, roughly 10 GiB, +with the defaults), plus one budget per slot reserved during its grace window, +plus the copy of the frame each writer is encoding. The inbound byte budgets +limit how fast that memory can fill, not how much of it can be held. Size +`limits.max_connections` and `outbound_queue_max_bytes` together against the +memory available to the relay, and cap the process itself (for example with a +container memory limit): the relay enforces no aggregate bound. ## Running diff --git a/crates/relay/configuration/base.yaml b/crates/relay/configuration/base.yaml index 4b2d5c9..89ef4c9 100644 --- a/crates/relay/configuration/base.yaml +++ b/crates/relay/configuration/base.yaml @@ -18,6 +18,11 @@ connection: handshake_timeout_seconds: 10 write_timeout_seconds: 10 outbound_queue_capacity: 64 + # Bytes held for one peer and not yet written: its queue and backlog, or + # its buffer while it reconnects. At least max_incoming_frame_bytes. + outbound_queue_max_bytes: 1048576 + # A frame not written this long after it was queued drops its slow reader. + outbound_queue_timeout_seconds: 30 ping_interval_seconds: 30 pong_timeout_seconds: 10 idle_timeout_seconds: 45 diff --git a/crates/relay/src/config.rs b/crates/relay/src/config.rs index 7e7cdcb..28c9498 100644 --- a/crates/relay/src/config.rs +++ b/crates/relay/src/config.rs @@ -34,6 +34,13 @@ pub const DEFAULT_WRITE_TIMEOUT: Duration = Duration::from_secs(10); /// Default number of frames queued for a connected peer. pub const DEFAULT_OUTBOUND_QUEUE_CAPACITY: usize = 64; +/// Default bytes held for one peer and not yet written: four maximal frames. +pub const DEFAULT_OUTBOUND_QUEUE_MAX_BYTES: usize = 4 * DEFAULT_MAX_INCOMING_FRAME_BYTES; + +/// Default longest time a frame may wait in a peer's outbound queue before +/// the peer is dropped as too slow to read. +pub const DEFAULT_OUTBOUND_QUEUE_TIMEOUT: Duration = Duration::from_secs(30); + /// Runtime limits for the relay connection handler. #[derive(Debug, Clone, PartialEq, Eq)] pub struct NativeRelayTransportConfig { @@ -70,6 +77,15 @@ pub struct NativeRelayTransportConfig { /// Maximum frames queued for a connected peer that is not reading. pub outbound_queue_capacity: usize, + /// Maximum bytes of frames held for one peer and not yet written: its + /// outbound queue and the backlog waiting to enter it, or its buffer + /// while it reconnects. A frame that does not fit is dropped. + pub outbound_queue_max_bytes: usize, + + /// Longest time a frame may wait in a peer's outbound queue, write + /// included, before the peer is dropped as too slow to read. + pub outbound_queue_timeout: Duration, + /// Maximum accepted lifetime of an auth token. pub max_token_lifetime: Duration, } @@ -89,6 +105,8 @@ impl NativeRelayTransportConfig { handshake_timeout: DEFAULT_HANDSHAKE_TIMEOUT, write_timeout: DEFAULT_WRITE_TIMEOUT, outbound_queue_capacity: DEFAULT_OUTBOUND_QUEUE_CAPACITY, + outbound_queue_max_bytes: DEFAULT_OUTBOUND_QUEUE_MAX_BYTES, + outbound_queue_timeout: DEFAULT_OUTBOUND_QUEUE_TIMEOUT, max_token_lifetime: DEFAULT_MAX_TOKEN_LIFETIME, } } @@ -162,6 +180,21 @@ impl NativeRelayTransportConfig { self } + /// Set the bytes that may be held for one peer and not yet written. + /// + /// Below the encoded size of a frame, that frame is always dropped; the + /// configuration file rejects a value below the incoming frame cap. + pub fn with_outbound_queue_max_bytes(mut self, bytes: usize) -> Self { + self.outbound_queue_max_bytes = bytes; + self + } + + /// Set the longest time a frame may wait in a peer's outbound queue. + pub fn with_outbound_queue_timeout(mut self, timeout: Duration) -> Self { + self.outbound_queue_timeout = timeout; + self + } + /// Set the maximum accepted auth-token lifetime. pub fn with_max_token_lifetime(mut self, lifetime: Duration) -> Self { self.max_token_lifetime = lifetime; diff --git a/crates/relay/src/configuration.rs b/crates/relay/src/configuration.rs index f76570e..64d4190 100644 --- a/crates/relay/src/configuration.rs +++ b/crates/relay/src/configuration.rs @@ -47,6 +47,12 @@ pub enum ConfigurationError { burst: u32, frame: u32, }, + + #[error( + "connection.outbound_queue_max_bytes ({bytes}) must be at least \ + connection.max_incoming_frame_bytes ({frame})" + )] + OutboundBytesBelowFrameCap { bytes: usize, frame: usize }, } /// Relay settings. @@ -118,6 +124,14 @@ pub struct ConnectionSettings { /// Frames queued for a connected peer that is not reading. pub outbound_queue_capacity: usize, + /// Bytes of frames held for one peer and not yet written: its queue and + /// backlog, or its buffer while it reconnects. + pub outbound_queue_max_bytes: usize, + + /// Time a frame may wait in the outbound queue before its slow reader + /// is dropped. + pub outbound_queue_timeout_seconds: u64, + /// Interval between relay keepalive pings. pub ping_interval_seconds: u64, @@ -208,6 +222,10 @@ impl Settings { "connection.write_timeout_seconds", self.connection.write_timeout_seconds, ), + ( + "connection.outbound_queue_timeout_seconds", + self.connection.outbound_queue_timeout_seconds, + ), ( "connection.ping_interval_seconds", self.connection.ping_interval_seconds, @@ -239,6 +257,10 @@ impl Settings { "connection.outbound_queue_capacity", self.connection.outbound_queue_capacity, ), + ( + "connection.outbound_queue_max_bytes", + self.connection.outbound_queue_max_bytes, + ), ("limits.max_connections", self.limits.max_connections), ( "limits.max_connections_per_identity", @@ -272,6 +294,16 @@ impl Settings { } } + // The relay re-encodes a frame no longer than the text it arrived as, + // so a budget that holds the largest incoming frame can queue any + // frame the relay accepts. + if self.connection.outbound_queue_max_bytes < self.connection.max_incoming_frame_bytes { + return Err(ConfigurationError::OutboundBytesBelowFrameCap { + bytes: self.connection.outbound_queue_max_bytes, + frame: self.connection.max_incoming_frame_bytes, + }); + } + let frame = self.connection.max_incoming_frame_bytes; let frame = u32::try_from(frame) .map_err(|_| ConfigurationError::FrameAboveByteCounter { frame })?; @@ -341,6 +373,10 @@ impl From<&Settings> for NativeRelayTransportConfig { handshake_timeout: Duration::from_secs(settings.connection.handshake_timeout_seconds), write_timeout: Duration::from_secs(settings.connection.write_timeout_seconds), outbound_queue_capacity: settings.connection.outbound_queue_capacity, + outbound_queue_max_bytes: settings.connection.outbound_queue_max_bytes, + outbound_queue_timeout: Duration::from_secs( + settings.connection.outbound_queue_timeout_seconds, + ), max_token_lifetime: Duration::from_secs(settings.auth.max_token_lifetime_seconds), } } @@ -369,3 +405,83 @@ impl From<&RateSettings> for RateLimit { } } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::config::{DEFAULT_OUTBOUND_QUEUE_MAX_BYTES, DEFAULT_OUTBOUND_QUEUE_TIMEOUT}; + + type TestResult = Result<(), Box>; + + /// The shipped `base.yaml` with the local overlay, without environment + /// overrides. + fn shipped() -> Result { + let settings = config::Config::builder() + .add_source(config::File::from_str( + include_str!("../configuration/base.yaml"), + config::FileFormat::Yaml, + )) + .add_source(config::File::from_str( + include_str!("../configuration/local.yaml"), + config::FileFormat::Yaml, + )) + .build()? + .try_deserialize::()?; + Ok(settings) + } + + #[test] + fn shipped_outbound_defaults_are_valid_and_match_the_transport() -> TestResult { + let settings = shipped()?; + settings.validate()?; + + let transport = NativeRelayTransportConfig::from(&settings); + assert_eq!( + transport.outbound_queue_max_bytes, + DEFAULT_OUTBOUND_QUEUE_MAX_BYTES + ); + assert_eq!( + transport.outbound_queue_timeout, + DEFAULT_OUTBOUND_QUEUE_TIMEOUT + ); + Ok(()) + } + + #[test] + fn outbound_budget_below_the_frame_cap_is_rejected() -> TestResult { + let mut settings = shipped()?; + let frame = settings.connection.max_incoming_frame_bytes; + + settings.connection.outbound_queue_max_bytes = frame - 1; + assert!(matches!( + settings.validate(), + Err(ConfigurationError::OutboundBytesBelowFrameCap { .. }) + )); + + settings.connection.outbound_queue_max_bytes = frame; + settings.validate()?; + Ok(()) + } + + #[test] + fn zero_outbound_budget_or_timeout_is_rejected() -> TestResult { + let mut settings = shipped()?; + settings.connection.outbound_queue_max_bytes = 0; + assert!(matches!( + settings.validate(), + Err(ConfigurationError::Zero { + setting: "connection.outbound_queue_max_bytes" + }) + )); + + let mut settings = shipped()?; + settings.connection.outbound_queue_timeout_seconds = 0; + assert!(matches!( + settings.validate(), + Err(ConfigurationError::Zero { + setting: "connection.outbound_queue_timeout_seconds" + }) + )); + Ok(()) + } +} diff --git a/crates/relay/src/connection.rs b/crates/relay/src/connection.rs index 291a8f2..5d76028 100644 --- a/crates/relay/src/connection.rs +++ b/crates/relay/src/connection.rs @@ -1,9 +1,6 @@ use std::{sync::Arc, time::Duration}; -use futures_util::{ - SinkExt, StreamExt, - stream::{SplitSink, SplitStream}, -}; +use futures_util::{Sink, SinkExt, StreamExt, stream::SplitStream}; use thiserror::Error; use tokio::{ net::TcpStream, @@ -25,7 +22,7 @@ use crate::{ auth::{Identity, IdentityError, unix_now}, config::NativeRelayTransportConfig, limits::{Bucket, InboundBytes, Limiter, RelayLimitsConfig}, - session::{Registration, Role, Session, SessionError}, + session::{Outbound, Registration, Role, Session, SessionError}, }; use wallet_session_core::{ did::DidKey, @@ -113,8 +110,8 @@ pub struct Connection { query: ConnectionQuery, identity: Identity, session: Arc, - outbound: mpsc::Receiver, - self_tx: mpsc::Sender, + outbound: mpsc::Receiver, + self_tx: mpsc::Sender, config: NativeRelayTransportConfig, limiter: Arc, } @@ -177,8 +174,8 @@ impl Connection { query: ConnectionQuery, identity: Identity, session: Arc, - outbound: mpsc::Receiver, - self_tx: mpsc::Sender, + outbound: mpsc::Receiver, + self_tx: mpsc::Sender, config: NativeRelayTransportConfig, limiter: Arc, ) -> Self { @@ -304,14 +301,25 @@ async fn stop_writer(writer: &mut JoinHandle<()>) { } } -async fn write_loop( - mut sink: SplitSink, - mut rx: mpsc::Receiver, +/// Write queued frames and keepalive pings until the queue closes, the +/// connection shuts down, or a write misses its deadline. +/// +/// Besides the per-write timeout, a frame must be written within the queue +/// timeout of entering the queue, so a reader that takes each write just in +/// time but never catches up is dropped as well. The frames behind the one +/// being written are checked when the writer reaches them, so a keepalive +/// ping written in between can delay that by up to one write timeout. A frame +/// holds its bytes in the peer's budget until it has been written. +async fn write_loop( + mut sink: S, + mut rx: mpsc::Receiver, mut shutdown: oneshot::Receiver>, config: NativeRelayTransportConfig, ping_deadline: watch::Sender>, backlog_ready: Arc, -) { +) where + S: Sink + Unpin, +{ let mut ping = tokio::time::interval_at( tokio::time::Instant::now() + config.ping_interval, config.ping_interval, @@ -323,15 +331,19 @@ async fn write_loop( loop { tokio::select! { received = rx.recv() => match received { - Some(frame) => { + Some(queued) => { + let frame = &queued.frame; let session_ready = matches!(&frame.body, RelayFrameBody::SessionReady); - let Ok(text) = serde_json::to_string(&frame) else { + let Ok(text) = serde_json::to_string(frame) else { break; }; log::trace!("relay -> topic={} frame={text}", frame.topic); - if !send_with_timeout(&mut sink, Message::text(text), config.write_timeout).await { + let message = Message::text(text); + if !send_queued(&mut sink, message, queued.queued_at, &config).await { break; } + // Written: its bytes go back to the peer's budget. + drop(queued); // Once the peer has been told it is ready, the backlog // held for it may follow, and every later write makes room // in the queue for more of it. @@ -424,11 +436,10 @@ fn charge_control_frame( } /// Send one message, returning whether the socket is still usable. -async fn send_with_timeout( - sink: &mut SplitSink, - message: Message, - write_timeout: Duration, -) -> bool { +async fn send_with_timeout(sink: &mut S, message: Message, write_timeout: Duration) -> bool +where + S: Sink + Unpin, +{ match tokio::time::timeout(write_timeout, sink.send(message)).await { Ok(Ok(())) => true, Ok(Err(_)) => false, @@ -439,6 +450,42 @@ async fn send_with_timeout( } } +/// Send a frame that entered the queue at `queued_at`, returning whether the +/// socket is still usable. +/// +/// The write gets the write timeout, cut short so that the frame is written +/// within the queue timeout of entering the queue or not at all. A frame +/// already past it is not attempted, even if the socket would take it now. +async fn send_queued( + sink: &mut S, + message: Message, + queued_at: tokio::time::Instant, + config: &NativeRelayTransportConfig, +) -> bool +where + S: Sink + Unpin, +{ + let left = config + .outbound_queue_timeout + .saturating_sub(queued_at.elapsed()); + if left >= config.write_timeout { + return send_with_timeout(sink, message, config.write_timeout).await; + } + let sent = if left.is_zero() { + None + } else { + tokio::time::timeout(left, sink.send(message)).await.ok() + }; + match sent { + Some(Ok(())) => true, + Some(Err(_)) => false, + None => { + log::warn!("outbound queue is not draining; closing connection"); + false + } + } +} + /// Shared, read-only state a connection's read loop works against. struct ReadContext<'a> { query: &'a ConnectionQuery, @@ -457,7 +504,7 @@ struct BacklogDelivery { /// The slot's channel. It is weak because the slot's handle must stay the /// only strong sender, so a replaced connection closes its queue - tx: mpsc::WeakSender, + tx: mpsc::WeakSender, } async fn read_loop( @@ -768,6 +815,8 @@ mod tests { use super::*; use crate::auth::test_utils::{self, signing_key}; + use crate::limits::OutboundBudget; + use tokio::time::timeout; use wallet_session_core::message::Topic; const AUDIENCE: &str = "wss://relay.example.com"; @@ -957,4 +1006,134 @@ mod tests { stop_writer(&mut writer).await; assert!(writer.await.is_err_and(|error| error.is_cancelled())); } + + /// Ids of the frames a test reader accepted. + type Accepted = Arc>>; + + /// A reader that takes `per_message` to accept each message, and the ids + /// of the frames it accepted. + fn slow_reader(per_message: Duration) -> (Accepted, impl Sink + Unpin) { + let accepted = Arc::new(std::sync::Mutex::new(Vec::new())); + let sink = futures_util::sink::unfold( + Arc::clone(&accepted), + move |accepted, message: Message| async move { + tokio::time::sleep(per_message).await; + if let Message::Text(text) = &message + && let Ok(frame) = serde_json::from_str::(text) + { + accepted.lock().expect("accepted ids").push(frame.id); + } + Ok::<_, std::convert::Infallible>(accepted) + }, + ); + (accepted, Box::pin(sink)) + } + + /// Run a writer over `sink` as a connection does. It runs until the + /// returned sender shuts it down or is dropped. + fn spawn_writer( + sink: S, + rx: mpsc::Receiver, + config: NativeRelayTransportConfig, + ) -> (oneshot::Sender>, JoinHandle<()>) + where + S: Sink + Unpin + Send + 'static, + { + let (shutdown, shutdown_rx) = oneshot::channel(); + let (ping_deadline, _) = watch::channel(None); + let writer = tokio::spawn(write_loop( + sink, + rx, + shutdown_rx, + config, + ping_deadline, + Arc::new(Notify::new()), + )); + (shutdown, writer) + } + + /// Writer settings for these tests, with no keepalive ping in the way. + fn writer_config() -> NativeRelayTransportConfig { + NativeRelayTransportConfig::new(AUDIENCE) + .with_ping_interval(Duration::from_secs(3600)) + .with_write_timeout(Duration::from_secs(10)) + .with_outbound_queue_timeout(Duration::from_secs(15)) + } + + /// A message frame entering the queue now, holding 100 bytes of `budget`. + fn queued( + budget: &Arc, + id: &str, + ) -> Result> { + let frame = RelayFrame { + v: RELAY_FRAME_VERSION, + id: id.to_owned(), + topic: topic()?, + body: RelayFrameBody::Message { + payload: "ciphertext".to_owned(), + }, + }; + Ok(Outbound::new(frame, budget.try_charge(100))) + } + + #[tokio::test(start_paused = true)] + async fn writer_drops_a_reader_that_never_catches_up_with_its_queue() -> TestResult { + let budget = Arc::new(OutboundBudget::new(1024)); + let (tx, rx) = mpsc::channel(8); + for i in 0..6 { + tx.try_send(queued(&budget, &format!("m{i}"))?)?; + } + + // Every write takes 4s, well inside the 10s write timeout, but the + // six frames take 24s to get through and none may wait over 15s. The + // fourth is still being written at 15s: that write is cut short. + let (accepted, sink) = slow_reader(Duration::from_secs(4)); + let start = tokio::time::Instant::now(); + let (_shutdown, writer) = spawn_writer(sink, rx, writer_config()); + timeout(Duration::from_secs(60), writer).await??; + + let queue_timeout = Duration::from_secs(15); + let window = queue_timeout..=queue_timeout + WRITER_SHUTDOWN_GRACE; + let elapsed = start.elapsed(); + assert!( + window.contains(&elapsed), + "the writer must stop once a frame has waited the queue timeout, not {elapsed:?}" + ); + assert_eq!(*accepted.lock().expect("accepted ids"), ["m0", "m1", "m2"]); + + drop(tx); + assert_eq!(budget.held(), 0, "frames left behind give their bytes back"); + Ok(()) + } + + #[tokio::test(start_paused = true)] + async fn writer_keeps_a_slow_reader_that_keeps_up() -> TestResult { + let budget = Arc::new(OutboundBudget::new(1024)); + let (tx, rx) = mpsc::channel(8); + let (accepted, sink) = slow_reader(Duration::from_secs(4)); + let (shutdown, writer) = spawn_writer(sink, rx, writer_config()); + + // The last of these waits 8s and is written by 12s. + for i in 0..3 { + tx.try_send(queued(&budget, &format!("m{i}"))?)?; + } + // A frame queued long after the queue emptied waits only from then. + tokio::time::sleep(Duration::from_secs(30)).await; + tx.try_send(queued(&budget, "m3")?)?; + tokio::time::sleep(Duration::from_secs(10)).await; + + assert!( + !writer.is_finished(), + "a reader that keeps up must not be dropped" + ); + assert_eq!( + *accepted.lock().expect("accepted ids"), + ["m0", "m1", "m2", "m3"] + ); + assert_eq!(budget.held(), 0, "written frames give their bytes back"); + + let _ = shutdown.send(None); + timeout(Duration::from_secs(1), writer).await??; + Ok(()) + } } diff --git a/crates/relay/src/health.rs b/crates/relay/src/health.rs index 0fb4957..fb217a8 100644 --- a/crates/relay/src/health.rs +++ b/crates/relay/src/health.rs @@ -115,7 +115,7 @@ mod tests { use super::*; use crate::{ - limits::{RateLimit, RelayLimitsConfig}, + limits::{OutboundBudget, RateLimit, RelayLimitsConfig}, session::{PeerHandle, Role}, }; @@ -177,6 +177,7 @@ mod tests { PeerHandle { identity: identity.clone(), tx, + budget: Arc::new(OutboundBudget::new(1024)), }, permit, &RelayFrame { diff --git a/crates/relay/src/limits.rs b/crates/relay/src/limits.rs index defeb1b..7eb670b 100644 --- a/crates/relay/src/limits.rs +++ b/crates/relay/src/limits.rs @@ -2,7 +2,7 @@ use std::{ collections::HashMap, sync::{ Arc, Mutex, MutexGuard, - atomic::{AtomicU64, Ordering}, + atomic::{AtomicU64, AtomicUsize, Ordering}, }, time::Duration, }; @@ -345,6 +345,72 @@ impl Drop for ConnectionPermit { } } +/// Bytes of frames held for one peer and not yet written to it. +/// +/// A connection's budget covers its outbound queue, the frame its writer is +/// writing, and the backlog waiting to enter the queue; a reserved slot has +/// one of its own for the frames buffered while its peer reconnects. Bytes +/// are the frame's JSON text, which is what the writer sends. +#[derive(Debug)] +pub struct OutboundBudget { + capacity: usize, + held: AtomicUsize, +} + +impl OutboundBudget { + /// An empty budget of `capacity` bytes. + pub fn new(capacity: usize) -> Self { + Self { + capacity, + held: AtomicUsize::new(0), + } + } + + /// Bytes this budget may hold at once. + pub fn capacity(&self) -> usize { + self.capacity + } + + /// Bytes currently charged to this budget. + pub fn held(&self) -> usize { + self.held.load(Ordering::Relaxed) + } + + /// Charge `bytes`, if they fit beside what is already held. + pub(crate) fn try_charge(self: &Arc, bytes: usize) -> Option { + self.held + .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |held| { + held.checked_add(bytes) + .filter(|total| *total <= self.capacity) + }) + .ok()?; + Some(OutboundCharge { + budget: Arc::clone(self), + bytes, + }) + } +} + +/// Bytes charged to an [`OutboundBudget`]; given back when dropped. +#[derive(Debug)] +pub(crate) struct OutboundCharge { + budget: Arc, + bytes: usize, +} + +impl OutboundCharge { + /// Bytes this charge holds. + pub(crate) fn bytes(&self) -> usize { + self.bytes + } +} + +impl Drop for OutboundCharge { + fn drop(&mut self) { + self.budget.held.fetch_sub(self.bytes, Ordering::Relaxed); + } +} + #[cfg(test)] mod tests { use super::*; @@ -610,6 +676,40 @@ mod tests { ); } + #[test] + fn outbound_budget_refuses_what_does_not_fit_and_takes_back_on_drop() { + let budget = Arc::new(OutboundBudget::new(100)); + + let first = budget.try_charge(60); + assert!(first.is_some()); + assert!( + budget.try_charge(41).is_none(), + "a charge must fit beside what is held" + ); + let second = budget.try_charge(40); + assert!(second.is_some(), "a charge may fill the budget exactly"); + assert_eq!(budget.held(), 100); + + drop(first); + assert_eq!(budget.held(), 40, "a dropped charge gives its bytes back"); + assert!(budget.try_charge(60).is_some()); + assert_eq!(budget.held(), 40, "an unbound charge is given back at once"); + } + + #[test] + fn outbound_budget_refuses_a_charge_larger_than_itself() { + let budget = Arc::new(OutboundBudget::new(100)); + assert!(budget.try_charge(101).is_none()); + assert!(budget.try_charge(usize::MAX).is_none()); + + let _held = budget.try_charge(1); + assert!( + budget.try_charge(usize::MAX).is_none(), + "an overflowing total must be refused, not wrapped" + ); + assert_eq!(budget.held(), 1); + } + #[tokio::test(start_paused = true)] async fn the_global_budget_still_caps_the_sum_of_shares() { let limiter = Limiter::new(RelayLimitsConfig { diff --git a/crates/relay/src/registry.rs b/crates/relay/src/registry.rs index 80b2300..6dbcbdc 100644 --- a/crates/relay/src/registry.rs +++ b/crates/relay/src/registry.rs @@ -7,7 +7,7 @@ use tokio::sync::mpsc; use wallet_session_core::{frame::RelayFrame, message::Topic}; use crate::limits::ConnectionPermit; -use crate::session::{PeerHandle, Registration, Role, Session, SessionError}; +use crate::session::{Outbound, PeerHandle, Registration, Role, Session, SessionError}; /// Directory of live sessions, keyed by topic. /// @@ -104,7 +104,7 @@ impl Registry { &self, topic: &Topic, role: Role, - tx: &mpsc::WeakSender, + tx: &mpsc::WeakSender, grace: Duration, ) -> Result { let sessions = self.lock()?; @@ -169,7 +169,7 @@ mod tests { use super::*; use tokio::sync::mpsc; - use crate::limits::{Limiter, RelayLimitsConfig}; + use crate::limits::{Limiter, OutboundBudget, RelayLimitsConfig}; /// A permit for `identity`, from a limiter that is never shared: these /// tests do not exercise the per-identity cap. @@ -197,12 +197,13 @@ mod tests { /// Outbound queue capacity used by peer handles in these tests. const PEER_QUEUE: usize = 8; - fn peer(seed: u8) -> (PeerHandle, mpsc::Receiver) { + fn peer(seed: u8) -> (PeerHandle, mpsc::Receiver) { let (tx, rx) = mpsc::channel(PEER_QUEUE); ( PeerHandle { identity: wallet_session_core::did::DidKey::from_ed25519_bytes([seed; 32]), tx, + budget: Arc::new(OutboundBudget::new(1024 * 1024)), }, rx, ) @@ -405,8 +406,8 @@ mod tests { )?; let mut kinds = Vec::new(); - while let Ok(frame) = proposer_rx.try_recv() { - kinds.push(match &frame.body { + while let Ok(queued) = proposer_rx.try_recv() { + kinds.push(match &queued.frame.body { RelayFrameBody::SessionReady => "ready", RelayFrameBody::PeerDisconnected => "disconnected", RelayFrameBody::Message { .. } => "message", diff --git a/crates/relay/src/server.rs b/crates/relay/src/server.rs index e023b80..f487639 100644 --- a/crates/relay/src/server.rs +++ b/crates/relay/src/server.rs @@ -11,7 +11,7 @@ use tokio_tungstenite::tungstenite::protocol::{CloseFrame, frame::coding::CloseC use crate::{ config::NativeRelayTransportConfig, connection::{Connection, ConnectionError, RELAY_FRAME_VERSION, RelayStream, random_id}, - limits::Limiter, + limits::{Limiter, OutboundBudget}, registry::{Registry, RegistryError}, session::{PeerHandle, Registration, SessionError}, }; @@ -94,6 +94,7 @@ pub async fn handle( let peer = PeerHandle { identity: identity.did.clone(), tx: tx.clone(), + budget: Arc::new(OutboundBudget::new(config.outbound_queue_max_bytes)), }; let ready = RelayFrame { @@ -678,6 +679,47 @@ mod tests { Ok(()) } + #[tokio::test] + async fn reconnect_buffer_keeps_the_newest_frames_its_byte_budget_holds() -> TestResult { + let topic = topic()?; + let payload = "a".repeat(1000); + // Each frame is re-encoded as the client sent it: room for two, not three. + let encoded = serde_json::to_string(&message(&topic, "m1", &payload))?; + let addr = serve_with( + NativeRelayTransportConfig::new(AUDIENCE) + .with_session_grace(Duration::from_secs(2)) + .with_outbound_queue_max_bytes(2 * encoded.len()), + ) + .await?; + + let (mut proposer, responder) = ready_pair(addr, &topic).await?; + drop(responder); + sleep(Duration::from_millis(100)).await; + for id in ["m1", "m2", "m3"] { + send(&mut proposer, &message(&topic, id, &payload)).await?; + } + // Let the relay buffer them before the responder returns. + sleep(Duration::from_millis(200)).await; + + let mut responder = connect(addr, "responder", &topic, None).await?; + assert!(matches!( + read_frame(&mut responder).await?.body, + RelayFrameBody::SessionReady + )); + let mut redelivered = Vec::new(); + while let Ok(Ok(frame)) = + timeout(Duration::from_millis(300), read_frame(&mut responder)).await + { + redelivered.push(frame.id); + } + assert_eq!( + redelivered, + ["m2", "m3"], + "the oldest buffered frame makes room for the newest" + ); + Ok(()) + } + #[tokio::test] async fn buffering_disabled_drops_frames_while_disconnected() -> TestResult { let addr = serve_with( @@ -1832,10 +1874,12 @@ mod tests { interval: Duration::from_secs(1), }; // The write deadline is far beyond the test, so only the token can end - // the stuck connection in time. + // the stuck connection in time. The byte budget holds the whole + // backlog, so it outlasts the socket buffers. let addr = serve_with_limits( NativeRelayTransportConfig::new(AUDIENCE) .with_outbound_queue_capacity(2) + .with_outbound_queue_max_bytes(32 * 1024 * 1024) .with_write_timeout(Duration::from_secs(60)), RelayLimitsConfig { inbound_message_rate: unlimited, diff --git a/crates/relay/src/session.rs b/crates/relay/src/session.rs index ea555f3..fd30728 100644 --- a/crates/relay/src/session.rs +++ b/crates/relay/src/session.rs @@ -1,6 +1,6 @@ -use crate::limits::ConnectionPermit; +use crate::limits::{ConnectionPermit, OutboundBudget, OutboundCharge}; use std::collections::VecDeque; -use std::sync::{Mutex, MutexGuard}; +use std::sync::{Arc, Mutex, MutexGuard}; use thiserror::Error; use tokio::sync::mpsc; @@ -16,7 +16,61 @@ pub struct PeerHandle { pub identity: DidKey, /// Channel to the peer's writer task. - pub tx: mpsc::Sender, + pub tx: mpsc::Sender, + + /// Bytes this connection may hold in its queue and backlog. It belongs + /// to the connection: frames a reconnect or a replacement inherits are + /// charged to the new connection's budget instead. + pub budget: Arc, +} + +/// A frame in a peer's outbound queue. +#[derive(Debug)] +pub struct Outbound { + /// The frame to write. + pub frame: RelayFrame, + + /// When the frame entered the queue. + pub queued_at: Instant, + + /// The frame's share of its peer's outbound budget, given back once the + /// writer is done with it. Relay control frames carry none: they are + /// small, the queue capacity bounds their number, and a full budget must + /// not drop a `session_ready` or `peer_disconnected`. + _charge: Option, +} + +impl Outbound { + /// A frame entering the queue now, holding `charge`. + pub(crate) fn new(frame: RelayFrame, charge: Option) -> Self { + Self { + frame, + queued_at: Instant::now(), + _charge: charge, + } + } +} + +/// Bytes of `frame`'s JSON text, which is what the writer sends, or `None` +/// if it cannot be encoded. +fn encoded_len(frame: &RelayFrame) -> Option { + /// Counts the bytes written to it. + struct Counter(usize); + + impl std::io::Write for Counter { + fn write(&mut self, buf: &[u8]) -> std::io::Result { + self.0 = self.0.saturating_add(buf.len()); + Ok(buf.len()) + } + + fn flush(&mut self) -> std::io::Result<()> { + Ok(()) + } + } + + let mut counter = Counter(0); + serde_json::to_writer(&mut counter, frame).ok()?; + Some(counter.0) } /// Outcome of binding a connection to a topic slot. @@ -62,13 +116,16 @@ enum Slot { permit: Option, /// Frames not yet handed to the outbound queue. They keep the expiry - /// they were buffered with. + /// they were buffered with, and are charged to the peer's budget. backlog: VecDeque, }, Disconnected { identity: DidKey, deadline: Instant, buffered: VecDeque, + + /// What `buffered` holds, as large as a connection's budget. + budget: Arc, }, } @@ -87,6 +144,9 @@ impl Slot { struct BufferedFrame { frame: RelayFrame, expires_at: Instant, + + /// Its encoded size, charged to the budget of the slot holding it. + charge: OutboundCharge, } /// Drop the frames whose TTL has passed. @@ -94,6 +154,69 @@ fn prune_expired(frames: &mut VecDeque, now: Instant) { frames.retain(|entry| entry.expires_at > now); } +/// Charge `bytes` to `budget` for a frame joining `frames`, dropping the +/// oldest of them if that makes room (spec §3: on overflow the oldest frame +/// is discarded). +/// +/// Every frame in `frames` is charged to `budget`. When the frame would not +/// fit even with all of them gone, because the rest of the budget is held by +/// the outbound queue, nothing is dropped and `None` is returned. +fn make_room( + frames: &mut VecDeque, + budget: &Arc, + bytes: usize, +) -> Option { + // Charges are only taken under the session lock, which the caller holds, + // so the queue's share can only shrink while this runs. + let held_here: usize = frames.iter().map(|entry| entry.charge.bytes()).sum(); + let queued = budget.held().saturating_sub(held_here); + if queued.saturating_add(bytes) > budget.capacity() { + return None; + } + loop { + if let Some(charge) = budget.try_charge(bytes) { + return Some(charge); + } + frames.pop_front()?; + } +} + +/// Move `frames` to `budget`, their new owner's, keeping the newest that fit. +/// +/// The budget they were charged to belonged to the reservation or the +/// connection being taken over. +fn recharge(frames: &mut VecDeque, budget: &Arc) { + let mut kept = VecDeque::with_capacity(frames.len()); + while let Some(mut entry) = frames.pop_back() { + match budget.try_charge(entry.charge.bytes()) { + Some(charge) => { + entry.charge = charge; + kept.push_front(entry); + } + None => { + log::debug!( + "dropping {} inherited frame(s) over the outbound budget", + frames.len() + 1 + ); + frames.clear(); + } + } + } + *frames = kept; +} + +/// Hand a frame to a connected peer's queue, dropping it if the queue is full. +fn enqueue(peer: &PeerHandle, outbound: Outbound) -> Result<(), SessionError> { + match peer.tx.try_send(outbound) { + Ok(()) => Ok(()), + Err(mpsc::error::TrySendError::Full(_)) => { + log::debug!("dropping frame for peer whose outbound queue is full"); + Ok(()) + } + Err(mpsc::error::TrySendError::Closed(_)) => Err(SessionError::NoCounterpart), + } +} + /// The two role slots for one topic. #[derive(Default)] struct Slots { @@ -180,6 +303,9 @@ impl Session { /// claimed by another identity keeps no frames, and the frames its previous /// owner sent to the counterpart are discarded. /// + /// The backlog is charged to the new connection's budget, newest first; a + /// replaced connection keeps only the frames already in its own queue. + /// /// `permit` is the connection's per-identity permit. A registration that /// replaces this identity's own live connection inherits that connection's /// permit instead, so it is accepted with `None`: the per-identity cap must @@ -241,6 +367,7 @@ impl Session { None => (VecDeque::new(), permit), }; prune_expired(&mut backlog, now); + recharge(&mut backlog, &peer.budget); if !backlog.is_empty() { log::debug!( "redelivering {} buffered frame(s) to {}", @@ -251,11 +378,18 @@ impl Session { let both_ready = slots.is_connected(role.counterpart()); if both_ready { - if peer.tx.try_send(ready.clone()).is_err() { + if peer + .tx + .try_send(Outbound::new(ready.clone(), None)) + .is_err() + { log::warn!("could not enqueue session_ready for {}", peer.identity); } if let Some(Slot::Connected { peer: other, .. }) = slots.get(role.counterpart()) - && other.tx.try_send(ready.clone()).is_err() + && other + .tx + .try_send(Outbound::new(ready.clone(), None)) + .is_err() { log::warn!("could not enqueue session_ready for {}", other.identity); } @@ -279,11 +413,12 @@ impl Session { /// Expired frames are dropped, never handed over. What does not fit stays /// in the backlog, in order, for a later call once the writer has made /// room. Only the connection owning `tx` may drain the slot, so a replaced - /// connection cannot consume its successor's backlog. + /// connection cannot consume its successor's backlog. A frame keeps its + /// charge as it moves: the queue and the backlog share one budget. pub(crate) fn drain_backlog( &self, role: Role, - tx: &mpsc::Sender, + tx: &mpsc::Sender, ) -> Result { let mut slots = self.lock()?; let Some(Slot::Connected { peer, backlog, .. }) = slots.get_mut(role) else { @@ -308,7 +443,7 @@ impl Session { match tx.try_reserve() { Ok(permit) => { if let Some(entry) = backlog.pop_front() { - permit.send(entry.frame); + permit.send(Outbound::new(entry.frame, Some(entry.charge))); delivered += 1; } } @@ -327,10 +462,11 @@ impl Session { /// /// Only acts if the slot still belongs to the connection owning `tx`, so a /// replaced connection tearing down cannot release its successor's slot. + /// The reservation buffers within a budget as large as the connection's. pub(crate) fn disconnect( &self, role: Role, - tx: &mpsc::WeakSender, + tx: &mpsc::WeakSender, grace: Duration, ) -> Result { let mut slots = self.lock()?; @@ -344,11 +480,13 @@ impl Session { } let identity = peer.identity.clone(); + let budget = Arc::new(OutboundBudget::new(peer.budget.capacity())); *slots.get_mut(role) = Some(Slot::Disconnected { identity, deadline: Instant::now() + grace, buffered: VecDeque::new(), + budget, }); Ok(true) @@ -377,7 +515,10 @@ impl Session { slots.discard_held_frames(role.counterpart()); if let Some(Slot::Connected { peer, .. }) = slots.get(role.counterpart()) - && peer.tx.try_send(notification.clone()).is_err() + && peer + .tx + .try_send(Outbound::new(notification.clone(), None)) + .is_err() { log::warn!( "dropping peer_disconnected for {}: outbound queue unavailable", @@ -393,6 +534,13 @@ impl Session { /// /// Buffered frames expire after `buffer_ttl`; a zero TTL drops them immediately. /// A frame queued behind a connected peer's backlog expires the same way. + /// + /// Each frame is charged its encoded size to the budget of the slot that + /// holds it, so a peer that reads slowly has frames dropped rather than + /// relay memory grown (spec §3). A full buffer or backlog drops its oldest + /// frames to make room, by count and by bytes; a frame that cannot fit + /// beside what the outbound queue holds is dropped, like one that finds + /// the queue full. pub(crate) fn route( &self, from: Role, @@ -400,6 +548,10 @@ impl Session { max_buffered: usize, buffer_ttl: Duration, ) -> Result<(), SessionError> { + let Some(bytes) = encoded_len(&frame) else { + log::debug!("dropping a frame that cannot be encoded"); + return Ok(()); + }; let mut slots = self.lock()?; let target = from.counterpart(); let now = Instant::now(); @@ -410,33 +562,47 @@ impl Session { while backlog.len() >= max_buffered.max(1) { backlog.pop_front(); } - backlog.push_back(BufferedFrame { - frame, - expires_at: now + buffer_ttl, - }); - return Ok(()); - } - return match peer.tx.try_send(frame) { - Ok(()) => Ok(()), - Err(mpsc::error::TrySendError::Full(_)) => { - log::debug!("dropping frame for peer whose outbound queue is full"); - Ok(()) + let Some(charge) = make_room(backlog, &peer.budget, bytes) else { + log::debug!("dropping frame for peer whose outbound byte budget is full"); + return Ok(()); + }; + if !backlog.is_empty() { + backlog.push_back(BufferedFrame { + frame, + expires_at: now + buffer_ttl, + charge, + }); + return Ok(()); } - Err(mpsc::error::TrySendError::Closed(_)) => Err(SessionError::NoCounterpart), + // Room was made by dropping the whole backlog: nothing is left + // ahead of the frame, and one left on an empty backlog would + // never be drained. + return enqueue(peer, Outbound::new(frame, Some(charge))); + } + let Some(charge) = peer.budget.try_charge(bytes) else { + log::debug!("dropping frame for peer whose outbound byte budget is full"); + return Ok(()); }; + return enqueue(peer, Outbound::new(frame, Some(charge))); } match slots.get_mut(target) { - Some(Slot::Disconnected { buffered, .. }) => { + Some(Slot::Disconnected { + buffered, budget, .. + }) => { if max_buffered > 0 && !buffer_ttl.is_zero() { prune_expired(buffered, now); if buffered.len() >= max_buffered { buffered.pop_front(); } - buffered.push_back(BufferedFrame { - frame, - expires_at: now + buffer_ttl, - }); + match make_room(buffered, budget, bytes) { + Some(charge) => buffered.push_back(BufferedFrame { + frame, + expires_at: now + buffer_ttl, + charge, + }), + None => log::debug!("dropping frame larger than the reconnect buffer"), + } } Ok(()) } @@ -448,7 +614,6 @@ impl Session { #[cfg(test)] mod tests { use super::*; - use std::sync::Arc; use crate::connection::RELAY_FRAME_VERSION; use crate::limits::{Limiter, RelayLimitsConfig}; @@ -465,21 +630,44 @@ mod tests { /// Outbound queue capacity used by peer handles in these tests. const PEER_QUEUE: usize = 8; + /// Outbound byte budget of peer handles in tests that do not exercise it. + const PEER_BUDGET: usize = 1024 * 1024; + fn topic() -> Result { Topic::new([1u8; 32]) } - fn peer(seed: u8) -> (PeerHandle, mpsc::Receiver) { - let (tx, rx) = mpsc::channel(PEER_QUEUE); + fn peer(seed: u8) -> (PeerHandle, mpsc::Receiver) { + peer_with(seed, PEER_QUEUE, PEER_BUDGET) + } + + /// A peer handle with a queue of `frames` and a budget of `bytes`. + fn peer_with(seed: u8, frames: usize, bytes: usize) -> (PeerHandle, mpsc::Receiver) { + let (tx, rx) = mpsc::channel(frames); ( PeerHandle { identity: DidKey::from_ed25519_bytes([seed; 32]), tx, + budget: Arc::new(OutboundBudget::new(bytes)), }, rx, ) } + /// Encoded size of a frame, as the writer will send it. + fn size(frame: &RelayFrame) -> usize { + encoded_len(frame).expect("encodable frame") + } + + /// Ids of the frames waiting in a queue, taking them out as a writer would. + fn take_ids(rx: &mut mpsc::Receiver) -> Vec { + let mut ids = Vec::new(); + while let Ok(queued) = rx.try_recv() { + ids.push(queued.frame.id); + } + ids + } + /// Disconnect whichever connection currently owns `role`'s slot. fn disconnect_current( session: &Session, @@ -530,6 +718,29 @@ mod tests { } } + /// Route a message frame with `id` from the proposer to the responder. + fn route_id(session: &Session, id: &str) -> Result<(), SessionError> { + let frame = frame_with_id(&topic().expect("topic"), id); + session.route(Role::Proposer, frame, MAX_BUFFERED, BUFFER_TTL) + } + + /// A ready pair whose responder went away while `ids` were sent to it, + /// with the proposer's queue. + fn responder_away_with( + ids: &[&str], + ) -> Result<(Session, mpsc::Receiver), SessionError> { + let session = Session::default(); + let (proposer, proposer_rx) = peer(1); + let (responder, _) = peer(2); + register(&session, Role::Proposer, proposer)?; + register(&session, Role::Responder, responder)?; + disconnect_current(&session, Role::Responder, Duration::from_secs(10))?; + for id in ids { + route_id(&session, id)?; + } + Ok((session, proposer_rx)) + } + /// Ids a freshly registered peer receives: queued frames, then the backlog. /// /// Drains in steps like a connection does, emptying the queue in between @@ -537,15 +748,13 @@ mod tests { async fn received_ids( session: &Session, role: Role, - tx: &mpsc::Sender, - rx: &mut mpsc::Receiver, + tx: &mpsc::Sender, + rx: &mut mpsc::Receiver, ) -> Result, SessionError> { let mut ids = Vec::new(); loop { let drain = session.drain_backlog(role, tx)?; - while let Ok(frame) = rx.try_recv() { - ids.push(frame.id); - } + ids.extend(take_ids(rx)); if drain.finished { return Ok(ids); } @@ -611,10 +820,10 @@ mod tests { session.route(Role::Proposer, frame(&topic()?), MAX_BUFFERED, BUFFER_TTL)?; assert!(matches!( - responder_rx.try_recv()?.body, + responder_rx.try_recv()?.frame.body, RelayFrameBody::SessionReady )); - let received = responder_rx.try_recv()?; + let received = responder_rx.try_recv()?.frame; assert!(matches!(received.body, RelayFrameBody::Message { .. })); Ok(()) } @@ -665,7 +874,7 @@ mod tests { // frames to the disconnected peer are buffered, not dropped session.route(Role::Proposer, frame(&topic()?), MAX_BUFFERED, BUFFER_TTL)?; - assert_eq!(responder_rx.try_recv()?.id, "ready"); + assert_eq!(responder_rx.try_recv()?.frame.id, "ready"); assert!(responder_rx.try_recv().is_err()); // same identity reconnects within grace and receives session_ready, @@ -772,13 +981,13 @@ mod tests { register(&session, Role::Proposer, proposer)?; register(&session, Role::Responder, responder)?; disconnect_current(&session, Role::Responder, Duration::from_secs(1))?; - assert_eq!(proposer_rx.try_recv()?.id, "ready"); + assert_eq!(proposer_rx.try_recv()?.frame.id, "ready"); tokio::time::advance(Duration::from_secs(2)).await; let notice = notification(); assert!(session.expire(Role::Responder, ¬ice)?.is_some()); - let queued = proposer_rx.try_recv()?; + let queued = proposer_rx.try_recv()?.frame; assert!( matches!(queued.body, RelayFrameBody::PeerDisconnected), "expected the disconnect in the survivor's queue, got {:?}", @@ -837,8 +1046,8 @@ mod tests { BUFFER_TTL, )?; - assert_eq!(responder_rx.try_recv()?.id, "ready"); - assert_eq!(responder_rx.try_recv()?.id, "m1"); + assert_eq!(responder_rx.try_recv()?.frame.id, "ready"); + assert_eq!(responder_rx.try_recv()?.frame.id, "m1"); Ok(()) } @@ -887,7 +1096,7 @@ mod tests { MAX_BUFFERED, BUFFER_TTL, )?; - assert_eq!(responder_rx.try_recv()?.id, "live"); + assert_eq!(responder_rx.try_recv()?.frame.id, "live"); Ok(()) } @@ -926,10 +1135,7 @@ mod tests { "session_ready holds one slot of the queue" ); - let mut ids = Vec::new(); - while let Ok(frame) = responder_rx.try_recv() { - ids.push(frame.id); - } + let mut ids = take_ids(&mut responder_rx); ids.extend(received_ids(&session, Role::Responder, &tx, &mut responder_rx).await?); let mut expected = vec!["ready".to_owned()]; @@ -978,7 +1184,7 @@ mod tests { let (old, mut old_rx) = peer(2); register(&session, Role::Proposer, proposer)?; register(&session, Role::Responder, old)?; - assert_eq!(proposer_rx.try_recv()?.id, "ready"); + assert_eq!(proposer_rx.try_recv()?.frame.id, "ready"); // Same identity, new connection, while the old one is still registered. let (new, mut new_rx) = peer(2); @@ -989,7 +1195,7 @@ mod tests { )); // The counterpart sees the peer come back. - assert_eq!(proposer_rx.try_recv()?.id, "ready"); + assert_eq!(proposer_rx.try_recv()?.frame.id, "ready"); // Routing now reaches the new connection only. session.route( @@ -1033,8 +1239,8 @@ mod tests { MAX_BUFFERED, BUFFER_TTL, )?; - assert_eq!(owner_rx.try_recv()?.id, "ready"); - assert_eq!(owner_rx.try_recv()?.id, "m1"); + assert_eq!(owner_rx.try_recv()?.frame.id, "ready"); + assert_eq!(owner_rx.try_recv()?.frame.id, "m1"); Ok(()) } @@ -1486,4 +1692,233 @@ mod tests { ); Ok(()) } + + #[test] + fn encoded_len_is_the_text_the_writer_sends() -> TestResult { + let frame = RelayFrame { + id: "quote \" backslash \\ newline \n control \u{1}".to_owned(), + body: RelayFrameBody::Message { + payload: "é \u{7f} \u{1f} \u{1F600}".to_owned(), + }, + ..frame(&topic()?) + }; + assert_eq!( + encoded_len(&frame), + Some(serde_json::to_string(&frame)?.len()) + ); + Ok(()) + } + + #[test] + fn frames_over_the_byte_budget_are_dropped_until_written() -> TestResult { + let one = size(&frame_with_id(&topic()?, "m1")); + let session = Session::default(); + let (proposer, _proposer_rx) = peer(1); + // Room in the queue for many frames, but bytes for two. + let (responder, mut responder_rx) = peer_with(2, PEER_QUEUE, 2 * one); + let budget = Arc::clone(&responder.budget); + register(&session, Role::Proposer, proposer)?; + register(&session, Role::Responder, responder)?; + + for id in ["m1", "m2", "m3"] { + route_id(&session, id)?; + } + assert_eq!(budget.held(), 2 * one, "session_ready is not charged"); + assert_eq!(take_ids(&mut responder_rx), ["ready", "m1", "m2"]); + + // Written frames give their bytes back. + assert_eq!(budget.held(), 0); + route_id(&session, "m4")?; + assert_eq!(take_ids(&mut responder_rx), ["m4"]); + Ok(()) + } + + #[tokio::test(start_paused = true)] + async fn a_full_byte_budget_never_refuses_relay_control_frames() -> TestResult { + let one = size(&frame_with_id(&topic()?, "m1")); + let session = Session::default(); + let (proposer, _proposer_rx) = peer(1); + let (responder, mut responder_rx) = peer_with(2, PEER_QUEUE, one); + let budget = Arc::clone(&responder.budget); + register(&session, Role::Proposer, proposer)?; + register(&session, Role::Responder, responder)?; + route_id(&session, "m1")?; + assert_eq!(budget.held(), budget.capacity()); + + // The proposer goes away for good, and another identity takes its role. + disconnect_current(&session, Role::Proposer, Duration::from_secs(1))?; + tokio::time::advance(Duration::from_secs(2)).await; + assert!(session.expire(Role::Proposer, ¬ification())?.is_some()); + let (other, _other_rx) = peer(3); + register(&session, Role::Proposer, other)?; + + assert_eq!( + take_ids(&mut responder_rx), + ["ready", "m1", "gone", "ready"], + "control frames bypass the byte budget" + ); + Ok(()) + } + + #[tokio::test(start_paused = true)] + async fn queue_and_backlog_share_one_byte_budget() -> TestResult { + let one = size(&frame_with_id(&topic()?, "m1")); + let (session, _proposer_rx) = responder_away_with(&["m1", "m2", "m3", "m4"])?; + + // It returns with queue room for session_ready and one frame, and + // bytes for four frames. + let (responder, mut responder_rx) = peer_with(2, 2, 4 * one); + let tx = responder.tx.clone(); + let budget = Arc::clone(&responder.budget); + register(&session, Role::Responder, responder)?; + assert_eq!(budget.held(), 4 * one, "charged on reconnect"); + let drain = session.drain_backlog(Role::Responder, &tx)?; + assert_eq!(drain.delivered, 1); + assert_eq!(budget.held(), 4 * one, "kept in the queue"); + + // A new frame needs room: the oldest backlog frame goes, not the one + // already queued. + route_id(&session, "m5")?; + assert_eq!(budget.held(), 4 * one); + assert_eq!( + received_ids(&session, Role::Responder, &tx, &mut responder_rx).await?, + ["ready", "m1", "m3", "m4", "m5"] + ); + assert_eq!(budget.held(), 0, "written frames give their bytes back"); + Ok(()) + } + + #[tokio::test(start_paused = true)] + async fn a_frame_too_large_beside_the_queue_leaves_the_backlog_alone() -> TestResult { + let topic = topic()?; + let one = size(&frame_with_id(&topic, "m1")); + let (session, _proposer_rx) = responder_away_with(&["m1", "m2"])?; + + let (responder, mut responder_rx) = peer_with(2, 2, 3 * one); + let tx = responder.tx.clone(); + register(&session, Role::Responder, responder)?; + let drain = session.drain_backlog(Role::Responder, &tx)?; + assert_eq!(drain.delivered, 1); + + // With one frame queued, this one cannot fit even if the backlog went. + let large = RelayFrame { + body: RelayFrameBody::Message { + payload: "a".repeat(2 * one), + }, + ..frame_with_id(&topic, "large") + }; + session.route(Role::Proposer, large, MAX_BUFFERED, BUFFER_TTL)?; + assert_eq!( + received_ids(&session, Role::Responder, &tx, &mut responder_rx).await?, + ["ready", "m1", "m2"] + ); + Ok(()) + } + + #[tokio::test(start_paused = true)] + async fn a_frame_that_displaces_the_whole_backlog_is_queued_directly() -> TestResult { + let one = size(&frame_with_id(&topic()?, "m1")); + let (session, _proposer_rx) = responder_away_with(&["m1", "m2"])?; + + let (responder, mut responder_rx) = peer_with(2, 2, 2 * one); + let tx = responder.tx.clone(); + register(&session, Role::Responder, responder)?; + let drain = session.drain_backlog(Role::Responder, &tx)?; + assert_eq!(drain.delivered, 1); + + // The writer takes session_ready, which frees a queue slot but no bytes. + assert_eq!(responder_rx.try_recv()?.frame.id, "ready"); + + // Room for m3 means dropping m2, the whole backlog. A frame left on an + // empty backlog would never be drained, so m3 goes to the queue. + route_id(&session, "m3")?; + let drain = session.drain_backlog(Role::Responder, &tx)?; + assert_eq!( + drain, + Drain { + delivered: 0, + finished: true, + } + ); + assert_eq!(take_ids(&mut responder_rx), ["m1", "m3"]); + Ok(()) + } + + #[tokio::test(start_paused = true)] + async fn reconnect_buffer_is_capped_by_bytes_dropping_the_oldest() -> TestResult { + let one = size(&frame_with_id(&topic()?, "m1")); + let session = Session::default(); + let (proposer, _proposer_rx) = peer(1); + let (responder, _) = peer_with(2, PEER_QUEUE, 2 * one); + register(&session, Role::Proposer, proposer)?; + register(&session, Role::Responder, responder)?; + disconnect_current(&session, Role::Responder, Duration::from_secs(10))?; + for id in ["m1", "m2", "m3"] { + route_id(&session, id)?; + } + let reserved = match session.lock()?.get(Role::Responder) { + Some(Slot::Disconnected { budget, .. }) => budget.held(), + _ => 0, + }; + assert_eq!( + reserved, + 2 * one, + "the reservation is as large as the connection's budget" + ); + + let (responder, mut responder_rx) = peer(2); + let tx = responder.tx.clone(); + let budget = Arc::clone(&responder.budget); + register(&session, Role::Responder, responder)?; + assert_eq!(budget.held(), 2 * one, "moved to the new budget"); + assert_eq!( + received_ids(&session, Role::Responder, &tx, &mut responder_rx).await?, + ["ready", "m2", "m3"] + ); + Ok(()) + } + + #[tokio::test(start_paused = true)] + async fn a_reconnect_keeps_the_newest_frames_its_budget_holds() -> TestResult { + let one = size(&frame_with_id(&topic()?, "m1")); + let (session, _proposer_rx) = responder_away_with(&["m1", "m2", "m3"])?; + + // A connection with a smaller budget than the reservation it reclaims. + let (responder, mut responder_rx) = peer_with(2, PEER_QUEUE, 2 * one); + let tx = responder.tx.clone(); + register(&session, Role::Responder, responder)?; + assert_eq!( + received_ids(&session, Role::Responder, &tx, &mut responder_rx).await?, + ["ready", "m2", "m3"] + ); + Ok(()) + } + + #[tokio::test(start_paused = true)] + async fn replacement_takes_the_backlog_onto_its_own_budget() -> TestResult { + let one = size(&frame_with_id(&topic()?, "m1")); + let (session, _proposer_rx) = responder_away_with(&["m1", "m2", "m3"])?; + + // The first reconnect queues session_ready and one frame. + let (first, _first_rx) = peer_with(2, 2, PEER_BUDGET); + let first_tx = first.tx.clone(); + let first_budget = Arc::clone(&first.budget); + register(&session, Role::Responder, first)?; + let drain = session.drain_backlog(Role::Responder, &first_tx)?; + assert_eq!(drain.delivered, 1); + assert_eq!(first_budget.held(), 3 * one); + + // Its replacement takes the backlog; the queued frame stays behind. + let (second, mut second_rx) = peer(2); + let second_tx = second.tx.clone(); + let second_budget = Arc::clone(&second.budget); + register(&session, Role::Responder, second)?; + assert_eq!(first_budget.held(), one); + assert_eq!(second_budget.held(), 2 * one); + assert_eq!( + received_ids(&session, Role::Responder, &second_tx, &mut second_rx).await?, + ["ready", "m2", "m3"] + ); + Ok(()) + } } diff --git a/deployment/configs/relay/base.yaml.template b/deployment/configs/relay/base.yaml.template index 7e775de..f931605 100644 --- a/deployment/configs/relay/base.yaml.template +++ b/deployment/configs/relay/base.yaml.template @@ -23,6 +23,11 @@ connection: handshake_timeout_seconds: 10 write_timeout_seconds: 10 outbound_queue_capacity: 64 + # Bytes held for one peer and not yet written: its queue and backlog, or + # its buffer while it reconnects. At least max_incoming_frame_bytes. + outbound_queue_max_bytes: 1048576 + # A frame not written this long after it was queued drops its slow reader. + outbound_queue_timeout_seconds: 30 ping_interval_seconds: 30 pong_timeout_seconds: 10 idle_timeout_seconds: 45 From 09eca9cff1176c1de101ed0bc40ec1f8afe99e98 Mon Sep 17 00:00:00 2001 From: Kyryl R Date: Wed, 28 Oct 2026 11:52:21 +0200 Subject: [PATCH 3/9] Charge WebSocket control frames to shared inbound budgets --- crates/relay/src/connection.rs | 67 ++++++++++----- crates/relay/src/server.rs | 143 +++++++++++++++++++++++++++++++-- 2 files changed, 185 insertions(+), 25 deletions(-) diff --git a/crates/relay/src/connection.rs b/crates/relay/src/connection.rs index 5d76028..4fdfa61 100644 --- a/crates/relay/src/connection.rs +++ b/crates/relay/src/connection.rs @@ -21,7 +21,7 @@ use wallet_session_core::{codec, connection::ConnectionQuery}; use crate::{ auth::{Identity, IdentityError, unix_now}, config::NativeRelayTransportConfig, - limits::{Bucket, InboundBytes, Limiter, RelayLimitsConfig}, + limits::{Bucket, InboundBytes, Limiter}, session::{Outbound, Registration, Role, Session, SessionError}, }; use wallet_session_core::{ @@ -41,7 +41,7 @@ const WRITER_SHUTDOWN_GRACE: Duration = Duration::from_millis(100); /// Consecutive frames a connection may have rejected by its own inbound rate /// limits before it is disconnected. -const MAX_CONSECUTIVE_RATE_DROPS: u32 = 32; +pub(crate) const MAX_CONSECUTIVE_RATE_DROPS: u32 = 32; /// Maximum accepted query-string length before parsing. pub const MAX_QUERY_BYTES: usize = 8 * 1024; @@ -401,36 +401,61 @@ fn spend_redelivery_credit(credit: &mut u32) -> bool { true } -/// Charge an inbound control frame to the connection's inbound budget. +/// Charge an inbound control frame like a text frame: to the connection's +/// inbound budget, then to the identity's share of the relay-wide budget and +/// to the budget itself. /// -/// The frame is dropped when the budget is exhausted, and a peer that keeps -/// exceeding it is disconnected like any other flooder. Returns `Err` only to -/// disconnect the peer. +/// The frame is dropped when a budget is exhausted, and a peer that keeps +/// exceeding its own or its identity's budget is disconnected like any other +/// flooder. An exhausted relay-wide budget alone drops the frame without +/// counting towards that, so other identities cannot get this peer +/// disconnected. Returns `Err` only to disconnect the peer. fn charge_control_frame( bytes: usize, message_bucket: &mut Bucket, byte_bucket: &mut Bucket, - limits: &RelayLimitsConfig, + limiter: &Limiter, + identity: &DidKey, consecutive_drops: &mut u32, topic: &Topic, ) -> Result<(), ConnectionError> { + let limits = limiter.limits(); let now = tokio::time::Instant::now(); - let within_limits = message_bucket.allow_n(1, limits.inbound_message_rate, now) + let within_connection_limits = message_bucket.allow_n(1, limits.inbound_message_rate, now) && byte_bucket.allow_n( u32::try_from(bytes).unwrap_or(u32::MAX), limits.inbound_byte_rate, now, ); - if within_limits { - *consecutive_drops = 0; + if !within_connection_limits { + *consecutive_drops += 1; + log::debug!("dropping control frame on topic {topic}: inbound rate limit exceeded"); + if *consecutive_drops > MAX_CONSECUTIVE_RATE_DROPS { + log::warn!("peer on topic {topic} keeps exceeding its inbound rate limit; dropping"); + return Err(ConnectionError::RateLimited); + } return Ok(()); } - *consecutive_drops += 1; - log::debug!("dropping control frame on topic {topic}: inbound rate limit exceeded"); - if *consecutive_drops > MAX_CONSECUTIVE_RATE_DROPS { - log::warn!("peer on topic {topic} keeps exceeding its inbound rate limit; dropping"); - return Err(ConnectionError::RateLimited); + match limiter.allow_inbound_bytes(identity, bytes) { + InboundBytes::Allowed => *consecutive_drops = 0, + InboundBytes::ShareExhausted => { + *consecutive_drops += 1; + log::debug!( + "dropping control frame on topic {topic}: identity share of the relay budget exhausted" + ); + if *consecutive_drops > MAX_CONSECUTIVE_RATE_DROPS { + log::warn!( + "peer on topic {topic} keeps exceeding its share of the relay budget; dropping" + ); + return Err(ConnectionError::RateLimited); + } + } + InboundBytes::GlobalExhausted => { + log::debug!( + "dropping control frame on topic {topic}: relay inbound byte budget exhausted" + ); + } } Ok(()) } @@ -660,14 +685,17 @@ async fn read_loop( return Err(ConnectionError::UnexpectedBinary); } // Keepalive pongs and other control frames count as activity, but - // they are charged to the same budget as data frames so a peer - // cannot hold the connection open with unmetered traffic. + // they are charged to the same budgets as data frames, the shared + // ones included, so a peer cannot hold the connection open with + // unmetered traffic and identities cannot exceed the relay-wide + // budget with it. Message::Ping(payload) | Message::Pong(payload) => { charge_control_frame( payload.len(), &mut message_bucket, &mut byte_bucket, - &limits, + limiter, + identity, &mut consecutive_drops, &query.topic, )?; @@ -678,7 +706,8 @@ async fn read_loop( frame.payload().len(), &mut message_bucket, &mut byte_bucket, - &limits, + limiter, + identity, &mut consecutive_drops, &query.topic, )?; diff --git a/crates/relay/src/server.rs b/crates/relay/src/server.rs index f487639..1684326 100644 --- a/crates/relay/src/server.rs +++ b/crates/relay/src/server.rs @@ -185,6 +185,7 @@ pub enum ServerError { mod tests { use super::*; use crate::auth::test_utils::{self, signing_key}; + use crate::connection::MAX_CONSECUTIVE_RATE_DROPS; use crate::limits::{Limiter, RateLimit, RelayLimitsConfig}; use futures_util::{SinkExt, StreamExt}; use std::net::SocketAddr; @@ -384,15 +385,17 @@ mod tests { async fn serve_with_limits( config: NativeRelayTransportConfig, limits: RelayLimitsConfig, + ) -> Result> { + serve_with_limiter(config, Arc::new(Limiter::new(limits))).await + } + + async fn serve_with_limiter( + config: NativeRelayTransportConfig, + limiter: Arc, ) -> Result> { let listener = TcpListener::bind("127.0.0.1:0").await?; let addr = listener.local_addr()?; - tokio::spawn(serve( - listener, - Arc::new(Registry::new()), - Arc::new(Limiter::new(limits)), - config, - )); + tokio::spawn(serve(listener, Arc::new(Registry::new()), limiter, config)); Ok(addr) } @@ -2106,6 +2109,134 @@ mod tests { Ok(()) } + /// Wait until `counter` reaches `expected`, for at most two seconds, and + /// return what it reads then. + async fn wait_for_count(counter: impl Fn() -> u64, expected: u64) -> u64 { + for _ in 0..100 { + if counter() >= expected { + break; + } + sleep(Duration::from_millis(20)).await; + } + counter() + } + + /// Send `count` pings carrying `payload` bytes each. + async fn send_pings(client: &mut Client, count: u32, payload: u32) -> TestResult { + let payload = Bytes::from(vec![b'p'; usize::try_from(payload)?]); + for _ in 0..count { + client.send(Message::Ping(payload.clone())).await?; + } + Ok(()) + } + + /// Bytes of the pings these tests send. + const PING_PAYLOAD: u32 = 100; + + /// Limits under which only the shared byte budgets can refuse a ping. + fn ping_limits() -> RelayLimitsConfig { + let roomy = RateLimit { + burst: 1024 * 1024, + interval: Duration::from_secs(3600), + }; + RelayLimitsConfig { + inbound_message_rate: RateLimit { + burst: 1000, + interval: Duration::from_secs(3600), + }, + inbound_byte_rate: roomy, + identity_inbound_byte_rate: roomy, + global_inbound_byte_rate: roomy, + ..RelayLimitsConfig::default() + } + } + + /// Pings spend the relay-wide inbound budget like text frames, and its + /// exhaustion drops them without disconnecting an identity that is within + /// its own limits. + #[tokio::test] + async fn control_frames_spend_the_global_budget_and_spare_others() -> TestResult { + // Room for ten pings, refilling far slower than the test runs. + let limiter = Arc::new(Limiter::new(RelayLimitsConfig { + global_inbound_byte_rate: RateLimit { + burst: 10 * PING_PAYLOAD, + interval: Duration::from_secs(3600), + }, + ..ping_limits() + })); + let config = NativeRelayTransportConfig::new(AUDIENCE); + let addr = serve_with_limiter(config, Arc::clone(&limiter)).await?; + + // One identity spends the whole relay-wide budget on pings: the + // eleventh no longer fits it. + let mut hog = connect_auth(addr, &token_with(7), "proposer", &topic()?, None).await?; + send_pings(&mut hog, 11, PING_PAYLOAD).await?; + let payload = u64::from(PING_PAYLOAD); + let dropped = wait_for_count(|| limiter.dropped_global_bytes(), payload).await; + assert_eq!( + dropped, payload, + "ten pings must fit the relay-wide budget and the eleventh must not" + ); + + // An unrelated identity within its own limits has its pings dropped + // for the exhausted relay budget, more often than the consecutive-drop + // limit allows, and is not disconnected for it. + let other = Topic::new([2u8; 32])?; + let mut victim = connect_auth(addr, &token_with(8), "proposer", &other, None).await?; + let flood = MAX_CONSECUTIVE_RATE_DROPS + 8; + send_pings(&mut victim, flood, PING_PAYLOAD).await?; + let expected = payload * (1 + u64::from(flood)); + let dropped = wait_for_count(|| limiter.dropped_global_bytes(), expected).await; + assert_eq!( + dropped, expected, + "pings past the exhausted relay budget must be dropped and counted" + ); + assert_eq!(limiter.dropped_share_bytes(), 0); + + for client in [&mut victim, &mut hog] { + assert!( + timeout(Duration::from_millis(500), expect_closed(client)) + .await + .is_err(), + "an exhausted relay-wide budget alone must not disconnect anyone" + ); + } + Ok(()) + } + + /// Pings spend the sender's identity share, and an identity that keeps + /// exceeding it with them is disconnected like a text flooder. + #[tokio::test] + async fn control_frames_over_the_identity_share_disconnect_the_sender() -> TestResult { + // Room for one ping, refilling far slower than the test runs. + let limiter = Arc::new(Limiter::new(RelayLimitsConfig { + identity_inbound_byte_rate: RateLimit { + burst: PING_PAYLOAD, + interval: Duration::from_secs(3600), + }, + ..ping_limits() + })); + let config = NativeRelayTransportConfig::new(AUDIENCE); + let addr = serve_with_limiter(config, Arc::clone(&limiter)).await?; + + // Every ping after the first is dropped, and once the drops pass the + // consecutive-drop limit the sender is disconnected. Sends may fail + // after that. + let mut hog = connect(addr, "proposer", &topic()?, None).await?; + let _ = send_pings(&mut hog, MAX_CONSECUTIVE_RATE_DROPS + 8, PING_PAYLOAD).await; + assert!( + expect_closed(&mut hog).await, + "an identity flooding control frames over its share must be disconnected" + ); + let limit = u64::from(PING_PAYLOAD) * u64::from(MAX_CONSECUTIVE_RATE_DROPS); + assert!( + limiter.dropped_share_bytes() > limit, + "the dropped pings must be counted against the identity share" + ); + assert_eq!(limiter.dropped_global_bytes(), 0); + Ok(()) + } + #[tokio::test] async fn frame_over_incoming_size_cap_closes_connection() -> TestResult { let addr = serve_with( From 34b76af23607457ca03fc41f5bb2f1187655a43d Mon Sep 17 00:00:00 2001 From: Kyryl R Date: Wed, 28 Oct 2026 12:15:01 +0200 Subject: [PATCH 4/9] Redact relay topics and bound malformed-frame logging --- crates/core/src/message.rs | 61 +++++ crates/relay/README.md | 6 + crates/relay/src/connection.rs | 155 ++++++++---- crates/relay/src/main.rs | 9 +- crates/relay/src/registry.rs | 17 +- crates/relay/src/server.rs | 124 +++++++++- crates/relay/src/telemetry.rs | 432 +++++++++++++++++++++++++++++++++ 7 files changed, 740 insertions(+), 64 deletions(-) diff --git a/crates/core/src/message.rs b/crates/core/src/message.rs index c858fbc..e460d26 100644 --- a/crates/core/src/message.rs +++ b/crates/core/src/message.rs @@ -3,6 +3,7 @@ use crate::did::DidKey; use crate::error::Error; use crate::payload::Payload; use serde::{Deserialize, Serialize}; +use sha2::{Digest, Sha256}; use std::fmt; use std::str::FromStr; @@ -162,6 +163,17 @@ impl Topic { } Ok(Self(bytes)) } + + /// A short tag that tells topics apart in logs without disclosing them. + /// + /// A topic is a bearer secret, so logs carry this tag instead: the first + /// four bytes of SHA-256 over a fixed domain separator and the topic, in + /// hex. It is the same in every process, and nothing short of searching + /// the topic space recovers a topic from it. The digest is only computed + /// when the tag is formatted. + pub fn log_tag(&self) -> TopicLogTag<'_> { + TopicLogTag(self) + } } impl fmt::Display for Topic { @@ -170,6 +182,37 @@ impl fmt::Display for Topic { } } +/// Log-safe stand-in for a [`Topic`]; see [`Topic::log_tag`]. +#[derive(Clone, Copy)] +pub struct TopicLogTag<'a>(&'a Topic); + +impl TopicLogTag<'_> { + /// Domain separator, so the tag is not a plain hash of the topic. + const DOMAIN: &'static [u8] = b"wallet-session/topic-log-tag/v1"; + + /// Digest bytes kept. + const LEN: usize = 4; +} + +impl fmt::Display for TopicLogTag<'_> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + let digest = Sha256::new() + .chain_update(Self::DOMAIN) + .chain_update(&self.0.0) + .finalize(); + digest[..Self::LEN] + .iter() + .try_for_each(|byte| write!(f, "{byte:02x}")) + } +} + +/// Same as [`fmt::Display`]: a derived `Debug` would print the topic. +impl fmt::Debug for TopicLogTag<'_> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + fmt::Display::fmt(self, f) + } +} + impl FromStr for Topic { type Err = Error; @@ -321,6 +364,24 @@ mod tests { assert_eq!(topic.to_string().parse::().unwrap(), topic); } + #[test] + fn topic_log_tag_is_short_stable_and_hides_the_topic() { + let topic = Topic::new([7u8; 32]).unwrap(); + let tag = topic.log_tag().to_string(); + + assert_eq!(tag.len(), 8); + assert!( + tag.bytes() + .all(|byte| matches!(byte, b'0'..=b'9' | b'a'..=b'f')) + ); + assert_eq!(topic.clone().log_tag().to_string(), tag); + assert_eq!(format!("{:?}", topic.log_tag()), tag); + assert_ne!(Topic::new([8u8; 32]).unwrap().log_tag().to_string(), tag); + + // The topic's own encoding is untouched. + assert_eq!(topic.to_string(), codec::encode(&[7u8; 32])); + } + #[test] fn ciphertext_base64url_round_trip() { let data = Ciphertext::new(vec![1, 2, 3, 4, 5]); diff --git a/crates/relay/README.md b/crates/relay/README.md index 097c154..21673b0 100644 --- a/crates/relay/README.md +++ b/crates/relay/README.md @@ -75,6 +75,12 @@ container memory limit): the relay enforces no aggregate bound. ## Running +Relay logs use a short topic tag instead of the bearer topic. Invalid-frame +records report the error class and location without echoing client content. +Stdout is written by a separate thread with an 8 MiB queue; when the sink +falls behind, whole records are dropped and a count is reported after it +recovers. Configure retention and collection for the deployment separately. + From this directory, since the configuration is read from the working directory: diff --git a/crates/relay/src/connection.rs b/crates/relay/src/connection.rs index 4fdfa61..8e7bca8 100644 --- a/crates/relay/src/connection.rs +++ b/crates/relay/src/connection.rs @@ -1,6 +1,7 @@ -use std::{sync::Arc, time::Duration}; +use std::{fmt, sync::Arc, time::Duration}; use futures_util::{Sink, SinkExt, StreamExt, stream::SplitStream}; +use serde_json::error::Category; use thiserror::Error; use tokio::{ net::TcpStream, @@ -161,7 +162,7 @@ impl Connection { let websocket = handshake.map_err(ConnectionError::from)?; log::debug!( "handshake accepted for topic {} as {:?} ({})", - query.topic, + query.topic.log_tag(), query.role, identity.did ); @@ -197,20 +198,17 @@ impl Connection { registration: Registration, ) -> Result<(), ConnectionError> { let mut websocket = websocket; + let topic = self.query.topic.log_tag(); log::debug!( - "serving peer {} on topic {} as {:?}", + "serving peer {} on topic {topic} as {:?}", self.identity.did, - self.query.topic, self.query.role ); let both_ready = match registration { Registration::BothReady => true, Registration::Waiting => false, Registration::RoleTaken(role) => { - log::warn!( - "rejecting connection for topic {}: {role:?} role already taken", - self.query.topic - ); + log::warn!("rejecting connection for topic {topic}: {role:?} role already taken"); let frame = CloseFrame { code: CloseCode::Policy, reason: "role already taken".into(), @@ -242,8 +240,7 @@ impl Connection { let config = &self.config; if both_ready { log::info!( - "session ready for topic {}, role {:?}", - self.query.topic, + "session ready for topic {topic}, role {:?}", self.query.role ); } @@ -337,7 +334,12 @@ async fn write_loop( let Ok(text) = serde_json::to_string(frame) else { break; }; - log::trace!("relay -> topic={} frame={text}", frame.topic); + log::trace!( + "relay -> topic={} type={} ({} bytes)", + frame.topic.log_tag(), + frame.body.relay_only_body().unwrap_or("message"), + text.len() + ); let message = Message::text(text); if !send_queued(&mut sink, message, queued.queued_at, &config).await { break; @@ -419,6 +421,7 @@ fn charge_control_frame( consecutive_drops: &mut u32, topic: &Topic, ) -> Result<(), ConnectionError> { + let topic = topic.log_tag(); let limits = limiter.limits(); let now = tokio::time::Instant::now(); let within_connection_limits = message_bucket.allow_n(1, limits.inbound_message_rate, now) @@ -548,6 +551,7 @@ async fn read_loop( identity, expires_at, } = *context; + let topic = query.topic.log_tag(); let limits = limiter.limits(); let now = tokio::time::Instant::now(); @@ -567,24 +571,20 @@ async fn read_loop( biased; _ = &mut *writer => { - log::debug!("connection writer stopped for topic {}; closing", query.topic); + log::debug!("connection writer stopped for topic {topic}; closing"); break; } _ = wait_pong_deadline(ping_deadline_rx) => { - log::debug!("peer on topic {} missed a keepalive pong; dropping", query.topic); + log::debug!("peer on topic {topic} missed a keepalive pong; dropping"); return Err(ConnectionError::PongTimeout); } _ = tokio::time::sleep_until(expires_at) => { - log::debug!( - "peer on topic {} has an expired token; closing", - query.topic - ); + log::debug!("peer on topic {topic} has an expired token; closing"); return Err(ConnectionError::TokenExpired); } _ = tokio::time::sleep_until(idle_deadline) => { log::debug!( - "peer on topic {} idle for {:?}; dropping", - query.topic, + "peer on topic {topic} idle for {:?}; dropping", config.idle_timeout ); return Err(ConnectionError::IdleTimeout); @@ -630,14 +630,10 @@ async fn read_loop( ); if !within_connection_limits { consecutive_drops += 1; - log::debug!( - "dropping frame on topic {}: inbound rate limit exceeded", - query.topic - ); + log::debug!("dropping frame on topic {topic}: inbound rate limit exceeded"); if consecutive_drops > MAX_CONSECUTIVE_RATE_DROPS { log::warn!( - "peer on topic {} keeps exceeding its inbound rate limit; dropping", - query.topic + "peer on topic {topic} keeps exceeding its inbound rate limit; dropping" ); return Err(ConnectionError::RateLimited); } @@ -649,13 +645,11 @@ async fn read_loop( InboundBytes::ShareExhausted => { consecutive_drops += 1; log::debug!( - "dropping frame on topic {}: identity share of the relay budget exhausted", - query.topic + "dropping frame on topic {topic}: identity share of the relay budget exhausted" ); if consecutive_drops > MAX_CONSECUTIVE_RATE_DROPS { log::warn!( - "peer on topic {} keeps exceeding its share of the relay budget; dropping", - query.topic + "peer on topic {topic} keeps exceeding its share of the relay budget; dropping" ); return Err(ConnectionError::RateLimited); } @@ -663,25 +657,25 @@ async fn read_loop( } InboundBytes::GlobalExhausted => { log::debug!( - "dropping frame on topic {}: relay inbound byte budget exhausted", - query.topic + "dropping frame on topic {topic}: relay inbound byte budget exhausted" ); continue; } } - log::trace!("relay <- topic={} frame={text}", query.topic); - serde_json::from_str::(&text).inspect_err(|error| { - log::warn!("invalid relay frame on topic {}: {error}", query.topic); + // Neither record carries the frame itself: it holds the topic, + // and the client picks its content. + log::trace!("relay <- topic={topic} ({} bytes)", text.len()); + serde_json::from_str::(&text).map_err(|error| { + let rejection = FrameRejection::new(&error, text.len()); + log::warn!("invalid relay frame on topic {topic}: {rejection}"); + ConnectionError::Frame(rejection) })? } Message::Close(_) => break, // The relay protocol is JSON text only. Message::Binary(_) => { - log::warn!( - "binary frame on topic {} is not part of the relay protocol", - query.topic - ); + log::warn!("binary frame on topic {topic} is not part of the relay protocol"); return Err(ConnectionError::UnexpectedBinary); } // Keepalive pongs and other control frames count as activity, but @@ -717,25 +711,20 @@ async fn read_loop( if frame.v != RELAY_FRAME_VERSION { log::warn!( - "unsupported relay frame version {} on topic {}", - frame.v, - query.topic + "unsupported relay frame version {} on topic {topic}", + frame.v ); return Err(ConnectionError::InvalidVersion(frame.v)); } if frame.topic != query.topic { log::warn!( - "frame topic {} does not match subscription {}", - frame.topic, - query.topic + "frame topic {} does not match subscription {topic}", + frame.topic.log_tag() ); return Err(ConnectionError::TopicMismatch); } if let Some(kind) = frame.body.relay_only_body() { - log::warn!( - "client sent relay-only {kind} frame on topic {}", - query.topic - ); + log::warn!("client sent relay-only {kind} frame on topic {topic}"); return Err(ConnectionError::ClientControlFrame(kind)); } @@ -747,10 +736,7 @@ async fn read_loop( ) { Ok(()) => {} Err(SessionError::NoCounterpart) => { - log::trace!( - "dropping frame on topic {}: counterpart not connected", - query.topic - ); + log::trace!("dropping frame on topic {topic}: counterpart not connected"); } Err(error) => return Err(error.into()), } @@ -794,7 +780,7 @@ pub enum ConnectionError { WebSocket(#[from] tokio_tungstenite::tungstenite::Error), #[error("malformed relay frame: {0}")] - Frame(#[from] serde_json::Error), + Frame(FrameRejection), #[error(transparent)] Session(#[from] SessionError), @@ -839,6 +825,47 @@ pub enum ConnectionError { Random(#[from] getrandom::Error), } +/// What may be logged about a text frame that failed to parse: the kind of +/// error, where it occurred and the frame's size, never the frame's content. +/// +/// The rendered `serde_json` error quotes the offending value, escaped, so +/// logging it would let a client write its frame into the log several times +/// over. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct FrameRejection { + category: Category, + line: usize, + column: usize, + bytes: usize, +} + +impl FrameRejection { + fn new(error: &serde_json::Error, bytes: usize) -> Self { + Self { + category: error.classify(), + line: error.line(), + column: error.column(), + bytes, + } + } +} + +impl fmt::Display for FrameRejection { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + let category = match self.category { + Category::Io => "i/o", + Category::Syntax => "syntax", + Category::Data => "data", + Category::Eof => "end of input", + }; + write!( + f, + "{category} error at line {} column {} ({} bytes)", + self.line, self.column, self.bytes + ) + } +} + #[cfg(test)] mod tests { use super::*; @@ -890,6 +917,28 @@ mod tests { ); } + #[test] + fn frame_rejection_is_short_and_independent_of_the_frame() { + let text = format!(r#"{{"v":"{}"}}"#, "\u{7f}".repeat(200_000)); + let Err(error) = serde_json::from_str::(&text) else { + panic!("a string version must be rejected"); + }; + // serde_json's own rendering quotes the value: this is what it replaces. + assert!(error.to_string().len() > 1_000_000); + + let rejection = FrameRejection::new(&error, text.len()); + let rendered = ConnectionError::Frame(rejection).to_string(); + assert_eq!( + rendered, + format!( + "malformed relay frame: data error at line 1 column {} ({} bytes)", + error.column(), + text.len() + ) + ); + assert!(rendered.len() < 100); + } + #[test] fn proposer_params_parse() -> TestResult { let topic = topic()?; diff --git a/crates/relay/src/main.rs b/crates/relay/src/main.rs index 59901e2..671e666 100644 --- a/crates/relay/src/main.rs +++ b/crates/relay/src/main.rs @@ -3,7 +3,8 @@ //! Configuration comes from `configuration/*.yaml`: `base.yaml` plus the //! overlay named by `APP_ENVIRONMENT`, with environment variables overriding //! both. The relay speaks plain WebSocket and expects TLS to be terminated in -//! front of it. Logs are bunyan JSON on stdout. +//! front of it. Logs are bunyan JSON on stdout, written by a thread of their +//! own: when stdout falls behind, records are dropped rather than waited for. use std::sync::Arc; @@ -16,12 +17,14 @@ use wallet_session_relay::{ limits::Limiter, registry::Registry, server::serve, - telemetry::{SERVICE_NAME, get_subscriber, init_subscriber}, + telemetry::{SERVICE_NAME, get_subscriber, init_subscriber, non_blocking}, }; #[tokio::main] async fn main() -> anyhow::Result<()> { - let subscriber = get_subscriber(SERVICE_NAME.to_owned(), "info".to_owned(), std::io::stdout); + let (log_writer, _log_guard) = + non_blocking(std::io::stdout()).context("failed to start the log writer")?; + let subscriber = get_subscriber(SERVICE_NAME.to_owned(), "info".to_owned(), log_writer); init_subscriber(subscriber); let settings = diff --git a/crates/relay/src/registry.rs b/crates/relay/src/registry.rs index 6dbcbdc..1f08f7e 100644 --- a/crates/relay/src/registry.rs +++ b/crates/relay/src/registry.rs @@ -49,14 +49,15 @@ impl Registry { max_sessions: usize, ) -> Result<(Arc, Registration), RegistryError> { let mut sessions = self.lock()?; + let tag = topic.log_tag(); let created = !sessions.contains_key(topic); if created && sessions.len() >= max_sessions { - log::warn!("rejecting new session for topic {topic}: session limit reached"); + log::warn!("rejecting new session for topic {tag}: session limit reached"); return Err(RegistryError::TooManySessions); } if created { - log::debug!("created session for topic {topic}"); + log::debug!("created session for topic {tag}"); } let session = sessions.entry(topic.clone()).or_default().clone(); @@ -70,7 +71,7 @@ impl Registry { } })?; - log::debug!("registered peer {identity} on topic {topic} as {role:?}"); + log::debug!("registered peer {identity} on topic {tag} as {role:?}"); Ok((session, registration)) } @@ -115,7 +116,10 @@ impl Registry { let reserved = session.disconnect(role, tx, grace)?; if reserved { - log::debug!("reserved {role:?} slot on topic {topic} for {grace:?}"); + log::debug!( + "reserved {role:?} slot on topic {} for {grace:?}", + topic.log_tag() + ); } Ok(reserved) @@ -139,11 +143,12 @@ impl Registry { match session.expire(role, notification)? { Some(expiry) => { + let tag = topic.log_tag(); if expiry.empty { sessions.remove(topic); - log::debug!("removed empty session for topic {topic}"); + log::debug!("removed empty session for topic {tag}"); } - log::debug!("expired {role:?} reservation on topic {topic}"); + log::debug!("expired {role:?} reservation on topic {tag}"); Ok(true) } None => Ok(false), diff --git a/crates/relay/src/server.rs b/crates/relay/src/server.rs index 1684326..dd9e891 100644 --- a/crates/relay/src/server.rs +++ b/crates/relay/src/server.rs @@ -14,6 +14,7 @@ use crate::{ limits::{Limiter, OutboundBudget}, registry::{Registry, RegistryError}, session::{PeerHandle, Registration, SessionError}, + telemetry::bounded, }; use wallet_session_core::frame::{RelayFrame, RelayFrameBody}; @@ -40,7 +41,8 @@ pub async fn serve( tokio::spawn(async move { if let Err(error) = handle(stream, registry, limiter, config).await { - log::debug!("relay connection ended with error: {error}"); + // Errors can quote what the client sent. + log::debug!("relay connection ended with error: {}", bounded(&error)); } }); } @@ -149,7 +151,10 @@ pub async fn handle( tokio::spawn(async move { sleep(grace).await; if let Err(error) = registry.expire(¬ify_topic, role, ¬ification) { - log::warn!("could not expire the reservation on topic {notify_topic}: {error}"); + log::warn!( + "could not expire the reservation on topic {}: {error}", + notify_topic.log_tag() + ); } }); } @@ -187,6 +192,7 @@ mod tests { use crate::auth::test_utils::{self, signing_key}; use crate::connection::MAX_CONSECUTIVE_RATE_DROPS; use crate::limits::{Limiter, RateLimit, RelayLimitsConfig}; + use crate::telemetry::test_utils::{CapturedLogs, capture}; use futures_util::{SinkExt, StreamExt}; use std::net::SocketAddr; use tokio::time::timeout; @@ -2250,4 +2256,118 @@ mod tests { assert!(expect_closed(&mut proposer).await); Ok(()) } + + /// Wait for a captured record that contains `needle`. + async fn logged(logs: &CapturedLogs, needle: &str) -> bool { + for _ in 0..200 { + if logs.records().iter().any(|record| record.contains(needle)) { + return true; + } + sleep(Duration::from_millis(10)).await; + } + false + } + + #[tokio::test] + async fn logs_tag_topics_instead_of_writing_them() -> TestResult { + // At trace, so that every record the relay can write is checked. + let (logs, _capture) = capture("trace"); + let addr = serve_with( + NativeRelayTransportConfig::new(AUDIENCE).with_session_grace(Duration::from_secs(2)), + ) + .await?; + + let topic = topic()?; + let (mut proposer, mut responder) = ready_pair(addr, &topic).await?; + send(&mut proposer, &message(&topic, "m1", "hello")).await?; + read_frame(&mut responder).await?; + + // Warnings that name topics: a role already taken, and a frame + // stamped with another topic. + let mut second = connect_auth(addr, &token_with(8), "proposer", &topic, None).await?; + assert!(expect_closed(&mut second).await); + let other = Topic::new([8u8; 32])?; + send(&mut proposer, &message(&other, "rogue", "cross-topic")).await?; + assert!(expect_closed(&mut proposer).await); + + let tag = topic.log_tag(); + for needle in [ + format!("session ready for topic {tag}"), + format!("rejecting connection for topic {tag}"), + format!( + "frame topic {} does not match subscription {tag}", + other.log_tag() + ), + ] { + assert!(logged(&logs, &needle).await, "nothing logged {needle:?}"); + } + + let records = logs.records(); + for raw in [ + topic.to_string(), + format!("{topic:?}"), + other.to_string(), + format!("{other:?}"), + ] { + assert!( + records.iter().all(|record| !record.contains(&raw)), + "a record carries the raw topic {raw}" + ); + } + Ok(()) + } + + /// Send a frame that serde_json would quote, escaped, into its error and + /// check the records the relay writes about it at `filter`. + async fn invalid_frame_records_are_short_and_content_free(filter: &str) -> TestResult { + let (logs, _capture) = capture(filter); + let addr = serve_with(NativeRelayTransportConfig::new(AUDIENCE)).await?; + let topic = topic()?; + let (mut proposer, _responder) = ready_pair(addr, &topic).await?; + + let frame = format!(r#"{{"v":"{}"}}"#, "\u{7f}".repeat(200_000)); + proposer.send(Message::text(frame)).await?; + assert!( + expect_closed(&mut proposer).await, + "an invalid frame closes the connection" + ); + + assert!(logged(&logs, "invalid relay frame on topic").await); + if filter != "info" { + assert!( + logged( + &logs, + "relay connection ended with error: malformed relay frame" + ) + .await + ); + } + for record in logs.records() { + assert!( + record.len() < 1024, + "a {}-byte record at {filter}", + record.len() + ); + assert!( + !record.contains('\u{7f}') && !record.contains("\\u{7f}"), + "a record at {filter} carries the frame" + ); + } + Ok(()) + } + + #[tokio::test] + async fn invalid_frame_records_are_short_at_info() -> TestResult { + invalid_frame_records_are_short_and_content_free("info").await + } + + #[tokio::test] + async fn invalid_frame_records_are_short_at_debug() -> TestResult { + invalid_frame_records_are_short_and_content_free("debug").await + } + + #[tokio::test] + async fn invalid_frame_records_are_short_at_trace() -> TestResult { + invalid_frame_records_are_short_and_content_free("trace").await + } } diff --git a/crates/relay/src/telemetry.rs b/crates/relay/src/telemetry.rs index 4f55914..d3ee7ae 100644 --- a/crates/relay/src/telemetry.rs +++ b/crates/relay/src/telemetry.rs @@ -1,7 +1,15 @@ +use std::collections::VecDeque; +use std::fmt::{self, Write as _}; +use std::io::{self, Write}; +use std::sync::{Arc, Condvar, Mutex, MutexGuard, PoisonError, mpsc}; +use std::thread; +use std::time::Duration; + use tracing::Subscriber; use tracing::subscriber::set_global_default; use tracing_bunyan_formatter::{BunyanFormattingLayer, JsonStorageLayer}; use tracing_log::LogTracer; +use tracing_subscriber::filter::filter_fn; use tracing_subscriber::fmt::MakeWriter; use tracing_subscriber::layer::SubscriberExt; use tracing_subscriber::{EnvFilter, Registry}; @@ -9,7 +17,20 @@ use tracing_subscriber::{EnvFilter, Registry}; /// Service name written into every log record. pub const SERVICE_NAME: &str = "wallet-session-relay"; +/// Bytes of log records [`non_blocking`] holds while its sink catches up. +pub const LOG_BUFFER_BYTES: usize = 8 * 1024 * 1024; + +/// How long dropping a [`LogGuard`] waits for held records to be written. +const LOG_FLUSH_TIMEOUT: Duration = Duration::from_secs(1); + +/// Longest rendering, in bytes, of a client-influenced value in a log record. +pub(crate) const MAX_LOGGED_DISPLAY_BYTES: usize = 256; + /// Build the subscriber: one bunyan JSON record per event, written to `sink`. +/// +/// `RUST_LOG` overrides `env_filter`. Either way `tungstenite` logs at debug +/// at most: at trace it writes out every frame it reads or writes, and with +/// them the topics they carry. pub fn get_subscriber( name: String, env_filter: String, @@ -20,9 +41,27 @@ where { let env_filter = EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new(env_filter)); + subscriber(name, env_filter, sink) +} + +/// [`get_subscriber`] with an explicit filter. +fn subscriber( + name: String, + env_filter: EnvFilter, + sink: Sink, +) -> impl Subscriber + Send + Sync +where + Sink: for<'a> MakeWriter<'a> + Send + Sync + 'static, +{ + let no_frame_dumps = filter_fn(|metadata| { + let target = metadata.target(); + let tungstenite = target == "tungstenite" || target.starts_with("tungstenite::"); + !(tungstenite && *metadata.level() == tracing::Level::TRACE) + }); let formatting_layer = BunyanFormattingLayer::new(name, sink); Registry::default() .with(env_filter) + .with(no_frame_dumps) .with(JsonStorageLayer) .with(formatting_layer) } @@ -32,3 +71,396 @@ pub fn init_subscriber(subscriber: impl Subscriber + Send + Sync) { LogTracer::init().expect("Failed to set logger"); set_global_default(subscriber).expect("Failed to set subscriber"); } + +/// Render `value` for a log record, cut to [`MAX_LOGGED_DISPLAY_BYTES`] and +/// with control characters replaced, so a value that echoes client input +/// cannot make the record long. +pub(crate) fn bounded(value: T) -> Bounded { + Bounded(value) +} + +/// See [`bounded`]. +pub(crate) struct Bounded(T); + +impl fmt::Display for Bounded { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + let mut out = Truncating { + out: f, + left: MAX_LOGGED_DISPLAY_BYTES, + cut: false, + }; + let written = write!(out, "{}", self.0); + match written { + Err(_) if out.cut => out.out.write_str("..."), + result => result, + } + } +} + +/// Passes characters through until `left` bytes are used up, then fails, +/// which also stops formatting the rest of the value. +struct Truncating<'a, 'b> { + out: &'a mut fmt::Formatter<'b>, + left: usize, + cut: bool, +} + +impl fmt::Write for Truncating<'_, '_> { + fn write_str(&mut self, s: &str) -> fmt::Result { + for c in s.chars() { + // The JSON encoding would write each of these six bytes long. + let c = if c.is_control() { + char::REPLACEMENT_CHARACTER + } else { + c + }; + let Some(left) = self.left.checked_sub(c.len_utf8()) else { + self.cut = true; + return Err(fmt::Error); + }; + self.left = left; + self.out.write_char(c)?; + } + Ok(()) + } +} + +/// Start a log writer that hands records to a thread writing them to +/// `sink`, so the threads that log never wait for the sink. +/// +/// Up to [`LOG_BUFFER_BYTES`] of records wait for the sink. A record that +/// does not fit is dropped, and the number dropped is logged once the sink +/// has caught up. Dropping the guard writes out what is held, waiting up to +/// a second. +pub fn non_blocking(sink: W) -> io::Result<(NonBlocking, LogGuard)> +where + W: Write + Send + 'static, +{ + NonBlocking::spawn(sink, LOG_BUFFER_BYTES) +} + +/// Log writer returned by [`non_blocking`]. +#[derive(Clone)] +pub struct NonBlocking { + shared: Arc, +} + +/// Stops the [`non_blocking`] writer once it has written out what it holds. +#[must_use = "dropping the guard stops the log writer"] +pub struct LogGuard { + shared: Arc, + stopped: mpsc::Receiver<()>, +} + +/// One log record on its way to the writer, queued whole when dropped. +pub struct Record<'a> { + shared: &'a Shared, + bytes: Vec, + oversized: bool, +} + +struct Shared { + queue: Mutex, + ready: Condvar, + capacity: usize, +} + +#[derive(Default)] +struct Queue { + /// Records waiting for the sink, oldest first. + records: VecDeque>, + + /// Bytes of the waiting records and of the one being written. + bytes: usize, + + /// Records dropped since the writer last reported them. + dropped: u64, + + /// Set by the guard: write out what is held, then stop. + closing: bool, +} + +/// What the writer thread does next. +enum Next { + Write(Vec), + ReportDropped(u64), +} + +impl NonBlocking { + fn spawn(sink: W, capacity: usize) -> io::Result<(Self, LogGuard)> + where + W: Write + Send + 'static, + { + let shared = Arc::new(Shared { + queue: Mutex::default(), + ready: Condvar::new(), + capacity, + }); + let (stopped_tx, stopped) = mpsc::channel(); + thread::Builder::new() + .name("log-writer".to_owned()) + .spawn({ + let shared = Arc::clone(&shared); + move || { + write_records(&shared, sink); + let _ = stopped_tx.send(()); + } + })?; + let writer = Self { + shared: Arc::clone(&shared), + }; + Ok((writer, LogGuard { shared, stopped })) + } +} + +impl<'a> MakeWriter<'a> for NonBlocking { + type Writer = Record<'a>; + + fn make_writer(&'a self) -> Self::Writer { + Record { + shared: &self.shared, + bytes: Vec::new(), + oversized: false, + } + } +} + +impl Write for Record<'_> { + fn write(&mut self, buf: &[u8]) -> io::Result { + // A record that can never fit is dropped without copying the rest. + if self.bytes.len().saturating_add(buf.len()) > self.shared.capacity { + self.oversized = true; + self.bytes = Vec::new(); + } + if !self.oversized { + self.bytes.extend_from_slice(buf); + } + Ok(buf.len()) + } + + fn flush(&mut self) -> io::Result<()> { + Ok(()) + } +} + +impl Drop for Record<'_> { + fn drop(&mut self) { + if self.oversized { + self.shared.drop_record(); + } else if !self.bytes.is_empty() { + self.shared.push(std::mem::take(&mut self.bytes)); + } + } +} + +impl Drop for LogGuard { + fn drop(&mut self) { + self.shared.lock().closing = true; + self.shared.ready.notify_one(); + let _ = self.stopped.recv_timeout(LOG_FLUSH_TIMEOUT); + } +} + +impl Shared { + fn lock(&self) -> MutexGuard<'_, Queue> { + // Nothing that can panic runs under the lock. + self.queue.lock().unwrap_or_else(PoisonError::into_inner) + } + + /// Queue `record` if the buffer has room for it, or drop it. + fn push(&self, record: Vec) { + let mut queue = self.lock(); + if record.len() <= self.capacity.saturating_sub(queue.bytes) { + queue.bytes += record.len(); + queue.records.push_back(record); + } else { + queue.dropped += 1; + } + drop(queue); + self.ready.notify_one(); + } + + fn drop_record(&self) { + self.lock().dropped += 1; + self.ready.notify_one(); + } +} + +/// Write records to `sink` as they arrive, until the guard is dropped and +/// nothing is held any more. The lock is never held while writing. +fn write_records(shared: &Shared, mut sink: W) { + loop { + let next = { + let mut queue = shared.lock(); + loop { + if let Some(record) = queue.records.pop_front() { + break Next::Write(record); + } + if queue.dropped > 0 { + break Next::ReportDropped(std::mem::take(&mut queue.dropped)); + } + if queue.closing { + return; + } + queue = shared + .ready + .wait(queue) + .unwrap_or_else(PoisonError::into_inner); + } + }; + match next { + Next::Write(record) => { + // A failed write loses this record only. + let _ = sink.write_all(&record).and_then(|()| sink.flush()); + shared.lock().bytes -= record.len(); + } + // The queue is empty, so the report has room. + Next::ReportDropped(dropped) => { + tracing::warn!(dropped, "log output fell behind; dropped log records"); + } + } + } +} + +#[cfg(test)] +pub(crate) mod test_utils { + use super::*; + use tracing::subscriber::DefaultGuard; + + /// Log output written to memory. + #[derive(Clone, Default)] + pub(crate) struct CapturedLogs(Arc>>); + + impl CapturedLogs { + /// The lines written so far: one JSON record each. + pub(crate) fn records(&self) -> Vec { + let bytes = self.0.lock().unwrap_or_else(PoisonError::into_inner); + String::from_utf8_lossy(&bytes) + .lines() + .map(str::to_owned) + .collect() + } + } + + impl Write for CapturedLogs { + fn write(&mut self, buf: &[u8]) -> io::Result { + self.0 + .lock() + .unwrap_or_else(PoisonError::into_inner) + .extend_from_slice(buf); + Ok(buf.len()) + } + + fn flush(&mut self) -> io::Result<()> { + Ok(()) + } + } + + impl<'a> MakeWriter<'a> for CapturedLogs { + type Writer = Self; + + fn make_writer(&'a self) -> Self::Writer { + self.clone() + } + } + + /// Capture the records logged on this thread through the relay's own + /// subscriber, with `filter` in place of `RUST_LOG`, until the guard is + /// dropped. + pub(crate) fn capture(filter: &str) -> (CapturedLogs, DefaultGuard) { + // `log` records reach the subscriber through the global logger. + let _ = LogTracer::init(); + let logs = CapturedLogs::default(); + let subscriber = subscriber( + SERVICE_NAME.to_owned(), + EnvFilter::new(filter), + logs.clone(), + ); + (logs, tracing::subscriber::set_default(subscriber)) + } +} + +#[cfg(test)] +mod tests { + use super::test_utils::CapturedLogs; + use super::*; + + #[test] + fn bounded_cuts_long_values_and_replaces_control_characters() { + assert_eq!(bounded("short value").to_string(), "short value"); + assert_eq!(bounded("a\u{7f}b\nc").to_string(), "a\u{fffd}b\u{fffd}c"); + + let long = bounded("x".repeat(100_000)).to_string(); + assert_eq!(long.len(), MAX_LOGGED_DISPLAY_BYTES + "...".len()); + assert!(long.ends_with("...")); + + // Characters are never split. + let wide = bounded("€".repeat(MAX_LOGGED_DISPLAY_BYTES)).to_string(); + assert!(wide.len() <= MAX_LOGGED_DISPLAY_BYTES + "...".len()); + assert!(wide.trim_end_matches("...").chars().all(|c| c == '€')); + } + + fn log(writer: &NonBlocking, record: &str) { + let _ = writer.make_writer().write_all(record.as_bytes()); + } + + #[test] + fn non_blocking_writer_writes_records_in_order() -> io::Result<()> { + let written = CapturedLogs::default(); + let (writer, guard) = NonBlocking::spawn(written.clone(), 1024)?; + + log(&writer, "one\n"); + log(&writer, "two\n"); + drop(guard); + + assert_eq!(written.records(), ["one", "two"]); + Ok(()) + } + + /// A sink whose writes wait until `release` is dropped, or 10s. + struct HeldSink { + written: CapturedLogs, + release: mpsc::Receiver<()>, + } + + impl Write for HeldSink { + fn write(&mut self, buf: &[u8]) -> io::Result { + let _ = self.release.recv_timeout(Duration::from_secs(10)); + self.written.write(buf) + } + + fn flush(&mut self) -> io::Result<()> { + Ok(()) + } + } + + #[test] + fn non_blocking_writer_drops_records_instead_of_waiting_for_the_sink() -> io::Result<()> { + let written = CapturedLogs::default(); + let (release, held) = mpsc::channel(); + let sink = HeldSink { + written: written.clone(), + release: held, + }; + let (writer, guard) = NonBlocking::spawn(sink, 100)?; + + let start = std::time::Instant::now(); + // The sink holds this one, and its bytes count until it is written. + log(&writer, &format!("{}\n", "a".repeat(59))); + // No room beside it. + log(&writer, &format!("{}\n", "b".repeat(59))); + // Room for this one. + log(&writer, &format!("{}\n", "c".repeat(29))); + // Never fits. + log(&writer, &format!("{}\n", "d".repeat(200))); + assert!( + start.elapsed() < Duration::from_secs(5), + "logging must not wait for the sink" + ); + + drop(release); + drop(guard); + assert_eq!(written.records(), ["a".repeat(59), "c".repeat(29)]); + Ok(()) + } +} From 04239065374cdf81e70c3e1b051b6562ed25d324 Mon Sep 17 00:00:00 2001 From: Kyryl R Date: Wed, 28 Oct 2026 12:42:25 +0200 Subject: [PATCH 5/9] Enforce source and paired-traffic inbound fairness --- crates/relay/README.md | 107 +++++- crates/relay/configuration/base.yaml | 12 + crates/relay/src/config.rs | 16 +- crates/relay/src/configuration.rs | 250 +++++++++++++- crates/relay/src/connection.rs | 182 +++++++--- crates/relay/src/health.rs | 48 +-- crates/relay/src/lib.rs | 1 + crates/relay/src/limits.rs | 356 ++++++++++++++++---- crates/relay/src/server.rs | 322 +++++++++++++++++- crates/relay/src/session.rs | 51 +++ crates/relay/src/source.rs | 341 +++++++++++++++++++ deployment/configs/relay/base.yaml.template | 13 + 12 files changed, 1534 insertions(+), 165 deletions(-) create mode 100644 crates/relay/src/source.rs diff --git a/crates/relay/README.md b/crates/relay/README.md index 21673b0..e0fcaca 100644 --- a/crates/relay/README.md +++ b/crates/relay/README.md @@ -27,11 +27,13 @@ See `configuration/base.yaml` for every key. The ones that matter most: | --- | --- | | `application.host` / `application.port` | WebSocket listener | | `application.audience` | `aud` required in connection tokens (no default in production) | +| `application.trusted_proxies` | peers allowed to name the client in `X-Real-IP` (see [Inbound fairness](#inbound-fairness)) | | `health.host` / `health.port` | probe listener, separate from the WebSocket port | | `auth.max_token_lifetime_seconds` | longest remaining lifetime accepted in a token | | `session.grace_seconds` | how long a disconnected peer's slot is reserved | | `session.buffer_max_messages` / `session.buffer_ttl_seconds` | frames buffered for it | -| `limits.*` | connection caps, per-identity/per-connection rates, and each identity's share of the global inbound byte budget | +| `limits.*` | connection caps, per-identity/per-connection rates, and each identity's and source's share of the global inbound byte budget | +| `limits.global_paired_reserve_percent` | part of the global budget only paired sessions may spend | | `connection.*` | frame size, queue capacity, and the handshake, write, ping, pong and idle timeouts | | `connection.outbound_queue_max_bytes` | bytes held for one peer and not yet written (see [Outbound memory](#outbound-memory)) | | `connection.outbound_queue_timeout_seconds` | longest a frame may wait in a peer's outbound queue before that peer is dropped | @@ -41,7 +43,103 @@ zero timeouts, a blank audience, a health address equal to the application address, a byte-rate burst below the maximum incoming frame (no frame could ever fit it), an outbound byte budget below the maximum incoming frame, or a maximum incoming frame above `u32::MAX`. Every byte bucket and outbound budget -must hold a whole frame. +must hold a whole frame, and so must the part of the global burst outside the +paired reserve, which must be below 100 percent. A trusted proxy range of +length 0 (`0.0.0.0/0`, `::/0`) is refused: it would let every peer choose its +own source. + +### Inbound fairness + +Every inbound text, ping and pong frame is charged, before it is parsed, to its +connection's own message and byte buckets (`inbound_message_rate`, +`inbound_byte_rate`). A +connection that keeps exceeding them is disconnected. What happens next +depends on whether anything can receive the frame: + +- **Paired**: the counterpart is connected. The frame is charged to its + identity's share (`identity_inbound_byte_rate`), its source's share + (`source_inbound_byte_rate`) and the relay-wide budget + (`global_inbound_byte_rate`), then parsed and forwarded. +- **Buffered**: the counterpart is within its reconnect grace window and + buffering is enabled. The frame will be delivered when it returns, so it is + charged like a paired one, except that it cannot spend the paired reserve. +- **No receiver**: there is no counterpart, or it is reconnecting while + buffering is disabled. The frame is dropped without being parsed and without + spending any shared budget: identities cannot drain the relay-wide budget + through sessions that have nobody on the other side. + +WebSocket pings and pongs are for the relay, not the counterpart. They are +charged to the same budgets whether or not a counterpart is connected, but +without one they cannot spend the paired reserve either. Refusing one still +counts as activity for the idle and pong deadlines. + +The shares and what exceeding them does: + +- The **identity share** (default 1 MiB/s) caps one `did:key`. A connection + that keeps exceeding it is disconnected. +- The **source share** (default 1 MiB/s, no larger than the identity share) + caps every identity connecting from one source together, since identities + cost nothing to mint. Exhausting it only drops the frame: other clients + behind the same address must not be able to get a peer disconnected. +- The **global budget** (default 16 MiB/s) caps the relay. Exhausting it only + drops the frame, so other identities cannot get a peer disconnected. + `global_paired_reserve_percent` (default 25) of its burst can be spent only + by paired frames: buffered frames and control frames without a counterpart + are refused once the budget is down to the reserve, so they cannot starve + sessions whose peers are talking to each other. + +A source is the TCP peer's IPv4 address, or the /64 holding its IPv6 address +(IPv4-mapped peers count as IPv4). When the peer is in +`application.trusted_proxies`, the source is instead the address it writes in +`X-Real-IP`. The header must appear once and hold exactly one address, +otherwise the proxy itself is the source. From any other peer the header is +ignored, so a client that reaches the relay directly cannot pick its source; +the first time that happens the relay logs a warning, since it also means a +proxy in front is missing from the list. `X-Forwarded-For` is never read. +`trusted_proxies` takes a YAML list or one comma-separated string, so +`APPLICATION__TRUSTED_PROXIES=172.20.0.5/32` sets it from the environment. + +Compatibility and defaults: + +- `application.trusted_proxies` is optional and empty by default: every + connection is accounted to its TCP peer. `limits.source_inbound_byte_rate` + and `limits.global_paired_reserve_percent` are required, like the other + limits; a configuration file written for an earlier version must add them. +- A relay behind a proxy that is not listed accounts every client to the + proxy's address, so all of them share one source share (1 MiB/s by + default). List the proxy, or raise the source share. +- Clients behind one NAT address, or one IPv6 /64, share a source share. +- A frame sent before the counterpart connects, or while it is reconnecting + with buffering disabled, is no longer parsed: an invalid or relay-only frame + sent then is dropped silently instead of closing the connection. It is + checked as before once a counterpart is present. Clients should wait for + `session_ready` before sending, as the protocol already requires. +- `/ready` reports `dropped_source_bytes` beside the other drop counters. + +Topology limits: + +- The source share is only as scarce as addresses are. A client with many IPv4 + addresses, or with routed IPv6 space (a /48 holds 65,536 /64s), gets one + share for each; the global budget and the paired reserve remain the bound + on what such a client can take. As for identities, at most 65,536 sources + are tracked at once; past that, idle ones are pruned first and then + arbitrary ones are forgotten, starting again with a full share. +- Trusting a proxy is only safe when the proxy overwrites `X-Real-IP` and + nothing else reachable from a trusted address can connect to the relay. The + compose deployment trusts the private ranges Docker assigns its networks + from, because the relay port is not published and nginx is its only peer + there (see `deployment/configs/relay/base.yaml.template`). Publishing the + relay port, attaching other containers to its network, or running a + relay-port listener on the Docker host breaks that assumption: narrow the + list to nginx's address. +- nginx resolves the client address it forwards from `X-Forwarded-For` only + for its own `TRUSTED_PROXY_CIDR`, so the source behind a load balancer or + CDN is only as accurate as that setting. nginx's own per-address limits are + keyed on the full address, not the /64. +- A relay without a proxy listens directly (`configuration/production.yaml`), + and its sources are the TCP peers. Nothing in the relay can tell apart + clients that share an address. +- Budgets are per process. Each relay instance has its own. ### Outbound memory @@ -94,8 +192,9 @@ APP_ENVIRONMENT=production APPLICATION__AUDIENCE=wss://relay.example.com \ `GET /health` on the health port is liveness; `GET /ready` is readiness and reports `connections`, `sessions`, and the monotonic `dropped_share_bytes` / -`dropped_global_bytes` counters. A rising drop counter means the relay is -shedding inbound traffic, so alert on its rate of change. +`dropped_source_bytes` / `dropped_global_bytes` counters. A rising drop counter +means the relay is shedding inbound traffic, so alert on its rate of change. +Frames dropped because nothing could receive them are not counted. ## Deployment diff --git a/crates/relay/configuration/base.yaml b/crates/relay/configuration/base.yaml index 89ef4c9..ffd3630 100644 --- a/crates/relay/configuration/base.yaml +++ b/crates/relay/configuration/base.yaml @@ -1,5 +1,10 @@ application: port: 8080 + # Peers allowed to name the client they forward in X-Real-IP, as addresses + # or address/length ranges. Empty: the header is ignored and every + # connection is accounted to its TCP peer. List a proxy only if it sets + # X-Real-IP itself and nothing else can reach this port. + trusted_proxies: [] health: port: 8081 @@ -43,6 +48,13 @@ limits: identity_inbound_byte_rate: burst: 1048576 interval_seconds: 1 + # Shared by every identity connecting from one address (an IPv6 /64). + source_inbound_byte_rate: + burst: 1048576 + interval_seconds: 1 global_inbound_byte_rate: burst: 16777216 interval_seconds: 1 + # Share of the global burst, in percent, that only frames between two + # connected peers may spend. + global_paired_reserve_percent: 25 diff --git a/crates/relay/src/config.rs b/crates/relay/src/config.rs index 28c9498..9a95df4 100644 --- a/crates/relay/src/config.rs +++ b/crates/relay/src/config.rs @@ -1,8 +1,8 @@ -use std::time::Duration; +use std::{sync::Arc, time::Duration}; use tokio_tungstenite::tungstenite::protocol::WebSocketConfig; -use crate::auth::DEFAULT_MAX_TOKEN_LIFETIME; +use crate::{auth::DEFAULT_MAX_TOKEN_LIFETIME, source::IpPrefix}; /// Default maximum WebSocket text or binary frame size accepted from a client. pub const DEFAULT_MAX_INCOMING_FRAME_BYTES: usize = 256 * 1024; @@ -88,6 +88,11 @@ pub struct NativeRelayTransportConfig { /// Maximum accepted lifetime of an auth token. pub max_token_lifetime: Duration, + + /// Peers allowed to name the client they forward in `X-Real-IP`. The + /// header is ignored from every other peer, which is then the source its + /// connection is accounted to. Empty by default: no peer is trusted. + pub trusted_proxies: Arc<[IpPrefix]>, } impl NativeRelayTransportConfig { @@ -108,6 +113,7 @@ impl NativeRelayTransportConfig { outbound_queue_max_bytes: DEFAULT_OUTBOUND_QUEUE_MAX_BYTES, outbound_queue_timeout: DEFAULT_OUTBOUND_QUEUE_TIMEOUT, max_token_lifetime: DEFAULT_MAX_TOKEN_LIFETIME, + trusted_proxies: Arc::from([]), } } @@ -201,6 +207,12 @@ impl NativeRelayTransportConfig { self } + /// Set the peers allowed to name the client they forward in `X-Real-IP`. + pub fn with_trusted_proxies(mut self, proxies: impl Into>) -> Self { + self.trusted_proxies = proxies.into(); + self + } + /// Get WebSocketConfig from this config. pub fn websocket_config(&self) -> WebSocketConfig { WebSocketConfig::default() diff --git a/crates/relay/src/configuration.rs b/crates/relay/src/configuration.rs index 64d4190..b8f176d 100644 --- a/crates/relay/src/configuration.rs +++ b/crates/relay/src/configuration.rs @@ -1,15 +1,20 @@ use std::{ + fmt, net::{IpAddr, SocketAddr}, path::Path, time::Duration, }; -use serde::Deserialize; +use serde::{ + Deserialize, Deserializer, + de::{self, SeqAccess, Visitor}, +}; use thiserror::Error; use crate::{ config::NativeRelayTransportConfig, limits::{RateLimit, RelayLimitsConfig}, + source::IpPrefix, }; pub const CONFIGURATION_DIRECTORY: &str = "configuration"; @@ -53,6 +58,21 @@ pub enum ConfigurationError { connection.max_incoming_frame_bytes ({frame})" )] OutboundBytesBelowFrameCap { bytes: usize, frame: usize }, + + #[error("limits.global_paired_reserve_percent ({0}) must be below 100")] + ReserveTooLarge(u8), + + #[error( + "limits.global_inbound_byte_rate.burst outside the paired reserve ({unreserved}) \ + must be at least connection.max_incoming_frame_bytes ({frame})" + )] + UnreservedBelowFrameCap { unreserved: u32, frame: u32 }, + + #[error( + "application.trusted_proxies must not contain {0}: it would let every peer \ + name its own source" + )] + TrustsEveryPeer(IpPrefix), } /// Relay settings. @@ -77,6 +97,44 @@ pub struct ApplicationSettings { /// Audience expected in connection tokens (`aud`). pub audience: String, + + /// Peers allowed to name the client they forward in `X-Real-IP`, as + /// addresses or `address/length` ranges. A list, or one comma-separated + /// string so that an environment variable can set it. Empty when absent: + /// the header is then ignored from every peer. + #[serde(default, deserialize_with = "address_ranges")] + pub trusted_proxies: Vec, +} + +/// Read a list of address ranges, or one string of comma-separated ranges. +fn address_ranges<'de, D: Deserializer<'de>>(deserializer: D) -> Result, D::Error> { + struct Ranges; + + impl<'de> Visitor<'de> for Ranges { + type Value = Vec; + + fn expecting(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.write_str("a list of address ranges, or a string of comma-separated ranges") + } + + fn visit_str(self, joined: &str) -> Result { + joined + .split(',') + .filter(|range| !range.trim().is_empty()) + .map(|range| range.parse().map_err(E::custom)) + .collect() + } + + fn visit_seq>(self, mut ranges: A) -> Result { + let mut parsed = Vec::new(); + while let Some(range) = ranges.next_element()? { + parsed.push(range); + } + Ok(parsed) + } + } + + deserializer.deserialize_any(Ranges) } /// HTTP listener used by probes. @@ -166,8 +224,16 @@ pub struct LimitsSettings { /// Inbound byte rate one identity may consume from the relay-wide budget. pub identity_inbound_byte_rate: RateSettings, + /// Inbound byte rate the identities connecting from one address (an IPv6 + /// /64) may consume together from the relay-wide budget. + pub source_inbound_byte_rate: RateSettings, + /// Inbound byte rate across all connections. pub global_inbound_byte_rate: RateSettings, + + /// Percentage of the global burst only frames between two connected peers + /// may spend. + pub global_paired_reserve_percent: u8, } /// Token bucket parameters. @@ -203,6 +269,14 @@ impl Settings { self.application_address(), )); } + if let Some(range) = self + .application + .trusted_proxies + .iter() + .find(|range| range.length() == 0) + { + return Err(ConfigurationError::TrustsEveryPeer(*range)); + } for (setting, value) in [ ( @@ -284,6 +358,10 @@ impl Settings { "limits.identity_inbound_byte_rate", &self.limits.identity_inbound_byte_rate, ), + ( + "limits.source_inbound_byte_rate", + &self.limits.source_inbound_byte_rate, + ), ( "limits.global_inbound_byte_rate", &self.limits.global_inbound_byte_rate, @@ -293,6 +371,11 @@ impl Settings { return Err(ConfigurationError::Zero { setting }); } } + if self.limits.global_paired_reserve_percent >= 100 { + return Err(ConfigurationError::ReserveTooLarge( + self.limits.global_paired_reserve_percent, + )); + } // The relay re-encodes a frame no longer than the text it arrived as, // so a budget that holds the largest incoming frame can queue any @@ -312,6 +395,10 @@ impl Settings { "limits.identity_inbound_byte_rate", self.limits.identity_inbound_byte_rate.burst, ), + ( + "limits.source_inbound_byte_rate", + self.limits.source_inbound_byte_rate.burst, + ), ( "limits.inbound_byte_rate", self.limits.inbound_byte_rate.burst, @@ -329,6 +416,16 @@ impl Settings { }); } } + // Frames buffered for a reconnecting peer cannot spend the reserve, + // so what is left of the burst must still hold a whole frame. + let limits = RelayLimitsConfig::from(self); + let unreserved = limits + .global_inbound_byte_rate + .burst + .saturating_sub(limits.paired_reserve_bytes()); + if unreserved < frame { + return Err(ConfigurationError::UnreservedBelowFrameCap { unreserved, frame }); + } Ok(()) } @@ -378,6 +475,7 @@ impl From<&Settings> for NativeRelayTransportConfig { settings.connection.outbound_queue_timeout_seconds, ), max_token_lifetime: Duration::from_secs(settings.auth.max_token_lifetime_seconds), + trusted_proxies: settings.application.trusted_proxies.clone().into(), } } } @@ -392,7 +490,9 @@ impl From<&Settings> for RelayLimitsConfig { inbound_message_rate: (&settings.limits.inbound_message_rate).into(), inbound_byte_rate: (&settings.limits.inbound_byte_rate).into(), identity_inbound_byte_rate: (&settings.limits.identity_inbound_byte_rate).into(), + source_inbound_byte_rate: (&settings.limits.source_inbound_byte_rate).into(), global_inbound_byte_rate: (&settings.limits.global_inbound_byte_rate).into(), + global_paired_reserve_percent: settings.limits.global_paired_reserve_percent, } } } @@ -463,6 +563,154 @@ mod tests { Ok(()) } + /// The deployment templates, with the placeholders the container + /// entrypoint substitutes filled in. + fn deployed() -> Result { + let render = |template: &str| { + template + .replace("${RELAY_PORT}", "8080") + .replace("${HEALTH_PORT}", "8081") + .replace("${APPLICATION_AUDIENCE}", "wss://relay.example.com") + .replace("${MAX_TOKEN_LIFETIME_SECONDS}", "600") + }; + let base = include_str!("../../../deployment/configs/relay/base.yaml.template"); + let production = include_str!("../../../deployment/configs/relay/production.yaml.template"); + let settings = config::Config::builder() + .add_source(config::File::from_str( + &render(base), + config::FileFormat::Yaml, + )) + .add_source(config::File::from_str( + &render(production), + config::FileFormat::Yaml, + )) + .build()? + .try_deserialize::()?; + Ok(settings) + } + + fn ip(text: &str) -> IpAddr { + text.parse().expect("test address") + } + + #[test] + fn shipped_fairness_defaults_match_the_limiter_and_trust_no_proxy() -> TestResult { + let settings = shipped()?; + settings.validate()?; + + let limits = RelayLimitsConfig::from(&settings); + assert_eq!(limits, RelayLimitsConfig::default()); + assert!( + limits.source_inbound_byte_rate.burst <= limits.identity_inbound_byte_rate.burst, + "a source's share must default to no more than an identity's" + ); + assert!(settings.application.trusted_proxies.is_empty()); + assert!( + NativeRelayTransportConfig::from(&settings) + .trusted_proxies + .is_empty() + ); + Ok(()) + } + + #[test] + fn the_deployment_trusts_only_the_compose_network() -> TestResult { + let deployed = deployed()?; + deployed.validate()?; + + let shipped = shipped()?; + assert_eq!(deployed.limits, shipped.limits); + assert_eq!(deployed.connection, shipped.connection); + assert_eq!(deployed.session, shipped.session); + + let trusted = &deployed.application.trusted_proxies; + let trusts = |peer: &str| trusted.iter().any(|range| range.contains(ip(peer))); + for nginx in ["172.18.0.3", "192.168.16.2", "10.10.0.5"] { + assert!(trusts(nginx), "a compose network address {nginx} is nginx"); + } + for client in ["203.0.113.9", "2001:db8::1", "127.0.0.1"] { + assert!(!trusts(client), "{client} must not name its own source"); + } + Ok(()) + } + + #[test] + fn trusted_proxies_read_from_a_list_or_one_joined_string() -> TestResult { + let with = |value: &str| -> Result { + Ok(config::Config::builder() + .add_source(config::File::from_str( + include_str!("../configuration/base.yaml"), + config::FileFormat::Yaml, + )) + .add_source(config::File::from_str( + include_str!("../configuration/local.yaml"), + config::FileFormat::Yaml, + )) + .set_override("application.trusted_proxies", value)? + .build()? + .try_deserialize::()?) + }; + + let settings = with("10.0.0.5, fd00::/8")?; + assert_eq!( + settings.application.trusted_proxies, + ["10.0.0.5/32".parse::()?, "fd00::/8".parse()?] + ); + assert!(with("")?.application.trusted_proxies.is_empty()); + assert!( + with("10.0.0.0/33").is_err(), + "an invalid range must fail the configuration, not be skipped" + ); + Ok(()) + } + + #[test] + fn fairness_settings_that_break_the_boundary_are_rejected() -> TestResult { + for everyone in ["0.0.0.0/0", "::/0"] { + let mut settings = shipped()?; + settings.application.trusted_proxies = vec![everyone.parse()?]; + assert!(matches!( + settings.validate(), + Err(ConfigurationError::TrustsEveryPeer(_)) + )); + } + + let mut settings = shipped()?; + settings.limits.global_paired_reserve_percent = 100; + assert!(matches!( + settings.validate(), + Err(ConfigurationError::ReserveTooLarge(100)) + )); + + // 1% of the shipped 16 MiB burst is less than one maximal frame. + settings.limits.global_paired_reserve_percent = 99; + assert!(matches!( + settings.validate(), + Err(ConfigurationError::UnreservedBelowFrameCap { .. }) + )); + + let mut settings = shipped()?; + settings.limits.source_inbound_byte_rate.burst = + u32::try_from(settings.connection.max_incoming_frame_bytes)? - 1; + assert!(matches!( + settings.validate(), + Err(ConfigurationError::BurstBelowFrameCap { + setting: "limits.source_inbound_byte_rate", + .. + }) + )); + + let mut settings = shipped()?; + settings.limits.source_inbound_byte_rate.interval_seconds = 0; + assert!(matches!( + settings.validate(), + Err(ConfigurationError::Zero { + setting: "limits.source_inbound_byte_rate" + }) + )); + Ok(()) + } + #[test] fn zero_outbound_budget_or_timeout_is_rejected() -> TestResult { let mut settings = shipped()?; diff --git a/crates/relay/src/connection.rs b/crates/relay/src/connection.rs index 8e7bca8..9f8054f 100644 --- a/crates/relay/src/connection.rs +++ b/crates/relay/src/connection.rs @@ -22,13 +22,13 @@ use wallet_session_core::{codec, connection::ConnectionQuery}; use crate::{ auth::{Identity, IdentityError, unix_now}, config::NativeRelayTransportConfig, - limits::{Bucket, InboundBytes, Limiter}, - session::{Outbound, Registration, Role, Session, SessionError}, + limits::{Bucket, InboundBytes, Limiter, Traffic}, + session::{Counterpart, Outbound, Registration, Role, Session, SessionError}, + source::Source, }; use wallet_session_core::{ did::DidKey, frame::{RelayFrame, RelayFrameBody}, - message::Topic, }; /// Server-side relay WebSocket stream. @@ -103,6 +103,16 @@ pub(crate) fn random_id() -> Result { Ok(codec::encode(&bytes)) } +/// The client behind an accepted handshake. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct Client { + /// Identity proven by the connection token. + pub identity: Identity, + + /// Address the connection's inbound bytes are accounted to. + pub source: Source, +} + /// A single peer connection bound to one topic slot. /// /// A connection only depends on its [`Session`]; peer registration and @@ -110,6 +120,7 @@ pub(crate) fn random_id() -> Result { pub struct Connection { query: ConnectionQuery, identity: Identity, + source: Source, session: Arc, outbound: mpsc::Receiver, self_tx: mpsc::Sender, @@ -119,21 +130,31 @@ pub struct Connection { impl Connection { /// Perform the WebSocket handshake and parse the connection parameters. + /// + /// The client's source is the TCP peer, or the address a peer listed in + /// the trusted proxies names in `X-Real-IP` (see [`Source::resolve`]). #[allow(clippy::result_large_err)] pub async fn handshake( stream: TcpStream, config: NativeRelayTransportConfig, - ) -> Result<(ConnectionQuery, Identity, RelayStream), ConnectionError> { + ) -> Result<(ConnectionQuery, Client, RelayStream), ConnectionError> { + let peer = stream + .peer_addr() + .map_err(ConnectionError::PeerAddress)? + .ip(); let audience = config.audience.as_str(); let now = unix_now(); - let mut parsed: Option> = None; + let mut parsed: Option> = None; let request = accept_hdr_async_with_config( stream, |request: &Request, response: Response| { let params = parse_connection(request.uri(), audience, now, config.max_token_lifetime); let accepted = params.is_ok(); - parsed = Some(params); + parsed = Some(params.map(|(query, identity)| { + let source = Source::resolve(peer, request.headers(), &config.trusted_proxies); + (query, Client { identity, source }) + })); if accepted { Ok(response) } else { @@ -150,7 +171,7 @@ impl Connection { } }; - let (query, identity) = match parsed { + let (query, client) = match parsed { Some(Ok(parsed)) => parsed, Some(Err(error)) => return Err(error), None => match handshake.err() { @@ -164,16 +185,16 @@ impl Connection { "handshake accepted for topic {} as {:?} ({})", query.topic.log_tag(), query.role, - identity.did + client.identity.did ); - Ok((query, identity, websocket)) + Ok((query, client, websocket)) } /// Build a connection for a peer that has been registered on `session`. pub fn new( query: ConnectionQuery, - identity: Identity, + client: Client, session: Arc, outbound: mpsc::Receiver, self_tx: mpsc::Sender, @@ -182,7 +203,8 @@ impl Connection { ) -> Self { Self { query, - identity, + identity: client.identity, + source: client.source, session, outbound, self_tx, @@ -259,6 +281,7 @@ impl Connection { config, limiter: &self.limiter, identity: &self.identity.did, + source: self.source, expires_at, }; let result = read_loop( @@ -404,25 +427,23 @@ fn spend_redelivery_credit(credit: &mut u32) -> bool { } /// Charge an inbound control frame like a text frame: to the connection's -/// inbound budget, then to the identity's share of the relay-wide budget and -/// to the budget itself. +/// inbound budget, then to the shared budgets (see [`charge_shared`]). /// -/// The frame is dropped when a budget is exhausted, and a peer that keeps -/// exceeding its own or its identity's budget is disconnected like any other -/// flooder. An exhausted relay-wide budget alone drops the frame without -/// counting towards that, so other identities cannot get this peer -/// disconnected. Returns `Err` only to disconnect the peer. +/// Control frames are for the relay, not the counterpart, so they are charged +/// whether or not one is connected. Without a connected counterpart they are +/// [`Traffic::Unpaired`] and cannot spend the share reserved for paired +/// traffic. The frame is dropped when a budget is exhausted, and a peer that +/// keeps exceeding its own budget is disconnected like any other flooder. +/// Returns `Err` only to disconnect the peer. fn charge_control_frame( bytes: usize, message_bucket: &mut Bucket, byte_bucket: &mut Bucket, - limiter: &Limiter, - identity: &DidKey, + context: &ReadContext<'_>, consecutive_drops: &mut u32, - topic: &Topic, + topic: &impl fmt::Display, ) -> Result<(), ConnectionError> { - let topic = topic.log_tag(); - let limits = limiter.limits(); + let limits = context.limiter.limits(); let now = tokio::time::Instant::now(); let within_connection_limits = message_bucket.allow_n(1, limits.inbound_message_rate, now) && byte_bucket.allow_n( @@ -440,12 +461,54 @@ fn charge_control_frame( return Ok(()); } - match limiter.allow_inbound_bytes(identity, bytes) { - InboundBytes::Allowed => *consecutive_drops = 0, + let traffic = match context.session.counterpart(context.query.role)? { + Counterpart::Connected => Traffic::Paired, + Counterpart::Reserved | Counterpart::Absent => Traffic::Unpaired, + }; + charge_shared( + context, + bytes, + traffic, + consecutive_drops, + "control frame", + topic, + )?; + Ok(()) +} + +/// Charge `bytes` the connection's own budget has accepted to the shared +/// budgets: the identity's and the source's shares of the relay-wide budget, +/// and the budget itself. Returns whether the frame may go on, or `Err` to +/// disconnect the peer. +/// +/// A peer that keeps exceeding its identity's share is disconnected like one +/// that keeps exceeding its own budget. An exhausted source share or +/// relay-wide budget only drops the frame, without counting towards that: +/// other identities, behind the same address or anywhere else, could +/// otherwise get this peer disconnected. +fn charge_shared( + context: &ReadContext<'_>, + bytes: usize, + traffic: Traffic, + consecutive_drops: &mut u32, + kind: &str, + topic: &impl fmt::Display, +) -> Result { + let ReadContext { + limiter, + identity, + source, + .. + } = *context; + match limiter.allow_inbound_bytes(identity, &source, bytes, traffic) { + InboundBytes::Allowed => { + *consecutive_drops = 0; + return Ok(true); + } InboundBytes::ShareExhausted => { *consecutive_drops += 1; log::debug!( - "dropping control frame on topic {topic}: identity share of the relay budget exhausted" + "dropping {kind} on topic {topic}: identity share of the relay budget exhausted" ); if *consecutive_drops > MAX_CONSECUTIVE_RATE_DROPS { log::warn!( @@ -454,13 +517,16 @@ fn charge_control_frame( return Err(ConnectionError::RateLimited); } } - InboundBytes::GlobalExhausted => { + InboundBytes::SourceExhausted => { log::debug!( - "dropping control frame on topic {topic}: relay inbound byte budget exhausted" + "dropping {kind} on topic {topic}: source share of the relay budget exhausted" ); } + InboundBytes::GlobalExhausted => { + log::debug!("dropping {kind} on topic {topic}: relay inbound byte budget exhausted"); + } } - Ok(()) + Ok(false) } /// Send one message, returning whether the socket is still usable. @@ -521,6 +587,7 @@ struct ReadContext<'a> { config: &'a NativeRelayTransportConfig, limiter: &'a Limiter, identity: &'a DidKey, + source: Source, expires_at: tokio::time::Instant, } @@ -548,10 +615,13 @@ async fn read_loop( session, config, limiter, - identity, expires_at, + .. } = *context; let topic = query.topic.log_tag(); + // A frame for a reconnecting counterpart is buffered only when buffering + // is on; otherwise nothing will receive it, as with no counterpart at all. + let buffering = config.session_buffer_max_messages > 0 && !config.session_buffer_ttl.is_zero(); let limits = limiter.limits(); let now = tokio::time::Instant::now(); @@ -640,27 +710,28 @@ async fn read_loop( continue; } - match limiter.allow_inbound_bytes(identity, text.len()) { - InboundBytes::Allowed => consecutive_drops = 0, - InboundBytes::ShareExhausted => { - consecutive_drops += 1; - log::debug!( - "dropping frame on topic {topic}: identity share of the relay budget exhausted" - ); - if consecutive_drops > MAX_CONSECUTIVE_RATE_DROPS { - log::warn!( - "peer on topic {topic} keeps exceeding its share of the relay budget; dropping" - ); - return Err(ConnectionError::RateLimited); - } - continue; - } - InboundBytes::GlobalExhausted => { - log::debug!( - "dropping frame on topic {topic}: relay inbound byte budget exhausted" - ); + // Only a frame something can receive, its counterpart or the + // buffer kept for it while it reconnects, spends the shared + // budgets and is parsed. Any other frame is dropped here, at + // the cost of this connection's own budget alone. + let traffic = match session.counterpart(query.role)? { + Counterpart::Connected => Traffic::Paired, + Counterpart::Reserved if buffering => Traffic::Unpaired, + Counterpart::Reserved | Counterpart::Absent => { + log::trace!("dropping frame on topic {topic}: counterpart not connected"); continue; } + }; + let charged = charge_shared( + context, + text.len(), + traffic, + &mut consecutive_drops, + "frame", + &topic, + )?; + if !charged { + continue; } // Neither record carries the frame itself: it holds the topic, @@ -688,10 +759,9 @@ async fn read_loop( payload.len(), &mut message_bucket, &mut byte_bucket, - limiter, - identity, + context, &mut consecutive_drops, - &query.topic, + &topic, )?; continue; } @@ -700,10 +770,9 @@ async fn read_loop( frame.payload().len(), &mut message_bucket, &mut byte_bucket, - limiter, - identity, + context, &mut consecutive_drops, - &query.topic, + &topic, )?; continue; } @@ -803,6 +872,9 @@ pub enum ConnectionError { #[error("websocket handshake timed out")] HandshakeTimeout, + #[error("peer address unavailable: {0}")] + PeerAddress(std::io::Error), + #[error("handshake query is too large")] QueryTooLarge, diff --git a/crates/relay/src/health.rs b/crates/relay/src/health.rs index fb217a8..c9ffc0f 100644 --- a/crates/relay/src/health.rs +++ b/crates/relay/src/health.rs @@ -32,6 +32,7 @@ pub struct ReadinessResponse { connections: usize, sessions: usize, dropped_share_bytes: u64, + dropped_source_bytes: u64, dropped_global_bytes: u64, } @@ -77,6 +78,7 @@ async fn ready(State(state): State) -> impl IntoResponse { connections: state.limiter.active_connections(), sessions, dropped_share_bytes: state.limiter.dropped_share_bytes(), + dropped_source_bytes: state.limiter.dropped_source_bytes(), dropped_global_bytes: state.limiter.dropped_global_bytes(), }), ) @@ -96,7 +98,7 @@ async fn ready(State(state): State) -> impl IntoResponse { #[cfg(test)] mod tests { - use std::time::Duration; + use std::{net::IpAddr, time::Duration}; use axum::{ body::{Body, to_bytes}, @@ -115,8 +117,9 @@ mod tests { use super::*; use crate::{ - limits::{OutboundBudget, RateLimit, RelayLimitsConfig}, + limits::{InboundBytes, OutboundBudget, RateLimit, RelayLimitsConfig, Traffic}, session::{PeerHandle, Role}, + source::Source, }; type TestResult = Result<(), Box>; @@ -200,6 +203,7 @@ mod tests { "connections": 1, "sessions": 1, "dropped_share_bytes": 0, + "dropped_source_bytes": 0, "dropped_global_bytes": 0, }) ); @@ -214,6 +218,10 @@ mod tests { burst: 10, interval: Duration::from_secs(3600), }, + source_inbound_byte_rate: RateLimit { + burst: 20, + interval: Duration::from_secs(3600), + }, global_inbound_byte_rate: RateLimit { burst: 100, interval: Duration::from_secs(3600), @@ -221,28 +229,21 @@ mod tests { ..RelayLimitsConfig::default() })); - let hog = DidKey::from_ed25519_bytes([1u8; 32]); - assert_eq!( - limiter.allow_inbound_bytes(&hog, 10), - crate::limits::InboundBytes::Allowed - ); - assert_eq!( - limiter.allow_inbound_bytes(&hog, 1), - crate::limits::InboundBytes::ShareExhausted - ); - - for seed in 2..11 { + let shared = Source::of(IpAddr::from([192, 0, 2, 1])); + let own = |seed: u8| Source::of(IpAddr::from([198, 51, 100, seed])); + let charge = |seed: u8, source: Source, bytes: usize| { let identity = DidKey::from_ed25519_bytes([seed; 32]); - assert_eq!( - limiter.allow_inbound_bytes(&identity, 10), - crate::limits::InboundBytes::Allowed - ); + limiter.allow_inbound_bytes(&identity, &source, bytes, Traffic::Paired) + }; + + assert_eq!(charge(1, shared, 10), InboundBytes::Allowed); + assert_eq!(charge(1, shared, 1), InboundBytes::ShareExhausted); + assert_eq!(charge(2, shared, 10), InboundBytes::Allowed); + assert_eq!(charge(3, shared, 10), InboundBytes::SourceExhausted); + for seed in 4..12 { + assert_eq!(charge(seed, own(seed), 10), InboundBytes::Allowed); } - let last = DidKey::from_ed25519_bytes([11u8; 32]); - assert_eq!( - limiter.allow_inbound_bytes(&last, 10), - crate::limits::InboundBytes::GlobalExhausted - ); + assert_eq!(charge(12, own(12), 10), InboundBytes::GlobalExhausted); let (status, body) = get(registry, limiter, READY_PATH).await; assert_eq!(status, StatusCode::OK); @@ -253,6 +254,7 @@ mod tests { "connections": 0, "sessions": 0, "dropped_share_bytes": 1, + "dropped_source_bytes": 10, "dropped_global_bytes": 10, }) ); @@ -342,7 +344,7 @@ mod tests { assert!(ready.contains("content-type: application/json"), "{ready}"); assert!( ready.ends_with( - r#"{"status":"ok","connections":0,"sessions":0,"dropped_share_bytes":0,"dropped_global_bytes":0}"# + r#"{"status":"ok","connections":0,"sessions":0,"dropped_share_bytes":0,"dropped_source_bytes":0,"dropped_global_bytes":0}"# ), "{ready}" ); diff --git a/crates/relay/src/lib.rs b/crates/relay/src/lib.rs index 11cbbcf..cc73595 100644 --- a/crates/relay/src/lib.rs +++ b/crates/relay/src/lib.rs @@ -7,4 +7,5 @@ pub mod limits; pub mod registry; pub mod server; pub mod session; +pub mod source; pub mod telemetry; diff --git a/crates/relay/src/limits.rs b/crates/relay/src/limits.rs index 7eb670b..23849c5 100644 --- a/crates/relay/src/limits.rs +++ b/crates/relay/src/limits.rs @@ -1,5 +1,6 @@ use std::{ collections::HashMap, + hash::Hash, sync::{ Arc, Mutex, MutexGuard, atomic::{AtomicU64, AtomicUsize, Ordering}, @@ -10,7 +11,9 @@ use std::{ use tokio::time::Instant; use wallet_session_core::did::DidKey; -/// Maximum number of identities tracked by the connect rate limiter. +use crate::source::Source; + +/// Maximum number of identities or sources tracked by one keyed rate limiter. const MAX_RATE_KEYS: usize = 65_536; /// Token-bucket parameters. @@ -26,16 +29,35 @@ pub struct RateLimit { /// Outcome of charging the relay-wide inbound byte budget. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum InboundBytes { - /// Accepted against both the identity share and the global budget. + /// Accepted against the identity and source shares and the global budget. Allowed, /// The identity already consumed its share of the budget. ShareExhausted, - /// The identity is within its share, but the global budget is dry. + /// The identity is within its share, but the identities connecting from + /// its source address have together consumed the source's share. + SourceExhausted, + + /// Both shares have room, but the global budget is dry, or down to the + /// reserve for paired traffic. GlobalExhausted, } +/// What the inbound bytes being charged are for. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum Traffic { + /// Frames between two connected peers. They may spend the whole relay-wide + /// budget, the share reserved for them included. + Paired, + + /// Anything else the relay spends work on: frames buffered for a peer + /// that is reconnecting, and control frames on a connection whose + /// counterpart is not connected. They are refused once the relay-wide + /// budget is down to the paired reserve. + Unpaired, +} + /// Static relay resource limits. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub struct RelayLimitsConfig { @@ -62,8 +84,20 @@ pub struct RelayLimitsConfig { /// [`RelayLimitsConfig::global_inbound_byte_rate`], but one cannot. pub identity_inbound_byte_rate: RateLimit, + /// Inbound byte rate the identities connecting from one source address + /// may consume together from the relay-wide budget. Identities cost + /// nothing to mint and addresses do, so this is what bounds one client + /// spreading its traffic across many identities. No larger than the + /// identity share by default. + pub source_inbound_byte_rate: RateLimit, + /// Inbound byte rate across all connections. pub global_inbound_byte_rate: RateLimit, + + /// Percentage of the global burst that only [`Traffic::Paired`] bytes may + /// spend, so that one-sided sessions cannot take the last of the budget + /// from peers that are talking to each other. + pub global_paired_reserve_percent: u8, } impl Default for RelayLimitsConfig { @@ -88,14 +122,28 @@ impl Default for RelayLimitsConfig { burst: 1024 * 1024, interval: Duration::from_secs(1), }, + source_inbound_byte_rate: RateLimit { + burst: 1024 * 1024, + interval: Duration::from_secs(1), + }, global_inbound_byte_rate: RateLimit { burst: 16 * 1024 * 1024, interval: Duration::from_secs(1), }, + global_paired_reserve_percent: 25, } } } +impl RelayLimitsConfig { + /// Bytes of the global burst only [`Traffic::Paired`] bytes may spend. + pub fn paired_reserve_bytes(&self) -> u32 { + let percent = u64::from(self.global_paired_reserve_percent.min(100)); + let reserve = u64::from(self.global_inbound_byte_rate.burst) * percent / 100; + u32::try_from(reserve).unwrap_or(u32::MAX) + } +} + /// Lock a mutex, recovering the guard if it was poisoned. fn lock(mutex: &Mutex) -> MutexGuard<'_, T> { match mutex.lock() { @@ -132,6 +180,12 @@ impl Bucket { /// Refill and consume `cost` tokens, returning whether the attempt is allowed. pub(crate) fn allow_n(&mut self, cost: u32, rate: RateLimit, now: Instant) -> bool { + self.allow_n_above(cost, 0, rate, now) + } + + /// Refill and consume `cost` tokens if at least `floor` remain afterwards, + /// returning whether the attempt is allowed. + fn allow_n_above(&mut self, cost: u32, floor: u32, rate: RateLimit, now: Instant) -> bool { let cost = f64::from(cost.clamp(1, rate.burst.max(1))); let elapsed = now.saturating_duration_since(self.last); let interval = rate.interval.as_secs_f64().max(f64::EPSILON); @@ -139,7 +193,7 @@ impl Bucket { self.tokens = (self.tokens + refill).min(f64::from(rate.burst)); self.last = now; - if self.tokens >= cost { + if self.tokens - cost >= f64::from(floor) { self.tokens -= cost; true } else { @@ -150,14 +204,23 @@ impl Bucket { /// Per-key token buckets, with the time of the last full prune. #[derive(Debug)] -struct BucketMap { - buckets: HashMap, +struct BucketMap { + buckets: HashMap, last_prune: Instant, } -fn bucket_for<'a>( - state: &'a mut BucketMap, - identity: &DidKey, +impl BucketMap { + fn new(now: Instant) -> Self { + Self { + buckets: HashMap::new(), + last_prune: now, + } + } +} + +fn bucket_for<'a, K: Eq + Hash + Clone>( + state: &'a mut BucketMap, + key: &K, rate: RateLimit, now: Instant, ) -> &'a mut Bucket { @@ -169,15 +232,15 @@ fn bucket_for<'a>( state.last_prune = now; } while state.buckets.len() >= MAX_RATE_KEYS { - let Some(key) = state.buckets.keys().next().cloned() else { + let Some(evicted) = state.buckets.keys().next().cloned() else { break; }; - state.buckets.remove(&key); + state.buckets.remove(&evicted); } } state .buckets - .entry(identity.clone()) + .entry(key.clone()) .or_insert_with(|| Bucket::full(rate.burst, now)) } @@ -187,10 +250,12 @@ pub struct Limiter { limits: RelayLimitsConfig, total: Mutex, per_identity: Mutex>, - connect_rate: Mutex, - identity_bytes: Mutex, + connect_rate: Mutex>, + identity_bytes: Mutex>, + source_bytes: Mutex>, inbound_bytes: Mutex, dropped_share_bytes: AtomicU64, + dropped_source_bytes: AtomicU64, dropped_global_bytes: AtomicU64, } @@ -202,16 +267,12 @@ impl Limiter { limits, total: Mutex::new(0), per_identity: Mutex::new(HashMap::new()), - connect_rate: Mutex::new(BucketMap { - buckets: HashMap::new(), - last_prune: now, - }), - identity_bytes: Mutex::new(BucketMap { - buckets: HashMap::new(), - last_prune: now, - }), + connect_rate: Mutex::new(BucketMap::new(now)), + identity_bytes: Mutex::new(BucketMap::new(now)), + source_bytes: Mutex::new(BucketMap::new(now)), inbound_bytes: Mutex::new(Bucket::full(limits.global_inbound_byte_rate.burst, now)), dropped_share_bytes: AtomicU64::new(0), + dropped_source_bytes: AtomicU64::new(0), dropped_global_bytes: AtomicU64::new(0), } } @@ -262,12 +323,26 @@ impl Limiter { }) } - /// Consume `bytes` from the identity's share of the relay-wide budget and - /// from the budget itself. - pub fn allow_inbound_bytes(&self, identity: &DidKey, bytes: usize) -> InboundBytes { + /// Consume `bytes` sent by `identity` from `source` from the identity's + /// and the source's shares of the relay-wide budget, and from the budget + /// itself. + /// + /// The identity share is checked first, so a sender over its own share + /// hears so even when its source is exhausted as well. Nothing is spent + /// unless every budget has room: a refused charge is given back to the + /// shares that had already accepted it. [`Traffic::Unpaired`] bytes are + /// refused once the global budget is down to the paired reserve. + pub fn allow_inbound_bytes( + &self, + identity: &DidKey, + source: &Source, + bytes: usize, + traffic: Traffic, + ) -> InboundBytes { let cost = u32::try_from(bytes).unwrap_or(u32::MAX); let now = Instant::now(); let share = self.limits.identity_inbound_byte_rate; + let source_share = self.limits.source_inbound_byte_rate; if cost > share.burst { self.dropped_share_bytes @@ -275,18 +350,38 @@ impl Limiter { return InboundBytes::ShareExhausted; } - let mut state = lock(&self.identity_bytes); - let bucket = bucket_for(&mut state, identity, share, now); - if !bucket.allow_n(cost, share, now) { + // Locks are always taken in this order: identities, sources, global. + let mut identities = lock(&self.identity_bytes); + let identity_bucket = bucket_for(&mut identities, identity, share, now); + if !identity_bucket.allow_n(cost, share, now) { self.dropped_share_bytes .fetch_add(u64::from(cost), Ordering::Relaxed); return InboundBytes::ShareExhausted; } - let within_global = - lock(&self.inbound_bytes).allow_n(cost, self.limits.global_inbound_byte_rate, now); + let mut sources = lock(&self.source_bytes); + let source_bucket = bucket_for(&mut sources, source, source_share, now); + // A cost above the burst would be charged as only the burst. + if cost > source_share.burst || !source_bucket.allow_n(cost, source_share, now) { + identity_bucket.refund(cost, share); + self.dropped_source_bytes + .fetch_add(u64::from(cost), Ordering::Relaxed); + return InboundBytes::SourceExhausted; + } + + let floor = match traffic { + Traffic::Paired => 0, + Traffic::Unpaired => self.limits.paired_reserve_bytes(), + }; + let within_global = lock(&self.inbound_bytes).allow_n_above( + cost, + floor, + self.limits.global_inbound_byte_rate, + now, + ); if !within_global { - bucket.refund(cost, share); + identity_bucket.refund(cost, share); + source_bucket.refund(cost, source_share); self.dropped_global_bytes .fetch_add(u64::from(cost), Ordering::Relaxed); return InboundBytes::GlobalExhausted; @@ -299,6 +394,12 @@ impl Limiter { self.dropped_share_bytes.load(Ordering::Relaxed) } + /// Bytes dropped because the identities behind one source address + /// exceeded the source's share. + pub fn dropped_source_bytes(&self) -> u64 { + self.dropped_source_bytes.load(Ordering::Relaxed) + } + /// Bytes dropped because the relay-wide budget was exhausted. pub fn dropped_global_bytes(&self) -> u64 { self.dropped_global_bytes.load(Ordering::Relaxed) @@ -415,10 +516,23 @@ impl Drop for OutboundCharge { mod tests { use super::*; + use std::net::IpAddr; + fn key(seed: u8) -> DidKey { DidKey::from_ed25519_bytes([seed; 32]) } + /// A source address of its own for each `seed`. + fn source(seed: u8) -> Source { + Source::of(IpAddr::from([192, 0, 2, seed])) + } + + /// Charge paired bytes from the source the tests share unless they say + /// otherwise. + fn paired(limiter: &Limiter, identity: &DidKey, bytes: usize) -> InboundBytes { + limiter.allow_inbound_bytes(identity, &source(1), bytes, Traffic::Paired) + } + fn limits() -> RelayLimitsConfig { RelayLimitsConfig { max_connections: 2, @@ -574,24 +688,15 @@ mod tests { }); let identity = key(1); + assert_eq!(paired(&limiter, &identity, 600), InboundBytes::Allowed); assert_eq!( - limiter.allow_inbound_bytes(&identity, 600), - InboundBytes::Allowed - ); - assert_eq!( - limiter.allow_inbound_bytes(&identity, 600), + paired(&limiter, &identity, 600), InboundBytes::GlobalExhausted ); - assert_eq!( - limiter.allow_inbound_bytes(&identity, 400), - InboundBytes::Allowed - ); + assert_eq!(paired(&limiter, &identity, 400), InboundBytes::Allowed); tokio::time::advance(Duration::from_secs(10)).await; - assert_eq!( - limiter.allow_inbound_bytes(&identity, 1_000), - InboundBytes::Allowed - ); + assert_eq!(paired(&limiter, &identity, 1_000), InboundBytes::Allowed); assert_eq!(limiter.dropped_global_bytes(), 600); assert_eq!(limiter.dropped_share_bytes(), 0); } @@ -612,17 +717,11 @@ mod tests { let hog = key(1); let other = key(2); - assert_eq!( - limiter.allow_inbound_bytes(&hog, 100), - InboundBytes::Allowed - ); - assert_eq!( - limiter.allow_inbound_bytes(&hog, 1), - InboundBytes::ShareExhausted - ); + assert_eq!(paired(&limiter, &hog, 100), InboundBytes::Allowed); + assert_eq!(paired(&limiter, &hog, 1), InboundBytes::ShareExhausted); assert_eq!( - limiter.allow_inbound_bytes(&other, 100), + paired(&limiter, &other, 100), InboundBytes::Allowed, "another identity keeps its own share" ); @@ -647,30 +746,27 @@ mod tests { let other = key(2); let oversized = key(3); - assert_eq!(limiter.allow_inbound_bytes(&hog, 4), InboundBytes::Allowed); - assert_eq!( - limiter.allow_inbound_bytes(&other, 4), - InboundBytes::GlobalExhausted - ); + assert_eq!(paired(&limiter, &hog, 4), InboundBytes::Allowed); + assert_eq!(paired(&limiter, &other, 4), InboundBytes::GlobalExhausted); assert_eq!(limiter.dropped_global_bytes(), 4); assert_eq!(limiter.dropped_share_bytes(), 0); tokio::time::advance(Duration::from_secs(1)).await; assert_eq!( - limiter.allow_inbound_bytes(&other, 4), + paired(&limiter, &other, 4), InboundBytes::Allowed, "a globally rejected frame must not spend the identity share" ); tokio::time::advance(Duration::from_secs(1)).await; assert_eq!( - limiter.allow_inbound_bytes(&oversized, 5), + paired(&limiter, &oversized, 5), InboundBytes::ShareExhausted, "a frame larger than the share burst can never fit it" ); assert_eq!(limiter.dropped_share_bytes(), 5); assert_eq!( - limiter.allow_inbound_bytes(&oversized, 4), + paired(&limiter, &oversized, 4), InboundBytes::Allowed, "the rejected oversized frame must not have spent the share" ); @@ -724,14 +820,146 @@ mod tests { ..RelayLimitsConfig::default() }); + assert_eq!(paired(&limiter, &key(1), 1_000), InboundBytes::Allowed); assert_eq!( - limiter.allow_inbound_bytes(&key(1), 1_000), + paired(&limiter, &key(2), 1_000), + InboundBytes::GlobalExhausted + ); + assert_eq!(limiter.dropped_global_bytes(), 1_000); + } + + #[tokio::test(start_paused = true)] + async fn a_source_share_caps_the_identities_behind_one_address() { + let roomy = RateLimit { + burst: 1_000_000, + interval: Duration::from_secs(60), + }; + let limiter = Limiter::new(RelayLimitsConfig { + identity_inbound_byte_rate: roomy, + source_inbound_byte_rate: RateLimit { + burst: 1_000, + interval: Duration::from_secs(60), + }, + global_inbound_byte_rate: roomy, + ..RelayLimitsConfig::default() + }); + + // Fresh identities from one address spend one share between them. + for seed in 1..=4 { + assert_eq!( + limiter.allow_inbound_bytes(&key(seed), &source(1), 250, Traffic::Paired), + InboundBytes::Allowed + ); + } + assert_eq!( + limiter.allow_inbound_bytes(&key(5), &source(1), 1, Traffic::Paired), + InboundBytes::SourceExhausted, + "a fifth identity from the same address finds the share spent" + ); + assert_eq!(limiter.dropped_source_bytes(), 1); + assert_eq!(limiter.dropped_share_bytes(), 0); + assert_eq!(limiter.dropped_global_bytes(), 0); + + assert_eq!( + limiter.allow_inbound_bytes(&key(1), &source(2), 1_000, Traffic::Paired), + InboundBytes::Allowed, + "another address keeps its own share, whichever identity uses it" + ); + } + + #[tokio::test(start_paused = true)] + async fn a_source_refusal_spends_no_identity_share() { + let limiter = Limiter::new(RelayLimitsConfig { + identity_inbound_byte_rate: RateLimit { + burst: 200, + interval: Duration::from_secs(3600), + }, + source_inbound_byte_rate: RateLimit { + burst: 100, + interval: Duration::from_secs(3600), + }, + ..RelayLimitsConfig::default() + }); + + assert_eq!( + limiter.allow_inbound_bytes(&key(1), &source(1), 100, Traffic::Paired), InboundBytes::Allowed ); assert_eq!( - limiter.allow_inbound_bytes(&key(2), 1_000), + limiter.allow_inbound_bytes(&key(2), &source(1), 100, Traffic::Paired), + InboundBytes::SourceExhausted + ); + assert_eq!( + limiter.allow_inbound_bytes(&key(2), &source(2), 100, Traffic::Paired), + InboundBytes::Allowed, + "the refused frame must not have spent the identity share" + ); + + assert_eq!( + limiter.allow_inbound_bytes(&key(3), &source(3), 150, Traffic::Paired), + InboundBytes::SourceExhausted, + "a frame larger than the source burst can never fit it" + ); + assert_eq!( + limiter.allow_inbound_bytes(&key(3), &source(3), 100, Traffic::Paired), + InboundBytes::Allowed, + "the oversized frame must have spent neither share" + ); + assert_eq!(limiter.dropped_source_bytes(), 250); + } + + #[tokio::test(start_paused = true)] + async fn unpaired_bytes_cannot_spend_the_paired_reserve() { + let roomy = RateLimit { + burst: 1_000_000, + interval: Duration::from_secs(60), + }; + let limiter = Limiter::new(RelayLimitsConfig { + identity_inbound_byte_rate: roomy, + source_inbound_byte_rate: roomy, + global_inbound_byte_rate: RateLimit { + burst: 1_000, + interval: Duration::from_secs(3600), + }, + global_paired_reserve_percent: 25, + ..RelayLimitsConfig::default() + }); + + // However many identities and addresses one-sided traffic comes + // from, it stops where the reserve begins. + for seed in 1..=3 { + assert_eq!( + limiter.allow_inbound_bytes(&key(seed), &source(seed), 250, Traffic::Unpaired), + InboundBytes::Allowed + ); + } + assert_eq!( + limiter.allow_inbound_bytes(&key(4), &source(4), 1, Traffic::Unpaired), InboundBytes::GlobalExhausted ); - assert_eq!(limiter.dropped_global_bytes(), 1_000); + assert_eq!(limiter.dropped_global_bytes(), 1); + + assert_eq!( + limiter.allow_inbound_bytes(&key(5), &source(5), 250, Traffic::Paired), + InboundBytes::Allowed, + "paired traffic still has the reserve" + ); + assert_eq!( + limiter.allow_inbound_bytes(&key(5), &source(5), 1, Traffic::Paired), + InboundBytes::GlobalExhausted + ); + } + + #[test] + fn the_paired_reserve_is_a_share_of_the_global_burst() { + let mut limits = RelayLimitsConfig::default(); + assert_eq!(limits.paired_reserve_bytes(), 4 * 1024 * 1024); + + limits.global_paired_reserve_percent = 0; + assert_eq!(limits.paired_reserve_bytes(), 0); + + limits.global_inbound_byte_rate.burst = u32::MAX; + limits.global_paired_reserve_percent = 99; + assert_eq!(limits.paired_reserve_bytes(), 4_252_017_622); } } diff --git a/crates/relay/src/server.rs b/crates/relay/src/server.rs index dd9e891..fd44101 100644 --- a/crates/relay/src/server.rs +++ b/crates/relay/src/server.rs @@ -67,7 +67,8 @@ pub async fn handle( }; let grace = config.session_grace; - let (query, identity, mut websocket) = Connection::handshake(stream, config.clone()).await?; + let (query, client, mut websocket) = Connection::handshake(stream, config.clone()).await?; + let identity = &client.identity; if !limiter.allow_connect(&identity.did) { log::warn!("connect rate limit exceeded for {}", identity.did); @@ -121,17 +122,9 @@ pub async fn handle( }; let holds_slot = !matches!(registration, Registration::RoleTaken(_)); - let result = Connection::new( - query, - identity, - session, - rx, - tx, - config, - Arc::clone(&limiter), - ) - .run(websocket, registration) - .await; + let result = Connection::new(query, client, session, rx, tx, config, Arc::clone(&limiter)) + .run(websocket, registration) + .await; if holds_slot && registry @@ -192,13 +185,14 @@ mod tests { use crate::auth::test_utils::{self, signing_key}; use crate::connection::MAX_CONSECUTIVE_RATE_DROPS; use crate::limits::{Limiter, RateLimit, RelayLimitsConfig}; + use crate::source::{IpPrefix, REAL_IP_HEADER}; use crate::telemetry::test_utils::{CapturedLogs, capture}; use futures_util::{SinkExt, StreamExt}; use std::net::SocketAddr; use tokio::time::timeout; use tokio_tungstenite::{ MaybeTlsStream, WebSocketStream, connect_async, - tungstenite::{Bytes, Message}, + tungstenite::{Bytes, Message, client::IntoClientRequest, http::HeaderValue}, }; use wallet_session_core::frame::{RelayFrame, RelayFrameBody}; use wallet_session_core::message::Topic; @@ -1894,6 +1888,7 @@ mod tests { inbound_message_rate: unlimited, inbound_byte_rate: unlimited, identity_inbound_byte_rate: unlimited, + source_inbound_byte_rate: unlimited, global_inbound_byte_rate: unlimited, ..RelayLimitsConfig::default() }, @@ -2139,10 +2134,10 @@ mod tests { /// Bytes of the pings these tests send. const PING_PAYLOAD: u32 = 100; - /// Limits under which only the shared byte budgets can refuse a ping. - fn ping_limits() -> RelayLimitsConfig { + /// Limits under which only a budget a test tightens can refuse a frame. + fn roomy_limits() -> RelayLimitsConfig { let roomy = RateLimit { - burst: 1024 * 1024, + burst: 16 * 1024 * 1024, interval: Duration::from_secs(3600), }; RelayLimitsConfig { @@ -2152,11 +2147,22 @@ mod tests { }, inbound_byte_rate: roomy, identity_inbound_byte_rate: roomy, + source_inbound_byte_rate: roomy, global_inbound_byte_rate: roomy, ..RelayLimitsConfig::default() } } + /// Limits under which only the shared byte budgets can refuse a ping. + /// The pinging connections have no counterpart, so their pings may not + /// spend a paired reserve: there is none, and the whole budget is theirs. + fn ping_limits() -> RelayLimitsConfig { + RelayLimitsConfig { + global_paired_reserve_percent: 0, + ..roomy_limits() + } + } + /// Pings spend the relay-wide inbound budget like text frames, and its /// exhaustion drops them without disconnecting an identity that is within /// its own limits. @@ -2243,6 +2249,290 @@ mod tests { Ok(()) } + /// Connect as `auth`, sending `X-Real-IP: real_ip` if one is given. + async fn connect_claiming( + addr: SocketAddr, + auth: &str, + role: &str, + topic: &Topic, + real_ip: Option<&str>, + ) -> Result> { + let mut request = + format!("ws://{addr}/?auth={auth}&role={role}&topic={topic}").into_client_request()?; + if let Some(real_ip) = real_ip { + request + .headers_mut() + .insert(REAL_IP_HEADER, HeaderValue::from_str(real_ip)?); + } + let (client, _response) = connect_async(request).await?; + Ok(client) + } + + /// Ready session `n` of a test, on a topic and between two identities of + /// its own. Its proposer, which the tests send from, claims `real_ip` if + /// one is given. + async fn session_claiming( + addr: SocketAddr, + n: u8, + real_ip: Option<&str>, + ) -> Result<(Topic, (Client, Client)), Box> { + let topic = Topic::new([100 + n; 32])?; + let proposer = token_with(2 * n + 1); + let responder = token_with(2 * n + 2); + let mut proposer = connect_claiming(addr, &proposer, "proposer", &topic, real_ip).await?; + let mut responder = connect_claiming(addr, &responder, "responder", &topic, None).await?; + read_frame(&mut proposer).await?; + read_frame(&mut responder).await?; + Ok((topic, (proposer, responder))) + } + + /// Send `frame` from the first client of `pair` and report whether the + /// second receives it. + async fn reaches( + pair: &mut (Client, Client), + frame: &RelayFrame, + ) -> Result> { + send(&mut pair.0, frame).await?; + let received = timeout(Duration::from_millis(300), read_frame(&mut pair.1)).await; + Ok(matches!(received, Ok(Ok(RelayFrame { ref id, .. })) if *id == frame.id)) + } + + /// Bytes the relay charges for `frame`, the text the client sends. + fn charged_bytes(frame: &RelayFrame) -> Result> { + Ok(u32::try_from(serde_json::to_string(frame)?.len())?) + } + + /// A frame nothing can receive is dropped before it is parsed, and costs + /// its connection's own budget only: none of the shared budgets. + #[tokio::test] + async fn frames_nothing_can_receive_spend_no_shared_budget() -> TestResult { + let tight = RateLimit { + burst: 4096, + interval: Duration::from_secs(3600), + }; + let limiter = Arc::new(Limiter::new(RelayLimitsConfig { + identity_inbound_byte_rate: tight, + source_inbound_byte_rate: tight, + global_inbound_byte_rate: tight, + ..RelayLimitsConfig::default() + })); + let config = NativeRelayTransportConfig::new(AUDIENCE); + let addr = serve_with_limiter(config, Arc::clone(&limiter)).await?; + + // Five times what any shared budget holds, then a frame that would + // close the connection if it were parsed. + let topic = topic()?; + let payload = "a".repeat(1000); + let mut proposer = connect(addr, "proposer", &topic, None).await?; + for i in 0..20 { + send( + &mut proposer, + &message(&topic, &format!("x{i:02}"), &payload), + ) + .await?; + } + proposer.send(Message::text("not a relay frame")).await?; + sleep(Duration::from_millis(200)).await; + assert_eq!( + ( + limiter.dropped_share_bytes(), + limiter.dropped_source_bytes(), + limiter.dropped_global_bytes() + ), + (0, 0, 0), + "a frame without a receiver must not be charged to a shared budget" + ); + + // The connection survived, and every shared budget is still whole. + let responder = connect(addr, "responder", &topic, None).await?; + let mut pair = (proposer, responder); + read_frame(&mut pair.0).await?; + read_frame(&mut pair.1).await?; + for i in 0..3 { + let frame = message(&topic, &format!("m{i}"), &payload); + assert!( + reaches(&mut pair, &frame).await?, + "paired frame {i} was refused" + ); + } + Ok(()) + } + + /// The connection's own budget still meters frames nothing can receive, + /// and a peer flooding past it is disconnected. + #[tokio::test] + async fn frames_nothing_can_receive_still_spend_the_connection_budget() -> TestResult { + let addr = serve_with_limits( + NativeRelayTransportConfig::new(AUDIENCE), + tight_inbound_limits(), + ) + .await?; + let topic = topic()?; + let mut proposer = connect(addr, "proposer", &topic, None).await?; + + for i in 0..64 { + if send(&mut proposer, &message(&topic, &format!("m{i}"), "x")) + .await + .is_err() + { + break; + } + } + assert!( + expect_closed(&mut proposer).await, + "a one-sided flood over the connection's own limit must be disconnected" + ); + Ok(()) + } + + /// Traffic that does not travel between two connected peers, here frames + /// buffered for a peer that is away, cannot spend the paired reserve; and + /// paired frames refused for an exhausted relay budget disconnect no one. + #[tokio::test] + async fn paired_traffic_keeps_its_reserve_under_one_sided_load() -> TestResult { + let payload = "a".repeat(900); + let frame = charged_bytes(&message(&topic()?, "b0", &payload))?; + + // Three frames beside the reserve, and the reserve holds one more. + let limiter = Arc::new(Limiter::new(RelayLimitsConfig { + global_inbound_byte_rate: RateLimit { + burst: 4 * frame, + interval: Duration::from_secs(3600), + }, + global_paired_reserve_percent: 25, + ..roomy_limits() + })); + assert_eq!(limiter.limits().paired_reserve_bytes(), frame); + let config = NativeRelayTransportConfig::new(AUDIENCE); + let addr = serve_with_limiter(config, Arc::clone(&limiter)).await?; + + let (away, (mut sender, gone)) = session_claiming(addr, 0, None).await?; + drop(gone); + sleep(Duration::from_millis(100)).await; + for id in ["b0", "b1", "b2", "b3"] { + send(&mut sender, &message(&away, id, &payload)).await?; + } + let dropped = wait_for_count(|| limiter.dropped_global_bytes(), u64::from(frame)).await; + assert_eq!( + dropped, + u64::from(frame), + "three buffered frames fit beside the reserve, and the fourth does not" + ); + + let (talking, mut pair) = session_claiming(addr, 1, None).await?; + assert!( + reaches(&mut pair, &message(&talking, "v0", &payload)).await?, + "a paired frame must find the reserve untouched" + ); + assert!( + !reaches(&mut pair, &message(&talking, "v1", &payload)).await?, + "the reserve is a share of the budget, not more of it" + ); + assert_eq!(limiter.dropped_global_bytes(), 2 * u64::from(frame)); + + for client in [&mut sender, &mut pair.0] { + assert!( + timeout(Duration::from_millis(500), expect_closed(client)) + .await + .is_err(), + "an exhausted relay-wide budget alone must not disconnect anyone" + ); + } + Ok(()) + } + + /// Limits under which a source's share holds two of the frames these tests + /// send, and nothing else can refuse one. + fn source_limits(frame: u32) -> RelayLimitsConfig { + RelayLimitsConfig { + source_inbound_byte_rate: RateLimit { + burst: 2 * frame, + interval: Duration::from_secs(3600), + }, + ..roomy_limits() + } + } + + /// Without trusted proxies, the source is the TCP peer: identities + /// connecting from one address share its share, and a client cannot buy + /// a fresh one by naming another address in `X-Real-IP`. + #[tokio::test] + async fn a_direct_peer_cannot_name_its_own_source() -> TestResult { + let payload = "a".repeat(1000); + let frame = charged_bytes(&message(&topic()?, "m0", &payload))?; + let limiter = Arc::new(Limiter::new(source_limits(frame))); + let config = NativeRelayTransportConfig::new(AUDIENCE); + let addr = serve_with_limiter(config, Arc::clone(&limiter)).await?; + + let (topic, mut spender) = session_claiming(addr, 0, None).await?; + for id in ["m0", "m1"] { + assert!(reaches(&mut spender, &message(&topic, id, &payload)).await?); + } + + let (topic, mut same_address) = session_claiming(addr, 1, None).await?; + assert!( + !reaches(&mut same_address, &message(&topic, "s0", &payload)).await?, + "another identity from the same address shares the spent share" + ); + let (topic, mut spoofer) = session_claiming(addr, 2, Some("203.0.113.7")).await?; + assert!( + !reaches(&mut spoofer, &message(&topic, "s1", &payload)).await?, + "X-Real-IP from a peer that is not a trusted proxy must be ignored" + ); + assert_eq!(limiter.dropped_source_bytes(), 2 * u64::from(frame)); + assert_eq!(limiter.dropped_share_bytes(), 0); + + assert!( + timeout(Duration::from_millis(500), expect_closed(&mut spoofer.0)) + .await + .is_err(), + "an exhausted source share alone must not disconnect anyone" + ); + Ok(()) + } + + /// Behind a trusted proxy, the source is the client address the proxy + /// names: one per IPv4 address, one per IPv6 /64. + #[tokio::test] + async fn a_trusted_proxy_names_each_client_source() -> TestResult { + let payload = "a".repeat(1000); + let frame = charged_bytes(&message(&topic()?, "m0", &payload))?; + let limiter = Arc::new(Limiter::new(source_limits(frame))); + let config = NativeRelayTransportConfig::new(AUDIENCE) + .with_trusted_proxies(vec!["127.0.0.1".parse::()?]); + let addr = serve_with_limiter(config, Arc::clone(&limiter)).await?; + + // Each client spends its source's share, which another identity + // behind the same source then finds spent, and a neighbour does not. + let clients = [ + ("198.51.100.1", "198.51.100.1", "198.51.100.2"), + ("2001:db8:1:1::1", "2001:db8:1:1::2", "2001:db8:1:2::1"), + ]; + for (n, (spender, same_source, neighbour)) in (0u8..).step_by(3).zip(clients) { + let (topic, mut pair) = session_claiming(addr, n, Some(spender)).await?; + for id in ["m0", "m1"] { + assert!( + reaches(&mut pair, &message(&topic, id, &payload)).await?, + "{spender} has a share of its own" + ); + } + + let (topic, mut pair) = session_claiming(addr, n + 1, Some(same_source)).await?; + assert!( + !reaches(&mut pair, &message(&topic, "s0", &payload)).await?, + "{same_source} is the same source as {spender}" + ); + + let (topic, mut pair) = session_claiming(addr, n + 2, Some(neighbour)).await?; + assert!( + reaches(&mut pair, &message(&topic, "n0", &payload)).await?, + "{neighbour} is a source of its own" + ); + } + assert_eq!(limiter.dropped_source_bytes(), 2 * u64::from(frame)); + Ok(()) + } + #[tokio::test] async fn frame_over_incoming_size_cap_closes_connection() -> TestResult { let addr = serve_with( diff --git a/crates/relay/src/session.rs b/crates/relay/src/session.rs index fd30728..7b062ca 100644 --- a/crates/relay/src/session.rs +++ b/crates/relay/src/session.rs @@ -259,6 +259,20 @@ impl Slots { } } +/// Where a frame from one peer would go if it were routed now. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum Counterpart { + /// The counterpart is connected: the frame goes to its queue. + Connected, + + /// The counterpart's slot is reserved while it reconnects: the frame is + /// buffered for it, if buffering is enabled. + Reserved, + + /// There is no counterpart: the frame is dropped. + Absent, +} + /// Result of expiring a reservation. pub(crate) struct Expiry { /// Whether the session now has no peers left. @@ -284,6 +298,18 @@ impl Session { self.slots.lock().map_err(|_| SessionError::LockPoisoned) } + /// What a frame from `from` would reach if [`Session::route`]d now. + /// + /// The answer can be stale by the time a frame is routed; it is meant for + /// deciding what a frame costs before it is parsed, not where it goes. + pub(crate) fn counterpart(&self, from: Role) -> Result { + Ok(match self.lock()?.get(from.counterpart()) { + Some(Slot::Connected { .. }) => Counterpart::Connected, + Some(Slot::Disconnected { .. }) => Counterpart::Reserved, + None => Counterpart::Absent, + }) + } + /// Bind a peer to its role slot. /// /// A same-identity reservation is reclaimed. A same-identity live @@ -841,6 +867,31 @@ mod tests { Ok(()) } + #[tokio::test(start_paused = true)] + async fn counterpart_follows_the_other_slot() -> TestResult { + let session = Session::default(); + let (proposer, _) = peer(1); + register(&session, Role::Proposer, proposer)?; + assert_eq!(session.counterpart(Role::Proposer)?, Counterpart::Absent); + + let (responder, _) = peer(2); + register(&session, Role::Responder, responder)?; + assert_eq!(session.counterpart(Role::Proposer)?, Counterpart::Connected); + assert_eq!( + session.counterpart(Role::Responder)?, + Counterpart::Connected + ); + + let grace = Duration::from_secs(10); + disconnect_current(&session, Role::Responder, grace)?; + assert_eq!(session.counterpart(Role::Proposer)?, Counterpart::Reserved); + + tokio::time::advance(grace).await; + session.expire(Role::Responder, ¬ification())?; + assert_eq!(session.counterpart(Role::Proposer)?, Counterpart::Absent); + Ok(()) + } + #[test] fn route_drops_frames_when_outbound_queue_is_full() -> TestResult { let session = Session::default(); diff --git a/crates/relay/src/source.rs b/crates/relay/src/source.rs new file mode 100644 index 0000000..064228d --- /dev/null +++ b/crates/relay/src/source.rs @@ -0,0 +1,341 @@ +//! The address a connection's inbound bytes are accounted to. +//! +//! Identities cost nothing to mint, so the identity share alone cannot stop +//! one client from spending the relay-wide budget under many names. An address +//! is harder to multiply. It is the TCP peer's, unless that peer is a proxy +//! listed in `application.trusted_proxies`: then it is the client address the +//! proxy writes into `X-Real-IP`. The header is ignored from every other peer, +//! so a client that reaches the relay directly cannot choose its own source. + +use std::{ + fmt, + net::{IpAddr, Ipv4Addr, Ipv6Addr}, + str::FromStr, + sync::atomic::{AtomicBool, Ordering}, +}; + +use serde::{Deserialize, Deserializer}; +use thiserror::Error; +use tokio_tungstenite::tungstenite::http::HeaderMap; + +/// Header in which a trusted proxy names the client it forwards. +pub const REAL_IP_HEADER: &str = "x-real-ip"; + +/// Leading bits of an IPv6 address that make up one source. A subscriber is +/// normally assigned a whole /64, so keying on the full address would give one +/// client as many shares as it has addresses. +pub const IPV6_SOURCE_PREFIX: u8 = 64; + +/// Whether an ignored `X-Real-IP` has been reported, so that clients sending +/// one cannot fill the log with the warning. +static IGNORED_HEADER_REPORTED: AtomicBool = AtomicBool::new(false); + +/// A range of addresses, written `address/length` or as a single address. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct IpPrefix { + network: IpAddr, + length: u8, +} + +/// A range that is not an address or an `address/length` pair. +#[derive(Debug, Error)] +#[error("invalid address range {0:?}: expected an IP address or address/length")] +pub struct InvalidPrefix(String); + +impl IpPrefix { + /// The addresses sharing the first `length` bits of `address`, or `None` + /// if `length` is longer than the address. + /// + /// An IPv4-mapped IPv6 range is stored as the IPv4 range it maps, since + /// addresses are compared in their canonical form. + pub fn new(address: IpAddr, length: u8) -> Option { + let (address, length) = match address { + IpAddr::V6(v6) if length >= 96 => match v6.to_ipv4_mapped() { + Some(v4) => (IpAddr::V4(v4), length - 96), + None => (address, length), + }, + _ => (address, length), + }; + if length > bits(address) { + return None; + } + Some(Self { + network: mask(address, length), + length, + }) + } + + /// Number of leading bits the range fixes. + pub fn length(&self) -> u8 { + self.length + } + + /// Whether `address` lies in this range. + pub fn contains(&self, address: IpAddr) -> bool { + let address = address.to_canonical(); + address.is_ipv4() == self.network.is_ipv4() && mask(address, self.length) == self.network + } +} + +impl FromStr for IpPrefix { + type Err = InvalidPrefix; + + fn from_str(text: &str) -> Result { + let invalid = || InvalidPrefix(text.to_owned()); + let (address, length) = match text.trim().split_once('/') { + Some((address, length)) => (address, Some(length)), + None => (text.trim(), None), + }; + let address: IpAddr = address.parse().map_err(|_| invalid())?; + let length = match length { + Some(length) => length.parse().map_err(|_| invalid())?, + None => bits(address), + }; + Self::new(address, length).ok_or_else(invalid) + } +} + +impl fmt::Display for IpPrefix { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{}/{}", self.network, self.length) + } +} + +impl<'de> Deserialize<'de> for IpPrefix { + fn deserialize>(deserializer: D) -> Result { + let text = String::deserialize(deserializer)?; + text.parse().map_err(serde::de::Error::custom) + } +} + +/// Bits in an address of `address`'s family. +fn bits(address: IpAddr) -> u8 { + if address.is_ipv4() { 32 } else { 128 } +} + +/// `address` with every bit after the first `length` cleared. +fn mask(address: IpAddr, length: u8) -> IpAddr { + match address { + IpAddr::V4(v4) => { + let kept = u32::MAX + .checked_shl(32u32.saturating_sub(u32::from(length))) + .unwrap_or(0); + IpAddr::V4(Ipv4Addr::from(u32::from(v4) & kept)) + } + IpAddr::V6(v6) => { + let kept = u128::MAX + .checked_shl(128u32.saturating_sub(u32::from(length))) + .unwrap_or(0); + IpAddr::V6(Ipv6Addr::from(u128::from(v6) & kept)) + } + } +} + +/// What a connection's inbound bytes are accounted to: an IPv4 address, or +/// the /64 holding an IPv6 one. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub struct Source(IpAddr); + +impl Source { + /// The source `address` belongs to. + pub fn of(address: IpAddr) -> Self { + match address.to_canonical() { + v4 @ IpAddr::V4(_) => Self(v4), + v6 @ IpAddr::V6(_) => Self(mask(v6, IPV6_SOURCE_PREFIX)), + } + } + + /// The source of a connection from `peer` that sent `headers` with its + /// handshake. + /// + /// Only a peer in `trusted_proxies` may name the client in `X-Real-IP`. + /// From any other peer the header is ignored and the peer is the source. + /// A trusted proxy that does not name exactly one address is the source + /// itself, so the clients it forwards that way share one share rather + /// than escape it. + pub fn resolve(peer: IpAddr, headers: &HeaderMap, trusted_proxies: &[IpPrefix]) -> Self { + if !trusted_proxies.iter().any(|proxy| proxy.contains(peer)) { + if headers.contains_key(REAL_IP_HEADER) + && !IGNORED_HEADER_REPORTED.swap(true, Ordering::Relaxed) + { + log::warn!( + "ignoring X-Real-IP from a peer outside application.trusted_proxies; \ + if a proxy fronts the relay, list it there, or every client it \ + forwards shares one source budget" + ); + } + return Self::of(peer); + } + match real_ip(headers) { + Some(client) => Self::of(client), + None => { + log::debug!("trusted proxy named no single client address; accounting the proxy"); + Self::of(peer) + } + } + } +} + +/// The address in `X-Real-IP`, if the header appears once and holds nothing +/// but an address. +fn real_ip(headers: &HeaderMap) -> Option { + let mut values = headers.get_all(REAL_IP_HEADER).iter(); + let value = values.next()?; + if values.next().is_some() { + return None; + } + value.to_str().ok()?.trim().parse().ok() +} + +#[cfg(test)] +mod tests { + use super::*; + + use tokio_tungstenite::tungstenite::http::HeaderValue; + + type TestResult = Result<(), Box>; + + fn ip(text: &str) -> IpAddr { + text.parse().expect("test address") + } + + fn headers(real_ips: &[&str]) -> HeaderMap { + let mut headers = HeaderMap::new(); + for real_ip in real_ips { + headers.append( + REAL_IP_HEADER, + HeaderValue::from_str(real_ip).expect("header value"), + ); + } + headers + } + + fn proxies(ranges: &[&str]) -> Vec { + ranges + .iter() + .map(|range| range.parse().expect("test range")) + .collect() + } + + #[test] + fn ranges_parse_as_addresses_or_prefixes() -> TestResult { + let range: IpPrefix = "10.1.2.3/8".parse()?; + assert_eq!(range.to_string(), "10.0.0.0/8", "host bits are cleared"); + assert!(range.contains(ip("10.200.0.1"))); + assert!(!range.contains(ip("11.0.0.1"))); + + let single: IpPrefix = " 192.0.2.7 ".parse()?; + assert_eq!(single.length(), 32); + assert!(single.contains(ip("192.0.2.7"))); + assert!(!single.contains(ip("192.0.2.8"))); + + let v6: IpPrefix = "fd00::/8".parse()?; + assert!(v6.contains(ip("fd12::1"))); + assert!(!v6.contains(ip("fe80::1"))); + assert!( + !v6.contains(ip("10.0.0.1")), + "an IPv6 range holds no IPv4 address" + ); + + let mapped: IpPrefix = "::ffff:10.0.0.0/104".parse()?; + assert_eq!(mapped.to_string(), "10.0.0.0/8"); + + for invalid in ["", "proxy", "10.0.0.0/33", "fd00::/129", "10.0.0.0/", "/8"] { + assert!( + invalid.parse::().is_err(), + "{invalid:?} must be rejected" + ); + } + Ok(()) + } + + #[test] + fn mapped_peers_match_their_ipv4_range() -> TestResult { + let range: IpPrefix = "172.16.0.0/12".parse()?; + assert!(range.contains(ip("::ffff:172.18.0.2"))); + Ok(()) + } + + #[test] + fn ipv6_sources_are_their_64() { + assert_eq!( + Source::of(ip("2001:db8:1:2::1")), + Source::of(ip("2001:db8:1:2:ffff:ffff:ffff:ffff")), + "addresses in one /64 are one source" + ); + assert_ne!( + Source::of(ip("2001:db8:1:2::1")), + Source::of(ip("2001:db8:1:3::1")), + "neighbouring /64s are different sources" + ); + assert_eq!(Source::of(ip("2001:db8:1:2::1")).0, ip("2001:db8:1:2::")); + } + + #[test] + fn ipv4_sources_are_the_address() { + assert_ne!(Source::of(ip("192.0.2.1")), Source::of(ip("192.0.2.2"))); + assert_eq!( + Source::of(ip("::ffff:192.0.2.1")), + Source::of(ip("192.0.2.1")), + "an IPv4-mapped peer is its IPv4 source, not a /64" + ); + } + + #[test] + fn real_ip_from_a_direct_peer_is_ignored() { + let peer = ip("198.51.100.20"); + let spoofed = headers(&["203.0.113.9"]); + + assert_eq!(Source::resolve(peer, &spoofed, &[]), Source::of(peer)); + assert_eq!( + Source::resolve(peer, &spoofed, &proxies(&["10.0.0.0/8", "fd00::/8"])), + Source::of(peer), + "a peer outside the trusted ranges names no one but itself" + ); + } + + #[test] + fn real_ip_from_a_trusted_proxy_names_the_source() { + let trusted = proxies(&["172.16.0.0/12"]); + let proxy = ip("172.18.0.3"); + + assert_eq!( + Source::resolve(proxy, &headers(&["203.0.113.9"]), &trusted), + Source::of(ip("203.0.113.9")) + ); + assert_eq!( + Source::resolve(proxy, &headers(&["2001:db8:5:6::77"]), &trusted), + Source::of(ip("2001:db8:5:6::1")), + "a proxied IPv6 client is its /64 too" + ); + assert_eq!( + Source::resolve( + ip("::ffff:172.18.0.3"), + &headers(&["203.0.113.9"]), + &trusted + ), + Source::of(ip("203.0.113.9")), + "a dual-stack listener sees the proxy as an IPv4-mapped peer" + ); + } + + #[test] + fn a_trusted_proxy_naming_no_single_address_is_the_source() { + let trusted = proxies(&["127.0.0.1"]); + let proxy = ip("127.0.0.1"); + + for named in [ + headers(&[]), + headers(&["not-an-address"]), + headers(&["203.0.113.9, 198.51.100.1"]), + headers(&["203.0.113.9:443"]), + headers(&["203.0.113.9", "198.51.100.1"]), + ] { + assert_eq!( + Source::resolve(proxy, &named, &trusted), + Source::of(proxy), + "{named:?} must not name a source" + ); + } + } +} diff --git a/deployment/configs/relay/base.yaml.template b/deployment/configs/relay/base.yaml.template index f931605..8303216 100644 --- a/deployment/configs/relay/base.yaml.template +++ b/deployment/configs/relay/base.yaml.template @@ -6,6 +6,12 @@ application: port: ${RELAY_PORT} audience: '${APPLICATION_AUDIENCE}' + # nginx writes the client address into X-Real-IP; the relay believes it only + # from these peers. docker-compose.yml publishes no relay port, so the only + # peer on the compose network is nginx, whose address Docker assigns from + # one of these private ranges. If the relay port is ever published, or + # anything else joins its network, narrow this to nginx's address. + trusted_proxies: ['10.0.0.0/8', '172.16.0.0/12', '192.168.0.0/16'] health: port: ${HEALTH_PORT} @@ -48,6 +54,13 @@ limits: identity_inbound_byte_rate: burst: 1048576 interval_seconds: 1 + # Shared by every identity connecting from one address (an IPv6 /64). + source_inbound_byte_rate: + burst: 1048576 + interval_seconds: 1 global_inbound_byte_rate: burst: 16777216 interval_seconds: 1 + # Share of the global burst, in percent, that only frames between two + # connected peers may spend. + global_paired_reserve_percent: 25 From 300bbd15f665f0a76c56572304a86420653970fa Mon Sep 17 00:00:00 2001 From: Kyryl R Date: Wed, 28 Oct 2026 12:52:35 +0200 Subject: [PATCH 6/9] Meter WebSocket wire bytes before fragment parsing --- crates/relay/README.md | 51 +++ crates/relay/configuration/base.yaml | 6 + crates/relay/src/configuration.rs | 51 +++ crates/relay/src/connection.rs | 17 +- crates/relay/src/lib.rs | 1 + crates/relay/src/limits.rs | 73 ++++- crates/relay/src/meter.rs | 331 ++++++++++++++++++++ crates/relay/src/server.rs | 83 ++++- deployment/configs/relay/base.yaml.template | 6 + 9 files changed, 608 insertions(+), 11 deletions(-) create mode 100644 crates/relay/src/meter.rs diff --git a/crates/relay/README.md b/crates/relay/README.md index e0fcaca..5087dbc 100644 --- a/crates/relay/README.md +++ b/crates/relay/README.md @@ -34,6 +34,7 @@ See `configuration/base.yaml` for every key. The ones that matter most: | `session.buffer_max_messages` / `session.buffer_ttl_seconds` | frames buffered for it | | `limits.*` | connection caps, per-identity/per-connection rates, and each identity's and source's share of the global inbound byte budget | | `limits.global_paired_reserve_percent` | part of the global budget only paired sessions may spend | +| `limits.inbound_wire_byte_rate` | rate at which one connection's socket is read, framing and fragments included (see [Wire bytes](#wire-bytes)) | | `connection.*` | frame size, queue capacity, and the handshake, write, ping, pong and idle timeouts | | `connection.outbound_queue_max_bytes` | bytes held for one peer and not yet written (see [Outbound memory](#outbound-memory)) | | `connection.outbound_queue_timeout_seconds` | longest a frame may wait in a peer's outbound queue before that peer is dropped | @@ -141,6 +142,56 @@ Topology limits: clients that share an address. - Budgets are per process. Each relay instance has its own. +### Wire bytes + +The budgets above count what reaches the read loop: complete text messages +and control frames. The WebSocket layer below it does work on every frame, +and a message may arrive in any number of fragments, so a client could send a +text frame without FIN followed by empty continuation frames, six bytes each, +that never form a message and were never charged. + +Every connection's socket is therefore read through its own wire budget, +`limits.inbound_wire_byte_rate` (default 1 MiB per 2 seconds), before the +WebSocket layer parses anything. Every byte read counts once: the upgrade +request, frame headers and masks, control frames, the continuations of a +fragmented message, and the payloads of complete messages. + +- A read takes no more than the budget holds. While the budget is empty the + relay leaves the socket alone and tries again when it has refilled, after at + most a second, so the sender is slowed down by TCP flow control. Nothing is + dropped, and exceeding the wire budget never disconnects anyone. +- The token, pong and idle deadlines keep running meanwhile. A throttled read + only holds up this connection's own frames. +- The wire budget is separate from `inbound_byte_rate`: a complete message's + payload is counted once in each, never twice in one. The default leaves room + for the framing of a peer that stays within `inbound_byte_rate`, so only + overhead-heavy traffic, like a fragment flood, is slowed. Keep it above + `inbound_byte_rate`. +- `limits.inbound_wire_byte_rate` is required, like the other limits; a + configuration file written for an earlier version must add it. + +Residual limits: + +- The wire budget is per connection only. Wire bytes spend no identity, source + or global budget, because frames that reach nobody must not (see above), and + a relay-wide wire budget would let one-sided floods slow paired sessions' + reads. A client's wire throughput therefore scales with its connections, up + to `limits.max_connections` across identities, and only per-connection + limits in front of the relay (nginx `limit_conn`) bound how many one source + holds. +- The budget counts bytes, not frames. An empty continuation costs the relay + several times more work per byte than a large frame, so at the default rate + one connection streaming nothing but empty fragments is read at about + 87,000 frames a second. There is no limit on the number of fragments per + message: tokio-tungstenite has no frame-level read to enforce one. +- Fragments are not messages, so they do not reset the idle deadline; a + connection that streams them and answers the relay's pings lasts until its + token expires, read at its wire rate. +- A peer that sends faster than its wire budget queues its own pongs and + messages behind the bytes waiting for it. If nothing complete gets through + within `connection.pong_timeout_seconds` of a keepalive ping, the pong + deadline drops it. + ### Outbound memory Frames waiting for a peer are bounded per peer, in frames and in bytes, and diff --git a/crates/relay/configuration/base.yaml b/crates/relay/configuration/base.yaml index ffd3630..b9d3303 100644 --- a/crates/relay/configuration/base.yaml +++ b/crates/relay/configuration/base.yaml @@ -45,6 +45,12 @@ limits: inbound_byte_rate: burst: 524288 interval_seconds: 2 + # Every byte read from one connection's socket: framing, control frames and + # fragments included. Over it the socket is read more slowly, nothing is + # dropped. Keep it above inbound_byte_rate, which counts message payloads. + inbound_wire_byte_rate: + burst: 1048576 + interval_seconds: 2 identity_inbound_byte_rate: burst: 1048576 interval_seconds: 1 diff --git a/crates/relay/src/configuration.rs b/crates/relay/src/configuration.rs index b8f176d..5a5b557 100644 --- a/crates/relay/src/configuration.rs +++ b/crates/relay/src/configuration.rs @@ -221,6 +221,10 @@ pub struct LimitsSettings { /// Inbound byte rate per connection. pub inbound_byte_rate: RateSettings, + /// Rate at which one connection's socket is read, framing, control + /// frames and fragments included. + pub inbound_wire_byte_rate: RateSettings, + /// Inbound byte rate one identity may consume from the relay-wide budget. pub identity_inbound_byte_rate: RateSettings, @@ -354,6 +358,10 @@ impl Settings { &self.limits.inbound_message_rate, ), ("limits.inbound_byte_rate", &self.limits.inbound_byte_rate), + ( + "limits.inbound_wire_byte_rate", + &self.limits.inbound_wire_byte_rate, + ), ( "limits.identity_inbound_byte_rate", &self.limits.identity_inbound_byte_rate, @@ -403,6 +411,12 @@ impl Settings { "limits.inbound_byte_rate", self.limits.inbound_byte_rate.burst, ), + // A frame larger than this burst would still be read, only + // throttled part-way through. + ( + "limits.inbound_wire_byte_rate", + self.limits.inbound_wire_byte_rate.burst, + ), ( "limits.global_inbound_byte_rate", self.limits.global_inbound_byte_rate.burst, @@ -489,6 +503,7 @@ impl From<&Settings> for RelayLimitsConfig { connect_rate: (&settings.limits.connect_rate).into(), inbound_message_rate: (&settings.limits.inbound_message_rate).into(), inbound_byte_rate: (&settings.limits.inbound_byte_rate).into(), + inbound_wire_byte_rate: (&settings.limits.inbound_wire_byte_rate).into(), identity_inbound_byte_rate: (&settings.limits.identity_inbound_byte_rate).into(), source_inbound_byte_rate: (&settings.limits.source_inbound_byte_rate).into(), global_inbound_byte_rate: (&settings.limits.global_inbound_byte_rate).into(), @@ -711,6 +726,42 @@ mod tests { Ok(()) } + #[test] + fn wire_rate_is_shipped_above_the_payload_rate_and_validated() -> TestResult { + let settings = shipped()?; + let limits = RelayLimitsConfig::from(&settings); + assert!( + limits.inbound_wire_byte_rate.burst > limits.inbound_byte_rate.burst + && limits.inbound_wire_byte_rate.interval == limits.inbound_byte_rate.interval, + "the wire rate must leave room for the framing of payloads within their rate" + ); + assert_eq!( + RelayLimitsConfig::from(&deployed()?).inbound_wire_byte_rate, + limits.inbound_wire_byte_rate + ); + + let mut settings = shipped()?; + settings.limits.inbound_wire_byte_rate.interval_seconds = 0; + assert!(matches!( + settings.validate(), + Err(ConfigurationError::Zero { + setting: "limits.inbound_wire_byte_rate" + }) + )); + + let mut settings = shipped()?; + settings.limits.inbound_wire_byte_rate.burst = + u32::try_from(settings.connection.max_incoming_frame_bytes)? - 1; + assert!(matches!( + settings.validate(), + Err(ConfigurationError::BurstBelowFrameCap { + setting: "limits.inbound_wire_byte_rate", + .. + }) + )); + Ok(()) + } + #[test] fn zero_outbound_budget_or_timeout_is_rejected() -> TestResult { let mut settings = shipped()?; diff --git a/crates/relay/src/connection.rs b/crates/relay/src/connection.rs index 9f8054f..67cd71f 100644 --- a/crates/relay/src/connection.rs +++ b/crates/relay/src/connection.rs @@ -22,7 +22,8 @@ use wallet_session_core::{codec, connection::ConnectionQuery}; use crate::{ auth::{Identity, IdentityError, unix_now}, config::NativeRelayTransportConfig, - limits::{Bucket, InboundBytes, Limiter, Traffic}, + limits::{Bucket, InboundBytes, Limiter, RateLimit, Traffic}, + meter::MeteredStream, session::{Counterpart, Outbound, Registration, Role, Session, SessionError}, source::Source, }; @@ -31,8 +32,9 @@ use wallet_session_core::{ frame::{RelayFrame, RelayFrameBody}, }; -/// Server-side relay WebSocket stream. -pub(crate) type RelayStream = WebSocketStream; +/// Server-side relay WebSocket stream, read through the connection's wire +/// budget. +pub(crate) type RelayStream = WebSocketStream>; /// Relay frame protocol version accepted from clients. pub(crate) const RELAY_FRAME_VERSION: u8 = 1; @@ -133,15 +135,21 @@ impl Connection { /// /// The client's source is the TCP peer, or the address a peer listed in /// the trusted proxies names in `X-Real-IP` (see [`Source::resolve`]). + /// + /// Every byte read from `stream` from here on, the upgrade request + /// included, is admitted at `wire_rate` before tungstenite parses it (see + /// [`crate::meter`]). #[allow(clippy::result_large_err)] pub async fn handshake( stream: TcpStream, config: NativeRelayTransportConfig, + wire_rate: RateLimit, ) -> Result<(ConnectionQuery, Client, RelayStream), ConnectionError> { let peer = stream .peer_addr() .map_err(ConnectionError::PeerAddress)? .ip(); + let stream = MeteredStream::new(stream, wire_rate); let audience = config.audience.as_str(); let now = unix_now(); let mut parsed: Option> = None; @@ -1136,7 +1144,8 @@ mod tests { let config = NativeRelayTransportConfig::new(AUDIENCE) .with_handshake_timeout(Duration::from_millis(100)); - let outcome = Connection::handshake(server, config).await; + let wire_rate = crate::limits::RelayLimitsConfig::default().inbound_wire_byte_rate; + let outcome = Connection::handshake(server, config, wire_rate).await; assert!(matches!(outcome, Err(ConnectionError::HandshakeTimeout))); Ok(()) } diff --git a/crates/relay/src/lib.rs b/crates/relay/src/lib.rs index cc73595..b60f61b 100644 --- a/crates/relay/src/lib.rs +++ b/crates/relay/src/lib.rs @@ -4,6 +4,7 @@ pub mod configuration; pub mod connection; pub mod health; pub mod limits; +pub mod meter; pub mod registry; pub mod server; pub mod session; diff --git a/crates/relay/src/limits.rs b/crates/relay/src/limits.rs index 23849c5..7aae3b1 100644 --- a/crates/relay/src/limits.rs +++ b/crates/relay/src/limits.rs @@ -79,6 +79,16 @@ pub struct RelayLimitsConfig { /// Inbound byte rate per connection. pub inbound_byte_rate: RateLimit, + /// Rate at which one connection's socket is read, counting every byte + /// that crosses the wire before the WebSocket layer parses it: the + /// upgrade request, frame headers and masks, control frames and the + /// continuation frames of fragmented messages. A connection over it is + /// read more slowly, not disconnected. Keep it above + /// [`RelayLimitsConfig::inbound_byte_rate`], which counts only the + /// payloads of complete messages, so that framing never slows a peer + /// that stays within that. + pub inbound_wire_byte_rate: RateLimit, + /// Inbound byte rate one identity may consume from the relay-wide budget. /// Sized so that several identities together can exhaust /// [`RelayLimitsConfig::global_inbound_byte_rate`], but one cannot. @@ -118,6 +128,10 @@ impl Default for RelayLimitsConfig { burst: 512 * 1024, interval: Duration::from_secs(2), }, + inbound_wire_byte_rate: RateLimit { + burst: 1024 * 1024, + interval: Duration::from_secs(2), + }, identity_inbound_byte_rate: RateLimit { burst: 1024 * 1024, interval: Duration::from_secs(1), @@ -187,11 +201,7 @@ impl Bucket { /// returning whether the attempt is allowed. fn allow_n_above(&mut self, cost: u32, floor: u32, rate: RateLimit, now: Instant) -> bool { let cost = f64::from(cost.clamp(1, rate.burst.max(1))); - let elapsed = now.saturating_duration_since(self.last); - let interval = rate.interval.as_secs_f64().max(f64::EPSILON); - let refill = elapsed.as_secs_f64() / interval * f64::from(rate.burst); - self.tokens = (self.tokens + refill).min(f64::from(rate.burst)); - self.last = now; + self.refill(rate, now); if self.tokens - cost >= f64::from(floor) { self.tokens -= cost; @@ -200,6 +210,39 @@ impl Bucket { false } } + + /// Add what has accrued since the last update, up to the burst. + fn refill(&mut self, rate: RateLimit, now: Instant) { + let elapsed = now.saturating_duration_since(self.last); + let interval = rate.interval.as_secs_f64().max(f64::EPSILON); + let refill = elapsed.as_secs_f64() / interval * f64::from(rate.burst); + self.tokens = (self.tokens + refill).min(f64::from(rate.burst)); + self.last = now; + } + + /// Refill, and return the whole tokens available now. + pub(crate) fn available(&mut self, rate: RateLimit, now: Instant) -> u32 { + self.refill(rate, now); + // Truncates, and saturates at both ends. + self.tokens as u32 + } + + /// Spend `cost` of the tokens [`Bucket::available`] reported. + pub(crate) fn spend(&mut self, cost: u32) { + self.tokens -= f64::from(cost); + } + + /// Time from the last refill until the bucket holds `tokens`, or its whole + /// burst if that is smaller. + pub(crate) fn time_until(&self, tokens: u32, rate: RateLimit) -> Duration { + let burst = rate.burst.max(1); + let missing = f64::from(tokens.clamp(1, burst)) - self.tokens; + if missing <= 0.0 { + return Duration::ZERO; + } + let seconds = missing * rate.interval.as_secs_f64() / f64::from(burst); + Duration::try_from_secs_f64(seconds).unwrap_or(Duration::MAX) + } } /// Per-key token buckets, with the time of the last full prune. @@ -610,6 +653,26 @@ mod tests { assert!(bucket.allow_n(30, rate, Instant::now())); } + #[tokio::test(start_paused = true)] + async fn bucket_reports_what_it_holds_and_when_it_will_hold_more() { + let rate = RateLimit { + burst: 100, + interval: Duration::from_secs(10), + }; + let mut bucket = Bucket::full(rate.burst, Instant::now()); + + assert_eq!(bucket.available(rate, Instant::now()), 100); + assert_eq!(bucket.time_until(50, rate), Duration::ZERO); + bucket.spend(100); + assert_eq!(bucket.available(rate, Instant::now()), 0); + // 10 tokens per second, and never more than the burst to wait for. + assert_eq!(bucket.time_until(50, rate), Duration::from_secs(5)); + assert_eq!(bucket.time_until(1_000, rate), Duration::from_secs(10)); + + tokio::time::advance(Duration::from_millis(2_500)).await; + assert_eq!(bucket.available(rate, Instant::now()), 25); + } + #[tokio::test(start_paused = true)] async fn oversized_cost_is_clamped_to_burst() { let rate = RateLimit { diff --git a/crates/relay/src/meter.rs b/crates/relay/src/meter.rs new file mode 100644 index 0000000..61599c1 --- /dev/null +++ b/crates/relay/src/meter.rs @@ -0,0 +1,331 @@ +//! Wire-level admission for client sockets. +//! +//! The read loop charges complete messages to the inbound budgets, but +//! tungstenite does work on every frame below it: a fragmented message is +//! reassembled from continuation frames that never surface as messages, so a +//! client could stream empty continuations, six bytes each, with nothing ever +//! charged. [`MeteredStream`] sits between the socket and tungstenite and +//! admits bytes against the connection's wire budget before tungstenite sees +//! them: the upgrade request, every frame header and mask, control frames and +//! continuations included, each byte once. +//! +//! A read never takes more than the budget holds. While it is empty the +//! stream leaves the socket alone until the budget has refilled, so the +//! kernel buffers fill and TCP flow control slows the sender down to the +//! budget. Nothing is dropped and the meter never closes a connection: a +//! throttled read returns `Pending`, and the read loop goes on polling its +//! token, pong and idle deadlines meanwhile. Every wait is a timer of at most +//! a second, so a throttled read always resumes. +//! +//! The wire budget is the connection's own, separate from the message-level +//! budgets the read loop charges, per connection and shared. A complete +//! message's payload is therefore not charged twice to one bucket, and bytes +//! that reach no counterpart still spend no shared budget. + +use std::{ + io, + pin::Pin, + task::{Context, Poll, ready}, + time::Duration, +}; + +use tokio::{ + io::{AsyncRead, AsyncWrite, ReadBuf}, + time::{Instant, Sleep}, +}; + +use crate::limits::{Bucket, RateLimit}; + +/// Bytes a throttled read waits to have available, or the whole burst if that +/// is smaller, so a throttled stream is read in chunks rather than a few bytes +/// per wakeup. +const WAKE_BYTES: u32 = 4096; + +/// Shortest wait before a throttled read tries again. +const MIN_WAIT: Duration = Duration::from_millis(1); + +/// Longest wait before a throttled read tries again, however far off the +/// refill it waits for is. +const MAX_WAIT: Duration = Duration::from_secs(1); + +/// A socket whose reads are admitted against a wire byte budget. +#[derive(Debug)] +pub struct MeteredStream { + inner: S, + budget: Bucket, + rate: RateLimit, + /// Wakes a throttled read once the budget has refilled. + throttle: Option>>, +} + +/// What the wire budget allows a read now. +enum Admission { + /// Read at most this many bytes. + Read(u32), + + /// Read nothing before this. + WaitUntil(Instant), +} + +impl MeteredStream { + /// Admit reads from `inner` at `rate`, starting with a full burst. + pub(crate) fn new(inner: S, rate: RateLimit) -> Self { + Self { + inner, + budget: Bucket::full(rate.burst, Instant::now()), + rate, + throttle: None, + } + } + + fn admit(&mut self, now: Instant) -> Admission { + let available = self.budget.available(self.rate, now); + if available > 0 { + return Admission::Read(available); + } + let wait = self.budget.time_until(WAKE_BYTES, self.rate); + Admission::WaitUntil(now + wait.clamp(MIN_WAIT, MAX_WAIT)) + } +} + +impl AsyncRead for MeteredStream { + fn poll_read( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &mut ReadBuf<'_>, + ) -> Poll> { + let this = self.get_mut(); + if buf.remaining() == 0 { + return Poll::Ready(Ok(())); + } + + // Never an empty successful read while throttled: tungstenite would + // take that for the end of the stream. + let admitted = loop { + if let Some(throttle) = this.throttle.as_mut() { + ready!(throttle.as_mut().poll(cx)); + this.throttle = None; + } + match this.admit(Instant::now()) { + Admission::Read(bytes) => break bytes, + Admission::WaitUntil(deadline) => { + this.throttle = Some(Box::pin(tokio::time::sleep_until(deadline))); + } + } + }; + + // No more than was admitted, so tungstenite never parses a byte the + // budget has not paid for. + let limit = usize::try_from(admitted) + .unwrap_or(usize::MAX) + .min(buf.remaining()); + let read = { + let mut limited = ReadBuf::new(buf.initialize_unfilled_to(limit)); + ready!(Pin::new(&mut this.inner).poll_read(cx, &mut limited))?; + limited.filled().len() + }; + buf.advance(read); + this.budget.spend(u32::try_from(read).unwrap_or(u32::MAX)); + Poll::Ready(Ok(())) + } +} + +/// Writes are the relay's own and pass through unmetered. +impl AsyncWrite for MeteredStream { + fn poll_write( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &[u8], + ) -> Poll> { + Pin::new(&mut self.get_mut().inner).poll_write(cx, buf) + } + + fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + Pin::new(&mut self.get_mut().inner).poll_flush(cx) + } + + fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + Pin::new(&mut self.get_mut().inner).poll_shutdown(cx) + } + + fn poll_write_vectored( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + bufs: &[io::IoSlice<'_>], + ) -> Poll> { + Pin::new(&mut self.get_mut().inner).poll_write_vectored(cx, bufs) + } + + fn is_write_vectored(&self) -> bool { + self.inner.is_write_vectored() + } +} + +#[cfg(test)] +mod tests { + use futures_util::StreamExt; + use tokio::io::{AsyncReadExt, AsyncWriteExt, DuplexStream}; + use tokio_tungstenite::{ + WebSocketStream, + tungstenite::{Message, protocol::Role}, + }; + + use super::*; + + type TestResult = Result<(), Box>; + + const HOUR: Duration = Duration::from_secs(3600); + + /// First bytes of the frames these tests send: FIN bit and opcode. + const TEXT_FIRST: u8 = 0x01; + const CONTINUATION: u8 = 0x00; + const CONTINUATION_FINAL: u8 = 0x80; + const PING: u8 = 0x89; + const PONG: u8 = 0x8a; + + /// A client frame as it crosses the wire: masked, with a payload short + /// enough for the one-byte length. + fn frame(first: u8, payload: &[u8]) -> Vec { + const MASK: [u8; 4] = [0x11, 0x22, 0x33, 0x44]; + let length = u8::try_from(payload.len()).expect("short payload"); + assert!(length < 126, "short payload"); + let mut wire = vec![first, 0x80 | length]; + wire.extend_from_slice(&MASK); + wire.extend( + payload + .iter() + .zip(MASK.iter().cycle()) + .map(|(byte, mask)| byte ^ mask), + ); + wire + } + + /// A text message split into an empty first frame, `continuations` empty + /// continuation frames, and a final one carrying `text`. + fn fragmented(continuations: usize, text: &[u8]) -> Vec { + let mut wire = frame(TEXT_FIRST, b""); + for _ in 0..continuations { + wire.extend(frame(CONTINUATION, b"")); + } + wire.extend(frame(CONTINUATION_FINAL, text)); + wire + } + + /// The relay end of a socket, read through a meter at `rate`, and the + /// client end. + fn metered(rate: RateLimit) -> (MeteredStream, DuplexStream) { + let (client, relay) = tokio::io::duplex(1 << 20); + (MeteredStream::new(relay, rate), client) + } + + /// The relay's WebSocket over a metered socket. + async fn websocket( + stream: MeteredStream, + ) -> WebSocketStream> { + WebSocketStream::from_raw_socket(stream, Role::Server, None).await + } + + fn left(stream: &mut MeteredStream) -> u32 { + let rate = stream.rate; + stream.budget.available(rate, Instant::now()) + } + + #[tokio::test(start_paused = true)] + async fn every_wire_byte_is_charged_once() -> TestResult { + let rate = RateLimit { + burst: 1 << 20, + interval: HOUR, + }; + let (stream, mut client) = metered(rate); + let mut websocket = websocket(stream).await; + + // Empty continuations, with control frames between the fragments of + // the message as the protocol allows. + let mut wire = frame(TEXT_FIRST, b""); + for _ in 0..1000 { + wire.extend(frame(CONTINUATION, b"")); + } + wire.extend(frame(PING, b"ping")); + wire.extend(frame(PONG, b"")); + wire.extend(frame(CONTINUATION_FINAL, b"hello")); + client.write_all(&wire).await?; + + assert!(matches!(websocket.next().await, Some(Ok(Message::Ping(_))))); + assert!(matches!(websocket.next().await, Some(Ok(Message::Pong(_))))); + let Some(Ok(Message::Text(text))) = websocket.next().await else { + return Err("the fragments must make up one text message".into()); + }; + assert_eq!(text.as_str(), "hello"); + + // Headers, masks, empty frames and payloads, each byte once. + let wire = u32::try_from(wire.len())?; + assert_eq!(left(websocket.get_mut()), rate.burst - wire); + Ok(()) + } + + /// The bypass the Cloud audit reproduced, scaled down: empty continuation + /// frames were parsed without bound and charged nothing. + #[tokio::test(start_paused = true)] + async fn empty_continuations_are_throttled_to_the_wire_rate() -> TestResult { + let rate = RateLimit { + burst: 600, + interval: Duration::from_secs(1), + }; + let (stream, mut client) = metered(rate); + let mut websocket = websocket(stream).await; + + // 6,014 bytes: the 600-byte burst at once, then 600 bytes a second. + client.write_all(&fragmented(1000, b"hi")).await?; + let start = Instant::now(); + let Some(Ok(Message::Text(text))) = websocket.next().await else { + return Err("the fragments must make up one text message".into()); + }; + assert_eq!(text.as_str(), "hi"); + + let elapsed = start.elapsed(); + assert!( + (Duration::from_secs(9)..=Duration::from_secs(11)).contains(&elapsed), + "the message must take about ten seconds at the wire rate, took {elapsed:?}" + ); + Ok(()) + } + + #[tokio::test(start_paused = true)] + async fn a_read_takes_no_more_than_the_budget_holds() -> TestResult { + let rate = RateLimit { + burst: 10, + interval: Duration::from_secs(1), + }; + let (mut stream, mut client) = metered(rate); + client.write_all(&[7; 25]).await?; + + let mut buf = [0u8; 64]; + let start = Instant::now(); + assert_eq!(stream.read(&mut buf).await?, 10, "the burst, at once"); + assert_eq!(start.elapsed(), Duration::ZERO); + assert_eq!(stream.read(&mut buf).await?, 10, "a refill later"); + assert!(start.elapsed() >= Duration::from_secs(1)); + assert_eq!(stream.read(&mut buf).await?, 5, "what is left"); + + // The end of the stream still reads as such. + drop(client); + assert_eq!(stream.read(&mut buf).await?, 0); + Ok(()) + } + + #[tokio::test(start_paused = true)] + async fn writes_pass_through_uncharged() -> TestResult { + let rate = RateLimit { + burst: 10, + interval: HOUR, + }; + let (mut stream, mut client) = metered(rate); + stream.write_all(&[1; 100]).await?; + + let mut buf = [0u8; 100]; + client.read_exact(&mut buf).await?; + assert_eq!(buf, [1; 100]); + assert_eq!(left(&mut stream), rate.burst); + Ok(()) + } +} diff --git a/crates/relay/src/server.rs b/crates/relay/src/server.rs index fd44101..f554a38 100644 --- a/crates/relay/src/server.rs +++ b/crates/relay/src/server.rs @@ -67,7 +67,9 @@ pub async fn handle( }; let grace = config.session_grace; - let (query, client, mut websocket) = Connection::handshake(stream, config.clone()).await?; + let wire_rate = limiter.limits().inbound_wire_byte_rate; + let (query, client, mut websocket) = + Connection::handshake(stream, config.clone(), wire_rate).await?; let identity = &client.identity; if !limiter.allow_connect(&identity.did) { @@ -192,7 +194,15 @@ mod tests { use tokio::time::timeout; use tokio_tungstenite::{ MaybeTlsStream, WebSocketStream, connect_async, - tungstenite::{Bytes, Message, client::IntoClientRequest, http::HeaderValue}, + tungstenite::{ + Bytes, Message, + client::IntoClientRequest, + http::HeaderValue, + protocol::frame::{ + Frame, FrameHeader, + coding::{Data, OpCode}, + }, + }, }; use wallet_session_core::frame::{RelayFrame, RelayFrameBody}; use wallet_session_core::message::Topic; @@ -1823,6 +1833,7 @@ mod tests { RelayLimitsConfig { inbound_message_rate: unlimited, inbound_byte_rate: unlimited, + inbound_wire_byte_rate: unlimited, global_inbound_byte_rate: unlimited, ..RelayLimitsConfig::default() }, @@ -1887,6 +1898,7 @@ mod tests { RelayLimitsConfig { inbound_message_rate: unlimited, inbound_byte_rate: unlimited, + inbound_wire_byte_rate: unlimited, identity_inbound_byte_rate: unlimited, source_inbound_byte_rate: unlimited, global_inbound_byte_rate: unlimited, @@ -2146,6 +2158,7 @@ mod tests { interval: Duration::from_secs(3600), }, inbound_byte_rate: roomy, + inbound_wire_byte_rate: roomy, identity_inbound_byte_rate: roomy, source_inbound_byte_rate: roomy, global_inbound_byte_rate: roomy, @@ -2533,6 +2546,72 @@ mod tests { Ok(()) } + /// A data frame as a client hands it to its WebSocket, which masks it on + /// the way out. + fn raw_frame(opcode: Data, is_final: bool, payload: &[u8]) -> Message { + Message::Frame(Frame::from_payload( + FrameHeader { + is_final, + rsv1: false, + rsv2: false, + rsv3: false, + opcode: OpCode::Data(opcode), + mask: None, + }, + Bytes::copy_from_slice(payload), + )) + } + + /// The fragment bypass the Cloud audit reproduced, at a bounded scale. A + /// FIN=0 text frame and 100,000 empty continuations, six bytes each, were + /// parsed in milliseconds with nothing charged, because none of them is a + /// message. Every byte now passes the connection's wire budget before it + /// is parsed, so the flood is read at that rate, and the sender is slowed + /// down, not disconnected. + #[tokio::test] + async fn empty_continuation_flood_is_read_at_the_wire_rate() -> TestResult { + const CONTINUATIONS: usize = 10_000; + // The flood is about 60 KB, nearly four bursts. + let wire_rate = RateLimit { + burst: 16 * 1024, + interval: Duration::from_secs(1), + }; + let addr = serve_with_limits( + NativeRelayTransportConfig::new(AUDIENCE), + RelayLimitsConfig { + inbound_wire_byte_rate: wire_rate, + ..RelayLimitsConfig::default() + }, + ) + .await?; + let topic = topic()?; + let (mut proposer, mut responder) = ready_pair(addr, &topic).await?; + + let started = tokio::time::Instant::now(); + proposer.feed(raw_frame(Data::Text, false, b"")).await?; + for _ in 0..CONTINUATIONS { + proposer.feed(raw_frame(Data::Continue, false, b"")).await?; + } + let last = serde_json::to_string(&message(&topic, "f0", "after-fragments"))?; + proposer + .send(raw_frame(Data::Continue, true, last.as_bytes())) + .await?; + + // The message is whole, so every fragment was read, in order. + assert_eq!(read_frame(&mut responder).await?.id, "f0"); + // What the burst did not cover takes about 2.7 seconds. + let elapsed = started.elapsed(); + assert!( + elapsed >= Duration::from_secs(2), + "empty continuations must be read at the wire rate, took {elapsed:?}" + ); + + // The sender is still connected, and its next frame goes through. + send(&mut proposer, &message(&topic, "f1", "still-connected")).await?; + assert_eq!(read_frame(&mut responder).await?.id, "f1"); + Ok(()) + } + #[tokio::test] async fn frame_over_incoming_size_cap_closes_connection() -> TestResult { let addr = serve_with( diff --git a/deployment/configs/relay/base.yaml.template b/deployment/configs/relay/base.yaml.template index 8303216..8de83df 100644 --- a/deployment/configs/relay/base.yaml.template +++ b/deployment/configs/relay/base.yaml.template @@ -51,6 +51,12 @@ limits: inbound_byte_rate: burst: 524288 interval_seconds: 2 + # Every byte read from one connection's socket: framing, control frames and + # fragments included. Over it the socket is read more slowly, nothing is + # dropped. Keep it above inbound_byte_rate, which counts message payloads. + inbound_wire_byte_rate: + burst: 1048576 + interval_seconds: 2 identity_inbound_byte_rate: burst: 1048576 interval_seconds: 1 From 00a24391a7eb99f668b026a3cf9815b9c7225272 Mon Sep 17 00:00:00 2001 From: Kyryl R Date: Wed, 28 Oct 2026 12:55:51 +0200 Subject: [PATCH 7/9] Cap backlog redelivery credit and stabilize buffered-frame tests --- crates/relay/src/connection.rs | 7 ++++++- crates/relay/src/server.rs | 23 +++++++++++++++++------ 2 files changed, 23 insertions(+), 7 deletions(-) diff --git a/crates/relay/src/connection.rs b/crates/relay/src/connection.rs index 67cd71f..5e7dde9 100644 --- a/crates/relay/src/connection.rs +++ b/crates/relay/src/connection.rs @@ -673,8 +673,13 @@ async fn read_loop( backlog_pending = match backlog.tx.upgrade() { Some(tx) => { let drain = session.drain_backlog(query.role, &tx)?; + // A continuously refilled backlog must not bank an + // unbounded exemption for later inbound messages. + let credit_cap = u32::try_from(config.session_buffer_max_messages) + .unwrap_or(u32::MAX); redelivery_credit = redelivery_credit - .saturating_add(u32::try_from(drain.delivered).unwrap_or(u32::MAX)); + .saturating_add(u32::try_from(drain.delivered).unwrap_or(u32::MAX)) + .min(credit_cap); !drain.finished } // The slot let go of the queue: this connection was replaced. diff --git a/crates/relay/src/server.rs b/crates/relay/src/server.rs index f554a38..0dd7e68 100644 --- a/crates/relay/src/server.rs +++ b/crates/relay/src/server.rs @@ -602,8 +602,12 @@ mod tests { /// Drain a client until the relay closes it (or the connection errors). async fn expect_closed(client: &mut Client) -> bool { + expect_closed_within(client, Duration::from_secs(2)).await + } + + async fn expect_closed_within(client: &mut Client, wait: Duration) -> bool { loop { - match timeout(Duration::from_secs(2), client.next()).await { + match timeout(wait, client.next()).await { Ok(None) | Ok(Some(Err(_))) | Ok(Some(Ok(Message::Close(_)))) => return true, Ok(Some(Ok(_))) => {} Err(_) => return false, @@ -662,6 +666,12 @@ mod tests { drop(responder); let _ = timeout(SYNC_TIMEOUT, responder_handler).await??; + // Keep the three writes in order before the reconnect, rather than + // leaving the last one in the client's Nagle buffer. + if let MaybeTlsStream::Plain(stream) = proposer.get_mut() { + stream.set_nodelay(true)?; + } + // Three frames while the responder is gone, but only two may be kept. for (id, payload) in [("m1", "one"), ("m2", "two"), ("m3", "three")] { send(&mut proposer, &message(&topic, id, payload)).await?; @@ -2249,14 +2259,15 @@ mod tests { // after that. let mut hog = connect(addr, "proposer", &topic()?, None).await?; let _ = send_pings(&mut hog, MAX_CONSECUTIVE_RATE_DROPS + 8, PING_PAYLOAD).await; + let limit = u64::from(PING_PAYLOAD) * u64::from(MAX_CONSECUTIVE_RATE_DROPS); + let dropped = wait_for_count(|| limiter.dropped_share_bytes(), limit + 1).await; assert!( - expect_closed(&mut hog).await, - "an identity flooding control frames over its share must be disconnected" + dropped > limit, + "the relay must process enough dropped pings to reach the disconnect threshold" ); - let limit = u64::from(PING_PAYLOAD) * u64::from(MAX_CONSECUTIVE_RATE_DROPS); assert!( - limiter.dropped_share_bytes() > limit, - "the dropped pings must be counted against the identity share" + expect_closed_within(&mut hog, Duration::from_secs(10)).await, + "an identity flooding control frames over its share must be disconnected" ); assert_eq!(limiter.dropped_global_bytes(), 0); Ok(()) From 4a68d50089f006b714e77896424f42d3d9fefbb0 Mon Sep 17 00:00:00 2001 From: Kyryl R Date: Wed, 28 Oct 2026 13:23:53 +0200 Subject: [PATCH 8/9] Align relay topology guidance with ingress limits --- crates/relay/README.md | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/crates/relay/README.md b/crates/relay/README.md index 5087dbc..84441bd 100644 --- a/crates/relay/README.md +++ b/crates/relay/README.md @@ -135,11 +135,14 @@ Topology limits: list to nginx's address. - nginx resolves the client address it forwards from `X-Forwarded-For` only for its own `TRUSTED_PROXY_CIDR`, so the source behind a load balancer or - CDN is only as accurate as that setting. nginx's own per-address limits are - keyed on the full address, not the /64. + CDN is only as accurate as that setting. In the shipped deployment, nginx + keys its IPv6 per-source limits on the /64, like the relay. - A relay without a proxy listens directly (`configuration/production.yaml`), and its sources are the TCP peers. Nothing in the relay can tell apart - clients that share an address. + clients that share an address. Its wire-byte throttle is per connection; + put a per-source connection cap in front of a directly exposed relay so + free identities cannot multiply fragment-processing work up to the global + connection limit. - Budgets are per process. Each relay instance has its own. ### Wire bytes From b8ee61e5b57187e2da14abac3db1c46612803b19 Mon Sep 17 00:00:00 2001 From: Kyryl R Date: Wed, 28 Oct 2026 14:03:33 +0200 Subject: [PATCH 9/9] Preserve relay fairness and backlog across edge cases --- crates/relay/README.md | 17 ++- crates/relay/src/configuration.rs | 89 +++++++++++--- crates/relay/src/connection.rs | 105 ++++++++++++++++- crates/relay/src/limits.rs | 77 +++++++++---- crates/relay/src/session.rs | 186 +++++++++++++++++++++++++++--- crates/relay/src/source.rs | 105 +++++++++++++++++ crates/relay/src/telemetry.rs | 26 ++++- 7 files changed, 545 insertions(+), 60 deletions(-) diff --git a/crates/relay/README.md b/crates/relay/README.md index 84441bd..42529bc 100644 --- a/crates/relay/README.md +++ b/crates/relay/README.md @@ -45,9 +45,12 @@ address, a byte-rate burst below the maximum incoming frame (no frame could ever fit it), an outbound byte budget below the maximum incoming frame, or a maximum incoming frame above `u32::MAX`. Every byte bucket and outbound budget must hold a whole frame, and so must the part of the global burst outside the -paired reserve, which must be below 100 percent. A trusted proxy range of -length 0 (`0.0.0.0/0`, `::/0`) is refused: it would let every peer choose its -own source. +paired reserve, which must be below 100 percent. The wire budget may not be +below `inbound_byte_rate`, in burst or in rate. Trusted proxy ranges that +together hold every IPv4 or every IPv6 address (`0.0.0.0/0`, `::/0`, or +`0.0.0.0/1` with `128.0.0.0/1`) are refused: they would let every peer choose +its own source. Narrower ranges are taken as given, so a range that holds +client addresses lets those clients choose theirs. ### Inbound fairness @@ -123,8 +126,9 @@ Topology limits: addresses, or with routed IPv6 space (a /48 holds 65,536 /64s), gets one share for each; the global budget and the paired reserve remain the bound on what such a client can take. As for identities, at most 65,536 sources - are tracked at once; past that, idle ones are pruned first and then - arbitrary ones are forgotten, starting again with a full share. + are tracked at once. Past that, a share is forgotten only once it has + refilled completely (checked at most once a second); until one has, new + sources spend one share between them instead of displacing a tracked one. - Trusting a proxy is only safe when the proxy overwrites `X-Real-IP` and nothing else reachable from a trusted address can connect to the relay. The compose deployment trusts the private ranges Docker assigns its networks @@ -169,7 +173,8 @@ fragmented message, and the payloads of complete messages. payload is counted once in each, never twice in one. The default leaves room for the framing of a peer that stays within `inbound_byte_rate`, so only overhead-heavy traffic, like a fragment flood, is slowed. Keep it above - `inbound_byte_rate`. + `inbound_byte_rate`: the relay refuses to start with a wire burst or rate + below it. - `limits.inbound_wire_byte_rate` is required, like the other limits; a configuration file written for an earlier version must add it. diff --git a/crates/relay/src/configuration.rs b/crates/relay/src/configuration.rs index 5a5b557..c6350c7 100644 --- a/crates/relay/src/configuration.rs +++ b/crates/relay/src/configuration.rs @@ -14,7 +14,7 @@ use thiserror::Error; use crate::{ config::NativeRelayTransportConfig, limits::{RateLimit, RelayLimitsConfig}, - source::IpPrefix, + source::{IpPrefix, covered_family}, }; pub const CONFIGURATION_DIRECTORY: &str = "configuration"; @@ -69,7 +69,19 @@ pub enum ConfigurationError { UnreservedBelowFrameCap { unreserved: u32, frame: u32 }, #[error( - "application.trusted_proxies must not contain {0}: it would let every peer \ + "limits.inbound_wire_byte_rate ({wire_burst} bytes per {wire_seconds} s) must be at \ + least limits.inbound_byte_rate ({payload_burst} bytes per {payload_seconds} s), in \ + burst and in rate" + )] + WireRateBelowPayloadRate { + wire_burst: u32, + wire_seconds: u64, + payload_burst: u32, + payload_seconds: u64, + }, + + #[error( + "application.trusted_proxies must not cover {0}: it would let every peer \ name its own source" )] TrustsEveryPeer(IpPrefix), @@ -273,13 +285,8 @@ impl Settings { self.application_address(), )); } - if let Some(range) = self - .application - .trusted_proxies - .iter() - .find(|range| range.length() == 0) - { - return Err(ConfigurationError::TrustsEveryPeer(*range)); + if let Some(family) = covered_family(&self.application.trusted_proxies) { + return Err(ConfigurationError::TrustsEveryPeer(family)); } for (setting, value) in [ @@ -430,6 +437,22 @@ impl Settings { }); } } + // Every payload byte crosses the wire too, framing on top: a wire + // budget below the payload budget would throttle the reads, pongs + // included, of a peer within its limits. Rates are compared over the + // product of the two intervals. + let wire = &self.limits.inbound_wire_byte_rate; + let payload = &self.limits.inbound_byte_rate; + let wire_bytes = u128::from(wire.burst) * u128::from(payload.interval_seconds); + let payload_bytes = u128::from(payload.burst) * u128::from(wire.interval_seconds); + if wire.burst < payload.burst || wire_bytes < payload_bytes { + return Err(ConfigurationError::WireRateBelowPayloadRate { + wire_burst: wire.burst, + wire_seconds: wire.interval_seconds, + payload_burst: payload.burst, + payload_seconds: payload.interval_seconds, + }); + } // Frames buffered for a reconnecting peer cannot spend the reserve, // so what is left of the burst must still hold a whole frame. let limits = RelayLimitsConfig::from(self); @@ -681,13 +704,25 @@ mod tests { #[test] fn fairness_settings_that_break_the_boundary_are_rejected() -> TestResult { - for everyone in ["0.0.0.0/0", "::/0"] { + let everyone: [&[&str]; 4] = [ + &["0.0.0.0/0"], + &["::/0"], + &["10.0.0.0/8", "0.0.0.0/1", "128.0.0.0/1"], + &["::/1", "8000::/1"], + ]; + for ranges in everyone { let mut settings = shipped()?; - settings.application.trusted_proxies = vec![everyone.parse()?]; - assert!(matches!( - settings.validate(), - Err(ConfigurationError::TrustsEveryPeer(_)) - )); + settings.application.trusted_proxies = ranges + .iter() + .map(|range| range.parse()) + .collect::>()?; + assert!( + matches!( + settings.validate(), + Err(ConfigurationError::TrustsEveryPeer(_)) + ), + "{ranges:?} trusts every peer of a family" + ); } let mut settings = shipped()?; @@ -762,6 +797,30 @@ mod tests { Ok(()) } + #[test] + fn wire_rate_below_the_payload_rate_is_rejected() -> TestResult { + let below = |burst: u32, interval_seconds: u64| -> Result { + let mut settings = shipped()?; + settings.limits.inbound_wire_byte_rate = RateSettings { + burst, + interval_seconds, + }; + Ok(matches!( + settings.validate(), + Err(ConfigurationError::WireRateBelowPayloadRate { .. }) + )) + }; + // The shipped payload rate: 512 KiB per 2 seconds. + let payload = shipped()?.limits.inbound_byte_rate; + assert_eq!((payload.burst, payload.interval_seconds), (512 * 1024, 2)); + + assert!(below(512 * 1024 - 1, 2)?, "a smaller burst"); + assert!(below(1024 * 1024, 5)?, "a larger burst refilling slower"); + assert!(!below(512 * 1024, 2)?, "the payload rate itself"); + assert!(!below(1024 * 1024, 4)?, "a larger burst at the same rate"); + Ok(()) + } + #[test] fn zero_outbound_budget_or_timeout_is_rejected() -> TestResult { let mut settings = shipped()?; diff --git a/crates/relay/src/connection.rs b/crates/relay/src/connection.rs index 5e7dde9..133b74b 100644 --- a/crates/relay/src/connection.rs +++ b/crates/relay/src/connection.rs @@ -726,11 +726,13 @@ async fn read_loop( // Only a frame something can receive, its counterpart or the // buffer kept for it while it reconnects, spends the shared // budgets and is parsed. Any other frame is dropped here, at - // the cost of this connection's own budget alone. + // the cost of this connection's own budget alone. That budget + // accepted it, so it ends a run of rate-limit drops. let traffic = match session.counterpart(query.role)? { Counterpart::Connected => Traffic::Paired, Counterpart::Reserved if buffering => Traffic::Unpaired, Counterpart::Reserved | Counterpart::Absent => { + consecutive_drops = 0; log::trace!("dropping frame on topic {topic}: counterpart not connected"); continue; } @@ -956,8 +958,10 @@ mod tests { use super::*; use crate::auth::test_utils::{self, signing_key}; - use crate::limits::OutboundBudget; - use tokio::time::timeout; + use crate::limits::{OutboundBudget, RelayLimitsConfig}; + use crate::session::PeerHandle; + use tokio::{net::TcpListener, time::timeout}; + use tokio_tungstenite::connect_async; use wallet_session_core::message::Topic; const AUDIENCE: &str = "wss://relay.example.com"; @@ -1300,4 +1304,99 @@ mod tests { timeout(Duration::from_secs(1), writer).await??; Ok(()) } + + /// Limits under which only the connection's own byte budget can drop a + /// text frame. It holds `burst` bytes and does not refill within a test. + fn own_byte_budget(burst: u32) -> RelayLimitsConfig { + RelayLimitsConfig { + inbound_message_rate: RateLimit { + burst: 1000, + interval: Duration::from_secs(3600), + }, + inbound_byte_rate: RateLimit { + burst, + interval: Duration::from_secs(86_400), + }, + ..RelayLimitsConfig::default() + } + } + + /// Serve one proposer that has no counterpart under `limits`, have it send + /// `messages` and then close, and return how its connection ended. + async fn serve_alone( + limits: RelayLimitsConfig, + messages: Vec, + ) -> Result, Box> { + let listener = TcpListener::bind("127.0.0.1:0").await?; + let addr = listener.local_addr()?; + let now = unix_now(); + let auth = test_utils::token(&signing_key(1), AUDIENCE, now, now + 300); + let url = format!("ws://{addr}/?auth={auth}&role=proposer&topic={}", topic()?); + let client = tokio::spawn(connect_async(url)); + + let (stream, _) = listener.accept().await?; + let config = NativeRelayTransportConfig::new(AUDIENCE); + let wire_rate = limits.inbound_wire_byte_rate; + let (query, peer, websocket) = + Connection::handshake(stream, config.clone(), wire_rate).await?; + let (mut client, _) = client.await??; + + let limiter = Arc::new(Limiter::new(limits)); + let session = Arc::new(Session::default()); + let (tx, rx) = mpsc::channel(config.outbound_queue_capacity); + let handle = PeerHandle { + identity: peer.identity.did.clone(), + tx: tx.clone(), + budget: Arc::new(OutboundBudget::new(config.outbound_queue_max_bytes)), + }; + let ready = RelayFrame { + v: RELAY_FRAME_VERSION, + id: "ready".to_owned(), + topic: topic()?, + body: RelayFrameBody::SessionReady, + }; + let permit = limiter.acquire_identity(&peer.identity.did); + let registration = session.register(Role::Proposer, handle, &ready, permit)?; + let connection = Connection::new(query, peer, session, rx, tx, config, limiter); + let served = tokio::spawn(connection.run(websocket, registration)); + + // Once the relay has dropped the peer, sending fails. + for message in messages { + let _ = client.send(message).await; + } + let _ = client.close(None).await; + Ok(timeout(Duration::from_secs(5), served).await??) + } + + /// A frame nothing can receive is dropped unparsed, but one within the + /// connection's own budget still ends a run of rate-limit drops: only + /// drops in a row disconnect a peer, with or without a counterpart. + #[tokio::test] + async fn accepted_frames_without_a_counterpart_end_a_run_of_drops() -> TestResult { + const BURST: u32 = 1000; + let fits = Message::text("s"); + // Once a byte is spent, a frame of the whole burst never fits again. + let over = Message::text("o".repeat(usize::try_from(BURST)?)); + let tolerated = usize::try_from(MAX_CONSECUTIVE_RATE_DROPS)?; + + let mut interleaved = vec![fits.clone()]; + for _ in 0..2 { + interleaved.extend(vec![over.clone(); tolerated]); + interleaved.push(fits.clone()); + } + let outcome = serve_alone(own_byte_budget(BURST), interleaved).await?; + assert!( + outcome.is_ok(), + "drops broken up by accepted frames must not disconnect: {outcome:?}" + ); + + let mut in_a_row = vec![fits]; + in_a_row.extend(vec![over; tolerated + 1]); + let outcome = serve_alone(own_byte_budget(BURST), in_a_row).await?; + assert!( + matches!(outcome, Err(ConnectionError::RateLimited)), + "drops in a row must still disconnect: {outcome:?}" + ); + Ok(()) + } } diff --git a/crates/relay/src/limits.rs b/crates/relay/src/limits.rs index 7aae3b1..90e8dab 100644 --- a/crates/relay/src/limits.rs +++ b/crates/relay/src/limits.rs @@ -16,6 +16,10 @@ use crate::source::Source; /// Maximum number of identities or sources tracked by one keyed rate limiter. const MAX_RATE_KEYS: usize = 65_536; +/// Least time between two scans of a full keyed rate limiter for buckets +/// that have refilled, which bounds the work a stream of new keys can cause. +const MIN_PRUNE_SPACING: Duration = Duration::from_secs(1); + /// Token-bucket parameters. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub struct RateLimit { @@ -211,6 +215,13 @@ impl Bucket { } } + /// Refill, and return whether the bucket holds its whole burst again, as + /// a new bucket would. + fn is_full(&mut self, rate: RateLimit, now: Instant) -> bool { + self.refill(rate, now); + self.tokens >= f64::from(rate.burst) + } + /// Add what has accrued since the last update, up to the burst. fn refill(&mut self, rate: RateLimit, now: Instant) { let elapsed = now.saturating_duration_since(self.last); @@ -249,36 +260,43 @@ impl Bucket { #[derive(Debug)] struct BucketMap { buckets: HashMap, + + /// Shared by the keys that arrive while every tracked bucket is still + /// refilling, so they neither take a tracked key's place nor each get a + /// full bucket of their own. + overflow: Bucket, last_prune: Instant, } impl BucketMap { - fn new(now: Instant) -> Self { + fn new(rate: RateLimit, now: Instant) -> Self { Self { buckets: HashMap::new(), + overflow: Bucket::full(rate.burst, now), last_prune: now, } } } +/// The bucket of `key`, tracking it if there is room. +/// +/// A tracked key keeps its bucket for as long as the bucket is still +/// refilling: forgetting it would hand the key a full bucket again. A full +/// map makes room only by forgetting buckets that have refilled completely, +/// and while none has, a new key is charged to the shared overflow bucket. fn bucket_for<'a, K: Eq + Hash + Clone>( state: &'a mut BucketMap, key: &K, rate: RateLimit, now: Instant, ) -> &'a mut Bucket { - if state.buckets.len() >= MAX_RATE_KEYS { - if now.saturating_duration_since(state.last_prune) >= rate.interval { - state - .buckets - .retain(|_, bucket| now.saturating_duration_since(bucket.last) < rate.interval); + if state.buckets.len() >= MAX_RATE_KEYS && !state.buckets.contains_key(key) { + if now.saturating_duration_since(state.last_prune) >= MIN_PRUNE_SPACING { + state.buckets.retain(|_, bucket| !bucket.is_full(rate, now)); state.last_prune = now; } - while state.buckets.len() >= MAX_RATE_KEYS { - let Some(evicted) = state.buckets.keys().next().cloned() else { - break; - }; - state.buckets.remove(&evicted); + if state.buckets.len() >= MAX_RATE_KEYS { + return &mut state.overflow; } } state @@ -310,9 +328,9 @@ impl Limiter { limits, total: Mutex::new(0), per_identity: Mutex::new(HashMap::new()), - connect_rate: Mutex::new(BucketMap::new(now)), - identity_bytes: Mutex::new(BucketMap::new(now)), - source_bytes: Mutex::new(BucketMap::new(now)), + connect_rate: Mutex::new(BucketMap::new(limits.connect_rate, now)), + identity_bytes: Mutex::new(BucketMap::new(limits.identity_inbound_byte_rate, now)), + source_bytes: Mutex::new(BucketMap::new(limits.source_inbound_byte_rate, now)), inbound_bytes: Mutex::new(Bucket::full(limits.global_inbound_byte_rate.burst, now)), dropped_share_bytes: AtomicU64::new(0), dropped_source_bytes: AtomicU64::new(0), @@ -702,16 +720,35 @@ mod tests { } #[tokio::test(start_paused = true)] - async fn connect_rate_map_stays_bounded() { - let limiter = Limiter::new(RelayLimitsConfig::default()); + async fn a_full_rate_map_forgets_no_bucket_that_is_still_refilling() { + let interval = Duration::from_secs(10); + let limiter = Limiter::new(RelayLimitsConfig { + connect_rate: RateLimit { burst: 1, interval }, + ..RelayLimitsConfig::default() + }); - for seed in 0..MAX_RATE_KEYS as u32 + 128 { - limiter.allow_connect(&key_from_u32(seed)); + // Every tracked identity spends its only attempt. + for seed in 0..MAX_RATE_KEYS as u32 { + assert!(limiter.allow_connect(&key_from_u32(seed))); } - assert!(limiter.tracked_identities() <= MAX_RATE_KEYS); - // A fresh identity can still attempt a connection. + // Past the bound, new identities share one bucket instead of taking + // the place of a tracked one, which would come back with a full one. + tokio::time::advance(MIN_PRUNE_SPACING).await; assert!(limiter.allow_connect(&key_from_u32(u32::MAX))); + assert!(!limiter.allow_connect(&key_from_u32(u32::MAX - 1))); + assert_eq!(limiter.tracked_identities(), MAX_RATE_KEYS); + for seed in 0..MAX_RATE_KEYS as u32 { + assert!( + !limiter.allow_connect(&key_from_u32(seed)), + "identity {seed} got its attempt back before its bucket refilled" + ); + } + + // Buckets that have refilled make room for new identities again. + tokio::time::advance(interval).await; + assert!(limiter.allow_connect(&key_from_u32(u32::MAX - 2))); + assert_eq!(limiter.tracked_identities(), 1); } #[tokio::test(start_paused = true)] diff --git a/crates/relay/src/session.rs b/crates/relay/src/session.rs index 7b062ca..8014173 100644 --- a/crates/relay/src/session.rs +++ b/crates/relay/src/session.rs @@ -155,8 +155,8 @@ fn prune_expired(frames: &mut VecDeque, now: Instant) { } /// Charge `bytes` to `budget` for a frame joining `frames`, dropping the -/// oldest of them if that makes room (spec §3: on overflow the oldest frame -/// is discarded). +/// oldest of them if that makes room, by count below `max_frames` and then by +/// bytes (spec §3: on overflow the oldest frame is discarded). /// /// Every frame in `frames` is charged to `budget`. When the frame would not /// fit even with all of them gone, because the rest of the budget is held by @@ -165,6 +165,7 @@ fn make_room( frames: &mut VecDeque, budget: &Arc, bytes: usize, + max_frames: usize, ) -> Option { // Charges are only taken under the session lock, which the caller holds, // so the queue's share can only shrink while this runs. @@ -173,6 +174,9 @@ fn make_room( if queued.saturating_add(bytes) > budget.capacity() { return None; } + while frames.len() >= max_frames.max(1) { + frames.pop_front(); + } loop { if let Some(charge) = budget.try_charge(bytes) { return Some(charge); @@ -489,6 +493,11 @@ impl Session { /// Only acts if the slot still belongs to the connection owning `tx`, so a /// replaced connection tearing down cannot release its successor's slot. /// The reservation buffers within a budget as large as the connection's. + /// + /// The backlog the connection had not yet handed to its queue stays + /// buffered for the peer, with the expiry it was buffered with, and is + /// charged to the reservation's budget instead. What the queue held goes + /// with the connection. pub(crate) fn disconnect( &self, role: Role, @@ -497,7 +506,7 @@ impl Session { ) -> Result { let mut slots = self.lock()?; - let Some(Slot::Connected { peer, .. }) = slots.get(role) else { + let Some(Slot::Connected { peer, backlog, .. }) = slots.get_mut(role) else { return Ok(false); }; let owns_slot = tx.upgrade().is_some_and(|tx| peer.tx.same_channel(&tx)); @@ -507,11 +516,15 @@ impl Session { let identity = peer.identity.clone(); let budget = Arc::new(OutboundBudget::new(peer.budget.capacity())); + let mut buffered = std::mem::take(backlog); + let now = Instant::now(); + prune_expired(&mut buffered, now); + recharge(&mut buffered, &budget); *slots.get_mut(role) = Some(Slot::Disconnected { identity, - deadline: Instant::now() + grace, - buffered: VecDeque::new(), + deadline: now + grace, + buffered, budget, }); @@ -566,7 +579,7 @@ impl Session { /// relay memory grown (spec §3). A full buffer or backlog drops its oldest /// frames to make room, by count and by bytes; a frame that cannot fit /// beside what the outbound queue holds is dropped, like one that finds - /// the queue full. + /// the queue full, and displaces none of them. pub(crate) fn route( &self, from: Role, @@ -585,10 +598,7 @@ impl Session { if let Some(Slot::Connected { peer, backlog, .. }) = slots.get_mut(target) { prune_expired(backlog, now); if !backlog.is_empty() { - while backlog.len() >= max_buffered.max(1) { - backlog.pop_front(); - } - let Some(charge) = make_room(backlog, &peer.budget, bytes) else { + let Some(charge) = make_room(backlog, &peer.budget, bytes, max_buffered) else { log::debug!("dropping frame for peer whose outbound byte budget is full"); return Ok(()); }; @@ -618,10 +628,7 @@ impl Session { }) => { if max_buffered > 0 && !buffer_ttl.is_zero() { prune_expired(buffered, now); - if buffered.len() >= max_buffered { - buffered.pop_front(); - } - match make_room(buffered, budget, bytes) { + match make_room(buffered, budget, bytes, max_buffered) { Some(charge) => buffered.push_back(BufferedFrame { frame, expires_at: now + buffer_ttl, @@ -1866,6 +1873,66 @@ mod tests { Ok(()) } + /// A message frame whose payload alone is `bytes` long. + fn large_frame(topic: &Topic, bytes: usize) -> RelayFrame { + RelayFrame { + body: RelayFrameBody::Message { + payload: "a".repeat(bytes), + }, + ..frame_with_id(topic, "large") + } + } + + #[tokio::test(start_paused = true)] + async fn a_frame_too_large_beside_the_queue_leaves_a_full_backlog_alone() -> TestResult { + let topic = topic()?; + let one = size(&frame_with_id(&topic, "m1")); + let (session, _proposer_rx) = responder_away_with(&["m1", "m2", "m3"])?; + + let (responder, mut responder_rx) = peer_with(2, 2, 3 * one); + let tx = responder.tx.clone(); + register(&session, Role::Responder, responder)?; + let drain = session.drain_backlog(Role::Responder, &tx)?; + assert_eq!(drain.delivered, 1); + + // The backlog is at its count cap of two, and with m1 queued the + // frame cannot fit even if the backlog went: it must not cost m2. + session.route(Role::Proposer, large_frame(&topic, 2 * one), 2, BUFFER_TTL)?; + assert_eq!( + received_ids(&session, Role::Responder, &tx, &mut responder_rx).await?, + ["ready", "m1", "m2", "m3"] + ); + Ok(()) + } + + #[tokio::test(start_paused = true)] + async fn a_frame_larger_than_the_reconnect_buffer_displaces_nothing() -> TestResult { + let topic = topic()?; + let one = size(&frame_with_id(&topic, "m1")); + let session = Session::default(); + let (proposer, _proposer_rx) = peer(1); + let (responder, _) = peer_with(2, PEER_QUEUE, 2 * one); + register(&session, Role::Proposer, proposer)?; + register(&session, Role::Responder, responder)?; + disconnect_current(&session, Role::Responder, Duration::from_secs(10))?; + for id in ["m1", "m2"] { + session.route(Role::Proposer, frame_with_id(&topic, id), 2, BUFFER_TTL)?; + } + + // The buffer is full by count and by bytes, and the frame is larger + // than its whole budget: it is dropped alone. + session.route(Role::Proposer, large_frame(&topic, 2 * one), 2, BUFFER_TTL)?; + + let (responder, mut responder_rx) = peer(2); + let tx = responder.tx.clone(); + register(&session, Role::Responder, responder)?; + assert_eq!( + received_ids(&session, Role::Responder, &tx, &mut responder_rx).await?, + ["ready", "m1", "m2"] + ); + Ok(()) + } + #[tokio::test(start_paused = true)] async fn a_frame_that_displaces_the_whole_backlog_is_queued_directly() -> TestResult { let one = size(&frame_with_id(&topic()?, "m1")); @@ -1972,4 +2039,95 @@ mod tests { ); Ok(()) } + + /// Bytes held by the reservation on `role`'s slot, if it is reserved. + fn reserved_bytes(session: &Session, role: Role) -> Result { + Ok(match session.lock()?.get(role) { + Some(Slot::Disconnected { budget, .. }) => budget.held(), + _ => 0, + }) + } + + #[tokio::test(start_paused = true)] + async fn a_backlog_outlives_its_connection_in_the_reservation() -> TestResult { + let topic = topic()?; + let early = frame_with_id(&topic, "early"); + let late = frame_with_id(&topic, "late"); + let both = size(&early) + size(&late); + let session = Session::default(); + let (proposer, _proposer_rx) = peer(1); + let (responder, _) = peer(2); + register(&session, Role::Proposer, proposer)?; + register(&session, Role::Responder, responder)?; + disconnect_current(&session, Role::Responder, Duration::from_secs(10))?; + + // t=0: buffered while the responder is away, expires at t=5. + let ttl = Duration::from_secs(5); + session.route(Role::Proposer, early, MAX_BUFFERED, ttl)?; + + // t=2: the responder takes it back, and "late" queues behind it, + // expiring at t=7. + tokio::time::advance(Duration::from_secs(2)).await; + let (first, _first_rx) = peer(2); + let first_budget = Arc::clone(&first.budget); + register(&session, Role::Responder, first)?; + session.route(Role::Proposer, late, MAX_BUFFERED, ttl)?; + assert_eq!(first_budget.held(), both); + + // The connection ends, its queue timeout say, before it drained any + // of its backlog: the reservation takes the backlog and its bytes. + disconnect_current(&session, Role::Responder, Duration::from_secs(10))?; + assert_eq!( + first_budget.held(), + 0, + "the connection that ended holds no bytes" + ); + assert_eq!(reserved_bytes(&session, Role::Responder)?, both); + + // t=6: only "late" is still within its TTL. + tokio::time::advance(Duration::from_secs(4)).await; + let (second, mut second_rx) = peer(2); + let second_tx = second.tx.clone(); + register(&session, Role::Responder, second)?; + assert_eq!( + received_ids(&session, Role::Responder, &second_tx, &mut second_rx).await?, + ["ready", "late"], + "the backlog survives the connection, with the expiry it was buffered with" + ); + Ok(()) + } + + #[tokio::test(start_paused = true)] + async fn a_lapsed_reservation_keeps_its_backlog_from_another_identity() -> TestResult { + let (session, _proposer_rx) = responder_away_with(&["m1", "m2"])?; + let (first, _first_rx) = peer(2); + register(&session, Role::Responder, first)?; + disconnect_current(&session, Role::Responder, Duration::from_secs(5))?; + let reservation = match session.lock()?.get(Role::Responder) { + Some(Slot::Disconnected { budget, .. }) => Arc::clone(budget), + _ => return Err("the responder's slot must be reserved".into()), + }; + assert!( + reservation.held() > 0, + "the reservation took the backlog over" + ); + + // The grace window lapses well within the frames' TTL, and another + // identity takes the role. + tokio::time::advance(Duration::from_secs(6)).await; + let (other, mut other_rx) = peer(3); + let other_tx = other.tx.clone(); + register(&session, Role::Responder, other)?; + assert_eq!( + received_ids(&session, Role::Responder, &other_tx, &mut other_rx).await?, + ["ready"], + "frames kept for the previous owner must not reach another identity" + ); + assert_eq!( + reservation.held(), + 0, + "the discarded frames give back their bytes" + ); + Ok(()) + } } diff --git a/crates/relay/src/source.rs b/crates/relay/src/source.rs index 064228d..1260cae 100644 --- a/crates/relay/src/source.rs +++ b/crates/relay/src/source.rs @@ -75,6 +75,67 @@ impl IpPrefix { let address = address.to_canonical(); address.is_ipv4() == self.network.is_ipv4() && mask(address, self.length) == self.network } + + /// The first and last address of the range, as integers. + fn span(&self) -> (u128, u128) { + let (first, host_bits) = match self.network { + IpAddr::V4(v4) => ( + u128::from(u32::from(v4)), + 32u32.saturating_sub(u32::from(self.length)), + ), + IpAddr::V6(v6) => ( + u128::from(v6), + 128u32.saturating_sub(u32::from(self.length)), + ), + }; + let hosts = u128::MAX.checked_shr(128 - host_bits).unwrap_or(0); + (first, first | hosts) + } +} + +/// The whole of an address family, as a range, if `ranges` together hold +/// every address a peer of that family can connect from. Trusting them would +/// let every such peer name its own source, as trusting `0.0.0.0/0` would. +/// +/// Peers are compared in canonical form, so an IPv6 list need not hold the +/// IPv4-mapped block to take in every IPv6 peer. +pub fn covered_family(ranges: &[IpPrefix]) -> Option { + let v4 = ranges + .iter() + .filter(|range| range.network.is_ipv4()) + .map(IpPrefix::span); + if covers_all(v4, u128::from(u32::MAX)) { + return IpPrefix::new(Ipv4Addr::UNSPECIFIED.into(), 0); + } + + let mapped = u128::from(Ipv6Addr::new(0, 0, 0, 0, 0, 0xffff, 0, 0)); + let v6 = ranges + .iter() + .filter(|range| range.network.is_ipv6()) + .map(IpPrefix::span) + .chain([(mapped, mapped | u128::from(u32::MAX))]); + if covers_all(v6, u128::MAX) { + return IpPrefix::new(Ipv6Addr::UNSPECIFIED.into(), 0); + } + None +} + +/// Whether `spans`, inclusive, together hold every integer up to `last`. +fn covers_all(spans: impl Iterator, last: u128) -> bool { + let mut spans: Vec<_> = spans.collect(); + spans.sort_unstable(); + // The lowest integer none of the spans seen so far holds. + let mut next = 0u128; + for (first, end) in spans { + if first > next { + return false; + } + match end.checked_add(1) { + Some(after) if after <= last => next = next.max(after), + _ => return true, + } + } + false } impl FromStr for IpPrefix { @@ -249,6 +310,50 @@ mod tests { Ok(()) } + #[test] + fn ranges_that_together_hold_every_peer_are_found() -> TestResult { + let every_v4: IpPrefix = "0.0.0.0/0".parse()?; + let every_v6: IpPrefix = "::/0".parse()?; + assert_eq!(covered_family(&proxies(&["0.0.0.0/0"])), Some(every_v4)); + assert_eq!( + covered_family(&proxies(&["128.0.0.0/1", "fd00::/8", "0.0.0.0/1"])), + Some(every_v4), + "two halves trust every IPv4 peer as /0 does" + ); + assert_eq!( + covered_family(&proxies(&["::ffff:0.0.0.0/97", "128.0.0.0/1"])), + Some(every_v4), + "an IPv4-mapped range is its IPv4 range" + ); + assert_eq!( + covered_family(&proxies(&["8000::/1", "::/1"])), + Some(every_v6) + ); + + // Every range beside the IPv4-mapped block: no IPv6 peer is left out. + let mapped = u128::from(Ipv6Addr::new(0, 0, 0, 0, 0, 0xffff, 0, 0)); + let around = (0..96u8) + .map(|bit| { + let differs = Ipv6Addr::from(mapped ^ (1u128 << (127 - bit))); + IpPrefix::new(IpAddr::V6(differs), bit + 1) + }) + .collect::>>() + .ok_or("ranges around the mapped block")?; + assert_eq!(covered_family(&around), Some(every_v6)); + assert_eq!(covered_family(&around[1..]), None); + + let partial: [&[&str]; 4] = [ + &["10.0.0.0/8", "172.16.0.0/12", "192.168.0.0/16", "fd00::/8"], + &["0.0.0.0/1", "128.0.0.0/2", "192.0.0.0/3", "::/1"], + &["127.0.0.1", "::1"], + &[], + ]; + for ranges in partial { + assert_eq!(covered_family(&proxies(ranges)), None, "{ranges:?}"); + } + Ok(()) + } + #[test] fn mapped_peers_match_their_ipv4_range() -> TestResult { let range: IpPrefix = "172.16.0.0/12".parse()?; diff --git a/crates/relay/src/telemetry.rs b/crates/relay/src/telemetry.rs index d3ee7ae..3bdfc08 100644 --- a/crates/relay/src/telemetry.rs +++ b/crates/relay/src/telemetry.rs @@ -91,14 +91,15 @@ impl fmt::Display for Bounded { }; let written = write!(out, "{}", self.0); match written { - Err(_) if out.cut => out.out.write_str("..."), + _ if out.cut => out.out.write_str("..."), result => result, } } } /// Passes characters through until `left` bytes are used up, then fails, -/// which also stops formatting the rest of the value. +/// which also stops formatting the rest of the value. Nothing is passed +/// through after a cut, even to a value that writes on regardless. struct Truncating<'a, 'b> { out: &'a mut fmt::Formatter<'b>, left: usize, @@ -116,6 +117,7 @@ impl fmt::Write for Truncating<'_, '_> { }; let Some(left) = self.left.checked_sub(c.len_utf8()) else { self.cut = true; + self.left = 0; return Err(fmt::Error); }; self.left = left; @@ -400,6 +402,26 @@ mod tests { assert!(wide.trim_end_matches("...").chars().all(|c| c == '€')); } + /// Writes on after a write has failed, and reports success. + struct Careless(String); + + impl fmt::Display for Careless { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + let _ = f.write_str(&self.0); + let _ = f.write_str("tail"); + Ok(()) + } + } + + #[test] + fn bounded_ends_at_the_cut_whatever_the_value_does_next() { + // The cut falls inside the last of these, with bytes left over for + // the short writes that follow it. + let wide = "€".repeat(MAX_LOGGED_DISPLAY_BYTES / 3 + 1); + let kept = "€".repeat(MAX_LOGGED_DISPLAY_BYTES / 3); + assert_eq!(bounded(Careless(wide)).to_string(), format!("{kept}...")); + } + fn log(writer: &NonBlocking, record: &str) { let _ = writer.make_writer().write_all(record.as_bytes()); }