Merge branch 'file_size_limit' into 'master'

Add file size limit for downloads

Closes #8

See merge request !8

debian/config.toml

@@ -4,6 +4,7 @@ key = "changeme"
 server_name = "http://localhost/content/"
 port = 3030
 max_retries = 3
+max_size_bytes = 10_000_000
 
 [[formats]]
 name = "original"

src/bin/worker.rs

 use crate::config::Format;
-use bytes::Bytes;
 use futures::StreamExt;
 use image::imageops::FilterType;
 use image::io::Reader as ImageReader;
 use image::GenericImageView;
 use sparkplug::config;
+use sparkplug::download;
 use sqlx::PgPool;
 use std::fs::File;
 use std::io::prelude::*;
@@ -12,13 +12,6 @@ use std::io::Cursor;
 use std::path::Path;
 use tokio::task;
 
-async fn download_file(url: &str) -> Option<Bytes> {
-	let resp = reqwest::get(url).await.ok()?;
-	let content = resp.bytes().await.ok()?;
-
-	Some(content)
-}
-
 fn write_file(path: &Path, bytes: &[u8]) -> Option<()> {
 	let mut file = File::create(path).ok()?;
 	file.write_all(bytes).ok()?;
@@ -28,11 +21,12 @@ fn write_file(path: &Path, bytes: &[u8]) -> Option<()> {
 
 async fn process_file(
 	url: &str,
+	max_size: usize,
 	output_directory: String,
 	output_id: String,
 	formats: Vec<Format>,
 ) -> Option<String> {
-	let content = download_file(url).await?;
+	let content = download::download_file(url, max_size).await?;
 
 	let extension = task::spawn_blocking(move || {
 		let reader = ImageReader::new(Cursor::new(content))
@@ -71,6 +65,7 @@ async fn main() {
 	let config = config::load_config();
 	let image_directory = &config.image_directory;
 	let formats = &config.formats;
+	let max_size = &config.max_size_bytes;
 
 	let pool = PgPool::connect(&config.database_url)
 		.await
@@ -101,6 +96,7 @@ LIMIT 1000
 			|(row, pool)| async move {
 				match process_file(
 					&row.original_url,
+					*max_size,
 					image_directory.to_string(),
 					row.id.to_string(),
 					formats.clone(),

src/config.rs

@@ -10,6 +10,7 @@ struct TomlConfig {
 	pub server_name: String,
 	pub port: u16,
 	pub max_retries: u16,
+	pub max_size_bytes: usize,
 
 	pub formats: Vec<Format>,
 }
@@ -29,6 +30,7 @@ pub struct Config {
 	pub server_name: Url,
 	pub port: u16,
 	pub max_retries: i32,
+	pub max_size_bytes: usize,
 
 	pub formats: Vec<Format>,
 }
@@ -45,6 +47,7 @@ impl From<TomlConfig> for Config {
 				.expect("Invalid server_name - Not a valid URL"),
 			port: config.port,
 			max_retries: config.max_retries as i32,
+			max_size_bytes: config.max_size_bytes,
 
 			formats: config.formats,
 		}

src/download.rs

+pub async fn download_file(url: &str, max_size: usize) -> Option<Vec<u8>> {
+	let mut resp = reqwest::get(url).await.ok()?;
+
+	let mut body_bytes = vec![];
+	let mut len = 0;
+	while let Some(chunk) = resp.chunk().await.ok()? {
+		len += chunk.len();
+		if len > max_size {
+			return None;
+		}
+		body_bytes.extend(chunk);
+	}
+
+	Some(body_bytes)
+}

src/lib.rs

 pub mod config;
+pub mod download;
 pub mod ext;
 pub mod routes;