Keep the hello size limit until a connection is admitted

Fixes #1087

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

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<HubIn>, 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<B: Backend> Hub<B> {
 
     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<B: Backend> Hub<B> {
                     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<B: Backend> Hub<B> {
                 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<B: Backend> Hub<B> {
 
     /// 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;

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<Admission, String> {
+        let _ = self.0.recv();
+        Err("wrong token".into())
+    }
+
+    fn open(&mut self, _file: &str) -> Result<Server, String> {
+        Err("no".into())
+    }
+
+    fn closing(&self) -> Closing {
+        Closing::Ended
+    }
+}
+
+fn connect(addr: SocketAddr) -> tungstenite::WebSocket<TcpStream> {
+    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"));
+}