diff --git a/src/lib.rs b/src/lib.rs index dfe664b..3932a58 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -28,9 +28,6 @@ pub mod mime; /// Helper functions. pub mod util; -/// Custom middleware implementation. -pub mod middleware; - // Use macros from tracing crate. #[macro_use] extern crate tracing; diff --git a/src/main.rs b/src/main.rs index 2a75a73..4937852 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,11 +1,15 @@ +use actix_multipart::form::MultipartFormConfig; +use actix_multipart::MultipartError; use actix_web::middleware::Logger; use actix_web::web::Data; +use actix_web::Error; +use actix_web::HttpRequest; use actix_web::{App, HttpServer}; use awc::ClientBuilder; +use byte_unit::Byte; use hotwatch::notify::event::ModifyKind; use hotwatch::{Event, EventKind, Hotwatch}; use rustypaste::config::{Config, ServerConfig}; -use rustypaste::middleware::ContentLengthLimiter; use rustypaste::paste::PasteType; use rustypaste::server; use rustypaste::util; @@ -25,6 +29,11 @@ use tracing_subscriber::{ #[macro_use] extern crate tracing; +fn handle_multipart_error(err: MultipartError, _req: &HttpRequest) -> Error { + error!("Multipart error: {}", err); + Error::from(err) +} + /// Sets up the application. /// /// * loads the configuration @@ -172,10 +181,21 @@ async fn main() -> IoResult<()> { App::new() .app_data(Data::clone(&config)) .app_data(Data::new(http_client)) + .app_data( + MultipartFormConfig::default() + .total_limit( + Byte::parse_str(&server_config.max_content_length, true) + .expect("cannot parse byte") + .as_u64() + .try_into() + .unwrap(), + ) + .memory_limit(10 * 1024 * 1024) + .error_handler(handle_multipart_error), + ) .wrap(Logger::new( "%{r}a \"%r\" %s %b \"%{Referer}i\" \"%{User-Agent}i\" %T", )) - .wrap(ContentLengthLimiter::new(server_config.max_content_length)) .configure(server::configure_routes) }) .bind(&server_config.address)?; diff --git a/src/middleware.rs b/src/middleware.rs deleted file mode 100644 index 0c647a8..0000000 --- a/src/middleware.rs +++ /dev/null @@ -1,96 +0,0 @@ -use actix_web::dev::{forward_ready, Service, ServiceRequest, ServiceResponse, Transform}; -use actix_web::http::header::CONTENT_LENGTH; -use actix_web::http::StatusCode; -use actix_web::{body::EitherBody, Error}; -use actix_web::{HttpMessage, HttpResponseBuilder}; -use byte_unit::Byte; -use futures_util::{Future, TryStreamExt}; -use std::{ - future::{ready, Ready}, - pin::Pin, - rc::Rc, -}; - -/// Content length limiter middleware. -#[derive(Debug)] -pub struct ContentLengthLimiter { - // Maximum amount of bytes to allow. - max_bytes: Byte, -} - -impl ContentLengthLimiter { - /// Constructs a new instance. - pub fn new(max_bytes: Byte) -> Self { - Self { max_bytes } - } -} - -impl Transform for ContentLengthLimiter -where - S: Service, Error = Error> + 'static, - S::Future: 'static, - B: 'static, -{ - type Response = ServiceResponse>; - type Error = Error; - type Transform = ContentLengthLimiterMiddleware; - type InitError = (); - type Future = Ready>; - fn new_transform(&self, service: S) -> Self::Future { - ready(Ok(ContentLengthLimiterMiddleware { - service: Rc::new(service), - max_bytes: self.max_bytes, - })) - } -} - -/// Content length limiter middleware implementation. -#[derive(Debug)] -pub struct ContentLengthLimiterMiddleware { - service: Rc, - max_bytes: Byte, -} - -impl Service for ContentLengthLimiterMiddleware -where - S: Service, Error = Error> + 'static, - S::Future: 'static, - B: 'static, -{ - type Response = ServiceResponse>; - type Error = Error; - type Future = Pin>>>; - forward_ready!(service); - fn call(&self, mut request: ServiceRequest) -> Self::Future { - let service = Rc::clone(&self.service); - if let Some(content_length) = request - .headers() - .get(CONTENT_LENGTH) - .and_then(|v| v.to_str().ok()) - .and_then(|v| v.parse::().ok()) - { - if content_length > self.max_bytes { - warn!( - "Upload rejected due to exceeded limit. ({:-#} > {:-#})", - content_length, self.max_bytes - ); - return Box::pin(async move { - // drain the body due to https://github.com/actix/actix-web/issues/2695 - let mut payload = request.take_payload(); - while let Ok(Some(_)) = payload.try_next().await {} - Ok(request.into_response( - HttpResponseBuilder::new(StatusCode::PAYLOAD_TOO_LARGE) - .body("upload limit exceeded") - .map_into_right_body(), - )) - }); - } - } - Box::pin(async move { - service - .call(request) - .await - .map(ServiceResponse::map_into_left_body) - }) - } -} diff --git a/src/server.rs b/src/server.rs index 7e1173d..1926fb0 100644 --- a/src/server.rs +++ b/src/server.rs @@ -425,7 +425,6 @@ pub fn configure_routes(cfg: &mut web::ServiceConfig) { mod tests { use super::*; use crate::config::LandingPageConfig; - use crate::middleware::ContentLengthLimiter; use crate::random::{RandomURLConfig, RandomURLType}; use actix_web::body::MessageBody; use actix_web::body::{BodySize, BoxBody};