From 5e2b9f4a0ee1a1390e1af04fec97b7ee25d36833 Mon Sep 17 00:00:00 2001
From: Ian Jackson <ijackson@chiark.greenend.org.uk>
Date: Sun, 8 Aug 2021 15:47:57 +0100
Subject: [PATCH] refactor body limit, for reuse in server, allow to discard
 data

Signed-off-by: Ian Jackson <ijackson@chiark.greenend.org.uk>
---
 src/bin/client.rs |  2 +-
 src/utils.rs      | 13 +++++++++++++
 2 files changed, 14 insertions(+), 1 deletion(-)

diff --git a/src/bin/client.rs b/src/bin/client.rs
index 1f382bb..201ca8e 100644
--- a/src/bin/client.rs
+++ b/src/bin/client.rs
@@ -143,7 +143,7 @@ fn submit_request<'r, 'c:'r, C:HCC>(
       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
-        .context("fetching response body")?;
+        .discard_data().context("fetching response body")?;
 
       if ! status.is_success() {
         throw!(anyhow!("HTTP error status={} body={:?}",
diff --git a/src/utils.rs b/src/utils.rs
index 0890cd5..a254b7b 100644
--- a/src/utils.rs
+++ b/src/utils.rs
@@ -25,6 +25,19 @@ pub enum ReadLimitedError {
   Hyper(#[from] hyper::Error),
 }
 
+impl ReadLimitedError {
+  pub fn discard_data(&mut self) { match self {
+    ReadLimitedError::Truncated { sofar,.. } => { mem::take(sofar); },
+    _ => { },
+  } }
+}
+#[ext(pub)]
+impl<T> Result<T,ReadLimitedError> {
+  fn discard_data(self) -> Self {
+    self.map_err(|mut e| { e.discard_data(); e })
+  }
+}
+
 #[throws(ReadLimitedError)]
 pub async fn read_limited_bytes<S>(limit: usize, stream: &mut S) -> Box<[u8]>
 where S: futures::Stream<Item=Result<hyper::body::Bytes,hyper::Error>>
-- 
2.30.2