fix(mod_zip): implement better file size checks for geode files (#52)

Author: Fleeym

src/mod_zip.rs

@@ -12,6 +12,13 @@ use zip::ZipArchive;
 use crate::types::api::ApiError;
 
 pub fn extract_mod_logo(file: &mut ZipFile<Cursor<Bytes>>) -> Result<Vec<u8>, ApiError> {
+    const FIVE_MEGABYTES: u64 = 5 * 1000 * 1000;
+    if file.size() > FIVE_MEGABYTES {
+        return Err(ApiError::BadRequest(
+            "Logo size excedes max allowed size (5 MB)".into(),
+        ));
+    }
+
     let mut logo: Vec<u8> = Vec::with_capacity(file.size() as usize);
     file.read_to_end(&mut logo)
         .inspect_err(|e| log::error!("logo.png read fail: {}", e))
@@ -61,6 +68,13 @@ pub fn extract_mod_logo(file: &mut ZipFile<Cursor<Bytes>>) -> Result<Vec<u8>, Ap
 }
 
 pub fn validate_mod_logo(file: &mut ZipFile<Cursor<Bytes>>) -> Result<(), ApiError> {
+    const FIVE_MEGABYTES: u64 = 5 * 1000 * 1000;
+    if file.size() > FIVE_MEGABYTES {
+        return Err(ApiError::BadRequest(
+            "Logo size excedes max allowed size (5 MB)".into(),
+        ));
+    }
+
     let mut logo: Vec<u8> = Vec::with_capacity(file.size() as usize);
     file.read_to_end(&mut logo)
         .inspect_err(|e| log::error!("logo.png read fail: {}", e))
@@ -117,26 +131,44 @@ pub fn bytes_to_ziparchive(bytes: Bytes) -> Result<ZipArchive<Cursor<Bytes>>, Ap
 }
 
 async fn download(url: &str, limit_mb: u32) -> Result<Bytes, ApiError> {
-    let limit_bytes = limit_mb * 1_000_000;
-    let response = reqwest::get(url).await.map_err(|e| {
-        log::error!("Failed to fetch .geode: {}", e);
-        ApiError::BadRequest("Couldn't download .geode file".into())
-    })?;
+    let limit_bytes: u64 = limit_mb as u64 * 1_000_000;
+    let mut response = reqwest::get(url)
+        .await
+        .inspect_err(|e| log::error!("Failed to fetch .geode file: {e}"))
+        .or(Err(ApiError::BadRequest(
+            "Failed to fetch .geode file".into(),
+        )))?;
 
-    let len = response.content_length().ok_or(ApiError::BadRequest(
-        "Couldn't determine .geode file size".into(),
-    ))?;
+    // Check Content-Length, but the server can lie about this, so we'll also stream the file
+    // If the header is somehow unavailable, we'll just check the size when streaming
+    let content_length = response.content_length().unwrap_or(0);
 
-    if len > limit_bytes as u64 {
+    if content_length > limit_bytes {
+        let len_mb = content_length / 1_000_000;
         return Err(ApiError::BadRequest(format!(
-            "File size is too large, max {}MB",
-            limit_mb
+            "Mod file is too large ({} mb), max size is {} mb",
+            len_mb, limit_mb
         )));
     }
 
-    response
-        .bytes()
-        .await
-        .inspect_err(|e| log::error!("Failed to get bytes from .geode: {}", e))
-        .or(Err(ApiError::InternalError))
+    let mut data: Vec<u8> = Vec::with_capacity(content_length as usize);
+
+    let mut streamed: u64 = 0;
+    while let Some(chunk) = response.chunk().await.or(Err(ApiError::BadRequest(
+        "Failed to read .geode chunk".into(),
+    )))? {
+        streamed += chunk.len() as u64;
+
+        if streamed > limit_bytes {
+            let len_mb = streamed / 1_000_000;
+            return Err(ApiError::BadRequest(format!(
+                "Mod file is too large ({} mb), max size is {} mb",
+                len_mb, limit_mb
+            )));
+        }
+
+        data.extend_from_slice(&chunk);
+    }
+
+    Ok(Bytes::from(data))
 }