feat(file size limitation): add max_file_size configuration and stop …

…fetching if limit is exceeded
olxgroup-oss · Oct 11, 2024 · 520e116 · 520e116
1 parent 2deb06b
commit 520e116
Show file tree

Hide file tree

Showing 13 changed files with 109 additions and 39 deletions.
diff --git a/.github/workflows/pull-requests-tests.yml b/.github/workflows/pull-requests-tests.yml
@@ -55,7 +55,7 @@ jobs:
           libheif-dev=1.16.2-r0
           libimagequant-dev=4.2.0-r0
           libjpeg-turbo-dev=2.1.5.1-r3
-          libpng-dev=1.6.39-r3
+          libpng-dev=1.6.44-r0
           librsvg-dev=2.56.3-r0
           libwebp-dev=1.3.2-r0
           openssl-dev=3.1.7-r0

diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -29,7 +29,7 @@ serde = "1.0.195"
 serde_json = "1.0.111"
 serde_qs = "0.12.0"
 config = "0.13.3"
-reqwest = { version = "0.11.24", optional = true }
+reqwest = { version = "0.11.24", optional = true, features = ["stream"] }
 libvips = "1.5.0"
 rexif = "0.7.3"
 lazy_static = "1.4.0"

diff --git a/Dockerfile b/Dockerfile
@@ -16,7 +16,7 @@ RUN apk add --update --no-cache --repository https://dl-cdn.alpinelinux.org/alpi
     libheif-dev=1.16.2-r0 \
     libimagequant-dev=4.2.0-r0 \
     libjpeg-turbo-dev=2.1.5.1-r3 \
-    libpng-dev=1.6.39-r3 \
+    libpng-dev=1.6.44-r0 \
     librsvg-dev=2.56.3-r0 \
     libwebp-dev=1.3.2-r0 \
     openssl-dev=3.1.7-r0 \
@@ -61,7 +61,7 @@ RUN apk add --update --no-cache  \
       libheif=1.16.2-r0 \
       libimagequant=4.2.0-r0 \
       libjpeg-turbo=2.1.5.1-r3 \
-      libpng=1.6.39-r3 \
+      libpng=1.6.44-r0 \
       librsvg=2.56.3-r0 \
       libwebp=1.3.2-r0 \
       openssl=3.1.7-r0 \

diff --git a/README.md b/README.md
diff --git a/config/default.json b/config/default.json
@@ -10,5 +10,6 @@
   "s3_key": "root",
   "s3_secret": "rootpassword",
   "s3_endpoint": "http://localhost:2969/",
-  "s3_bucket": "dali-private"
+  "s3_bucket": "dali-private",
+  "max_file_size" : 41943040
 }
diff --git a/src/commons/config.rs b/src/commons/config.rs
@@ -6,7 +6,7 @@ use serde::Serialize;
 use std::env;
 use std::fmt;
 
-#[derive(Debug, Deserialize, Serialize)]
+#[derive(Debug, Deserialize, Serialize, Clone)]
 pub struct Configuration {
     pub app_port: u16,
     pub health_port: u16,
@@ -21,6 +21,7 @@ pub struct Configuration {
     pub s3_secret: Option<String>,
     pub s3_endpoint: Option<String>,
     pub s3_bucket: Option<String>,
+    pub max_file_size: Option<u32>,
 }
 
 impl fmt::Display for Configuration {

diff --git a/src/image_provider/mod.rs b/src/image_provider/mod.rs
@@ -22,7 +22,7 @@ pub struct ImageResponse {
 
 #[async_trait]
 pub trait ImageProvider: Send + Sync {
-    async fn get_file(&self, resource: &str) -> Result<ImageResponse, ImageProcessingError>;
+    async fn get_file(&self, resource: &str, config: &Configuration) -> Result<ImageResponse, ImageProcessingError>;
 }
 
 #[allow(unreachable_code)]

diff --git a/src/image_provider/reqwest.rs b/src/image_provider/reqwest.rs
@@ -1,15 +1,17 @@
 #[cfg(feature = "reqwest")]
 pub mod client {
+    use std::io::Write;
     use std::time::Duration;
 
     use async_trait::async_trait;
+    use futures::TryStreamExt;
     use log::error;
     use reqwest::{Client, Url};
 
     use crate::commons::config::Configuration;
     use crate::image_provider::ImageProcessingError::{
         ClientReturnedErrorStatusCode, ImageDownloadFailed, ImageDownloadTimedOut,
-        InvalidResourceUriProvided,
+        InvalidResourceUriProvided, FileSizeExceeded
     };
     use crate::image_provider::{ImageProvider, ImageResponse};
     use crate::routes::image::ImageProcessingError;
@@ -49,7 +51,7 @@ pub mod client {
 
     #[async_trait]
     impl ImageProvider for ReqwestImageProvider {
-        async fn get_file(&self, resource: &str) -> Result<ImageResponse, ImageProcessingError> {
+        async fn get_file(&self, resource: &str, config: &Configuration) -> Result<ImageResponse, ImageProcessingError> {
             let url = Url::parse(resource).map_err(|_| {
                 error!(
                     "the provided resource uri is not a valid http url: '{}'",
@@ -82,15 +84,36 @@ pub mod client {
                 })
                 .collect();
             if status.is_success() {
-                let bytes = response.bytes().await.map_err(|e| {
+                let mut stream = response.bytes_stream();
+                let mut total_bytes = 0;
+                let mut binary_payload: Vec<u8> = Vec::new();
+                while let Some(bytes) = stream.try_next().await.map_err(|e| {
                     error!(
                         "failed to read the binary payload of the image '{}'. error: {}",
                         resource, e
                     );
                     ImageDownloadFailed
-                })?;
+                })? {
+                    if let Some(max_size) = config.max_file_size {
+                        total_bytes += bytes.len() as u32;
+                        if total_bytes > max_size {
+                            error!(
+                                "the downloaded image '{}' exceeds the maximum allowed size of {} bytes",
+                                resource, max_size
+                            );
+                            return Err(FileSizeExceeded(max_size));
+                        }
+                    }
+                    binary_payload.write_all(&bytes).map_err(|e| {
+                        error!(
+                        "failed to read the response for the file '{}'. error: '{}'",
+                        resource, e
+                    );
+                        ImageDownloadFailed
+                    })?;
+                }
                 Ok(ImageResponse {
-                    bytes: bytes.to_vec(),
+                    bytes: binary_payload,
                     response_headers: headers,
                 })
             } else if status.is_client_error() {

diff --git a/src/image_provider/s3.rs b/src/image_provider/s3.rs
@@ -12,13 +12,14 @@ pub mod s3 {
     use async_trait::async_trait;
     use aws_config::{BehaviorVersion, Region};
     use aws_sdk_s3::error::ProvideErrorMetadata;
+    use futures::TryStreamExt;
 
     use crate::commons::config::Configuration;
     use crate::image_provider::ImageResponse;
     use crate::image_provider::{
         ImageProcessingError::{
             self, ClientReturnedErrorStatusCode, ImageDownloadFailed, ImageDownloadTimedOut,
-            InvalidResourceUriProvided,
+            InvalidResourceUriProvided, FileSizeExceeded,
         },
         ImageProvider,
     };
@@ -80,7 +81,7 @@ pub mod s3 {
 
     #[async_trait]
     impl ImageProvider for S3ImageProvider {
-        async fn get_file(&self, resource: &str) -> Result<ImageResponse, ImageProcessingError> {
+        async fn get_file(&self, resource: &str, config: &Configuration) -> Result<ImageResponse, ImageProcessingError> {
             if String::from(resource).is_empty() {
                 error!("the provided resource uri is empty");
                 return Err(InvalidResourceUriProvided(String::new()));
@@ -146,13 +147,24 @@ pub mod s3 {
                     .collect(),
             };
             let mut binary_payload: Vec<u8> = Vec::new();
+            let mut total_bytes = 0;
             while let Some(bytes) = result.body.try_next().await.map_err(|e| {
                 error!(
                     "failed to read the response for the file '{}'. error: '{}'",
                     resource, e
                 );
                 ImageDownloadFailed
             })? {
+                if let Some(max_size) = config.max_file_size {
+                    total_bytes += bytes.len() as u32;
+                    if total_bytes > max_size {
+                        error!(
+                            "the downloaded image '{}' exceeds the maximum allowed size of {} bytes",
+                            resource, max_size
+                        );
+                        return Err(FileSizeExceeded(max_size));
+                    }
+                }
                 binary_payload.write_all(&bytes).map_err(|e| {
                     error!(
                         "failed to read the response for the file '{}'. error: '{}'",
@@ -161,11 +173,10 @@ pub mod s3 {
                     ImageDownloadFailed
                 })?;
             }
-
             Ok(ImageResponse {
                 bytes: binary_payload,
                 response_headers: headers,
             })
         }
     }
-}
+}
diff --git a/src/main.rs b/src/main.rs
@@ -83,6 +83,7 @@ async fn start_management_server(config: &Configuration) {
 pub struct AppState {
     vips_app: Arc<VipsApp>,
     image_provider: Arc<Box<dyn ImageProvider>>,
+    config: Arc<Configuration>,
 }
 
 async fn measure_request_handling_duration(
@@ -108,6 +109,7 @@ async fn start_main_server(config: &Configuration) {
     let app_state = AppState {
         vips_app: Arc::new(create_vips_app(&config).unwrap()),
         image_provider: Arc::new(create_image_provider(&config).await),
+        config: Arc::new(config.clone()),
     };
 
     let app = Router::new()

diff --git a/src/routes/image.rs b/src/routes/image.rs
@@ -15,7 +15,9 @@ use thiserror::Error;
 
 use crate::{
     commons::{ImageFormat, ProcessImageRequest},
-    image_processor, AppState,
+    image_processor,
+    routes::metric::FILES_EXCEEDING_MAX_ALLOWED_SIZE,
+    AppState,
 };
 
 use super::metric::{FETCH_DURATION, INPUT_SIZE, OUTPUT_SIZE};
@@ -71,6 +73,8 @@ pub enum ImageProcessingError {
     ProcessingWorkerJoinError,
     #[error("the image processing with libvips has failed")]
     LibvipsProcessingFailed(libvips::error::Error),
+    #[error("the image exceeds the allowed size")]
+    FileSizeExceeded(u32),
 }
 
 impl IntoResponse for ImageProcessingError {
@@ -97,6 +101,12 @@ impl IntoResponse for ImageProcessingError {
                 StatusCode::BAD_REQUEST,
                 format!("The provided resource URI is not valid: '{}'", resource_uri)
             ),
+            ImageProcessingError::FileSizeExceeded(max_allowed_size) => {
+                FILES_EXCEEDING_MAX_ALLOWED_SIZE.inc();
+                (
+                StatusCode::BAD_REQUEST,
+                format!("The image exceeds the allowed size of {max_allowed_size} bytes. Please ensure the file size is within the permissible limit or adjust the configuration."),
+            )},
             _ => (
                 StatusCode::INTERNAL_SERVER_ERROR,
                 String::from("Something went wrong on our side."),
@@ -115,17 +125,20 @@ pub async fn process_image(
     State(AppState {
         vips_app,
         image_provider,
+        config,
     }): State<AppState>,
     ProcessImageRequestExtractor(params): ProcessImageRequestExtractor<ProcessImageRequest>,
 ) -> Result<Response<Body>, ImageProcessingError> {
     let now = SystemTime::now();
-    let main_img = image_provider.get_file(&params.image_address).await?;
+    let main_img = image_provider
+        .get_file(&params.image_address, &config)
+        .await?;
     let mut total_input_size = main_img.bytes.len();
 
     let watermarks_futures = params
         .watermarks
         .iter()
-        .map(|wm| image_provider.get_file(&wm.image_address));
+        .map(|wm| image_provider.get_file(&wm.image_address, &config));
     let watermarks = join_all(watermarks_futures)
         .await
         .into_iter()

diff --git a/src/routes/metric.rs b/src/routes/metric.rs
@@ -5,7 +5,9 @@ use axum::{
 };
 use lazy_static::lazy_static;
 use log::error;
-use prometheus::{register_histogram_vec, Encoder, HistogramVec, TextEncoder};
+use prometheus::{
+    register_histogram_vec, register_int_counter, Encoder, HistogramVec, IntCounter, TextEncoder,
+};
 use prometheus_static_metric::make_static_metric;
 
 make_static_metric! {
@@ -64,6 +66,11 @@ lazy_static! {
         &["format"]
     )
     .expect("Cannot register metric");
+    pub static ref FILES_EXCEEDING_MAX_ALLOWED_SIZE: IntCounter = register_int_counter!(
+        "dali_files_exceeding_max_allowed_size",
+        "Amount of files that were not processed due to exceeding the maximum allowed size"
+    )
+    .expect("Cannot register metric");
     pub static ref HTTP_DURATION: HttpRequestDuration =
         HttpRequestDuration::from(&HTTP_DURATION_VEC);
     pub static ref FETCH_DURATION: FetchRequestDuration =