hyperhive/hive-sock-client/src/request_line.rs

//! 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. A generous bound, not
/// tuned to one request: it was sized for an operator message from the agent web
/// UI (a body of up to 2 MiB, up to 4 MiB once JSON-escaped, with 4× headroom),
/// and no client sends anything near it now.
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<Duration>,
    /// Wait from the first byte to the newline.
    pub line: Option<Duration>,
}

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<String> {
        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<R: AsyncBufRead + Unpin>(
    reader: &mut R,
    buf: &mut Vec<u8>,
    deadlines: LineDeadlines,
) -> std::io::Result<RequestLine> {
    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<F: Future>(limit: Option<Duration>, fut: F) -> Option<F::Output> {
    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);
    }
}