implement custom payload

meilisearch-http/src/lib.rs

@@ -10,10 +10,13 @@ pub mod routes;
 #[cfg(all(not(debug_assertions), feature = "analytics"))]
 pub mod analytics;
 
+use std::{pin::Pin, task::{Context, Poll}};
+
 pub use self::data::Data;
+use futures::{Stream, future::{Ready, ready}};
 pub use option::Opt;
 
-use actix_web::{HttpResponse, web};
+use actix_web::{FromRequest, HttpRequest, dev, error::PayloadError, web};
 
 pub fn configure_data(config: &mut web::ServiceConfig, data: Data) {
     let http_payload_size_limit = data.http_payload_size_limit();
@@ -35,6 +38,7 @@ pub fn configure_data(config: &mut web::ServiceConfig, data: Data) {
 #[cfg(feature = "mini-dashboard")]
 pub fn dashboard(config: &mut web::ServiceConfig, enable_frontend: bool) {
     use actix_web_static_files::Resource;
+    use actix_web::HttpResponse;
 
     mod dashboard {
         include!(concat!(env!("OUT_DIR"), "/generated.rs"));
@@ -102,3 +106,57 @@ macro_rules! create_app {
             .wrap(middleware::NormalizePath::new(middleware::TrailingSlash::Trim))
     }};
 }
+
+pub struct Payload {
+    payload: dev::Payload,
+    limit: usize,
+}
+
+pub struct PayloadConfig {
+    limit: usize,
+}
+
+impl Default for PayloadConfig {
+    fn default() -> Self {
+        Self { limit: 256 * 1024  }
+    }
+}
+
+impl FromRequest for Payload {
+    type Config = PayloadConfig;
+
+    type Error = PayloadError;
+
+    type Future = Ready<Result<Payload, Self::Error>>;
+
+    #[inline]
+    fn from_request(req: &HttpRequest, payload: &mut dev::Payload) -> Self::Future {
+        let limit = req.app_data::<PayloadConfig>().map(|c| c.limit).unwrap_or(Self::Config::default().limit);
+        ready(Ok(Payload { payload: payload.take(), limit }))
+    }
+}
+
+impl Stream for Payload {
+    type Item = Result<web::Bytes, PayloadError>;
+
+    #[inline]
+    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+        match Pin::new(&mut self.payload).poll_next(cx) {
+            Poll::Ready(Some(result)) => {
+                match result {
+                    Ok(bytes) => {
+                        match self.limit.checked_sub(bytes.len()) {
+                            Some(new_limit) => {
+                                self.limit = new_limit;
+                                Poll::Ready(Some(Ok(bytes)))
+                            }
+                            None => Poll::Ready(Some(Err(PayloadError::Overflow))),
+                        }
+                    }
+                    x => Poll::Ready(Some(x)),
+                }
+            },
+            otherwise => otherwise,
+        }
+    }
+}