From ca8d67d571b2a1e990405125ed469da56e3d01e9 Mon Sep 17 00:00:00 2001
From: Ian Jackson <ijackson@chiark.greenend.org.uk>
Date: Sun, 15 Aug 2021 17:08:25 +0100
Subject: [PATCH] read_limited_bytes: take a `capacity` argument

Signed-off-by: Ian Jackson <ijackson@chiark.greenend.org.uk>
---
 src/bin/client.rs | 4 +++-
 src/bin/server.rs | 2 +-
 src/utils.rs      | 3 +++
 3 files changed, 7 insertions(+), 2 deletions(-)

diff --git a/src/bin/client.rs b/src/bin/client.rs
index c3beab3..559b1af 100644
--- a/src/bin/client.rs
+++ b/src/bin/client.rs
@@ -140,7 +140,9 @@ 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, default(), &mut resp).await
+      let resp = read_limited_bytes(
+        max_body, default(), 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 8c3c23e..b0bac2a 100644
--- a/src/bin/server.rs
+++ b/src/bin/server.rs
@@ -86,7 +86,7 @@ async fn handle(
 
     let mut body = req.into_body();
     let initial = match read_limited_bytes(
-      METADATA_MAX_LEN, default(), &mut body
+      METADATA_MAX_LEN, default(), default(), &mut body
     ).await {
       Ok(all) => all,
       Err(ReadLimitedError::Truncated { sofar,.. }) => sofar,
diff --git a/src/utils.rs b/src/utils.rs
index 4b345d0..c210a7d 100644
--- a/src/utils.rs
+++ b/src/utils.rs
@@ -40,12 +40,15 @@ impl<T> Result<T,ReadLimitedError> {
 
 #[throws(ReadLimitedError)]
 pub async fn read_limited_bytes<S>(limit: usize, initial: Box<[u8]>,
+                                   capacity: usize,
                                    stream: &mut S) -> Box<[u8]>
 where S: futures::Stream<Item=Result<hyper::body::Bytes,hyper::Error>>
          + Debug + Unpin,
       // we also require that the Stream is cancellation-safe
 {
   let mut accum = initial.into_vec();
+  let capacity = min(limit, capacity);
+  if capacity > accum.len() { accum.reserve(capacity - accum.len()); }
   while let Some(item) = stream.next().await {
     let b = item?;
     accum.extend(b);
-- 
2.30.2