diff --git a/crates/ui-egui/src/control.rs b/crates/ui-egui/src/control.rs index 742641330c..630d66190c 100644 --- a/crates/ui-egui/src/control.rs +++ b/crates/ui-egui/src/control.rs @@ -793,14 +793,27 @@ struct Limits { max_unauthenticated: usize, /// How long a write may block on a client that never reads its replies. write_timeout: std::time::Duration, + /// How long a connection refused after a reply is kept half-open, discarding input, so the + /// reply is not lost to a reset (see [`close_after_reply`]). + linger: std::time::Duration, + /// How many refused connections may linger at once. A lingering connection has given back + /// its place among [`Self::max_unauthenticated`], so this is what bounds their threads; + /// beyond it a refused connection is closed at once. + max_lingering: usize, } #[cfg(not(target_arch = "wasm32"))] -const LIMITS: Limits = - Limits { auth_deadline: std::time::Duration::from_secs(10), max_unauthenticated: 8, write_timeout: std::time::Duration::from_secs(30) }; +const LIMITS: Limits = Limits { + auth_deadline: std::time::Duration::from_secs(10), + max_unauthenticated: 8, + write_timeout: std::time::Duration::from_secs(30), + linger: std::time::Duration::from_secs(1), + max_lingering: 8, +}; -/// One of the [`Limits::max_unauthenticated`] places, given back when dropped: when the -/// connection authenticates, and on every way its thread can end. +/// One of a limited number of places, given back when dropped: a place among +/// [`Limits::max_unauthenticated`] (given back when the connection authenticates or is refused, +/// and on every way its thread can end), or among [`Limits::max_lingering`]. #[cfg(not(target_arch = "wasm32"))] struct Slot(std::sync::Arc); @@ -850,6 +863,38 @@ impl std::io::Read for Reading { } } +/// Close a connection whose last reply has just been written without losing that reply. +/// +/// Closing a socket with input still unread makes the kernel send a reset instead of a FIN, and +/// the reset can discard the reply before the client reads it (Linux does): a client still +/// sending when it is refused would never learn why. So send FIN first, then read and discard +/// whatever arrives until the client closes its side or `linger` passes, and only then close. +/// +/// Only [`Limits::max_lingering`] connections linger at once: `place` is one of those places, +/// taken by [`refuse`] before the reply was written, and held for the drain only. Without one +/// the connection is closed at once after the FIN, accepting the risk of a reset rather than +/// holding another thread. +#[cfg(not(target_arch = "wasm32"))] +fn close_after_reply(reading: &mut Reading, place: Option, linger: std::time::Duration) { + use std::io::Read; + let _ = reading.stream.shutdown(std::net::Shutdown::Write); + let Some(_place) = place else { return }; + reading.until = Some(std::time::Instant::now() + linger); + let mut sink = [0u8; 1024]; + // Each read is bounded by the time left, and none is made once `linger` has passed. + while matches!(reading.read(&mut sink), Ok(n) if n > 0) {} +} + +/// A connection is being refused before `auth`: give back its place among those waiting to +/// authenticate, so it does not keep a real client out while it lingers, and take a place to +/// linger in if one is free. Done before the reply is written, so a client that has read the +/// reply sees the places as they will be. +#[cfg(not(target_arch = "wasm32"))] +fn refuse(slot: &mut Option, lingering: &std::sync::Arc, limits: Limits) -> Option { + drop(slot.take()); + Slot::take(lingering, limits.max_lingering) +} + /// Listen on `127.0.0.1` (random port) and forward authenticated requests to `client`. /// /// Protocol: newline-delimited JSON-RPC 2.0. The first request on a connection must be @@ -868,6 +913,7 @@ fn serve_with(client: ControlClient, limits: Limits) -> std::io::Result std::io::Result std::io::Result { let secs = limits.auth_deadline.as_secs_f64(); let e = rpc_error(Value::Null, -32004, &format!("not authenticated within {secs} s")); - let _ = writeln!(write, "{e}").and_then(|_| write.flush()); + let place = refuse(&mut slot, &lingering, limits); + if writeln!(write, "{e}").and_then(|_| write.flush()).is_ok() { + close_after_reply(reader.get_mut(), place, limits.linger); + } break; } // End of stream or a failed read ends the connection, as before. Ok(LineRead::End) | Err(_) => break, Ok(LineRead::TooLong) => { let e = rpc_error(Value::Null, -32600, &format!("request line longer than {max} bytes before auth")); - let _ = writeln!(write, "{e}").and_then(|_| write.flush()); + let place = refuse(&mut slot, &lingering, limits); + if writeln!(write, "{e}").and_then(|_| write.flush()).is_ok() { + close_after_reply(reader.get_mut(), place, limits.linger); + } break; } } @@ -945,7 +998,13 @@ fn serve_with(client: ControlClient, limits: Limits) -> std::io::Result = (0..limits.max_unauthenticated).map(|_| connect(ep.port)).collect(); + for s in &others { + // A refused one has already been closed, so its write may fail; its reply says why. + let _ = writeln!(s.try_clone().unwrap(), "{auth}"); + let reply = next_line(&mut BufReader::new(s.try_clone().unwrap())).expect("a reply"); + assert_eq!(reply["result"]["ok"], true, "admitted while the refused one lingers: {reply}"); + } + assert!(started.elapsed() < limits.linger, "checked while it lingered: {:?}", started.elapsed()); + drop(refused); + } + + /// Only [`super::Limits::max_lingering`] refused connections linger at once: one more still + /// gets its reply, and is then closed at once instead of lingering. + #[cfg(not(target_arch = "wasm32"))] + #[test] + fn only_so_many_refused_connections_linger() { + use std::io::{BufReader, Write}; + // A linger far longer than the checks below take, even on a loaded machine. + let limits = super::Limits { linger: std::time::Duration::from_secs(5), max_lingering: 1, ..TEST_LIMITS }; + let (_control, client) = super::attach(&egui::Context::default()); + let ep = super::serve_with(client, limits).unwrap(); + let bad = r#"{"jsonrpc": "2.0", "id": 1, "method": "page_count"}"#; + // The first refused connection takes the only lingering place and keeps it. + let first = connect(ep.port); + writeln!(first.try_clone().unwrap(), "{bad}").unwrap(); + let mut r1 = BufReader::new(first.try_clone().unwrap()); + assert_eq!(next_line(&mut r1).expect("refused with a reason")["error"]["code"], -32001); + // The second still gets its reason (it has nothing unread, so no reset)... + let second = connect(ep.port); + let mut w2 = second.try_clone().unwrap(); + writeln!(w2, "{bad}").unwrap(); + let mut r2 = BufReader::new(second.try_clone().unwrap()); + assert_eq!(next_line(&mut r2).expect("refused with a reason")["error"]["code"], -32001); + // ...and is closed outright, not lingering: writing to it soon fails, where a lingering + // connection would read and discard for the whole of `linger`. + let started = std::time::Instant::now(); + let closed = loop { + if w2.write_all(b"x").is_err() { + break true; + } + if started.elapsed() > limits.linger { + break false; + } + std::thread::sleep(std::time::Duration::from_millis(20)); + }; + assert!(closed && started.elapsed() < limits.linger / 2, "it lingered: {:?}", started.elapsed()); + drop(first); + } + /// #1806: only so many connections may wait to authenticate at once. One more is refused /// on the spot; once a waiting one authenticates or times out, its place is free again. #[cfg(not(target_arch = "wasm32"))]