diff --git a/crates/collab-net/src/hub.rs b/crates/collab-net/src/hub.rs index 5cf3e519..fe621d98 100644 --- a/crates/collab-net/src/hub.rs +++ b/crates/collab-net/src/hub.rs @@ -402,6 +402,8 @@ enum HubIn { sock: TcpStream, /// Its size (in flight until handled). bytes: usize, + /// Told when the hub takes the connection in (dropped unsent when it doesn't). + admitted: SyncSender<()>, }, Msg { conn: u64, @@ -946,9 +948,13 @@ fn serve(mut stream: TcpStream, conn: u64, hub: &Sender, limits: NetLimit if reserved || !shared.in_flight.take(bytes, limits.max_in_flight) { end(Some(BUSY)); } else { - ws.set_config(|c| *c = config(max)); - let opened = sock.try_clone().is_ok_and(|s| hub.send(HubIn::Open { conn, ip, first, out: out.clone(), sock: s, bytes }).is_ok()); - if opened && ws.get_ref().set_read_timeout(Some(idle)).is_ok() { + let (admitted, verdict) = sync_channel(1); + let opened = + sock.try_clone().is_ok_and(|s| hub.send(HubIn::Open { conn, ip, first, out: out.clone(), sock: s, bytes, admitted }).is_ok()); + // Nothing more is read until the hub took the connection in (its token checked): + // only then may it send messages of the full size. + if opened && verdict.recv().is_ok() && ws.get_ref().set_read_timeout(Some(idle)).is_ok() { + ws.set_config(|c| *c = config(max)); read_loop(&mut ws, conn, hub, limits, max, &shared.in_flight, &end); let _ = hub.send(HubIn::Gone { conn }); } @@ -1074,6 +1080,8 @@ struct Pending { out: Line, sock: TcpStream, ip: IpAddr, + /// Told when it is taken in. + admitted: SyncSender<()>, } /// A connection the hub took in. @@ -1195,7 +1203,7 @@ impl Hub { fn handle(&mut self, m: HubIn) { match m { - HubIn::Open { conn, ip, first, out, sock, bytes } => { + HubIn::Open { conn, ip, first, out, sock, bytes, admitted } => { self.in_flight.give(bytes); // The operator (with the server key) gets in from anywhere; everyone else is // checked (blocked addresses, connection rates). @@ -1204,7 +1212,7 @@ impl Hub { log::info!("event=refused ip={ip} reason={why:?}"); return refuse(&out, &why); } - let c = Pending { out, sock, ip }; + let c = Pending { out, sock, ip, admitted }; match first { ClientMsg::Hello(hello) => self.join(conn, hello, c), ClientMsg::Manage(m) => self.manage(conn, &m, c), @@ -1399,6 +1407,7 @@ impl Hub { if let Some(c) = old.and_then(|old| self.conns.remove(&old)) { shutdown(&c.sock); } + let _ = c.admitted.try_send(()); self.conns .insert(conn, Conn { out: c.out, sock: c.sock, kind: Kind::Joined { file: file.clone(), id, link: admission.link }, ip: c.ip }); self.queue(&file, more); @@ -1444,6 +1453,7 @@ impl Hub { /// Take in connection `c` as `kind` and send it `first`. fn take_in(&mut self, conn: u64, c: Pending, kind: Kind, first: &ServerMsg) { + let _ = c.admitted.try_send(()); self.conns.insert(conn, Conn { out: c.out, sock: c.sock, kind, ip: c.ip }); self.answer(conn, first); } @@ -1747,3 +1757,7 @@ mod queue_tests; #[cfg(test)] #[path = "barred_tests.rs"] mod barred_tests; + +#[cfg(test)] +#[path = "preauth_tests.rs"] +mod preauth_tests; diff --git a/crates/collab-net/src/preauth_tests.rs b/crates/collab-net/src/preauth_tests.rs new file mode 100644 index 00000000..5ae91665 --- /dev/null +++ b/crates/collab-net/src/preauth_tests.rs @@ -0,0 +1,72 @@ +//! A connection whose hello the backend hasn't admitted yet is held to the hello's size limit: it +//! can't have a message of the full size read while its token is unchecked. + +use std::net::{SocketAddr, TcpListener, TcpStream}; +use std::sync::mpsc::{Receiver, channel}; +use std::time::{Duration, Instant}; + +use figcraft_collab::message::{self, ClientMsg, ServerMsg, Signal}; +use figcraft_collab::{Client, Server}; +use tungstenite::{Bytes, Message}; + +use super::{Admission, Backend, Closing, Hosting}; +use crate::net::NetLimits; + +/// A backend that refuses every hello, once it is told to. +struct Slow(Receiver<()>); + +impl Backend for Slow { + fn admit(&mut self, _file: &str, _token: &str) -> Result { + let _ = self.0.recv(); + Err("wrong token".into()) + } + + fn open(&mut self, _file: &str) -> Result { + Err("no".into()) + } + + fn closing(&self) -> Closing { + Closing::Ended + } +} + +fn connect(addr: SocketAddr) -> tungstenite::WebSocket { + let s = TcpStream::connect(addr).unwrap(); + s.set_read_timeout(Some(Duration::from_secs(10))).unwrap(); + tungstenite::client::client(format!("ws://{addr}/"), s).unwrap().0 +} + +#[test] +fn an_unadmitted_connection_cant_send_a_large_message() { + let limits = NetLimits { hello_timeout: Duration::from_secs(5), ..NetLimits::default() }; + let (release, wait) = channel(); + let mut hosting = Hosting::serve(TcpListener::bind("127.0.0.1:0").unwrap(), Slow(wait), limits).unwrap(); + let addr = hosting.addr(); + let sender = std::thread::spawn(move || { + let mut ws = connect(addr); + let hello = message::encode(&Client::hello("0123456789abcdef", "bad", "x", "")).unwrap(); + ws.send(Message::Binary(Bytes::from(hello))).unwrap(); + let big = ClientMsg::Signal(Signal { kind: Signal::CHAT.into(), text: "x".repeat(1 << 20) }); + let _ = ws.send(Message::Binary(Bytes::from(message::encode(&big).unwrap()))); + let mut refused = None; + while let Ok(m) = ws.read() { + if let Message::Binary(b) = m + && let Ok(ServerMsg::Refused { reason }) = message::decode(&b, 64 << 20) + { + refused = Some(reason); + } + } + refused + }); + let mut peak = 0; + let until = Instant::now() + Duration::from_secs(2); + while Instant::now() < until { + peak = peak.max(hosting.in_flight_bytes()); + std::thread::sleep(Duration::from_millis(2)); + } + release.send(()).unwrap(); + let refused = sender.join().unwrap(); + hosting.stop(); + assert!(peak <= limits.max_hello, "a message of {peak} bytes was read before the token was checked"); + assert_eq!(refused.as_deref(), Some("wrong token")); +}