diff --git a/src/routes/v3/collections.rs b/src/routes/v3/collections.rs index 473e47df..bfb17787 100644 --- a/src/routes/v3/collections.rs +++ b/src/routes/v3/collections.rs @@ -368,7 +368,7 @@ pub async fn collection_icon_edit( pool: web::Data, redis: web::Data, file_host: web::Data>, - mut payload: web::Payload, + payload: web::Payload, session_queue: web::Data, ) -> Result { if let Some(content_type) = crate::util::ext::get_image_content_type(&ext.ext) { @@ -403,8 +403,7 @@ pub async fn collection_icon_edit( } } - let bytes = - read_from_payload(&mut payload, 262144, "Icons must be smaller than 256KiB").await?; + let bytes = read_from_payload(payload, 262144, "Icons must be smaller than 256KiB").await?; let color = crate::util::img::get_color_from_img(&bytes)?; @@ -414,7 +413,7 @@ pub async fn collection_icon_edit( .upload_file( content_type, &format!("data/{}/{}.{}", collection_id, hash, ext.ext), - bytes.freeze(), + bytes, ) .await?; diff --git a/src/routes/v3/images.rs b/src/routes/v3/images.rs index 86b202ef..c2b77809 100644 --- a/src/routes/v3/images.rs +++ b/src/routes/v3/images.rs @@ -41,7 +41,7 @@ pub async fn images_add( req: HttpRequest, web::Query(data): web::Query, file_host: web::Data>, - mut payload: web::Payload, + payload: web::Payload, pool: web::Data, redis: web::Data, session_queue: web::Data, @@ -166,14 +166,14 @@ pub async fn images_add( // Upload the image to the file host let bytes = - read_from_payload(&mut payload, 1_048_576, "Icons must be smaller than 1MiB").await?; + read_from_payload(payload, 1_048_576, "Icons must be smaller than 1MiB").await?; let hash = sha1::Sha1::from(&bytes).hexdigest(); let upload_data = file_host .upload_file( content_type, &format!("data/cached_images/{}.{}", hash, data.ext), - bytes.freeze(), + bytes, ) .await?; diff --git a/src/routes/v3/oauth_clients.rs b/src/routes/v3/oauth_clients.rs index ad8cb6a4..6190746a 100644 --- a/src/routes/v3/oauth_clients.rs +++ b/src/routes/v3/oauth_clients.rs @@ -345,7 +345,7 @@ pub async fn oauth_client_icon_edit( pool: web::Data, redis: web::Data, file_host: web::Data>, - mut payload: web::Payload, + payload: web::Payload, session_queue: web::Data, ) -> Result { if let Some(content_type) = crate::util::ext::get_image_content_type(&ext.ext) { @@ -376,14 +376,13 @@ pub async fn oauth_client_icon_edit( } } - let bytes = - read_from_payload(&mut payload, 262144, "Icons must be smaller than 256KiB").await?; + let bytes = read_from_payload(payload, 262144, "Icons must be smaller than 256KiB").await?; let hash = sha1::Sha1::from(&bytes).hexdigest(); let upload_data = file_host .upload_file( content_type, &format!("data/{}/{}.{}", client_id, hash, ext.ext), - bytes.freeze(), + bytes, ) .await?; diff --git a/src/routes/v3/organizations.rs b/src/routes/v3/organizations.rs index 48b2867b..a3b955b8 100644 --- a/src/routes/v3/organizations.rs +++ b/src/routes/v3/organizations.rs @@ -923,7 +923,7 @@ pub async fn organization_icon_edit( pool: web::Data, redis: web::Data, file_host: web::Data>, - mut payload: web::Payload, + payload: web::Payload, session_queue: web::Data, ) -> Result { if let Some(content_type) = crate::util::ext::get_image_content_type(&ext.ext) { @@ -973,8 +973,7 @@ pub async fn organization_icon_edit( } } - let bytes = - read_from_payload(&mut payload, 262144, "Icons must be smaller than 256KiB").await?; + let bytes = read_from_payload(payload, 262144, "Icons must be smaller than 256KiB").await?; let color = crate::util::img::get_color_from_img(&bytes)?; @@ -984,7 +983,7 @@ pub async fn organization_icon_edit( .upload_file( content_type, &format!("data/{}/{}.{}", organization_id, hash, ext.ext), - bytes.freeze(), + bytes, ) .await?; diff --git a/src/routes/v3/project_creation.rs b/src/routes/v3/project_creation.rs index da984f89..3359977c 100644 --- a/src/routes/v3/project_creation.rs +++ b/src/routes/v3/project_creation.rs @@ -508,9 +508,7 @@ async fn project_create_inner( CreateError::InvalidIconFormat(file_extension.to_string()) })?; let url = format!("data/{project_id}/images/{hash}.{file_extension}"); - let upload_data = file_host - .upload_file(content_type, &url, data.freeze()) - .await?; + let upload_data = file_host.upload_file(content_type, &url, data).await?; uploaded_files.push(UploadedFile { file_id: upload_data.file_id, file_name: upload_data.file_name, @@ -955,7 +953,7 @@ async fn process_icon_upload( .upload_file( content_type, &format!("data/{id}/{hash}.{file_extension}"), - data.freeze(), + data, ) .await?; diff --git a/src/routes/v3/projects.rs b/src/routes/v3/projects.rs index 5053e7df..fe634883 100644 --- a/src/routes/v3/projects.rs +++ b/src/routes/v3/projects.rs @@ -1314,7 +1314,7 @@ pub async fn project_icon_edit( pool: web::Data, redis: web::Data, file_host: web::Data>, - mut payload: web::Payload, + payload: web::Payload, session_queue: web::Data, ) -> Result { if let Some(content_type) = crate::util::ext::get_image_content_type(&ext.ext) { @@ -1374,8 +1374,7 @@ pub async fn project_icon_edit( } } - let bytes = - read_from_payload(&mut payload, 262144, "Icons must be smaller than 256KiB").await?; + let bytes = read_from_payload(payload, 262144, "Icons must be smaller than 256KiB").await?; let color = crate::util::img::get_color_from_img(&bytes)?; @@ -1385,7 +1384,7 @@ pub async fn project_icon_edit( .upload_file( content_type, &format!("data/{}/{}.{}", project_id, hash, ext.ext), - bytes.freeze(), + bytes, ) .await?; @@ -1524,7 +1523,7 @@ pub async fn add_gallery_item( pool: web::Data, redis: web::Data, file_host: web::Data>, - mut payload: web::Payload, + payload: web::Payload, session_queue: web::Data, ) -> Result { if let Some(content_type) = crate::util::ext::get_image_content_type(&ext.ext) { @@ -1586,7 +1585,7 @@ pub async fn add_gallery_item( } let bytes = read_from_payload( - &mut payload, + payload, 5 * (1 << 20), "Gallery image exceeds the maximum of 5MiB.", ) @@ -1607,9 +1606,7 @@ pub async fn add_gallery_item( )); } - file_host - .upload_file(content_type, &url, bytes.freeze()) - .await?; + file_host.upload_file(content_type, &url, bytes).await?; let mut transaction = pool.begin().await?; diff --git a/src/routes/v3/users.rs b/src/routes/v3/users.rs index f2cb1629..c33ad396 100644 --- a/src/routes/v3/users.rs +++ b/src/routes/v3/users.rs @@ -464,7 +464,7 @@ pub async fn user_icon_edit( pool: web::Data, redis: web::Data, file_host: web::Data>, - mut payload: web::Payload, + payload: web::Payload, session_queue: web::Data, ) -> Result { if let Some(content_type) = crate::util::ext::get_image_content_type(&ext.ext) { @@ -499,14 +499,14 @@ pub async fn user_icon_edit( } let bytes = - read_from_payload(&mut payload, 2097152, "Icons must be smaller than 2MiB").await?; + read_from_payload(payload, 2097152, "Icons must be smaller than 2MiB").await?; let hash = sha1::Sha1::from(&bytes).hexdigest(); let upload_data = file_host .upload_file( content_type, &format!("user/{}/{}.{}", user_id, hash, ext.ext), - bytes.freeze(), + bytes, ) .await?; diff --git a/src/routes/v3/version_creation.rs b/src/routes/v3/version_creation.rs index db3defeb..47cc66d2 100644 --- a/src/routes/v3/version_creation.rs +++ b/src/routes/v3/version_creation.rs @@ -904,7 +904,6 @@ pub async fn upload_file( } } - let data = data.freeze(); let primary = (validation_result.is_passed() && version_files.iter().all(|x| !x.primary) && !ignore_primary) diff --git a/src/util/routes.rs b/src/util/routes.rs index 00bb288e..2c363509 100644 --- a/src/util/routes.rs +++ b/src/util/routes.rs @@ -1,34 +1,43 @@ use crate::routes::v3::project_creation::CreateError; use crate::routes::ApiError; use actix_multipart::Field; +use actix_web::http::header::CONTENT_LENGTH; use actix_web::web::Payload; -use bytes::BytesMut; +use bytes::{Bytes, BytesMut}; use futures::StreamExt; pub async fn read_from_payload( - payload: &mut Payload, + payload: Payload, cap: usize, err_msg: &'static str, -) -> Result { - let mut bytes = BytesMut::new(); - while let Some(item) = payload.next().await { - if bytes.len() >= cap { - return Err(ApiError::InvalidInput(String::from(err_msg))); - } else { - bytes.extend_from_slice(&item.map_err(|_| { - ApiError::InvalidInput("Unable to parse bytes in payload sent!".to_string()) - })?); - } - } - Ok(bytes) +) -> Result { + payload + .to_bytes_limited(cap) + .await + .map_err(|_| ApiError::InvalidInput(String::from(err_msg)))? + .map_err(|_| ApiError::InvalidInput("Unable to parse bytes in payload sent!".to_string())) } pub async fn read_from_field( field: &mut Field, cap: usize, err_msg: &'static str, -) -> Result { - let mut bytes = BytesMut::new(); +) -> Result { + /// Sensible default (32kB) for initial, bounded allocation when collecting body bytes. + const INITIAL_ALLOC_BYTES: usize = 32 * 1024; + + let capacity = match field.headers().get(&CONTENT_LENGTH) { + None => INITIAL_ALLOC_BYTES, + Some(len) => match len.to_str().ok().and_then(|len| len.parse::().ok()) { + None => INITIAL_ALLOC_BYTES, + Some(len) if len as usize > cap => { + return Err(CreateError::InvalidInput(String::from(err_msg))) + } + Some(len) => (len as usize).min(INITIAL_ALLOC_BYTES), + }, + }; + + let mut bytes = BytesMut::with_capacity(capacity); while let Some(chunk) = field.next().await { if bytes.len() >= cap { return Err(CreateError::InvalidInput(String::from(err_msg))); @@ -36,5 +45,5 @@ pub async fn read_from_field( bytes.extend_from_slice(&chunk?); } } - Ok(bytes) + Ok(bytes.freeze()) }