//! Bounded reader for the request line a socket *server* accepts. //! //! The JSON-line sockets share one framing rule: a request is one line. A //! peer that never sends the newline, or sends it a byte at a time, would pin //! a connection task and an unbounded buffer, so every server reads its //! request through [`read_request_line`]. use std::time::Duration; use tokio::io::{AsyncBufRead, AsyncBufReadExt, AsyncReadExt}; /// Longest request line accepted, newline included. The largest request a /// client sends is an `OperatorMsg` from the agent web UI's `/send` /// (`hive-agent/src/web_ui/actions.rs:21`), whose body axum's default limit /// caps at 2 MiB; JSON escaping can roughly double that to 4 MiB, and this is /// 4× that. pub const REQUEST_LINE_MAX_BYTES: usize = 16 * 1024 * 1024; /// Default for [`LineDeadlines::line`] and [`LineDeadlines::idle`] on the /// sockets that bound time. Local peers send a request in one write. pub const REQUEST_LINE_DEADLINE: Duration = Duration::from_secs(30); /// How long a server waits on a request line. `None` waits forever. #[derive(Debug, Clone, Copy)] pub struct LineDeadlines { /// Wait for the first byte of the line. Unbounded for connections that /// legitimately sit idle between requests. pub idle: Option, /// Wait from the first byte to the newline. pub line: Option, } impl LineDeadlines { /// No deadline at all; only the size cap applies. pub const NONE: Self = Self { idle: None, line: None, }; } /// What [`read_request_line`] found. #[derive(Debug, PartialEq, Eq)] pub enum RequestLine { Eof, Complete, /// [`REQUEST_LINE_MAX_BYTES`] arrived without a newline. TooLong, /// A [`LineDeadlines`] expired first. TimedOut, } impl RequestLine { /// The error text for a line that can't be served, `None` when it can. /// Every server replies with this before closing the connection, because /// the rest of the line is still unread. #[must_use] pub fn refusal(&self) -> Option { match self { Self::Eof | Self::Complete => None, Self::TooLong => Some(format!( "request too long: no newline within {REQUEST_LINE_MAX_BYTES} bytes; \ connection closed" )), Self::TimedOut => { Some("request timed out: no complete line in time; connection closed".to_owned()) } } } } /// Read one request line into `buf`, consuming at most /// [`REQUEST_LINE_MAX_BYTES`] from `reader`. /// /// # Errors /// The underlying read failed. pub async fn read_request_line( reader: &mut R, buf: &mut Vec, deadlines: LineDeadlines, ) -> std::io::Result { buf.clear(); match with_deadline(deadlines.idle, reader.fill_buf()).await { None => return Ok(RequestLine::TimedOut), Some(Err(e)) => return Err(e), Some(Ok([])) => return Ok(RequestLine::Eof), Some(Ok(_)) => {} } let mut limited = (&mut *reader).take(REQUEST_LINE_MAX_BYTES as u64); let Some(n) = with_deadline(deadlines.line, limited.read_until(b'\n', buf)).await else { return Ok(RequestLine::TimedOut); }; let n = n?; Ok(if n == REQUEST_LINE_MAX_BYTES && !buf.ends_with(b"\n") { RequestLine::TooLong } else { RequestLine::Complete }) } async fn with_deadline(limit: Option, fut: F) -> Option { match limit { Some(d) => tokio::time::timeout(d, fut).await.ok(), None => Some(fut.await), } } #[cfg(test)] mod tests { use tokio::io::{AsyncWriteExt, BufReader}; use super::*; #[tokio::test] async fn over_long_line_is_rejected_after_reading_only_the_bound() { let total = 2 * REQUEST_LINE_MAX_BYTES as u64; let mut reader = BufReader::new(tokio::io::repeat(b'x').take(total)); let mut buf = Vec::new(); let outcome = read_request_line(&mut reader, &mut buf, LineDeadlines::NONE) .await .expect("read"); // At most one `BufReader` fill past the bound left the source. let unread = reader.into_inner().limit(); let consumed = total - unread; assert!( consumed <= (REQUEST_LINE_MAX_BYTES + 8 * 1024) as u64, "consumed {consumed} of {total} bytes" ); assert_eq!(buf.len(), REQUEST_LINE_MAX_BYTES); assert_eq!(outcome, RequestLine::TooLong); } #[tokio::test] async fn complete_line_and_eof() { let mut reader = BufReader::new(&b"{\"a\":1}\n"[..]); let mut buf = Vec::new(); let d = LineDeadlines::NONE; assert_eq!( read_request_line(&mut reader, &mut buf, d).await.unwrap(), RequestLine::Complete ); assert_eq!(buf, b"{\"a\":1}\n"); assert_eq!( read_request_line(&mut reader, &mut buf, d).await.unwrap(), RequestLine::Eof ); } #[tokio::test] async fn sender_that_stalls_mid_line_times_out() { let (mut tx, rx) = tokio::io::duplex(64); tx.write_all(b"{\"partial\"").await.unwrap(); let mut reader = BufReader::new(rx); let mut buf = Vec::new(); let d = LineDeadlines { idle: None, line: Some(Duration::from_millis(50)), }; let outcome = read_request_line(&mut reader, &mut buf, d).await.unwrap(); assert_eq!(outcome, RequestLine::TimedOut); drop(tx); } #[tokio::test] async fn sender_that_never_starts_times_out_only_with_an_idle_deadline() { let (tx, rx) = tokio::io::duplex(64); let mut reader = BufReader::new(rx); let mut buf = Vec::new(); let d = LineDeadlines { idle: Some(Duration::from_millis(50)), line: None, }; let outcome = read_request_line(&mut reader, &mut buf, d).await.unwrap(); assert_eq!(outcome, RequestLine::TimedOut); // Without an idle deadline the same silence keeps waiting. let line_only = LineDeadlines { idle: None, line: Some(Duration::from_millis(50)), }; let waiting = tokio::time::timeout( Duration::from_millis(200), read_request_line(&mut reader, &mut buf, line_only), ) .await; assert!(waiting.is_err(), "idle wait must not be bounded by `line`"); drop(tx); } }