feat(server): implement middleware for limiting the content length (#53)

fixtures/test-server-payload-limit/test.sh

@@ -3,7 +3,7 @@
 setup() {
   touch emptyfile
   truncate -s 9KB smallfile
-  truncate -s 10KB normalfile
+  fallocate -l 10000 normalfile
   truncate -s 11KB bigfile
 }
 
@@ -15,7 +15,7 @@ run_test() {
   test "upload limit exceeded" = "$result"
 
   result=$(curl -s -F "file=@normalfile" localhost:8000)
-  test "upload limit exceeded" != "$result"
+  test "upload limit exceeded" = "$result"
 
   result=$(curl -s -F "file=@smallfile" localhost:8000)
   test "upload limit exceeded" != "$result"

src/lib.rs

@@ -28,6 +28,9 @@ pub mod mime;
 /// Helper functions.
 pub mod util;
 
+/// Custom middleware implementation.
+pub mod middleware;
+
 /// Environment variable for setting the configuration file path.
 pub const CONFIG_ENV: &str = "CONFIG";
 

src/main.rs

@@ -5,6 +5,7 @@ use actix_web::{App, HttpServer};
 use awc::ClientBuilder;
 use hotwatch::{Event, Hotwatch};
 use rustypaste::config::{Config, ServerConfig};
+use rustypaste::middleware::ContentLengthLimiter;
 use rustypaste::paste::PasteType;
 use rustypaste::server;
 use rustypaste::util;
@@ -154,6 +155,9 @@ async fn main() -> IoResult<()> {
             .app_data(Data::clone(&config))
             .app_data(Data::new(http_client))
             .wrap(Logger::default())
+            .wrap(ContentLengthLimiter::new(
+                server_config.max_content_length.get_bytes(),
+            ))
             .configure(server::configure_routes)
     })
     .bind(&server_config.address)?;
@@ -191,6 +195,9 @@ async fn actix_web(
                 .app_data(Data::clone(&config))
                 .app_data(Data::new(http_client))
                 .wrap(Logger::default())
+                .wrap(ContentLengthLimiter::new(
+                    server_config.max_content_length.get_bytes(),
+                ))
                 .configure(server::configure_routes),
         );
     };

src/middleware.rs

@@ -0,0 +1,101 @@
+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 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: u128,
+}
+
+impl ContentLengthLimiter {
+    /// Contructs a new instance.
+    pub fn new(limit: u128) -> Self {
+        Self { max_bytes: limit }
+    }
+}
+
+impl Default for ContentLengthLimiter {
+    fn default() -> Self {
+        Self {
+            max_bytes: 10 * 1024 * 1024,
+        }
+    }
+}
+
+impl<S, B> Transform<S, ServiceRequest> for ContentLengthLimiter
+where
+    S: Service<ServiceRequest, Response = ServiceResponse<B>, Error = Error> + 'static,
+    S::Future: 'static,
+    B: 'static,
+{
+    type Response = ServiceResponse<EitherBody<B>>;
+    type Error = Error;
+    type Transform = ContentLengthLimiterMiddleware<S>;
+    type InitError = ();
+    type Future = Ready<Result<Self::Transform, Self::InitError>>;
+    fn new_transform(&self, service: S) -> Self::Future {
+        ready(Ok(ContentLengthLimiterMiddleware {
+            service: Rc::new(service),
+            limit: self.max_bytes,
+        }))
+    }
+}
+
+/// Content length limiter middleware implementation.
+#[derive(Debug)]
+pub struct ContentLengthLimiterMiddleware<S> {
+    service: Rc<S>,
+    limit: u128,
+}
+
+impl<S, B> Service<ServiceRequest> for ContentLengthLimiterMiddleware<S>
+where
+    S: Service<ServiceRequest, Response = ServiceResponse<B>, Error = Error> + 'static,
+    S::Future: 'static,
+    B: 'static,
+{
+    type Response = ServiceResponse<EitherBody<B>>;
+    type Error = Error;
+    type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>>>>;
+    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::<u128>().ok())
+        {
+            if content_length > self.limit {
+                log::warn!("{} > {}", content_length, self.limit);
+                log::warn!("Upload rejected due to exceeded limit.");
+                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)
+        })
+    }
+}

src/server.rs

@@ -195,13 +195,6 @@ async fn upload(
             let mut bytes = Vec::<u8>::new();
             while let Some(chunk) = field.next().await {
                 bytes.append(&mut chunk?.to_vec());
-                let config = config
-                    .read()
-                    .map_err(|_| error::ErrorInternalServerError("cannot acquire config"))?;
-                if bytes.len() as u128 > config.server.max_content_length.get_bytes() {
-                    log::warn!("Upload rejected for {}", host);
-                    return Err(error::ErrorPayloadTooLarge("upload limit exceeded"));
-                }
             }
             if bytes.is_empty() {
                 log::warn!("{} sent zero bytes", host);
@@ -286,17 +279,16 @@ pub fn configure_routes(cfg: &mut web::ServiceConfig) {
 #[cfg(test)]
 mod tests {
     use super::*;
+    use crate::middleware::ContentLengthLimiter;
     use crate::random::{RandomURLConfig, RandomURLType};
-    use actix_web::body::BodySize;
     use actix_web::body::MessageBody;
-    use actix_web::dev::ServiceResponse;
+    use actix_web::body::{BodySize, BoxBody};
     use actix_web::error::Error;
     use actix_web::http::{header, StatusCode};
     use actix_web::test::{self, TestRequest};
     use actix_web::web::Data;
     use actix_web::App;
     use awc::ClientBuilder;
-    use byte_unit::Byte;
     use glob::glob;
     use std::path::PathBuf;
     use std::str;
@@ -324,8 +316,7 @@ mod tests {
             .set_payload(multipart_data)
     }
 
-    async fn assert_body(response: ServiceResponse, expected: &str) -> Result<(), Error> {
-        let body = response.into_body();
+    async fn assert_body(body: BoxBody, expected: &str) -> Result<(), Error> {
         if let BodySize::Sized(size) = body.size() {
             assert_eq!(size, expected.as_bytes().len() as u64);
             let body_bytes = actix_web::body::to_bytes(body).await?;
@@ -368,7 +359,7 @@ mod tests {
             .to_request();
         let response = test::call_service(&app, request).await;
         assert_eq!(StatusCode::OK, response.status());
-        assert_body(response, "landing page").await?;
+        assert_body(response.into_body(), "landing page").await?;
         Ok(())
     }
 
@@ -390,7 +381,7 @@ mod tests {
             .to_request();
         let response = test::call_service(&app, request).await;
         assert_eq!(StatusCode::UNAUTHORIZED, response.status());
-        assert_body(response, "unauthorized").await?;
+        assert_body(response.into_body(), "unauthorized").await?;
         Ok(())
     }
 
@@ -410,7 +401,7 @@ mod tests {
             .to_request();
         let response = test::call_service(&app, request).await;
         assert_eq!(StatusCode::FORBIDDEN, response.status());
-        assert_body(response, "endpoint is not exposed").await?;
+        assert_body(response.into_body(), "endpoint is not exposed").await?;
         Ok(())
     }
 
@@ -432,7 +423,7 @@ mod tests {
             .to_request();
         let response = test::call_service(&app, request).await;
         assert_eq!(StatusCode::OK, response.status());
-        assert_body(response, env!("CARGO_PKG_VERSION")).await?;
+        assert_body(response.into_body(), env!("CARGO_PKG_VERSION")).await?;
         Ok(())
     }
 
@@ -452,7 +443,7 @@ mod tests {
         let response =
             test::call_service(&app, get_multipart_request("", "", "").to_request()).await;
         assert_eq!(StatusCode::UNAUTHORIZED, response.status());
-        assert_body(response, "unauthorized").await?;
+        assert_body(response.into_body(), "unauthorized").await?;
 
         Ok(())
     }
@@ -463,6 +454,7 @@ mod tests {
             App::new()
                 .app_data(Data::new(RwLock::new(Config::default())))
                 .app_data(Data::new(Client::default()))
+                .wrap(ContentLengthLimiter::new(30000))
                 .configure(configure_routes),
         )
         .await;
@@ -473,7 +465,7 @@ mod tests {
         )
         .await;
         assert_eq!(StatusCode::PAYLOAD_TOO_LARGE, response.status());
-        assert_body(response, "upload limit exceeded").await?;
+        assert_body(response.into_body().boxed(), "upload limit exceeded").await?;
 
         Ok(())
     }
@@ -482,7 +474,6 @@ mod tests {
     async fn test_upload_file() -> Result<(), Error> {
         let mut config = Config::default();
         config.server.upload_path = env::current_dir()?;
-        config.server.max_content_length = Byte::from_bytes(100);
 
         let app = test::init_service(
             App::new()
@@ -500,14 +491,18 @@ mod tests {
         )
         .await;
         assert_eq!(StatusCode::OK, response.status());
-        assert_body(response, &format!("http://localhost:8080/{file_name}\n")).await?;
+        assert_body(
+            response.into_body(),
+            &format!("http://localhost:8080/{file_name}\n"),
+        )
+        .await?;
 
         let serve_request = TestRequest::get()
             .uri(&format!("/{file_name}"))
             .to_request();
         let response = test::call_service(&app, serve_request).await;
         assert_eq!(StatusCode::OK, response.status());
-        assert_body(response, &timestamp).await?;
+        assert_body(response.into_body(), &timestamp).await?;
 
         fs::remove_file(file_name)?;
         let serve_request = TestRequest::get()
@@ -526,7 +521,6 @@ mod tests {
 
         let mut config = Config::default();
         config.server.upload_path = PathBuf::from(&test_upload_dir);
-        config.server.max_content_length = Byte::from_bytes(100);
         config.paste.duplicate_files = Some(false);
         config.paste.random_url = RandomURLConfig {
             enabled: true,
@@ -571,7 +565,6 @@ mod tests {
     async fn test_upload_expiring_file() -> Result<(), Error> {
         let mut config = Config::default();
         config.server.upload_path = env::current_dir()?;
-        config.server.max_content_length = Byte::from_bytes(100);
 
         let app = test::init_service(
             App::new()
@@ -594,14 +587,18 @@ mod tests {
         )
         .await;
         assert_eq!(StatusCode::OK, response.status());
-        assert_body(response, &format!("http://localhost:8080/{file_name}\n")).await?;
+        assert_body(
+            response.into_body(),
+            &format!("http://localhost:8080/{file_name}\n"),
+        )
+        .await?;
 
         let serve_request = TestRequest::get()
             .uri(&format!("/{file_name}"))
             .to_request();
         let response = test::call_service(&app, serve_request).await;
         assert_eq!(StatusCode::OK, response.status());
-        assert_body(response, &timestamp).await?;
+        assert_body(response.into_body(), &timestamp).await?;
 
         thread::sleep(Duration::from_millis(40));
 
@@ -651,7 +648,11 @@ mod tests {
         )
         .await;
         assert_eq!(StatusCode::OK, response.status());
-        assert_body(response, &format!("http://localhost:8080/{file_name}\n")).await?;
+        assert_body(
+            response.into_body().boxed(),
+            &format!("http://localhost:8080/{file_name}\n"),
+        )
+        .await?;
 
         let serve_request = TestRequest::get()
             .uri(&format!("/{file_name}"))
@@ -681,7 +682,6 @@ mod tests {
     async fn test_upload_url() -> Result<(), Error> {
         let mut config = Config::default();
         config.server.upload_path = env::current_dir()?;
-        config.server.max_content_length = Byte::from_bytes(100);
 
         let app = test::init_service(
             App::new()
@@ -700,7 +700,7 @@ mod tests {
         )
         .await;
         assert_eq!(StatusCode::OK, response.status());
-        assert_body(response, "http://localhost:8080/url\n").await?;
+        assert_body(response.into_body(), "http://localhost:8080/url\n").await?;
 
         let serve_request = TestRequest::get().uri("/url").to_request();
         let response = test::call_service(&app, serve_request).await;
@@ -720,7 +720,6 @@ mod tests {
     async fn test_upload_oneshot() -> Result<(), Error> {
         let mut config = Config::default();
         config.server.upload_path = env::current_dir()?;
-        config.server.max_content_length = Byte::from_bytes(100);
 
         let app = test::init_service(
             App::new()
@@ -741,14 +740,18 @@ mod tests {
         )
         .await;
         assert_eq!(StatusCode::OK, response.status());
-        assert_body(response, &format!("http://localhost:8080/{file_name}\n")).await?;
+        assert_body(
+            response.into_body(),
+            &format!("http://localhost:8080/{file_name}\n"),
+        )
+        .await?;
 
         let serve_request = TestRequest::get()
             .uri(&format!("/{file_name}"))
             .to_request();
         let response = test::call_service(&app, serve_request).await;
         assert_eq!(StatusCode::OK, response.status());
-        assert_body(response, &timestamp).await?;
+        assert_body(response.into_body(), &timestamp).await?;
 
         let serve_request = TestRequest::get()
             .uri(&format!("/{file_name}"))
@@ -775,7 +778,6 @@ mod tests {
     async fn test_upload_oneshot_url() -> Result<(), Error> {
         let mut config = Config::default();
         config.server.upload_path = env::current_dir()?;
-        config.server.max_content_length = Byte::from_bytes(100);
 
         let oneshot_url_suffix = "oneshot_url";
 
@@ -802,7 +804,7 @@ mod tests {
         .await;
         assert_eq!(StatusCode::OK, response.status());
         assert_body(
-            response,
+            response.into_body(),
             &format!("http://localhost:8080/{}\n", oneshot_url_suffix),
         )
         .await?;