bump hbb_common: cover the message cap's lifecycle on WebSocket and WebRTC

Tests only, plus a doc note: the header-only refusal and the fragmented
reassembly bound on WebSocket, the lift on WebSocket and WebRTC, and the
contract that the bound is set before the first read of untrusted data -
which is where this branch sets it.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01EZ49AbZJYfm8NTp5yDPMab

src/rendezvous_mediator.rs

@@ -93,73 +93,10 @@ impl Drop for AnswererSlot {
         WEBRTC_ANSWERERS.fetch_sub(1, Ordering::AcqRel);
     }
 }
-
-/// Punches in flight: each holds a socket of its own and waits up to `CONNECT_TIMEOUT` for the
-/// peer, and neither request behind them needs authentication, anyone who knows this id can ask
-/// hbbs to have us open one. A place is given up the moment the peer's session is up, so it
-/// stands for that wait and nothing past it: the places turn over on their own, within
-/// `CONNECT_TIMEOUT` and the few seconds a punch's phases add to it, and at the limit an arrival
-/// is declined rather than an older punch cut short. Declined is the listen alone, never the
-/// reply: that carries the WebRTC answer and the v6 address as well, and those go on without a
-/// v4 punch.
-///
-/// Two pools, so that neither transport pays for the other's crowd: a connection costs a place in
-/// each, the controller's preferred request punching UDP and its TCP fallback request punching
-/// TCP, and the transports that lose the race hold theirs for the whole wait.
-struct PunchPool {
-    places: AtomicUsize,
-    max: usize,
-}
-
-/// UDP over v4 and v6. A request that carries a v6 address costs two, until the one the peer
-/// does not use times out.
-static UDP_PUNCHES: PunchPool = PunchPool::new(32);
-/// The TCP punch's listener, and the LAN listen a FetchLocalAddr opens, which is that listener
-/// again.
-static TCP_PUNCHES: PunchPool = PunchPool::new(32);
-
-impl PunchPool {
-    const fn new(max: usize) -> Self {
-        Self {
-            places: AtomicUsize::new(0),
-            max,
-        }
-    }
-
-    fn take(&'static self) -> Option<PunchSlot> {
-        self.places
-            .fetch_update(Ordering::AcqRel, Ordering::Acquire, |n| {
-                (n < self.max).then(|| n + 1)
-            })
-            .ok()
-            .map(|_| PunchSlot(self))
-    }
-
-    /// For a request that punches over both, v4 first: at the last place it is then the v6
-    /// punch that goes without, whatever order the two listens start in - v4 is the one the
-    /// peer can count on, v6 the one it may have no route for.
-    fn take_pair(&'static self, v4: bool, v6: bool) -> (Option<PunchSlot>, Option<PunchSlot>) {
-        let v4 = v4.then(|| self.take()).flatten();
-        let v6 = v6.then(|| self.take()).flatten();
-        (v4, v6)
-    }
-}
-
-/// One of a pool's places, given back on drop.
-pub(crate) struct PunchSlot(&'static PunchPool);
-
-impl Drop for PunchSlot {
-    fn drop(&mut self) {
-        self.0.places.fetch_sub(1, Ordering::AcqRel);
-    }
-}
 // The rendezvous ICE route is reachable without a prior punch and the peer decides how many
 // candidates it sends, so these sites would let someone else set how much this machine writes to
 // its log file. One line a minute each, carrying the suppressed count.
 const ICE_LOG_INTERVAL: std::time::Duration = std::time::Duration::from_secs(60);
-/// A declined punch is one line a minute, carrying the number it stands for, for the same
-/// reason: the peer decides how often it asks.
-const PUNCH_LOG_INTERVAL: std::time::Duration = std::time::Duration::from_secs(60);
 static UNKNOWN_ICE_SESSION_LOG: hbb_common::log_throttle::LogThrottle =
     hbb_common::log_throttle::LogThrottle::new(ICE_LOG_INTERVAL);
 static REJECTED_REMOTE_ICE_LOG: hbb_common::log_throttle::LogThrottle =
@@ -778,14 +715,7 @@ impl RendezvousMediator {
             fla.controlled_context.clone().into_option(),
         );
         if peer_addr_v6.port() > 0 && !relay {
-            socket_addr_v6 = start_ipv6(
+            socket_addr_v6 = start_ipv6(peer_addr_v6, addr, server.clone(), meta.clone()).await;
-                peer_addr_v6,
-                addr,
-                server.clone(),
-                meta.clone(),
-                UDP_PUNCHES.take(),
-            )
-            .await;
         }
         if is_ipv4(&self.addr) && !relay && !config::is_disable_tcp_listen() {
             if let Err(err) = self
@@ -828,17 +758,6 @@ impl RendezvousMediator {
     ) -> ResultType<()> {
         let peer_addr = AddrMangle::decode(&fla.socket_addr);
         log::debug!("Handle intranet from {:?}", peer_addr);
-        // The listen this opens waits for the peer like a TCP punch and is declined like one; the
-        // caller then relays, as it does when the listen fails for any other reason.
-        let Some(slot) = TCP_PUNCHES.take() else {
-            hbb_common::throttled_log!(
-                PUNCH_LOG_INTERVAL,
-                warn,
-                "declined a LAN listen: {} TCP punches already in flight",
-                TCP_PUNCHES.max
-            );
-            bail!("no place among the punches in flight");
-        };
         let mut socket = connect_tcp(&*self.host, CONNECT_TIMEOUT).await?;
         let local_addr = socket.local_addr();
         // we saw invalid local_addr while using proxy, local_addr.ip() == "::1"
@@ -856,7 +775,7 @@ impl RendezvousMediator {
         });
         let bytes = msg_out.write_to_bytes()?;
         socket.send_raw(bytes).await?;
-        crate::accept_connection(server.clone(), socket, peer_addr, true, meta, slot).await;
+        crate::accept_connection(server.clone(), socket, peer_addr, true, meta).await;
         Ok(())
     }
 
@@ -1094,29 +1013,9 @@ impl RendezvousMediator {
         } else {
             String::new()
         };
-        // Whether the v4 legs relay is known here, and decides whether a v4 place is taken at
+        if peer_addr_v6.port() > 0 && !relay {
-        // all: the relay branch below runs the whole session, and a place held across it would
+            socket_addr_v6 =
-        // let ordinary relay traffic use the pool up. The v6 punch is not relayed with them - a
+                start_ipv6(peer_addr_v6, peer_addr, server.clone(), meta.clone()).await;
-        // symmetric NAT on v4 says nothing about v6 - and its place is taken after the v4 one,
-        // or it could be the last and leave the punch the peer counts on with none.
-        let relay_v4 = ph.nat_type.enum_value() == Ok(NatType::SYMMETRIC)
-            || Config::get_nat_type() == NatType::SYMMETRIC as i32
-            || relay
-            || (config::is_disable_tcp_listen() && ph.udp_port <= 0);
-        let punch_udp = !relay_v4 && ph.udp_port > 0;
-        let punch_tcp = !relay_v4 && ph.udp_port <= 0 && ph.webrtc_sdp_offer.is_empty();
-        let punch_v6 = peer_addr_v6.port() > 0 && !relay;
-        let (slot_udp, slot_v6) = UDP_PUNCHES.take_pair(punch_udp, punch_v6);
-        let slot_tcp = punch_tcp.then(|| TCP_PUNCHES.take()).flatten();
-        if punch_v6 {
-            socket_addr_v6 = start_ipv6(
-                peer_addr_v6,
-                peer_addr,
-                server.clone(),
-                meta.clone(),
-                slot_v6,
-            )
-            .await;
         }
         let relay_server = self.get_relay_server(ph.relay_server);
         // for ensure, websocket go relay directly
@@ -1125,7 +1024,11 @@ impl RendezvousMediator {
         // than trusting this classification, so a direct WebRTC pair can still form on a
         // connection this branch has already called relay-only. Do not gate the answerer on
         // nat_type to make the two agree.
-        if relay_v4 {
+        if ph.nat_type.enum_value() == Ok(NatType::SYMMETRIC)
+            || Config::get_nat_type() == NatType::SYMMETRIC as i32
+            || relay
+            || (config::is_disable_tcp_listen() && ph.udp_port <= 0)
+        {
             let uuid = Uuid::new_v4().to_string();
             return self
                 .create_relay(
@@ -1155,7 +1058,7 @@ impl RendezvousMediator {
         };
         if ph.udp_port > 0 {
             peer_addr.set_port(ph.udp_port as u16);
-            self.punch_udp_hole(peer_addr, server, msg_punch, meta, slot_udp)
+            self.punch_udp_hole(peer_addr, server, msg_punch, meta)
                 .await?;
             return Ok(());
         }
@@ -1191,18 +1094,7 @@ impl RendezvousMediator {
         let local_addr = socket.local_addr();
         // The listener inside takes this address over, so the mediator's socket goes first.
         drop(socket);
-        // The reply went out above: declined is the listen alone, so the controller's TCP attempt
+        punch_tcp_until_connected(server, peer_addr, local_addr, meta).await;
-        // meets nothing and its v6 one goes on.
-        let Some(slot) = slot_tcp else {
-            hbb_common::throttled_log!(
-                PUNCH_LOG_INTERVAL,
-                warn,
-                "declined a TCP punch: {} TCP punches already in flight",
-                TCP_PUNCHES.max
-            );
-            return Ok(());
-        };
-        punch_tcp_until_connected(server, peer_addr, local_addr, meta, slot).await;
         Ok(())
     }
 
@@ -1212,26 +1104,12 @@ impl RendezvousMediator {
         server: ServerPtr,
         msg_punch: PunchHoleSent,
         meta: ConnectionMeta,
-        slot: Option<PunchSlot>,
     ) -> ResultType<()> {
         let mut msg_out = Message::new();
         msg_out.set_punch_hole_sent(msg_punch);
         let (socket, addr) = new_direct_udp_for(&self.host).await?;
         let data = msg_out.write_to_bytes()?;
         socket.send_to(&data, addr).await?;
-        // The reply is out, and with it the answer and the v6 address: declined is the listen
-        // alone, and the socket goes at once - a declined request is not worth one kept for its
-        // resends, and the controller re-asks on its own. Its v4 attempt at a mapping that no
-        // longer answers fails on its own while its other transports go on.
-        let Some(slot) = slot else {
-            hbb_common::throttled_log!(
-                PUNCH_LOG_INTERVAL,
-                warn,
-                "declined a UDP punch: {} UDP punches already in flight",
-                UDP_PUNCHES.max
-            );
-            return Ok(());
-        };
         let socket_cloned = socket.clone();
         tokio::spawn(async move {
             for _ in 0..2 {
@@ -1240,15 +1118,7 @@ impl RendezvousMediator {
                 socket.send_to(&data, addr).await.ok();
             }
         });
-        udp_nat_listen(
+        udp_nat_listen(socket_cloned.clone(), peer_addr, peer_addr, server, meta).await?;
-            socket_cloned.clone(),
-            peer_addr,
-            peer_addr,
-            server,
-            meta,
-            slot,
-        )
-        .await?;
         Ok(())
     }
 
@@ -1438,33 +1308,13 @@ async fn start_ipv6(
     peer_addr_v4: SocketAddr,
     server: ServerPtr,
     meta: ConnectionMeta,
-    slot: Option<PunchSlot>,
 ) -> bytes::Bytes {
-    // Declining leaves the v4 path to carry the connection, as it already does wherever this
-    // machine has no public IPv6 address.
-    let Some(slot) = slot else {
-        hbb_common::throttled_log!(
-            PUNCH_LOG_INTERVAL,
-            warn,
-            "declined an IPv6 punch: {} UDP punches already in flight",
-            UDP_PUNCHES.max
-        );
-        return Default::default();
-    };
     crate::test_ipv6().await;
     if let Some((socket, local_addr_v6)) = crate::get_ipv6_socket().await {
         let server = server.clone();
         tokio::spawn(async move {
             allow_err!(
-                udp_nat_listen(
+                udp_nat_listen(socket.clone(), peer_addr_v6, peer_addr_v4, server, meta).await
-                    socket.clone(),
-                    peer_addr_v6,
-                    peer_addr_v4,
-                    server,
-                    meta,
-                    slot
-                )
-                .await
             );
         });
         return local_addr_v6;
@@ -1478,7 +1328,6 @@ async fn udp_nat_listen(
     peer_addr_v4: SocketAddr,
     server: ServerPtr,
     meta: ConnectionMeta,
-    slot: PunchSlot,
 ) -> ResultType<()> {
     let tm = Instant::now();
     let socket_cloned = socket.clone();
@@ -1491,9 +1340,6 @@ async fn udp_nat_listen(
             init_packet,
         )
         .await?;
-        // The KCP session is up: from here it is a connection like any other and the connection
-        // layer's own limits apply to it, so the place goes back for the next punch.
-        drop(slot);
         crate::server::create_tcp_connection(server, stream.1, peer_addr_v4, true, meta).await?;
         Ok(())
     };
@@ -1537,15 +1383,11 @@ const PUNCH_GRACE: u64 = 3000;
 /// a hole that no longer exists. Punching again across the window in which the controller dials
 /// rebuilds it, and once the controller sits in SYN_SENT one of those punches meets its SYN and
 /// completes as a simultaneous open: a second way in, which a single punch never had.
-///
-/// `slot` is this punch's place among `TCP_PUNCHES`, given back the moment a connection is in
-/// hand and before the session runs on it: from there the connection layer's own limits apply.
 async fn punch_tcp_until_connected(
     server: ServerPtr,
     peer_addr: SocketAddr,
     local_addr: SocketAddr,
     meta: ConnectionMeta,
-    slot: PunchSlot,
 ) {
     use hbb_common::tcp::new_listener;
     // Shadows the module's `std::time::Instant`: the deadline is held against tokio's sleeps and
@@ -1572,7 +1414,6 @@ async fn punch_tcp_until_connected(
     });
     let Some(listener) = listener else {
         if let Some(stream) = punch.await {
-            drop(slot);
             serve_punched(server, stream, peer_addr, meta).await;
         }
         return;
@@ -1617,12 +1458,10 @@ async fn punch_tcp_until_connected(
         biased;
         Some(stream) = punch => stream,
         Some((stream, addr)) = accept => {
-            drop(slot);
             return accept_punched_connection(server, stream, addr, meta).await;
         }
         else => return,
     };
-    drop(slot);
     serve_punched(server, punched, peer_addr, meta).await;
 }
 
@@ -1730,137 +1569,12 @@ impl Drop for CheckIfResendPk {
 #[cfg(test)]
 mod tests {
     use super::{
-        connection_meta, mpsc, socket_client, tokio, udp_nat_listen, AnswererSlot, Arc, IceRoute,
+        mpsc, socket_client, tokio, AnswererSlot, IceRoute, ICE_DEDUP_WINDOW,
-        IntoTargetAddr, Ordering, PunchHoleSent, RendezvousMediator, RendezvousMessage,
+        MAX_PENDING_REMOTE_ICE, MAX_WEBRTC_ANSWERERS,
-        ICE_DEDUP_WINDOW, MAX_PENDING_REMOTE_ICE, MAX_WEBRTC_ANSWERERS, TCP_PUNCHES, UDP_PUNCHES,
     };
-    use hbb_common::{protobuf::Message as _, tcp::new_listener};
+    use hbb_common::tcp::new_listener;
     use std::net::SocketAddr;
 
-    // The pools are statics and the tests run in parallel: the ones that count them take turns.
-    static POOLS: std::sync::Mutex<()> = std::sync::Mutex::new(());
-
-    fn pools() -> std::sync::MutexGuard<'static, ()> {
-        POOLS.lock().unwrap_or_else(|e| e.into_inner())
-    }
-
-    // The places are the bound and they come back on drop: a punch past the limit is declined
-    // while the ones in flight are left alone, and the next one in is admitted only once a
-    // place has actually been given back.
-    #[test]
-    fn test_punch_slots_cap_and_release() {
-        let _pools = pools();
-        for pool in [&UDP_PUNCHES, &TCP_PUNCHES] {
-            let held: Vec<_> = (0..pool.max)
-                .map(|_| pool.take().expect("a place up to the limit"))
-                .collect();
-            assert!(pool.take().is_none(), "the limit is the limit");
-            drop(held);
-            assert!(pool.take().is_some(), "a place comes back on drop");
-        }
-    }
-
-    // Neither transport pays for the other's crowd: a full UDP pool leaves the TCP places alone.
-    #[test]
-    fn a_full_udp_pool_leaves_the_tcp_places_alone() {
-        let _pools = pools();
-        let held: Vec<_> = (0..UDP_PUNCHES.max)
-            .map(|_| UDP_PUNCHES.take().expect("a place up to the limit"))
-            .collect();
-        assert!(UDP_PUNCHES.take().is_none(), "the UDP pool is full");
-        assert!(
-            TCP_PUNCHES.take().is_some(),
-            "a TCP punch must not wait on the UDP pool"
-        );
-        drop(held);
-    }
-
-    // At the last place it is the v6 punch that goes without, never the v4 one the peer counts on.
-    #[test]
-    fn the_last_place_goes_to_the_v4_punch() {
-        let _pools = pools();
-        let held: Vec<_> = (0..UDP_PUNCHES.max - 1)
-            .map(|_| UDP_PUNCHES.take().expect("a place up to the last"))
-            .collect();
-        let (v4, v6) = UDP_PUNCHES.take_pair(true, true);
-        assert!(v4.is_some(), "the last place goes to the v4 punch");
-        assert!(v6.is_none(), "and the v6 punch goes without");
-        drop(held);
-    }
-
-    // Declined is the listen, never the reply: with no place the PunchHoleSent still goes out,
-    // carrying the answer and the v6 address the controller's other transports run on, and the
-    // call returns at once - a listen would have waited for the peer's probe and failed without.
-    #[tokio::test]
-    async fn a_declined_udp_punch_still_replies() {
-        let hbbs = tokio::net::UdpSocket::bind("127.0.0.1:0").await.unwrap();
-        let host = hbbs.local_addr().unwrap();
-        let mediator = RendezvousMediator {
-            addr: host.into_target_addr().unwrap(),
-            host: host.to_string(),
-            host_prefix: String::new(),
-            keep_alive: 0,
-        };
-        let msg_punch = PunchHoleSent {
-            webrtc_sdp_answer: "answer".to_owned(),
-            socket_addr_v6: bytes::Bytes::from_static(b"v6"),
-            ..Default::default()
-        };
-        mediator
-            .punch_udp_hole(
-                host,
-                crate::server::new_for_test(),
-                msg_punch,
-                connection_meta(None, None),
-                None,
-            )
-            .await
-            .expect("declined is not an error, and not a listen");
-        let mut buf = [0u8; 4096];
-        let (n, _) = hbb_common::timeout(3000, hbbs.recv_from(&mut buf))
-            .await
-            .expect("the reply must reach hbbs")
-            .unwrap();
-        let sent = RendezvousMessage::parse_from_bytes(&buf[..n]).unwrap();
-        let sent = sent.punch_hole_sent();
-        assert_eq!(sent.webrtc_sdp_answer, "answer");
-        assert_eq!(&sent.socket_addr_v6[..], b"v6");
-    }
-
-    // The place is the wait: a listen that ends without a peer gives it back on its own. The end
-    // a test can reach is the peer's probe never coming, which `punch_udp` gives up on.
-    #[test]
-    fn a_listen_that_finds_no_peer_gives_its_place_back() {
-        let _pools = pools();
-        let rt = tokio::runtime::Builder::new_current_thread()
-            .enable_all()
-            .build()
-            .unwrap();
-        rt.block_on(async {
-            let socket = tokio::net::UdpSocket::bind("127.0.0.1:0").await.unwrap();
-            let silent = tokio::net::UdpSocket::bind("127.0.0.1:0").await.unwrap();
-            let peer = silent.local_addr().unwrap();
-            let before = UDP_PUNCHES.places.load(Ordering::Acquire);
-            let slot = UDP_PUNCHES.take().expect("a place");
-            assert_eq!(UDP_PUNCHES.places.load(Ordering::Acquire), before + 1);
-            let listened = udp_nat_listen(
-                Arc::new(socket),
-                peer,
-                peer,
-                crate::server::new_for_test(),
-                connection_meta(None, None),
-                slot,
-            )
-            .await;
-            assert!(listened.is_err(), "no peer, no session");
-            assert_eq!(
-                UDP_PUNCHES.places.load(Ordering::Acquire),
-                before,
-                "the place must come back with the wait"
-            );
-        });
-    }
-
     // A SOCKS proxy makes `connect_tcp_local` dial the proxy and ignore the local address, so
     // nothing these two assert can hold. Read once, from the same global config production reads.
     fn proxied() -> bool {

src/server.rs

@@ -182,7 +182,6 @@ async fn accept_connection_(
     socket: Stream,
     secure: bool,
     meta: ConnectionMeta,
-    slot: crate::rendezvous_mediator::PunchSlot,
 ) -> ResultType<()> {
     let local_addr = socket.local_addr();
     drop(socket);
@@ -192,8 +191,6 @@ async fn accept_connection_(
     let listener = new_listener(local_addr, true).await?;
     log::info!("Server listening on: {}", &listener.local_addr()?);
     if let Ok((stream, addr)) = timeout(CONNECT_TIMEOUT, listener.accept()).await? {
-        // The peer is in: the place goes back before the session runs, as every punch's does.
-        drop(slot);
         stream.set_nodelay(true).ok();
         let stream_addr = stream.local_addr()?;
         create_tcp_connection(
@@ -224,6 +221,8 @@ pub async fn create_tcp_connection(
     let Some(unauthorized) = admit_unauthorized(id, addr.ip()) else {
         bail!("too many unauthenticated connections from {}", addr.ip());
     };
+    // Before the handshake, so its read is bounded too; lifted again at authorization.
+    stream.set_max_packet_length(MAX_UNAUTHORIZED_MESSAGE);
     tokio::select! {
         handshake = identity_handshake(&mut stream, secure) => handshake?,
         _ = unauthorized.evicted() => {
@@ -323,15 +322,14 @@ async fn identity_handshake(stream: &mut Stream, secure: bool) -> ResultType<()>
     Ok(())
 }
 
-pub(crate) async fn accept_connection(
+pub async fn accept_connection(
     server: ServerPtr,
     socket: Stream,
     peer_addr: SocketAddr,
     secure: bool,
     meta: ConnectionMeta,
-    slot: crate::rendezvous_mediator::PunchSlot,
 ) {
-    if let Err(err) = accept_connection_(server, socket, secure, meta, slot).await {
+    if let Err(err) = accept_connection_(server, socket, secure, meta).await {
         log::warn!("Failed to accept connection from {}: {}", peer_addr, err);
     }
 }

src/server/connection.rs

@@ -93,6 +93,14 @@ const MAX_UNAUTHORIZED_CONNS: usize = 64;
 /// of addresses passes it, and the bound above is what holds. Meaningful only while the
 /// address is the controller's own, which punch and relay messages carry today.
 const MAX_UNAUTHORIZED_CONNS_PER_ADDR: usize = 16;
+/// The largest message a connection may send before it authorizes. Until then a peer sends only
+/// a public key, a login request, a test delay and a close reason, none of which carries an
+/// unbounded field - a server hands the login request's avatar out as a URL, and only a custom
+/// client that inlines an image into the avatar option instead reaches this. Sized to the read
+/// buffer tungstenite allocates per WebSocket connection regardless, so there the cap costs
+/// nothing beyond a floor already paid; with MAX_UNAUTHORIZED_CONNS it holds them to 8 MiB in
+/// all, against the 1 GiB a single one could make us hold before.
+pub const MAX_UNAUTHORIZED_MESSAGE: usize = 128 * 1024;
 
 /// A place among the unauthorized connections, taken before the identity handshake and given
 /// back on drop: at authorization, or when the connection ends first. The count of live
@@ -1871,6 +1879,10 @@ impl Connection {
         if let Some(keep_alive) = self.prepare_terminal_login_for_authorization().await {
             return keep_alive;
         }
+        // Lifted here rather than below with the rest of authorization: a multiplexed tunnel
+        // narrows it again for its own framing (`port_forward_mux::cap_packet_size`), so that
+        // call has to come after this one, not before.
+        self.stream.set_max_packet_length(usize::MAX);
         if !self.connect_port_forward_if_needed().await {
             return false;
         }