Share expected_len calc in read_limited_body

Deduplicate the sensitive MAX_CONTENT_LENGTH restriction and calculation
based off of the Content-Length header.

Author: DanGould

payjoin-cli/src/app/v1.rs

@@ -93,18 +93,7 @@ impl AppTrait for App {
             payjoin::bitcoin::consensus::encode::serialize_hex(&fallback_tx)
         );
 
-        let expected_length = response
-            .headers()
-            .get("Content-Length")
-            .and_then(|val| val.to_str().ok())
-            .and_then(|s| s.parse::<usize>().ok())
-            .unwrap_or(MAX_CONTENT_LENGTH);
-
-        if expected_length > MAX_CONTENT_LENGTH {
-            return Err(anyhow!("Response body is too large: {} bytes", expected_length));
-        }
-
-        let body = read_limited_body(response.bytes_stream(), expected_length).await?;
+        let body = read_limited_body(&response.headers().clone(), response.bytes_stream()).await?;
 
         let psbt = ctx.process_response(&body).map_err(|e| {
             log::debug!("Error processing response: {e:?}");
@@ -312,26 +301,11 @@ impl App {
         req: Request<Incoming>,
     ) -> Result<Response<BoxBody<Bytes, hyper::Error>>, ReplyableError> {
         let (parts, body) = req.into_parts();
-        let headers = Headers(&parts.headers);
-
-        let expected_length = headers
-            .0
-            .get("Content-Length")
-            .and_then(|val| val.to_str().ok())
-            .and_then(|s| s.parse::<usize>().ok())
-            .unwrap_or(MAX_CONTENT_LENGTH);
-
-        if expected_length > MAX_CONTENT_LENGTH {
-            log::error!("Error: Content length exceeds max allowed");
-            return Err(Implementation(ImplementationError::from(
-                anyhow!("Content length too large: {expected_length}").into_boxed_dyn_error(),
-            )));
-        }
-
-        let body = read_limited_body(body.into_data_stream(), expected_length)
+        let body = read_limited_body(&parts.headers, body.into_data_stream())
             .await
             .map_err(|e| Implementation(ImplementationError::from(e.into_boxed_dyn_error())))?;
 
+        let headers = Headers(&parts.headers);
         let query_string = parts.uri.query().unwrap_or("");
         let proposal = UncheckedProposal::from_request(&body, query_string, headers)?;
 
@@ -407,11 +381,22 @@ impl App {
     }
 }
 
-pub async fn read_limited_body<S, E>(mut stream: S, expected_len: usize) -> Result<Vec<u8>>
+pub async fn read_limited_body<S, E>(headers: &hyper::HeaderMap, mut stream: S) -> Result<Vec<u8>>
 where
     S: Stream<Item = Result<Bytes, E>> + Unpin,
     E: std::error::Error + Send + Sync + 'static,
 {
+    let expected_len = headers
+        .get("Content-Length")
+        .and_then(|val| val.to_str().ok())
+        .and_then(|s| s.parse::<usize>().ok())
+        .unwrap_or(MAX_CONTENT_LENGTH);
+
+    if expected_len > MAX_CONTENT_LENGTH {
+        log::error!("Header-defined Content-Length exceeds max allowed: {expected_len}");
+        return Err(anyhow!("Header-defined Content-Length exceeds max allowed: {expected_len}"));
+    }
+
     let mut body = Vec::with_capacity(expected_len);
 
     while let Some(chunk) = stream.next().await {