From 91d35fcae001bdbf46c79d6403701a09fef41a12 Mon Sep 17 00:00:00 2001 From: yahya-nusseibeh <254622510+yahya-nusseibeh@users.noreply.github.com> Date: Wed, 4 Nov 2026 14:16:00 +0200 Subject: [PATCH] Relay: reserve a bridge slot at the match, give it back on every exit max_bridges compared against bridges already running, so joins waiting (up to attach_wait) for their host's attach all passed the check and a burst could exceed RELAY_MAX_BRIDGES. - BridgeSlot, taken under the lookup lock when a guest is matched with its host, counts toward max_bridges from then on; dropping it gives it back, so every way out of a join (no attach, a refused or forged attach, an error inside bridge(), a panic) releases it. - counts().bridges now includes matched joins awaiting attach. - tests/bridge_slots.rs: a waiting join holds its slot (a second guest gets Full); the count returns to 0 after a timeout, after a forged attach, and after a bridge closes; refusals never take one. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01SJ3stH5q6PgMAVUSmJqKDe --- crates/horizon-relay/src/lib.rs | 64 ++++-- crates/horizon-relay/tests/bridge_slots.rs | 217 +++++++++++++++++++++ 2 files changed, 262 insertions(+), 19 deletions(-) create mode 100644 crates/horizon-relay/tests/bridge_slots.rs diff --git a/crates/horizon-relay/src/lib.rs b/crates/horizon-relay/src/lib.rs index 36e7c7d247..c0b92e58b4 100644 --- a/crates/horizon-relay/src/lib.rs +++ b/crates/horizon-relay/src/lib.rs @@ -116,6 +116,8 @@ impl Config { #[derive(Clone, Debug, Default, PartialEq, Eq)] pub struct Counts { pub hosts: usize, + /// Bridges open, and joins matched with a host and waiting for its attach (each holds one of + /// `max_bridges` from the match on). pub bridges: usize, pub registrations: u64, pub joins: u64, @@ -485,8 +487,28 @@ fn serve_host( Ok(()) } -enum Lookup { - Found(Sender), +/// One of `max_bridges`, reserved when a guest is matched with its host and given back when it drops: +/// on every way out of the join (no attach, a refusal, an error, a panic) and at the bridge's end. A +/// reservation counts from the match, not from the bridge, so joins waiting for their host's attach +/// can't add up past the cap. +struct BridgeSlot<'a>(&'a AtomicUsize); +impl<'a> BridgeSlot<'a> { + fn take(count: &'a AtomicUsize, max: usize) -> Option { + if count.fetch_add(1, Ordering::AcqRel) >= max { + count.fetch_sub(1, Ordering::AcqRel); + return None; + } + Some(Self(count)) + } +} +impl Drop for BridgeSlot<'_> { + fn drop(&mut self) { + self.0.fetch_sub(1, Ordering::AcqRel); + } +} + +enum Lookup<'a> { + Found(Sender, BridgeSlot<'a>), Unknown, Ended, Version(String), @@ -521,18 +543,20 @@ fn join( .ended .get(&nameplate) .is_some_and(|t| now.saturating_duration_since(*t) < state.config.ended_kept); - let busy = state.bridges.load(Ordering::Relaxed) >= state.config.max_bridges; let lookup = match inner.hosts.get_mut(&nameplate) { None if ended => Lookup::Ended, None => Lookup::Unknown, Some(h) if h.protocol != protocol => Lookup::Version(h.game.clone()), - Some(h) if busy || h.pending >= state.config.max_pending_per_host => Lookup::Full, - Some(h) => { - h.pending += 1; - Lookup::Found(h.to_host.clone()) - } + Some(h) if h.pending >= state.config.max_pending_per_host => Lookup::Full, + Some(h) => match BridgeSlot::take(&state.bridges, state.config.max_bridges) { + Some(slot) => { + h.pending += 1; + Lookup::Found(h.to_host.clone(), slot) + } + None => Lookup::Full, + }, }; - if matches!(lookup, Lookup::Found(_)) { + if matches!(lookup, Lookup::Found(..)) { inner.pending.insert( session, Pending { @@ -543,8 +567,8 @@ fn join( } lookup }; - let to_host = match lookup { - Lookup::Found(tx) => tx, + let (to_host, slot) = match lookup { + Lookup::Found(tx, slot) => (tx, slot), Lookup::Unknown => { state.unknown.fetch_add(1, Ordering::Relaxed); locked(&state.limiter).unknown(peer.ip(), now); @@ -573,8 +597,11 @@ fn join( } } match attached { - Ok((host, host_peer)) => bridge(state, stream, peer, host, host_peer, session), - Err(_) => refuse(&mut stream, ErrorCode::HostUnreachable), + Ok((host, host_peer)) => bridge(state, stream, peer, host, host_peer, session, slot), + Err(_) => { + drop(slot); + refuse(&mut stream, ErrorCode::HostUnreachable) + } } } @@ -610,7 +637,8 @@ fn attach( } } -/// Tell both ends they are bridged, then copy each way until either ends. +/// Tell both ends they are bridged, then copy each way until either ends. `_slot` (the guest's +/// reservation of one of `max_bridges`) is given back when this returns, however it returns. fn bridge( state: &Arc, mut guest: TcpStream, @@ -618,9 +646,9 @@ fn bridge( mut host: TcpStream, host_peer: SocketAddr, session: u64, + _slot: BridgeSlot<'_>, ) -> io::Result<()> { - state.bridges.fetch_add(1, Ordering::Relaxed); - let result = (|| -> io::Result<()> { + (|| -> io::Result<()> { wire::send( &mut guest, &FromRelay::Bridged { @@ -646,9 +674,7 @@ fn bridge( copy(state, host, guest); let _ = up.join(); Ok(()) - })(); - state.bridges.fetch_sub(1, Ordering::Relaxed); - result + })() } fn copy(state: &State, mut from: TcpStream, mut to: TcpStream) { diff --git a/crates/horizon-relay/tests/bridge_slots.rs b/crates/horizon-relay/tests/bridge_slots.rs new file mode 100644 index 0000000000..a16a991a1c --- /dev/null +++ b/crates/horizon-relay/tests/bridge_slots.rs @@ -0,0 +1,217 @@ +//! `max_bridges` counts from the moment a guest is matched with its host, and every way out of a +//! join gives the slot back, spoken to the relay over raw sockets. +use horizon_friends::wire::{self, ErrorCode, FromRelay, ToRelay}; +use horizon_relay::{Config, Relay}; +use std::io::BufReader; +use std::net::{Shutdown, SocketAddr, TcpStream}; +use std::time::{Duration, Instant}; + +const PROTOCOL: i32 = 777; + +fn relay(max_bridges: usize) -> (Relay, SocketAddr) { + let (relay, address, _) = Relay::spawn_local(Config { + max_bridges, + per_address: false, + attach_wait: Duration::from_millis(400), + ..Config::default() + }) + .unwrap(); + (relay, address) +} + +fn connect(address: SocketAddr) -> TcpStream { + let s = TcpStream::connect(address).unwrap(); + s.set_read_timeout(Some(Duration::from_secs(5))).unwrap(); + s +} + +/// A registered host's control connection (reader, writer), its nameplate and token. +fn host(address: SocketAddr) -> (BufReader, TcpStream, String, String) { + let mut s = connect(address); + wire::send( + &mut s, + &ToRelay::Host { + v: wire::VERSION, + protocol: PROTOCOL, + game: "26.3".into(), + }, + ) + .unwrap(); + let mut reader = BufReader::new(s.try_clone().unwrap()); + match wire::recv_buffered::<_, FromRelay>(&mut reader).unwrap() { + FromRelay::Registered { + nameplate, token, .. + } => (reader, s, nameplate, token), + other => panic!("{other:?}"), + } +} + +/// A guest's join, sent; the stream to read the relay's answer from. +fn join(address: SocketAddr, nameplate: &str) -> TcpStream { + let mut s = connect(address); + wire::send( + &mut s, + &ToRelay::Join { + v: wire::VERSION, + nameplate: nameplate.into(), + protocol: PROTOCOL, + game: "26.3".into(), + }, + ) + .unwrap(); + s +} + +/// The next message that isn't a ping. +fn next(reader: &mut BufReader) -> FromRelay { + loop { + match wire::recv_buffered::<_, FromRelay>(reader).unwrap() { + FromRelay::Ping => continue, + m => return m, + } + } +} + +fn answer(s: &mut TcpStream) -> FromRelay { + wire::recv(s).unwrap() +} + +fn settles_to_zero(relay: &Relay) { + let deadline = Instant::now() + Duration::from_secs(5); + while relay.counts().bridges != 0 { + assert!( + Instant::now() < deadline, + "bridges stuck at {}", + relay.counts().bridges + ); + std::thread::sleep(Duration::from_millis(10)); + } +} + +#[test] +fn a_waiting_join_holds_its_slot_and_a_timeout_gives_it_back() { + let (relay, address) = relay(1); + let (mut control, _keep, nameplate, _) = host(address); + // The first guest is matched: the host is asked to attach, and doesn't. + let mut first = join(address, &nameplate); + assert!(matches!(next(&mut control), FromRelay::Join { .. })); + assert_eq!(relay.counts().bridges, 1, "reserved at the match"); + // A second guest meanwhile finds the relay full. + let mut second = join(address, &nameplate); + assert!(matches!( + answer(&mut second), + FromRelay::Error { + code: ErrorCode::Full, + .. + } + )); + // The first one's wait ends: refused, and the slot is back. + assert!(matches!( + answer(&mut first), + FromRelay::Error { + code: ErrorCode::HostUnreachable, + .. + } + )); + settles_to_zero(&relay); + // So the next guest is matched again. + let mut third = join(address, &nameplate); + assert!(matches!(next(&mut control), FromRelay::Join { .. })); + assert!(matches!( + answer(&mut third), + FromRelay::Error { + code: ErrorCode::HostUnreachable, + .. + } + )); + settles_to_zero(&relay); +} + +#[test] +fn a_bridge_gives_its_slot_back_when_either_end_closes() { + let (relay, address) = relay(1); + let (mut control, _keep, nameplate, token) = host(address); + let mut guest = join(address, &nameplate); + let FromRelay::Join { session } = next(&mut control) else { + panic!() + }; + let mut attach = connect(address); + wire::send( + &mut attach, + &ToRelay::Attach { + v: wire::VERSION, + session, + token, + }, + ) + .unwrap(); + assert!(matches!(answer(&mut guest), FromRelay::Bridged { .. })); + assert!(matches!(answer(&mut attach), FromRelay::Bridged { .. })); + assert_eq!(relay.counts().bridges, 1); + guest.shutdown(Shutdown::Both).unwrap(); + settles_to_zero(&relay); +} + +#[test] +fn a_wrong_attach_or_a_refusal_holds_no_slot() { + let (relay, address) = relay(1); + let (mut control, _keep, nameplate, _) = host(address); + // Refused before any match: an unknown code, another protocol. + let mut unknown = join(address, "ZZZZ"); + assert!(matches!( + answer(&mut unknown), + FromRelay::Error { + code: ErrorCode::UnknownCode, + .. + } + )); + let mut other = connect(address); + wire::send( + &mut other, + &ToRelay::Join { + v: wire::VERSION, + nameplate: nameplate.clone(), + protocol: PROTOCOL + 1, + game: "x".into(), + }, + ) + .unwrap(); + assert!(matches!( + answer(&mut other), + FromRelay::Error { + code: ErrorCode::VersionMismatch, + .. + } + )); + assert_eq!(relay.counts().bridges, 0); + // Matched, then an attach with the wrong token: refused, and the join's wait runs out. + let mut guest = join(address, &nameplate); + let FromRelay::Join { session } = next(&mut control) else { + panic!() + }; + let mut forged = connect(address); + wire::send( + &mut forged, + &ToRelay::Attach { + v: wire::VERSION, + session, + token: "0".repeat(32), + }, + ) + .unwrap(); + assert!(matches!( + answer(&mut forged), + FromRelay::Error { + code: ErrorCode::BadRequest, + .. + } + )); + assert!(matches!( + answer(&mut guest), + FromRelay::Error { + code: ErrorCode::HostUnreachable, + .. + } + )); + settles_to_zero(&relay); +}