diff --git a/src/server.rs b/src/server.rs index f39620f51..1675b28b3 100644 --- a/src/server.rs +++ b/src/server.rs @@ -118,6 +118,15 @@ pub struct Server { pub type ServerPtr = Arc>; pub type ServerPtrWeak = Weak>; +#[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,48 @@ 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 +317,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(()) } diff --git a/src/server/connection.rs b/src/server/connection.rs index 24f5bf537..7614705fe 100644 --- a/src/server/connection.rs +++ b/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,104 @@ 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 not authorized within this long of starting is closed, however alive it +/// keeps itself: a wrong password, a pending 2FA, an accept prompt or an admin-terminal +/// credential prompt still unanswered. The controller reconnects on its own and the prompt +/// comes back. A connection that says nothing at all goes at the 30 s idle timeout already. +const LOGIN_GRACE: Duration = Duration::from_secs(180); +/// Connections between accept and authorization, across every transport: the resource bound. +/// At this many a further arrival is refused and the oldest is told to go, one at a time. +const MAX_UNAUTHORIZED_CONNS: usize = 64; +/// Of those, how many one address may hold at once: a quarter of the room. A fairness cap +/// against the cheapest flood, one host with one address, not a security boundary: any pool +/// 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; + +/// 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, +} + +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 { + 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) { + 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. +async fn login_deadline(authorized: bool, started: Instant) { + if authorized { + return std::future::pending().await; + } + time::sleep_until(started + LOGIN_GRACE).await +} + lazy_static::lazy_static! { + // Connections between accept and authorization, oldest first; see admit_unauthorized. + static ref UNAUTHORIZED_CONNS: Mutex)>> = 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 +362,8 @@ pub struct Connection { port_forward_address: String, tx_to_cm: mpsc::UnboundedSender, authorized: bool, + // The place among the unauthorized connections; given back at authorization. + unauthorized_id: Option, require_2fa: Option, awaiting_2fa: bool, keyboard: bool, @@ -419,6 +519,7 @@ impl Connection { id: i32, server: super::ServerPtrWeak, meta: super::ConnectionMeta, + unauthorized: UnauthorizedID, ) { let super::ConnectionMeta { control_permissions, @@ -477,6 +578,7 @@ impl Connection { port_forward_address: "".to_owned(), tx_to_cm, authorized: false, + unauthorized_id: Some(unauthorized), 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), @@ -589,6 +691,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. @@ -624,6 +727,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 authorization. Neither + // fires once the connection has authorized. + _ = unauthorized_evicted(&conn.unauthorized_id) => { + conn.on_close("Timeout", true).await; + break; + } + _ = login_deadline(conn.authorized, started) => { + conn.on_close("Timeout", true).await; + break; + } Some(data) = rx_from_cm.recv() => { match data { ipc::Data::Authorize => { @@ -1761,6 +1875,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(); @@ -6963,6 +7078,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]