fix: process LengthLimitError

viz-rs · Dec 22, 2023 · bb7e87c · bb7e87c
1 parent 341850f
commit bb7e87c
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 11 deletions.
diff --git a/viz-core/src/body.rs b/viz-core/src/body.rs
@@ -98,12 +98,9 @@ impl HttpBody for Body {
     ) -> Poll<Option<Result<Frame<Self::Data>, Self::Error>>> {
         match self.get_mut() {
             Self::Empty => Poll::Ready(None),
-            Self::Full(inner) => Pin::new(inner).poll_frame(cx).map_err(Error::from),
-            Self::Incoming(inner) => Pin::new(inner).poll_frame(cx).map_err(Error::from),
-            Self::Boxed(inner) => Pin::new(inner)
-                .get_pin_mut()
-                .poll_frame(cx)
-                .map_err(Error::from),
+            Self::Full(inner) => Pin::new(inner).poll_frame(cx).map_err(Into::into),
+            Self::Incoming(inner) => Pin::new(inner).poll_frame(cx).map_err(Into::into),
+            Self::Boxed(inner) => Pin::new(inner).get_pin_mut().poll_frame(cx),
         }
     }
 

diff --git a/viz-core/src/request.rs b/viz-core/src/request.rs
@@ -223,8 +223,10 @@ impl RequestExt for Request {
     {
         self.headers()
             .get(key)
-            .and_then(|v| v.to_str().ok())
-            .and_then(|v| v.parse::<T>().ok())
+            .map(header::HeaderValue::to_str)
+            .and_then(Result::ok)
+            .map(str::parse)
+            .and_then(Result::ok)
     }
 
     fn header_typed<H>(&self) -> Option<H>
@@ -271,8 +273,16 @@ impl RequestExt for Request {
         self.incoming()?
             .collect()
             .await
+            .map_err(|err| {
+                if err.downcast_ref::<LengthLimitError>().is_some() {
+                    return PayloadError::TooLarge;
+                }
+                if let Ok(err) = err.downcast::<hyper::Error>() {
+                    return PayloadError::Hyper(err);
+                }
+                PayloadError::Read
+            })
             .map(Collected::to_bytes)
-            .map_err(|_| PayloadError::Read)
     }
 
     #[cfg(feature = "limits")]
@@ -369,7 +379,8 @@ impl RequestExt for Request {
             boundary,
             self.extensions()
                 .get::<std::sync::Arc<MultipartLimits>>()
-                .map(|ml| ml.as_ref().clone())
+                .map(AsRef::as_ref)
+                .cloned()
                 .unwrap_or_default(),
         ))
     }

diff --git a/viz-core/src/types/payload.rs b/viz-core/src/types/payload.rs
@@ -34,8 +34,8 @@ pub enum PayloadError {
     #[error("JSON serialize or deserialize faild, {0}")]
     Json(#[from] serde_json::Error),
 
-    #[cfg(any(feature = "form", feature = "query"))]
     /// 400
+    #[cfg(any(feature = "form", feature = "query"))]
     #[error("url decode failed, {0}")]
     UrlDecode(#[from] serde_urlencoded::de::Error),