read at most 265 MiB from mycelium per message

1 files changed, 36 insertions, 2 deletions

diff --git a/src/transport/mycelium/mod.rs b/src/transport/mycelium/mod.rs
index 495f75e..5062d4a 100644
--- a/src/transport/mycelium/mod.rs
+++ b/src/transport/mycelium/mod.rs
@@ -15,6 +15,11 @@ pub mod daemon;
 pub use daemon::{MyceliumSetup, MyceliumSpawnConfig};
 
 const MYCELIUM_POLL_TIMEOUT_SECS: u64 = 300;
+/// The largest inbound message read (its JSON, base64 payload included: about 190 MiB of
+/// payload). Anyone can send us anything; this bounds what one message costs before it
+/// is even parsed. Generous, because a sync message can carry a whole history until
+/// large deltas travel as blobs.
+const MAX_INBOUND_BYTES: usize = 256 << 20;
 /// Polls of the inbox in progress at once.
 const INBOX_POLLS: usize = 2;
 /// How long a request to our Mycelium may take, before what its body adds
@@ -371,8 +376,15 @@ async fn poll_inbox(
             continue;
         }
 
-        let raw_bytes = match response.bytes().await {
-            Ok(bytes) => bytes,
+        let raw_bytes = match read_capped(response, MAX_INBOUND_BYTES).await {
+            Ok(Some(bytes)) => bytes,
+            Ok(None) => {
+                tracing::warn!(
+                    limit_bytes = MAX_INBOUND_BYTES,
+                    "Dropping an inbound message larger than the limit"
+                );
+                continue;
+            }
             Err(error) => {
                 tracing::warn!(error = %error, "Failed to read raw Mycelium response stream");
                 tokio::time::sleep(Duration::from_millis(100)).await;
@@ -399,6 +411,28 @@ async fn poll_inbox(
     }
 }
 
+/// The body of `response`, unless it is larger than `limit` bytes (None): read no
+/// further than that, whatever the sender claims.
+async fn read_capped(
+    mut response: reqwest::Response,
+    limit: usize,
+) -> reqwest::Result<Option<bytes::Bytes>> {
+    if response
+        .content_length()
+        .is_some_and(|length| length > limit as u64)
+    {
+        return Ok(None);
+    }
+    let mut body = bytes::BytesMut::new();
+    while let Some(chunk) = response.chunk().await? {
+        if body.len() + chunk.len() > limit {
+            return Ok(None);
+        }
+        body.extend_from_slice(&chunk);
+    }
+    Ok(Some(body.freeze()))
+}
+
 /// A message as the Mycelium REST API takes it (the JSON body of an outbound message).
 /// Also the simulator's wire format (`transport::sim`), so tests go through the same
 /// encoding.