server: bound unauthenticated connections in number and in time

A connection that never logs in costs whatever its transport costs, for as
long as it keeps itself alive: the only limit was the 30s idle timeout, which
any message resets. Nothing bounded how many such connections one machine
holds, on any transport. The shape sshd_config answers with LoginGraceTime
and MaxStartups.

Every connection is admitted among the unauthorized ones before its identity
handshake, in create_tcp_connection, and holds that place until it
authorizes or ends: the count of live places is the bound, not a ledger
beside the connections. One address may hold four; a further connection from
it is refused before the handshake. With 64 held in all, a further arrival
is refused too, and the oldest connection is told to go, unless one is on
its way out already: the handshake is raced against that eviction and ends
at once, and the session loop has it as a branch of its select, so the place
opens as soon as the connection has actually gone and not on a timer tick.
The newcomer is not let in on a place still occupied; the controller retries
on its own with backoff, and by then the place is free. At most one
connection is ever on its way out, so a burst of refused arrivals clears no
more room than a single one, and the retry that takes the freed place counts
against its address's share: one address turns out at most as many
connections as it may hold.

Two deadlines, from the moment the connection starts, each a branch of the
session loop's select rather than a check on the TestDelay tick. A
controller sends its login request as soon as it has our Hash, so a
connection without one after 20s is closed. One with a login request but
still not authorized after 180s, a wrong password, a pending 2FA or an
accept prompt left unanswered, is closed too; the controller reconnects on
its own and the prompt comes back. Both close with the Timeout reason the
idle path uses.

The peer address is normalized with try_into_v4 before admission, the same
form Connection::start keys the whitelist on, so an IPv4 peer and its
IPv4-mapped IPv6 form are one address and not two shares.

The WebRTC answerer's slot keeps bounding peer connection setup up to the
open data channel; from there this covers it like every other transport.

Tests cover the registry and the live bound: an address over its share is
refused while others are admitted; at the limit the newcomer is refused, the
oldest is told to go, nobody else is while it is on its way out, and its
place frees only when it has; an address at the limit turns out no more
connections than its share and is then refused without evicting anyone; and
with the limit held by 64 connections stalled in the handshake, one more
arrival is refused while the oldest handshake ends at once and only then is
there a place again.

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

src/server.rs

@@ -118,6 +118,15 @@ pub struct Server {
 pub type ServerPtr = Arc<RwLock<Server>>;
 pub type ServerPtrWeak = Weak<RwLock<Server>>;
 
+#[cfg(test)]
+pub fn new_for_test() -> ServerPtr {
+    Arc::new(RwLock::new(Server {
+        connections: HashMap::new(),
+        services: HashMap::new(),
+        id_count: 1000,
+    }))
+}
+
 pub fn new() -> ServerPtr {
     let mut server = Server {
         connections: HashMap::new(),
@@ -204,7 +213,40 @@ pub async fn create_tcp_connection(
     meta: ConnectionMeta,
 ) -> ResultType<()> {
     let mut stream = stream;
+    // The address the connection layer keys on, whitelist and admission alike.
+    let addr = hbb_common::try_into_v4(addr);
     let id = server.write().unwrap().get_new_id();
+    // Admitted before the identity handshake, so a peer that stalls in it, or after it without
+    // logging in, holds its place the whole time; an address over its share is turned away.
+    let Some(unauthorized) = admit_unauthorized(id, addr.ip()) else {
+        bail!("too many unauthenticated connections from {}", addr.ip());
+    };
+    tokio::select! {
+        handshake = identity_handshake(&mut stream, secure) => handshake?,
+        _ = unauthorized.evicted() => {
+            bail!("evicted to make room for a newer unauthenticated connection");
+        }
+    }
+
+    #[cfg(target_os = "macos")]
+    {
+        use std::process::Command;
+        if let Ok(task) = Command::new("/usr/bin/caffeinate")
+            .arg("-u")
+            .arg("-t 5")
+            .spawn()
+        {
+            super::CHILD_PROCESS.lock().unwrap().push(task);
+        }
+        log::info!("wake up macos");
+    }
+    Connection::start(addr, stream, id, Arc::downgrade(&server), meta, unauthorized).await;
+    Ok(())
+}
+
+/// Our signed identity goes out and, when `secure`, the controller's reply keys `stream`.
+/// Separate so it can be raced against the connection's eviction.
+async fn identity_handshake(stream: &mut Stream, secure: bool) -> ResultType<()> {
     let (sk, pk) = Config::get_key_pair();
     if secure && pk.len() == sign::PUBLICKEYBYTES && sk.len() == sign::SECRETKEYBYTES {
         let mut sk_ = [0u8; sign::SECRETKEYBYTES];
@@ -267,19 +309,6 @@ pub async fn create_tcp_connection(
         }
     }
 
-    #[cfg(target_os = "macos")]
-    {
-        use std::process::Command;
-        if let Ok(task) = Command::new("/usr/bin/caffeinate")
-            .arg("-u")
-            .arg("-t 5")
-            .spawn()
-        {
-            super::CHILD_PROCESS.lock().unwrap().push(task);
-        }
-        log::info!("wake up macos");
-    }
-    Connection::start(addr, stream, id, Arc::downgrade(&server), meta).await;
     Ok(())
 }
 

src/server/connection.rs

@@ -55,15 +55,16 @@ use scrap::android::{call_main_service_key_event, call_main_service_pointer_inpu
 use scrap::camera;
 use serde_derive::Serialize;
 use serde_json::{json, value::Value};
-#[cfg(not(any(target_os = "android", target_os = "ios")))]
-use std::sync::atomic::Ordering;
 use std::{
     collections::HashSet,
-    net::Ipv6Addr,
+    net::{IpAddr, Ipv6Addr},
     num::NonZeroI64,
     path::PathBuf,
     str::FromStr,
-    sync::{atomic::AtomicI64, mpsc as std_mpsc},
+    sync::{
+        atomic::{AtomicBool, AtomicI64, Ordering},
+        mpsc as std_mpsc,
+    },
 };
 #[cfg(not(any(target_os = "android", target_os = "ios")))]
 use system_shutdown;
@@ -79,7 +80,111 @@ const FAILURE_IDX_ID_WHITELIST: usize = 2;
 // throttles enumeration harder; shorter limits collateral on whitelisted neighbours.
 const ID_WHITELIST_FAILURE_DECAY_MINUTES: i32 = 10;
 
+/// A connection that has not sent its login request within this long is closed. A controller
+/// sends one as soon as the handshake is done, so nothing legitimate waits; a peer that holds
+/// the connection open and says nothing is not a controller.
+const LOGIN_REQUEST_GRACE: Duration = Duration::from_secs(20);
+/// A connection that has sent a login request but not authorized within this long is closed,
+/// however alive it keeps itself: a wrong password, a pending 2FA or an accept prompt still
+/// unanswered. The controller reconnects on its own and the prompt comes back.
+const LOGIN_GRACE: Duration = Duration::from_secs(180);
+/// Connections between accept and authorization, across every transport. Beyond this many the
+/// oldest is closed to make room, so a flood of them holds a legitimate controller off only
+/// while newer ones keep arriving faster than a password is typed.
+const MAX_UNAUTHORIZED_CONNS: usize = 64;
+/// Of those, how many one address may hold at once; a further connection from that address is
+/// refused. Filling the global limit takes this many times more addresses.
+const MAX_UNAUTHORIZED_CONNS_PER_ADDR: usize = 4;
+
+/// 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
+/// guards is the bound; an evicted one is told to go and keeps its place until it has.
+pub struct UnauthorizedID {
+    id: i32,
+    shared: Arc<UnauthorizedShared>,
+}
+
+struct UnauthorizedShared {
+    evicted: AtomicBool,
+    notify: hbb_common::tokio::sync::Notify,
+}
+
+/// Admit a connection from `ip` among the unauthorized ones. `None` when that address already
+/// holds its share, or when the global limit is reached: then the oldest connection is told to
+/// go, unless one is on its way out already, and this one is refused rather than let in on a
+/// place that is still occupied. At most one connection is ever on its way out, so a burst of
+/// refused arrivals clears no more room than a single one. The controller retries on its own.
+pub fn admit_unauthorized(id: i32, ip: IpAddr) -> Option<UnauthorizedID> {
+    let mut conns = UNAUTHORIZED_CONNS.lock().unwrap();
+    if conns.iter().filter(|(_, held, _)| *held == ip).count() >= MAX_UNAUTHORIZED_CONNS_PER_ADDR {
+        return None;
+    }
+    if conns.len() >= MAX_UNAUTHORIZED_CONNS {
+        if let Some((_, _, oldest)) = conns.first() {
+            if !oldest.evicted.swap(true, Ordering::AcqRel) {
+                oldest.notify.notify_one();
+            }
+        }
+        return None;
+    }
+    let shared = Arc::new(UnauthorizedShared {
+        evicted: AtomicBool::new(false),
+        notify: hbb_common::tokio::sync::Notify::new(),
+    });
+    conns.push((id, ip, shared.clone()));
+    Some(UnauthorizedID { id, shared })
+}
+
+impl UnauthorizedID {
+    /// Whether this connection was told to go to make room for a newer one.
+    pub fn is_evicted(&self) -> bool {
+        self.shared.evicted.load(Ordering::Acquire)
+    }
+
+    /// Resolves once this connection is told to go; at once if it already was.
+    pub async fn evicted(&self) {
+        if self.is_evicted() {
+            return;
+        }
+        self.shared.notify.notified().await;
+    }
+}
+
+impl Drop for UnauthorizedID {
+    fn drop(&mut self) {
+        UNAUTHORIZED_CONNS
+            .lock()
+            .unwrap()
+            .retain(|(id, _, _)| *id != self.id);
+    }
+}
+
+/// Resolves when the connection holding `unauthorized` is evicted; never once it has
+/// authorized and given its place back.
+async fn unauthorized_evicted(unauthorized: &Option<UnauthorizedID>) {
+    match unauthorized {
+        Some(u) => u.evicted().await,
+        None => std::future::pending().await,
+    }
+}
+
+/// Resolves at the login deadline of a connection started at `started`; never once it has
+/// authorized. Until the login request arrives the shorter grace applies.
+async fn login_deadline(authorized: bool, login_request_seen: bool, started: Instant) {
+    if authorized {
+        return std::future::pending().await;
+    }
+    let grace = if login_request_seen {
+        LOGIN_GRACE
+    } else {
+        LOGIN_REQUEST_GRACE
+    };
+    time::sleep_until(started + grace).await
+}
+
 lazy_static::lazy_static! {
+    // Connections between accept and authorization, oldest first; see admit_unauthorized.
+    static ref UNAUTHORIZED_CONNS: Mutex<Vec<(i32, IpAddr, Arc<UnauthorizedShared>)>> = Default::default();
     // [0] password, [1] 2FA, [2] ID whitelist.
     // Bucket 2 is separate so its rejections do not touch the password / 2FA budgets. It is
     // decayed in `check_id_whitelist` and cleared on auth, never on a bare id match.
@@ -264,6 +369,9 @@ pub struct Connection {
     port_forward_address: String,
     tx_to_cm: mpsc::UnboundedSender<ipc::Data>,
     authorized: bool,
+    // The place among the unauthorized connections; given back at authorization.
+    unauthorized_id: Option<UnauthorizedID>,
+    login_request_seen: bool,
     require_2fa: Option<totp_rs::TOTP>,
     awaiting_2fa: bool,
     keyboard: bool,
@@ -425,6 +533,7 @@ impl Connection {
         id: i32,
         server: super::ServerPtrWeak,
         meta: super::ConnectionMeta,
+        unauthorized: UnauthorizedID,
     ) {
         let super::ConnectionMeta {
             control_permissions,
@@ -483,6 +592,8 @@ impl Connection {
             port_forward_address: "".to_owned(),
             tx_to_cm,
             authorized: false,
+            unauthorized_id: Some(unauthorized),
+            login_request_seen: false,
             keyboard: Self::permission(keys::OPTION_ENABLE_KEYBOARD, &control_permissions),
             clipboard: Self::permission(keys::OPTION_ENABLE_CLIPBOARD, &control_permissions),
             audio: Self::permission(keys::OPTION_ENABLE_AUDIO, &control_permissions),
@@ -598,6 +709,7 @@ impl Connection {
         let mut test_delay_timer =
             crate::rustdesk_interval(time::interval_at(Instant::now(), TEST_DELAY_TIMEOUT));
         let mut last_recv_time = Instant::now();
+        let started = Instant::now();
 
         // The connection type is not known until the login request arrives;
         // `on_message` picks the type-specific timeout then.
@@ -633,6 +745,17 @@ impl Connection {
             tokio::select! {
                 // biased; // video has higher priority // causing test_delay_timer failed while transferring big file
 
+                // Both end an unauthorized connection at once, not on the next timer tick:
+                // told to go to make room, or past the grace for its login request or its
+                // authorization. Neither fires once the connection has authorized.
+                _ = unauthorized_evicted(&conn.unauthorized_id) => {
+                    conn.on_close("Timeout", true).await;
+                    break;
+                }
+                _ = login_deadline(conn.authorized, conn.login_request_seen, started) => {
+                    conn.on_close("Timeout", true).await;
+                    break;
+                }
                 Some(data) = rx_from_cm.recv() => {
                     match data {
                         ipc::Data::Authorize => {
@@ -1792,6 +1915,7 @@ impl Connection {
             return false;
         }
         self.authorized = true;
+        self.unauthorized_id = None;
         // Releases the budget `check_id_whitelist` charges against this address: only a peer
         // that got this far proved more than a self-reported id.
         self.clear_id_whitelist_failures();
@@ -2758,6 +2882,7 @@ impl Connection {
         }
         // After handling CloseReason messages, proceed to process other message types
         if let Some(message::Union::LoginRequest(lr)) = msg.union {
+            self.login_request_seen = true;
             if !self.check_login_scope(&lr).await {
                 return false;
             }
@@ -6994,6 +7119,167 @@ mod test {
     #[allow(unused)]
     use super::*;
 
+    // The registry is process-global and the harness runs tests in parallel threads, so every
+    // test that admits connections holds this first; a poisoned lock is still a lock.
+    static UNAUTHORIZED_TESTS: std::sync::Mutex<()> = std::sync::Mutex::new(());
+
+    fn unauthorized_count() -> usize {
+        UNAUTHORIZED_CONNS.lock().unwrap().len()
+    }
+
+    #[test]
+    fn test_unauthorized_admission_is_per_address() {
+        let _serial = UNAUTHORIZED_TESTS.lock().unwrap_or_else(|e| e.into_inner());
+        let a: IpAddr = "203.0.113.1".parse().unwrap();
+        let b: IpAddr = "203.0.113.2".parse().unwrap();
+        let held: Vec<_> = (0..MAX_UNAUTHORIZED_CONNS_PER_ADDR as i32)
+            .map(|i| admit_unauthorized(1_000_000 + i, a).unwrap())
+            .collect();
+        assert!(admit_unauthorized(1_000_100, a).is_none());
+        assert!(
+            held.iter().all(|u| !u.is_evicted()),
+            "a refusal evicts nobody"
+        );
+        let other = admit_unauthorized(1_000_101, b).unwrap();
+        assert!(!other.is_evicted());
+        drop(held);
+        assert!(admit_unauthorized(1_000_102, a).is_some());
+    }
+
+    #[tokio::test]
+    async fn test_unauthorized_full_tells_the_oldest_to_go_and_frees_its_place_only_when_it_has() {
+        let _serial = UNAUTHORIZED_TESTS.lock().unwrap_or_else(|e| e.into_inner());
+        let mut held: Vec<_> = (0..MAX_UNAUTHORIZED_CONNS as i32)
+            .map(|i| {
+                let ip: IpAddr = format!("198.51.100.{}", i + 1).parse().unwrap();
+                admit_unauthorized(2_000_000 + i, ip).unwrap()
+            })
+            .collect();
+        // At the limit the newcomer is refused, the oldest is told to go, and the count does
+        // not move: the place is still occupied.
+        assert!(admit_unauthorized(2_000_999, "198.51.100.250".parse().unwrap()).is_none());
+        assert!(held[0].is_evicted());
+        assert!(held[1..].iter().all(|u| !u.is_evicted()));
+        assert_eq!(unauthorized_count(), MAX_UNAUTHORIZED_CONNS);
+        hbb_common::timeout(1000, held[0].evicted()).await.unwrap();
+        // A further arrival while that one is still on its way out tells nobody else to go: a
+        // burst of refused arrivals clears no more room than a single one.
+        assert!(admit_unauthorized(2_001_000, "198.51.100.251".parse().unwrap()).is_none());
+        assert!(held[1..].iter().all(|u| !u.is_evicted()));
+        // Only once an evicted connection has gone is there a place for a newcomer.
+        drop(held.remove(0));
+        assert_eq!(unauthorized_count(), MAX_UNAUTHORIZED_CONNS - 1);
+        let newcomer = admit_unauthorized(2_001_001, "198.51.100.252".parse().unwrap()).unwrap();
+        assert!(!newcomer.is_evicted());
+        assert_eq!(unauthorized_count(), MAX_UNAUTHORIZED_CONNS);
+        // Full again, the next arrival tells the connection now oldest to go.
+        assert!(admit_unauthorized(2_001_002, "198.51.100.253".parse().unwrap()).is_none());
+        assert!(held[0].is_evicted());
+        assert!(held[1..].iter().all(|u| !u.is_evicted()));
+        assert!(!newcomer.is_evicted());
+    }
+
+    // The per-address share holds at the limit too: an address can turn out at most that many
+    // connections, one per place it then takes, and is refused before any eviction from then on.
+    #[test]
+    fn test_unauthorized_full_one_address_turns_out_at_most_its_share() {
+        let _serial = UNAUTHORIZED_TESTS.lock().unwrap_or_else(|e| e.into_inner());
+        let mut held: Vec<_> = (0..MAX_UNAUTHORIZED_CONNS as i32)
+            .map(|i| {
+                let ip: IpAddr = format!("198.51.100.{}", i + 1).parse().unwrap();
+                admit_unauthorized(3_000_000 + i, ip).unwrap()
+            })
+            .collect();
+        let flooder: IpAddr = "203.0.113.9".parse().unwrap();
+        let mut taken = Vec::new();
+        for i in 0..MAX_UNAUTHORIZED_CONNS_PER_ADDR as i32 {
+            assert!(admit_unauthorized(3_001_000 + i, flooder).is_none());
+            assert!(held[0].is_evicted());
+            drop(held.remove(0));
+            taken.push(admit_unauthorized(3_002_000 + i, flooder).unwrap());
+        }
+        assert_eq!(unauthorized_count(), MAX_UNAUTHORIZED_CONNS);
+        assert!(admit_unauthorized(3_003_000, flooder).is_none());
+        assert!(held.iter().all(|u| !u.is_evicted()));
+        assert!(taken.iter().all(|u| !u.is_evicted()));
+    }
+
+    /// A loopback TCP connection as create_tcp_connection sees it, and the controller's end,
+    /// which never speaks: the connection stalls in the identity handshake.
+    async fn stalled_incoming() -> (Stream, Stream, SocketAddr) {
+        let listener = hbb_common::tcp::new_listener("127.0.0.1:0", false)
+            .await
+            .unwrap();
+        let host = listener.local_addr().unwrap().to_string();
+        let controller = hbb_common::socket_client::connect_tcp(host, 3000)
+            .await
+            .unwrap();
+        let (accepted, addr) = listener.accept().await.unwrap();
+        let served = Stream::Tcp(hbb_common::tcp::FramedStream::from(accepted, addr));
+        (served, controller, addr)
+    }
+
+    // Live connections, not bookkeeping: with the limit reached by connections stalled in the
+    // handshake, one more arrival is refused and the oldest handshake is ended at once, not on
+    // a timer tick, so the live count never exceeds the limit and a place opens only then.
+    #[tokio::test]
+    async fn test_unauthorized_limit_bounds_live_handshakes() {
+        let _serial = UNAUTHORIZED_TESTS.lock().unwrap_or_else(|e| e.into_inner());
+        let server = crate::server::new_for_test();
+        let mut controllers = Vec::new();
+        let mut handshakes = Vec::new();
+        for i in 0..MAX_UNAUTHORIZED_CONNS {
+            let (served, controller, _) = stalled_incoming().await;
+            controllers.push(controller);
+            // Each from an address of its own, so only the global limit is in play.
+            let addr: SocketAddr = format!("192.0.2.{}:1", i + 1).parse().unwrap();
+            let server = server.clone();
+            handshakes.push(tokio::spawn(async move {
+                crate::server::create_tcp_connection(server, served, addr, true, Default::default())
+                    .await
+            }));
+        }
+        for _ in 0..200 {
+            if unauthorized_count() == MAX_UNAUTHORIZED_CONNS {
+                break;
+            }
+            hbb_common::sleep(0.02).await;
+        }
+        assert_eq!(unauthorized_count(), MAX_UNAUTHORIZED_CONNS);
+        assert!(handshakes.iter().all(|h| !h.is_finished()));
+
+        let (served, _controller, _) = stalled_incoming().await;
+        let addr: SocketAddr = "192.0.2.200:1".parse().unwrap();
+        let refused = crate::server::create_tcp_connection(
+            server.clone(),
+            served,
+            addr,
+            true,
+            Default::default(),
+        )
+        .await;
+        assert!(refused.is_err());
+        let ended = hbb_common::timeout(2000, handshakes.remove(0)).await;
+        assert!(
+            matches!(ended, Ok(Ok(Err(_)))),
+            "the oldest handshake ends on eviction"
+        );
+        assert_eq!(unauthorized_count(), MAX_UNAUTHORIZED_CONNS - 1);
+        assert!(
+            handshakes.iter().all(|h| !h.is_finished()),
+            "only the oldest was ended"
+        );
+
+        drop(controllers);
+        for h in handshakes {
+            assert!(
+                matches!(hbb_common::timeout(3000, h).await, Ok(Ok(Err(_)))),
+                "a stalled handshake ends when its controller goes"
+            );
+        }
+        assert_eq!(unauthorized_count(), 0, "no handshake outlives the test");
+    }
+
     #[cfg(feature = "flutter")]
     #[cfg(not(any(target_os = "android", target_os = "ios")))]
     #[test]