Limit the amount of data that can be fetched from untrusted external HTTP resources

1 files changed, 33 insertions, 0 deletions

diff --git a/crates/common/src/lib.rs b/crates/common/src/lib.rs
index d602a1ca..e8dbf60e 100644
--- a/crates/common/src/lib.rs
+++ b/crates/common/src/lib.rs
@@ -30,6 +30,7 @@ use directory::{
     Principal, QueryBy, Type,
 };
 use expr::if_block::IfBlock;
+use futures::StreamExt;
 use listener::{
     blocked::{AllowedIps, BlockedIps},
     tls::TlsManager,
@@ -38,6 +39,7 @@ use mail_send::Credentials;
 
 use manager::webadmin::Resource;
 use parking_lot::Mutex;
+use reqwest::Response;
 use sieve::Sieve;
 use store::{
     write::{QueueClass, ValueClass},
@@ -415,3 +417,34 @@ impl Clone for Security {
         }
     }
 }
+
+pub trait HttpLimitResponse: Sync + Send {
+    fn bytes_with_limit(
+        self,
+        limit: usize,
+    ) -> impl std::future::Future<Output = reqwest::Result<Option<Vec<u8>>>> + Send;
+}
+
+impl HttpLimitResponse for Response {
+    async fn bytes_with_limit(self, limit: usize) -> reqwest::Result<Option<Vec<u8>>> {
+        if self
+            .content_length()
+            .map_or(false, |len| len as usize > limit)
+        {
+            return Ok(None);
+        }
+
+        let mut bytes = Vec::with_capacity(std::cmp::min(limit, 1024));
+        let mut stream = self.bytes_stream();
+
+        while let Some(chunk) = stream.next().await {
+            let chunk = chunk?;
+            if bytes.len() + chunk.len() > limit {
+                return Ok(None);
+            }
+            bytes.extend_from_slice(&chunk);
+        }
+
+        Ok(Some(bytes))
+    }
+}