feat(relay): bound frame memory globally and connections per IP (#959)

.changeset/relay-frame-memory-budget.md

@@ -0,0 +1,22 @@
+---
+---
+
+feat(relay): bound the relay's total queued frames and the sockets one address may hold
+
+The outbound queue was bounded per socket and the number of sockets was bounded by nothing:
+one address that opened sockets and stopped reading could hold 8 MiB for every socket it kept,
+measured as 384 MiB from the shipped attempt limits, and nothing gave a closed socket's queued
+bytes back. Two bounds now sit beside that one, both recorded in
+`docs/rewrite/decisions/0008-frame-memory-budget.md`:
+
+- `RelayConfig::frame_memory_budget_bytes` (256 MiB) is one process-wide budget every socket's
+  queued payload is counted against, admitted frame by frame at the sink and refunded when a
+  socket's last handle drops.
+- `RelayConfig::max_connections_per_ip` (64) caps the live sockets one address may hold, as a
+  gauge beside the attempt limiter rather than folded into it.
+
+Both refuse the way the policy comment on #959 asks: an immediate `1013` close with a fixed
+reason and its own counter, never a text frame and never a silent drop. See #959.
+
+The TypeScript tree is untouched: no `src/` or `test/` file changes, so the published package
+has no user-visible change.

crates/nostream-relay/src/connection_cap.rs

@@ -0,0 +1,168 @@
+//! The per-IP live-connection cap: a gauge, beside the attempt limiter and not part of it.
+//!
+//! `limits.rs` answers "how fast may this address open connections" over a sliding window;
+//! this answers "how many may it hold at once". The capture has neither question
+//! (`src/@types/settings.ts:163-167`: `ConnectionLimits` is `rateLimits`, `ipWhitelist` and
+//! `ipBlacklist`), so the cap is the port's own policy with no counterpart there (ADR 0008).
+//!
+//! Two things it deliberately is not. It is not a second `RateLimits` store: no window, no
+//! clock, no whitelist, nothing to expire — a count that only `Drop` changes. And it is not
+//! advisory: the connection task holds a [`ConnectionSlot`] for as long as the socket lives,
+//! so an address at the cap is refused before the NIP-42 challenge is sent rather than after
+//! the client has already been spoken to.
+
+use std::collections::HashMap;
+use std::sync::atomic::{AtomicU64, Ordering};
+use std::sync::{Arc, Mutex};
+
+/// The live sockets per source address, with the cap each is checked against (ADR 0008).
+#[derive(Debug)]
+pub struct ConnectionCap {
+    max_per_ip: usize,
+    live: Mutex<HashMap<String, usize>>,
+    refusals: AtomicU64,
+}
+
+impl ConnectionCap {
+    /// A cap of `max_per_ip` sockets per address, with nothing counted.
+    pub fn new(max_per_ip: usize) -> Self {
+        Self {
+            max_per_ip,
+            live: Mutex::new(HashMap::new()),
+            refusals: AtomicU64::new(0),
+        }
+    }
+
+    /// The most sockets one address may hold at once.
+    pub fn max_per_ip(&self) -> usize {
+        self.max_per_ip
+    }
+
+    /// How many sockets `ip` holds now.
+    pub fn live(&self, ip: &str) -> usize {
+        self.live
+            .lock()
+            .unwrap_or_else(|poisoned| poisoned.into_inner())
+            .get(ip)
+            .copied()
+            .unwrap_or(0)
+    }
+
+    /// How many addresses hold a socket: an address that releases its last one leaves no entry.
+    pub fn tracked_addresses(&self) -> usize {
+        self.live
+            .lock()
+            .unwrap_or_else(|poisoned| poisoned.into_inner())
+            .len()
+    }
+
+    /// How many sockets were refused because their address was at the cap.
+    pub fn refusals(&self) -> u64 {
+        self.refusals.load(Ordering::SeqCst)
+    }
+
+    /// Takes a slot for `ip`, or counts the refusal and returns `None` at the cap.
+    ///
+    /// The slot releases on `Drop`, so a caller that holds it for the whole connection
+    /// releases it on every exit — the normal one, a registration failure, or one added
+    /// later — without a decrement at each.
+    pub fn acquire(self: &Arc<Self>, ip: &str) -> Option<ConnectionSlot> {
+        let mut live = self
+            .live
+            .lock()
+            .unwrap_or_else(|poisoned| poisoned.into_inner());
+        let held = live.get(ip).copied().unwrap_or(0);
+
+        if held >= self.max_per_ip {
+            // Nothing is inserted for a refused address: an entry that must not be there
+            // would read as a socket held.
+            drop(live);
+            self.refusals.fetch_add(1, Ordering::SeqCst);
+
+            return None;
+        }
+
+        live.insert(ip.to_owned(), held + 1);
+
+        Some(ConnectionSlot {
+            cap: Arc::clone(self),
+            ip: ip.to_owned(),
+        })
+    }
+}
+
+/// One socket's slot in a [`ConnectionCap`], given back when it drops.
+#[derive(Debug)]
+pub struct ConnectionSlot {
+    cap: Arc<ConnectionCap>,
+    ip: String,
+}
+
+impl Drop for ConnectionSlot {
+    fn drop(&mut self) {
+        let mut live = self
+            .cap
+            .live
+            .lock()
+            .unwrap_or_else(|poisoned| poisoned.into_inner());
+
+        if let Some(held) = live.get_mut(&self.ip) {
+            *held = held.saturating_sub(1);
+
+            if *held == 0 {
+                live.remove(&self.ip);
+            }
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    #![allow(clippy::expect_used)]
+
+    use super::*;
+
+    #[test]
+    fn the_cap_counts_live_slots_and_releases_each_once() {
+        let cap = Arc::new(ConnectionCap::new(2));
+
+        let first = cap.acquire("127.0.0.1").expect("a slot below the cap");
+        let second = cap
+            .acquire("127.0.0.1")
+            .expect("the last slot below the cap");
+
+        assert_eq!(cap.live("127.0.0.1"), 2);
+        assert_eq!(cap.tracked_addresses(), 1);
+
+        assert!(cap.acquire("127.0.0.1").is_none());
+        assert_eq!(cap.refusals(), 1);
+        assert_eq!(cap.live("127.0.0.1"), 2, "a refusal holds no slot");
+
+        // A second address has its own count, and an address at the cap does not affect it.
+        let elsewhere = cap
+            .acquire("10.0.0.1")
+            .expect("another address is unaffected");
+
+        assert_eq!(cap.live("10.0.0.1"), 1);
+        assert_eq!(cap.tracked_addresses(), 2);
+
+        drop(first);
+
+        assert_eq!(cap.live("127.0.0.1"), 1);
+        assert_eq!(cap.tracked_addresses(), 2);
+
+        drop(second);
+
+        assert_eq!(cap.live("127.0.0.1"), 0);
+        assert_eq!(
+            cap.tracked_addresses(),
+            1,
+            "an address that releases its last slot leaves no entry"
+        );
+
+        drop(elsewhere);
+
+        assert_eq!(cap.tracked_addresses(), 0);
+        assert_eq!(cap.refusals(), 1);
+    }
+}

crates/nostream-relay/src/handlers/subscribe.rs

@@ -268,6 +268,26 @@ async fn fetch_and_send(mut read: InFlightRead, filters: &[Filter], context: &Ha
             return;
         }
 
+        // A refused frame is never queued, and the connection loop is about to close this
+        // socket whatever this read does. Encoding the rest of the stream would only discard
+        // it, and with a shared budget the loop can be a whole blocked write away from
+        // running, so the window is not short. The sink must not latch on a budget refusal —
+        // a frame that fits later is admitted — which is why the stop lives here.
+        //
+        // The read then parks instead of unwinding: it has nothing left to send, and parking
+        // keeps the close's promptness a property of the refusal signal rather than of this
+        // handler happening to finish. On the refusal path the loop's `abort_all` drops the
+        // parked task before the watch below ever resolves, so the `end_read` after it is the
+        // path a client `CLOSE` or a re-`REQ` takes, not the refusal's. It is the same call
+        // every other exit of this function makes, and it is safe beside the loop's own
+        // `close` because it removes the entry only while the generation it carries still
+        // matches (`connection.rs:290-297`): whichever of the two runs second does nothing.
+        if context.sink.refusal().is_some() {
+            let _ = read.cancellation().changed().await;
+            context.connection.end_read(read);
+            return;
+        }
+
         let record = tokio::select! {
             changed = read.cancellation().changed() => {
                 // An aborted read sends no frame at all, not even EOSE.
@@ -408,3 +428,105 @@ fn can_subscribe(
 
     None
 }
+
+#[cfg(test)]
+mod tests {
+    #![allow(clippy::unwrap_used, clippy::expect_used, clippy::panic)]
+
+    use std::time::Duration;
+
+    use nostream_core::clock::TestClock;
+    use nostream_core::event::{EventId, Pubkey, Signature};
+    use nostream_core::nip42::Nip42Session;
+
+    use super::*;
+    use crate::connection::Connection;
+    use crate::relay_key::RelayKeySource;
+    use crate::sink::{FrameBudget, FrameSink, Refusal};
+    use crate::testing::{Fixture, settings};
+
+    /// An event the fixture can hold, signed by `secret`.
+    fn signed_event(secret: &[u8; 32], created_at: i64) -> nostream_core::Event {
+        let mut event = nostream_core::Event {
+            id: EventId::from_byte_array([0; 32]),
+            pubkey: Pubkey::from_byte_array([0; 32]),
+            created_at,
+            kind: Kind(1),
+            tags: Vec::new(),
+            content: "stored".to_owned(),
+            sig: Signature::from_byte_array([0; 64]),
+        };
+
+        nostream_core::crypto::sign_event(secret, &mut event).expect("sign the test event");
+
+        event
+    }
+
+    #[tokio::test]
+    async fn a_refused_sink_stops_the_read_instead_of_encoding_the_rest_of_it() {
+        // Three stored events behind a per-item gate, and a 64-byte budget: every event frame
+        // is refused, so the read must neither take another record nor encode one. The gate is
+        // the measure — two permits left means one record was taken and the stream was not
+        // drained — and the queue is the other: a read that ran on to the end would leave the
+        // 12-byte `EOSE` in it, which does fit.
+        let fixture = Fixture::new(settings());
+        let gate = fixture.events.gate_reads_per_item();
+
+        for index in 0..3 {
+            fixture
+                .seed_event(signed_event(&[7u8; 32], 1_700_000_000 + index))
+                .await;
+        }
+
+        let budget = Arc::new(FrameBudget::new(64));
+        let (sink, mut frames) = FrameSink::channel(64 * 1024, Some(Arc::clone(&budget)));
+        let refused = sink.clone();
+        let connection = Arc::new(Connection::new(
+            "127.0.0.1".to_owned(),
+            Nip42Session::with_challenge("challenge", 0),
+        ));
+        let context = Arc::new(fixture.context(
+            Arc::new(TestClock::new(1_700_000_000_000)),
+            RelayKeySource::from_values(None, None),
+            Arc::clone(&connection),
+            sink,
+        ));
+
+        let message = SubscribeMessage {
+            subscription_id: "s".to_owned(),
+            filters: vec![Filter {
+                kinds: Some(vec![Kind(1)]),
+                ..Filter::default()
+            }],
+        };
+
+        let mut read = begin(message, &context).expect("the subscription is registered");
+        gate.add_permits(3);
+
+        assert!(
+            tokio::time::timeout(Duration::from_millis(200), &mut read)
+                .await
+                .is_err(),
+            "the read parks once its frame is refused instead of draining the stream"
+        );
+        assert_eq!(refused.refusal(), Some(Refusal::BudgetExhausted));
+        assert_eq!(
+            gate.available_permits(),
+            2,
+            "one record was taken and the rest of the stream was never asked for"
+        );
+        assert!(frames.try_recv().is_err(), "no EOSE reached the queue");
+        assert_eq!(refused.queued_bytes(), 0);
+        assert_eq!(
+            budget.refusals(),
+            1,
+            "one refusal for the socket, not one per event"
+        );
+
+        // The close the read is waiting for ends it.
+        connection.close();
+        tokio::time::timeout(Duration::from_secs(1), read)
+            .await
+            .expect("the connection's close ends the read");
+    }
+}

crates/nostream-relay/src/lib.rs

@@ -44,6 +44,7 @@
 
 pub mod broadcast;
 pub mod connection;
+pub mod connection_cap;
 pub mod error;
 pub mod handlers;
 pub mod hidden_content;