bound the request line on the host admin and in-agent sockets
hive-c0re's admin socket and hive-agent's in-agent socket (mode 0666) read the request with an unbounded read_line: no length cap and no deadline, so a peer that never sends a newline grows a buffer and pins a task forever. The bounded reader (16 MiB cap) moves from hive-c0re's socket_server into hive-sock-client, the one crate both daemons already share or can take, and gains optional deadlines for the first byte and for the rest of the line. Both sockets now use it and answer an over-long or stalled line with the same error reply the bounded sockets give, then close. The admin socket bounds only the started line (clients may idle between requests); the in-agent socket is single-shot and bounds silence too. The per-agent sockets keep their behaviour (no deadline).
This commit is contained in:
parent
88dd58e947
commit
1f214f9f01
10 changed files with 425 additions and 87 deletions
|
|
@ -107,7 +107,9 @@ Wire types + socket plumbing:
|
|||
spec. ⚠️ **This socket never leaves the container** — `hive-c0re` is not
|
||||
in the path, unlike `hive-core-agent-sock`.
|
||||
- **`hive-sock-client/`** — the shared JSON-line unix-socket client; retry
|
||||
is a policy value (`Retry::None` / `Retry::RideOutRestart`).
|
||||
is a policy value (`Retry::None` / `Retry::RideOutRestart`). Also holds
|
||||
the bounded request-line reader (`read_request_line`) the socket servers
|
||||
use.
|
||||
|
||||
Swarm-level:
|
||||
|
||||
|
|
|
|||
1
Cargo.lock
generated
1
Cargo.lock
generated
|
|
@ -1776,6 +1776,7 @@ dependencies = [
|
|||
"hive-log",
|
||||
"hive-priv-sock",
|
||||
"hive-sh4re",
|
||||
"hive-sock-client",
|
||||
"hive-types",
|
||||
"hmac 0.13.0",
|
||||
"indicatif",
|
||||
|
|
|
|||
|
|
@ -27,7 +27,8 @@ use std::sync::Arc;
|
|||
use anyhow::{Context, Result};
|
||||
use hive_agent_sock::{Request, Response};
|
||||
use hive_sh4re::inbox::{LooseEnd, saturating_age};
|
||||
use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
|
||||
use hive_sock_client::{LineDeadlines, REQUEST_LINE_DEADLINE, RequestLine, read_request_line};
|
||||
use tokio::io::{AsyncWriteExt, BufReader};
|
||||
use tokio::net::{UnixListener, UnixStream};
|
||||
use tokio::sync::Notify;
|
||||
|
||||
|
|
@ -146,18 +147,43 @@ async fn handle_conn(
|
|||
wake: &Notify,
|
||||
reminders: Option<&Reminders>,
|
||||
bus: &Bus,
|
||||
) -> Result<()> {
|
||||
serve_one(stream, REQUEST_DEADLINES, |req| {
|
||||
dispatch(req, store, wake, reminders, bus)
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
/// One request line, one response line. A single-shot connection, so silence
|
||||
/// is bounded too: the socket is mode 0666 and any local process can open it.
|
||||
const REQUEST_DEADLINES: LineDeadlines = LineDeadlines {
|
||||
idle: Some(REQUEST_LINE_DEADLINE),
|
||||
line: Some(REQUEST_LINE_DEADLINE),
|
||||
};
|
||||
|
||||
/// Read the request line under `deadlines`, answer it with `apply`, and write
|
||||
/// the response. An over-long or stalled line gets a `Response::Err`.
|
||||
async fn serve_one(
|
||||
stream: UnixStream,
|
||||
deadlines: LineDeadlines,
|
||||
apply: impl FnOnce(Request) -> Response,
|
||||
) -> Result<()> {
|
||||
let (read, mut write) = stream.into_split();
|
||||
let mut reader = BufReader::new(read);
|
||||
let mut line = String::new();
|
||||
if reader.read_line(&mut line).await? == 0 {
|
||||
let mut line = Vec::new();
|
||||
let outcome = read_request_line(&mut reader, &mut line, deadlines).await?;
|
||||
if outcome == RequestLine::Eof {
|
||||
return Ok(());
|
||||
}
|
||||
let resp = match serde_json::from_str::<Request>(line.trim()) {
|
||||
Ok(req) => dispatch(req, store, wake, reminders, bus),
|
||||
Err(e) => Response::Err {
|
||||
message: format!("bad request: {e}"),
|
||||
},
|
||||
let resp = if let Some(message) = outcome.refusal() {
|
||||
Response::Err { message }
|
||||
} else {
|
||||
match serde_json::from_slice::<Request>(line.trim_ascii()) {
|
||||
Ok(req) => apply(req),
|
||||
Err(e) => Response::Err {
|
||||
message: format!("bad request: {e}"),
|
||||
},
|
||||
}
|
||||
};
|
||||
let mut out = serde_json::to_string(&resp)?;
|
||||
out.push('\n');
|
||||
|
|
@ -464,6 +490,8 @@ fn to_loose_end(t: Todo) -> LooseEnd {
|
|||
mod tests {
|
||||
use super::*;
|
||||
use chrono::{DateTime, Duration, Utc};
|
||||
use hive_sock_client::REQUEST_LINE_MAX_BYTES;
|
||||
use tokio::io::AsyncBufReadExt;
|
||||
|
||||
/// `saturating_age`'s own arms are exact (`hive-sh4re::inbox`); what
|
||||
/// is only checkable here is *which* stored timestamp each projection
|
||||
|
|
@ -583,4 +611,70 @@ mod tests {
|
|||
other => panic!("expected a Todo loose end, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
/// Send `input` (then keep the write half open), return the one response line.
|
||||
async fn exchange(input: Vec<u8>, deadlines: LineDeadlines) -> (Response, Result<()>) {
|
||||
let (client, server) = UnixStream::pair().expect("socketpair");
|
||||
let handler = tokio::spawn(serve_one(server, deadlines, |_| Response::Ok));
|
||||
let (read, mut write) = client.into_split();
|
||||
let _writer = tokio::spawn(async move {
|
||||
let _ = write.write_all(&input).await;
|
||||
std::future::pending::<()>().await;
|
||||
});
|
||||
let mut reply = String::new();
|
||||
BufReader::new(read)
|
||||
.read_line(&mut reply)
|
||||
.await
|
||||
.expect("reply");
|
||||
let resp = serde_json::from_str(&reply).expect("response json");
|
||||
(resp, handler.await.expect("join"))
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn a_valid_request_is_applied() {
|
||||
let (resp, done) = exchange(b"{\"cmd\":\"compact\"}\n".to_vec(), REQUEST_DEADLINES).await;
|
||||
assert!(matches!(resp, Response::Ok), "{resp:?}");
|
||||
done.expect("served");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn an_over_long_line_gets_an_error_not_an_unbounded_buffer() {
|
||||
let (resp, done) =
|
||||
exchange(vec![b'x'; REQUEST_LINE_MAX_BYTES + 1], REQUEST_DEADLINES).await;
|
||||
assert!(
|
||||
matches!(&resp, Response::Err { message } if message.starts_with("request too long")),
|
||||
"{resp:?}"
|
||||
);
|
||||
done.expect("served");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn a_sender_that_stalls_mid_line_times_out() {
|
||||
let short = std::time::Duration::from_millis(50);
|
||||
let deadlines = LineDeadlines {
|
||||
idle: Some(short),
|
||||
line: Some(short),
|
||||
};
|
||||
let (resp, done) = exchange(b"{\"cmd\"".to_vec(), deadlines).await;
|
||||
assert!(
|
||||
matches!(&resp, Response::Err { message } if message.starts_with("request timed out")),
|
||||
"{resp:?}"
|
||||
);
|
||||
done.expect("served");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn a_sender_that_never_writes_times_out() {
|
||||
let short = std::time::Duration::from_millis(50);
|
||||
let deadlines = LineDeadlines {
|
||||
idle: Some(short),
|
||||
line: Some(short),
|
||||
};
|
||||
let (resp, done) = exchange(Vec::new(), deadlines).await;
|
||||
assert!(
|
||||
matches!(&resp, Response::Err { message } if message.starts_with("request timed out")),
|
||||
"{resp:?}"
|
||||
);
|
||||
done.expect("served");
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -39,6 +39,7 @@ hive-jobq.workspace = true
|
|||
hive-jobq-wire.workspace = true
|
||||
hive-log.workspace = true
|
||||
hive-priv-sock.workspace = true
|
||||
hive-sock-client.workspace = true
|
||||
hive-types.workspace = true
|
||||
libc.workspace = true
|
||||
listenfd = "1"
|
||||
|
|
|
|||
|
|
@ -4,7 +4,8 @@ use std::sync::Arc;
|
|||
use anyhow::{Context, Result};
|
||||
use hive_host_sock::{HostRequest, HostResponse, LifecycleScope};
|
||||
use hive_priv_sock::{InfraAction, InfraContainer};
|
||||
use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
|
||||
use hive_sock_client::{LineDeadlines, REQUEST_LINE_DEADLINE, RequestLine, read_request_line};
|
||||
use tokio::io::{AsyncWriteExt, BufReader};
|
||||
use tokio::net::{UnixListener, UnixStream};
|
||||
|
||||
use crate::actions;
|
||||
|
|
@ -53,17 +54,35 @@ pub async fn serve(socket: &Path, coord: Arc<Coordinator>) -> Result<()> {
|
|||
}
|
||||
|
||||
async fn handle(stream: UnixStream, coord: Arc<Coordinator>) -> Result<()> {
|
||||
handle_with(stream, coord, ADMIN_LINE_DEADLINES).await
|
||||
}
|
||||
|
||||
/// A client may idle between requests on one connection, so only the time to
|
||||
/// finish a started line is bounded.
|
||||
const ADMIN_LINE_DEADLINES: LineDeadlines = LineDeadlines {
|
||||
idle: None,
|
||||
line: Some(REQUEST_LINE_DEADLINE),
|
||||
};
|
||||
|
||||
async fn handle_with(
|
||||
stream: UnixStream,
|
||||
coord: Arc<Coordinator>,
|
||||
deadlines: LineDeadlines,
|
||||
) -> Result<()> {
|
||||
let (read, mut write) = stream.into_split();
|
||||
let mut reader = BufReader::new(read);
|
||||
let mut line = String::new();
|
||||
let mut line = Vec::new();
|
||||
|
||||
loop {
|
||||
line.clear();
|
||||
let n = reader.read_line(&mut line).await?;
|
||||
if n == 0 {
|
||||
let outcome = read_request_line(&mut reader, &mut line, deadlines).await?;
|
||||
if outcome == RequestLine::Eof {
|
||||
return Ok(());
|
||||
}
|
||||
let req = match serde_json::from_str::<HostRequest>(line.trim()) {
|
||||
if let Some(message) = outcome.refusal() {
|
||||
write_response(&mut write, &HostResponse::error(message)).await?;
|
||||
return Ok(());
|
||||
}
|
||||
let req = match serde_json::from_slice::<HostRequest>(line.trim_ascii()) {
|
||||
Ok(req) => req,
|
||||
Err(e) => {
|
||||
write_response(
|
||||
|
|
@ -828,3 +847,73 @@ fn finish_lifecycle(ok_items: Vec<String>, errors: &[String]) -> HostResponse {
|
|||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::time::Duration;
|
||||
|
||||
use hive_sock_client::REQUEST_LINE_MAX_BYTES;
|
||||
use tokio::io::{AsyncBufReadExt, AsyncWriteExt};
|
||||
|
||||
use super::*;
|
||||
|
||||
/// A real `Coordinator` over a throwaway sqlite dir; `handle_with` takes
|
||||
/// an `Arc<Coordinator>` and `open` touches nothing outside the db path.
|
||||
fn coordinator() -> (tempfile::TempDir, Arc<Coordinator>) {
|
||||
let dir = tempfile::tempdir().expect("tempdir");
|
||||
let coord = Coordinator::open(
|
||||
&dir.path().join("broker.sqlite"),
|
||||
crate::coordinator::HiveEnv::default(),
|
||||
crate::hive_stats::PriceTable::new(),
|
||||
1,
|
||||
)
|
||||
.expect("open coordinator");
|
||||
(dir, Arc::new(coord))
|
||||
}
|
||||
|
||||
async fn reply_error(client: &mut tokio::net::unix::OwnedReadHalf) -> String {
|
||||
let mut reply = String::new();
|
||||
BufReader::new(client)
|
||||
.read_line(&mut reply)
|
||||
.await
|
||||
.expect("reply");
|
||||
let resp: HostResponse = serde_json::from_str(&reply).expect("response json");
|
||||
resp.error.expect("error reply")
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn over_long_request_gets_an_error_and_the_connection_closes() {
|
||||
let (_dir, coord) = coordinator();
|
||||
let (client, server) = UnixStream::pair().expect("socketpair");
|
||||
let handler = tokio::spawn(handle_with(server, coord, ADMIN_LINE_DEADLINES));
|
||||
let (mut read, mut write) = client.into_split();
|
||||
let writer = tokio::spawn(async move {
|
||||
let chunk = vec![b'x'; 64 * 1024];
|
||||
for _ in 0..2 * REQUEST_LINE_MAX_BYTES / chunk.len() {
|
||||
if write.write_all(&chunk).await.is_err() {
|
||||
return;
|
||||
}
|
||||
}
|
||||
});
|
||||
let message = reply_error(&mut read).await;
|
||||
assert!(message.starts_with("request too long"), "{message}");
|
||||
handler.await.expect("join").expect("clean close");
|
||||
writer.await.expect("join");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn stalled_request_gets_an_error_and_the_connection_closes() {
|
||||
let (_dir, coord) = coordinator();
|
||||
let (client, server) = UnixStream::pair().expect("socketpair");
|
||||
let deadlines = LineDeadlines {
|
||||
idle: None,
|
||||
line: Some(Duration::from_millis(50)),
|
||||
};
|
||||
let handler = tokio::spawn(handle_with(server, coord, deadlines));
|
||||
let (mut read, mut write) = client.into_split();
|
||||
write.write_all(b"{\"partial\"").await.expect("write");
|
||||
let message = reply_error(&mut read).await;
|
||||
assert!(message.starts_with("request timed out"), "{message}");
|
||||
handler.await.expect("join").expect("clean close");
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -10,7 +10,8 @@ use std::sync::Arc;
|
|||
use anyhow::{Context, Result};
|
||||
use hive_core_agent_sock::{Request, Response};
|
||||
use hive_sh4re::inbox::Message;
|
||||
use tokio::io::{AsyncBufRead, AsyncBufReadExt, AsyncReadExt, AsyncWriteExt, BufReader};
|
||||
use hive_sock_client::{LineDeadlines, RequestLine, read_request_line};
|
||||
use tokio::io::{AsyncWriteExt, BufReader};
|
||||
use tokio::net::{UnixListener, UnixStream};
|
||||
use tokio::task::JoinHandle;
|
||||
|
||||
|
|
@ -130,39 +131,6 @@ fn exit_on_dead_listener(agent: &str, e: &std::io::Error) -> ! {
|
|||
std::process::exit(1)
|
||||
}
|
||||
|
||||
/// Longest request line accepted on the agent sockets, 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.
|
||||
const REQUEST_LINE_MAX_BYTES: usize = 16 * 1024 * 1024;
|
||||
|
||||
enum RequestLine {
|
||||
Eof,
|
||||
Complete,
|
||||
TooLong,
|
||||
}
|
||||
|
||||
/// Read one request line into `buf`, consuming at most
|
||||
/// [`REQUEST_LINE_MAX_BYTES`] from `reader`.
|
||||
async fn read_request_line<R: AsyncBufRead + Unpin>(
|
||||
reader: &mut R,
|
||||
buf: &mut Vec<u8>,
|
||||
) -> std::io::Result<RequestLine> {
|
||||
buf.clear();
|
||||
let n = (&mut *reader)
|
||||
.take(REQUEST_LINE_MAX_BYTES as u64)
|
||||
.read_until(b'\n', buf)
|
||||
.await?;
|
||||
Ok(if n == 0 {
|
||||
RequestLine::Eof
|
||||
} else if n == REQUEST_LINE_MAX_BYTES && !buf.ends_with(b"\n") {
|
||||
RequestLine::TooLong
|
||||
} else {
|
||||
RequestLine::Complete
|
||||
})
|
||||
}
|
||||
|
||||
/// Whether `e`'s root cause is the peer closing the connection.
|
||||
/// `downcast_ref` searches the whole error chain (anyhow's own
|
||||
/// `object_downcast`, not just the outermost type), so this finds a
|
||||
|
|
@ -177,7 +145,7 @@ fn is_broken_pipe(e: &anyhow::Error) -> bool {
|
|||
})
|
||||
}
|
||||
|
||||
/// Answer each request line on `stream` until EOF. An over-long line gets a
|
||||
/// Answer each request line on `stream` until EOF. An over-long (or timed-out) line gets a
|
||||
/// `Response::Err` and then the connection is closed, because the rest of the
|
||||
/// line is still unread.
|
||||
async fn serve(stream: UnixStream, agent: String, coord: Arc<Coordinator>) -> Result<()> {
|
||||
|
|
@ -185,18 +153,14 @@ async fn serve(stream: UnixStream, agent: String, coord: Arc<Coordinator>) -> Re
|
|||
let mut reader = BufReader::new(read);
|
||||
let mut line = Vec::new();
|
||||
loop {
|
||||
let (resp, too_long) = match read_request_line(&mut reader, &mut line).await? {
|
||||
RequestLine::Eof => return Ok(()),
|
||||
RequestLine::TooLong => (
|
||||
Response::Err {
|
||||
message: format!(
|
||||
"request too long: no newline within {REQUEST_LINE_MAX_BYTES} bytes; \
|
||||
connection closed"
|
||||
),
|
||||
},
|
||||
true,
|
||||
),
|
||||
RequestLine::Complete => match serde_json::from_slice::<Request>(&line) {
|
||||
let outcome = read_request_line(&mut reader, &mut line, LineDeadlines::NONE).await?;
|
||||
if outcome == RequestLine::Eof {
|
||||
return Ok(());
|
||||
}
|
||||
let refusal = outcome.refusal();
|
||||
let (resp, too_long) = match refusal {
|
||||
Some(message) => (Response::Err { message }, true),
|
||||
None => match serde_json::from_slice::<Request>(&line) {
|
||||
Ok(req) => (dispatch(&req, &agent, &coord).await, false),
|
||||
Err(e) => (
|
||||
Response::Err {
|
||||
|
|
@ -211,9 +175,7 @@ async fn serve(stream: UnixStream, agent: String, coord: Arc<Coordinator>) -> Re
|
|||
write.write_all(payload.as_bytes()).await?;
|
||||
write.flush().await?;
|
||||
if too_long {
|
||||
anyhow::bail!(
|
||||
"`{agent}` sent a request line over {REQUEST_LINE_MAX_BYTES} bytes; closed"
|
||||
);
|
||||
anyhow::bail!("`{agent}` sent an over-long or stalled request line; closed");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -602,6 +564,9 @@ pub(crate) fn handle_send(
|
|||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use hive_sock_client::REQUEST_LINE_MAX_BYTES;
|
||||
use tokio::io::AsyncBufReadExt;
|
||||
|
||||
use super::*;
|
||||
|
||||
/// A real `Coordinator` over a throwaway sqlite dir. `serve` takes an
|
||||
|
|
@ -636,25 +601,6 @@ mod tests {
|
|||
);
|
||||
}
|
||||
|
||||
#[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)
|
||||
.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!(matches!(outcome, RequestLine::TooLong));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn serve_answers_an_over_long_line_with_an_error_and_closes() {
|
||||
let (_dir, coord) = coordinator();
|
||||
|
|
|
|||
|
|
@ -1,5 +1,6 @@
|
|||
# `hive-sock-client` — the one JSON-line-over-unix-socket client shared by
|
||||
# every daemon that talks to a hyperhive socket. Protocol-agnostic on
|
||||
# every daemon that talks to a hyperhive socket, plus the bounded
|
||||
# request-line reader the socket servers share. Protocol-agnostic on
|
||||
# purpose: it is generic over the request/response types, so it works for
|
||||
# the host-served control socket and the harness's in-agent socket alike.
|
||||
# Deliberately NOT folded into one of the `*-sock` wire-type crates —
|
||||
|
|
|
|||
|
|
@ -1,7 +1,8 @@
|
|||
# hive-sock-client
|
||||
|
||||
One JSON-line-over-unix-socket client, shared by every daemon that speaks
|
||||
to a hyperhive socket.
|
||||
to a hyperhive socket, plus the bounded request-line reader the socket
|
||||
servers share.
|
||||
|
||||
The wire protocol is the same everywhere: connect, write one line of JSON,
|
||||
read one line of JSON back. Before this crate existed, five daemons each
|
||||
|
|
@ -23,6 +24,15 @@ in-agent socket both use it, with their own wire-type crates.
|
|||
on. The drain is not optional: without it the server's write-back lands
|
||||
on a closed socket.
|
||||
|
||||
## Server side
|
||||
|
||||
`read_request_line` is the one bounded reader for the request line a
|
||||
socket server accepts: at most `REQUEST_LINE_MAX_BYTES`, with optional
|
||||
`LineDeadlines` (wait for the first byte, wait for the newline). A server
|
||||
answers `RequestLine::refusal()` and closes. `hive-c0re` uses it on the
|
||||
host admin socket and the per-agent sockets; `hive-agent` on
|
||||
the in-agent socket.
|
||||
|
||||
## Retry is a policy value, not a fork
|
||||
|
||||
`Retry::RideOutRestart` backs off 2/4/8/16/30s (60s total), sized to ride
|
||||
|
|
|
|||
|
|
@ -1,5 +1,6 @@
|
|||
//! JSON-line-over-unix-socket client, shared by every daemon that talks to
|
||||
//! a hyperhive socket.
|
||||
//! a hyperhive socket, plus [`read_request_line`], the bounded request-line
|
||||
//! reader its servers share.
|
||||
//!
|
||||
//! The wire protocol is identical everywhere — connect, write one line of
|
||||
//! JSON, read one line of JSON back — so this crate is generic over the
|
||||
|
|
@ -27,6 +28,12 @@
|
|||
//! choice and stays at the call site — it is not a property of the
|
||||
//! transport.
|
||||
|
||||
mod request_line;
|
||||
|
||||
pub use request_line::{
|
||||
LineDeadlines, REQUEST_LINE_DEADLINE, REQUEST_LINE_MAX_BYTES, RequestLine, read_request_line,
|
||||
};
|
||||
|
||||
use std::path::Path;
|
||||
use std::time::Duration;
|
||||
|
||||
|
|
|
|||
187
hive-sock-client/src/request_line.rs
Normal file
187
hive-sock-client/src/request_line.rs
Normal file
|
|
@ -0,0 +1,187 @@
|
|||
//! 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<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);
|
||||
}
|
||||
}
|
||||
Loading…
Reference in a new issue