Watch
0
0
Fork
You've already forked hyperhive
0

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:
atlas 2026-10-16 11:36:14 +02:00 • committed by mara
commit 1f214f9f01
10 changed files with 425 additions and 87 deletions

View file

@ -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"

View file

@ -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");
}
}

View file

@ -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();