Limit push endpoint response size

* Does not trust `Content-Length`.
* Will not allocate more than 64 KiB plus one frame, even on chunked
  responses.

Author: niklasf

src/clients/hyper_client.rs

@@ -66,18 +66,13 @@ impl WebPushClient for HyperWebPushClient {
         let response_status = response.status();
         trace!("Response status: {}", response_status);
 
-        let content_length: usize = response
-            .headers()
-            .get(CONTENT_LENGTH)
-            .and_then(|s| s.to_str().ok())
-            .and_then(|s| s.parse().ok())
-            .unwrap_or(0);
-
-        let mut body: Vec<u8> = Vec::with_capacity(content_length);
         let mut chunks = response.into_body();
-
+        let mut body = Vec::new();
         while let Some(chunk) = chunks.data().await {
             body.extend(&chunk?);
+            if body.len() > MAX_RESPONSE_SIZE {
+                return Err(WebPushError::ResponseTooLarge);
+            }
         }
         trace!("Body: {:?}", body);
 

src/clients/isahc_client.rs

@@ -1,10 +1,10 @@
 use async_trait::async_trait;
 use futures_lite::AsyncReadExt;
-use http::header::{CONTENT_LENGTH, RETRY_AFTER};
+use http::header::RETRY_AFTER;
 use isahc::HttpClient;
 
 use crate::clients::request_builder;
-use crate::clients::WebPushClient;
+use crate::clients::{WebPushClient, MAX_RESPONSE_SIZE};
 use crate::error::{RetryAfter, WebPushError};
 use crate::message::WebPushMessage;
 
@@ -67,21 +67,16 @@ impl WebPushClient for IsahcWebPushClient {
         let response_status = response.status();
         trace!("Response status: {}", response_status);
 
-        let content_length: usize = response
-            .headers()
-            .get(CONTENT_LENGTH)
-            .and_then(|s| s.to_str().ok())
-            .and_then(|s| s.parse().ok())
-            .unwrap_or(0);
-
-        let mut body: Vec<u8> = Vec::with_capacity(content_length);
-        let mut chunks = response.into_body();
-
-        chunks
+        let mut body = Vec::new();
+        if response
+            .into_body()
+            .take(MAX_RESPONSE_SIZE as u64 + 1)
             .read_to_end(&mut body)
-            .await
-            .map_err(|_| WebPushError::InvalidResponse)?;
-
+            .await?
+            > MAX_RESPONSE_SIZE
+        {
+            return Err(WebPushError::ResponseTooLarge);
+        }
         trace!("Body: {:?}", body);
 
         trace!("Body text: {:?}", std::str::from_utf8(&body));

src/clients/mod.rs

@@ -14,6 +14,8 @@ pub mod hyper_client;
 #[cfg(feature = "isahc-client")]
 pub mod isahc_client;
 
+const MAX_RESPONSE_SIZE: usize = 64 * 1024;
+
 /// An async client for sending the notification payload.
 /// Other features, such as thread safety, may vary by implementation.
 #[async_trait]

src/error.rs

@@ -64,6 +64,8 @@ pub enum WebPushError {
     InvalidResponse,
     /// A claim had invalid data
     InvalidClaims,
+    /// Response from push endpoint was too large
+    ResponseTooLarge,
     Other(ErrorInfo),
 }
 
@@ -134,6 +136,7 @@ impl WebPushError {
             WebPushError::Io(_) => "io_error",
             WebPushError::Other(_) => "other",
             WebPushError::InvalidClaims => "invalidClaims",
+            WebPushError::ResponseTooLarge => "response_too_large",
         }
     }
 }
@@ -162,6 +165,7 @@ impl fmt::Display for WebPushError {
             WebPushError::InvalidCryptoKeys => write!(f, "request has invalid cryptographic keys"),
             WebPushError::Other(info) => write!(f, "other: {}", info),
             WebPushError::InvalidClaims => write!(f, "at least one jwt claim was invalid"),
+            WebPushError::ResponseTooLarge => write!(f, "response from push endpoint was too large"),
         }
     }
 }