From 29d55c57efac7745daa364d15e18aa4ba00945a6 Mon Sep 17 00:00:00 2001 From: Ian Jackson Date: Sun, 15 Aug 2021 17:04:16 +0100 Subject: [PATCH] read_limited_bytes: take an `initial` argument Signed-off-by: Ian Jackson --- src/bin/client.rs | 2 +- src/bin/server.rs | 4 +++- src/utils.rs | 5 +++-- 3 files changed, 7 insertions(+), 4 deletions(-) diff --git a/src/bin/client.rs b/src/bin/client.rs index 2497780..c3beab3 100644 --- a/src/bin/client.rs +++ b/src/bin/client.rs @@ -140,7 +140,7 @@ fn submit_request<'r, 'c:'r, C:HCC>( let status = resp.status(); let mut resp = resp.into_body(); let max_body = c.ic.max_batch_down.sat() + MAX_OVERHEAD; - let resp = read_limited_bytes(max_body, &mut resp).await + let resp = read_limited_bytes(max_body, default(), &mut resp).await .discard_data().context("fetching response body")?; if ! status.is_success() { diff --git a/src/bin/server.rs b/src/bin/server.rs index 8c0d996..8c3c23e 100644 --- a/src/bin/server.rs +++ b/src/bin/server.rs @@ -85,7 +85,9 @@ async fn handle( }; let mut body = req.into_body(); - let initial = match read_limited_bytes(METADATA_MAX_LEN, &mut body).await { + let initial = match read_limited_bytes( + METADATA_MAX_LEN, default(), &mut body + ).await { Ok(all) => all, Err(ReadLimitedError::Truncated { sofar,.. }) => sofar, Err(ReadLimitedError::Hyper(e)) => throw!(e), diff --git a/src/utils.rs b/src/utils.rs index 5c10752..4b345d0 100644 --- a/src/utils.rs +++ b/src/utils.rs @@ -39,12 +39,13 @@ impl Result { } #[throws(ReadLimitedError)] -pub async fn read_limited_bytes(limit: usize, stream: &mut S) -> Box<[u8]> +pub async fn read_limited_bytes(limit: usize, initial: Box<[u8]>, + stream: &mut S) -> Box<[u8]> where S: futures::Stream> + Debug + Unpin, // we also require that the Stream is cancellation-safe { - let mut accum = vec![]; + let mut accum = initial.into_vec(); while let Some(item) = stream.next().await { let b = item?; accum.extend(b); -- 2.30.2