limit response body size

[hippotat.git] / src / utils.rs

diff --git a/src/utils.rs b/src/utils.rs
index 37984e5ba32ee6128575ec9c062fb4c90d15a9b9..2271fd47cb2b0a97c7136bacd58d4ddbd63d1b8f 100644 (file)
--- a/src/utils.rs
+++ b/src/utils.rs
@@ -16,6 +16,23 @@ impl<T,E> Result<T,E> where AE: From<E> {
   }
 }
 
+#[throws(AE)]
+pub async fn read_limited_body<S,E>(limit: usize, mut stream: S) -> Box<[u8]>
+where S: futures::Stream<Item=Result<hyper::body::Bytes,E>> + Unpin,
+      // we also require that the Stream is cancellation-safe
+      E: std::error::Error + Sync + Send + 'static,
+{
+  let mut accum = vec![];
+  while let Some(item) = stream.next().await {
+    let b = item.context("HTTP error fetching response body")?;
+    if accum.len() + b.len() > limit {
+      throw!(anyhow!("maximum response body size {} exceeded", limit));
+    }
+    accum.extend(b);
+  }
+  accum.into()
+}
+
 use sha2::Digest as _;
 
 type HmacH = sha2::Sha256;