Implement MessageBody for Cow<'static, {[u8], str}> (#2959)

actix · Jan 6, 2023 · cfd40b4 · cfd40b4
1 parent 08c2cdf
commit cfd40b4
Show file tree

Hide file tree

Showing 2 changed files with 69 additions and 1 deletion.
diff --git a/actix-http/CHANGES.md b/actix-http/CHANGES.md
@@ -2,6 +2,7 @@
 
 ## Unreleased - 2022-xx-xx
 ### Added
+- Implement `MessageBody` for `Cow<'static, str>` and `Cow<'static, [u8]>`. [#2959]
 - Implement `MessageBody` for `&mut B` where `B: MessageBody + Unpin`. [#2868]
 - Implement `MessageBody` for `Pin<B>` where `B::Target: MessageBody`. [#2868]
 - Automatic h2c detection via new service finalizer `HttpService::tcp_auto_h2c()`. [#2957]
@@ -17,6 +18,7 @@
 ### Performance
 - Improve overall performance of operations on `Extensions`. [#2890]
 
+[#2959]: https://github.com/actix/actix-web/pull/2959
 [#2868]: https://github.com/actix/actix-web/pull/2868
 [#2890]: https://github.com/actix/actix-web/pull/2890
 [#2957]: https://github.com/actix/actix-web/pull/2957

diff --git a/actix-http/src/body/message_body.rs b/actix-http/src/body/message_body.rs
@@ -120,7 +120,7 @@ pub trait MessageBody {
 }
 
 mod foreign_impls {
-    use std::ops::DerefMut;
+    use std::{borrow::Cow, ops::DerefMut};
 
     use super::*;
 
@@ -324,6 +324,39 @@ mod foreign_impls {
         }
     }
 
+    impl MessageBody for Cow<'static, [u8]> {
+        type Error = Infallible;
+
+        #[inline]
+        fn size(&self) -> BodySize {
+            BodySize::Sized(self.len() as u64)
+        }
+
+        #[inline]
+        fn poll_next(
+            self: Pin<&mut Self>,
+            _cx: &mut Context<'_>,
+        ) -> Poll<Option<Result<Bytes, Self::Error>>> {
+            if self.is_empty() {
+                Poll::Ready(None)
+            } else {
+                let bytes = match mem::take(self.get_mut()) {
+                    Cow::Borrowed(b) => Bytes::from_static(b),
+                    Cow::Owned(b) => Bytes::from(b),
+                };
+                Poll::Ready(Some(Ok(bytes)))
+            }
+        }
+
+        #[inline]
+        fn try_into_bytes(self) -> Result<Bytes, Self> {
+            match self {
+                Cow::Borrowed(b) => Ok(Bytes::from_static(b)),
+                Cow::Owned(b) => Ok(Bytes::from(b)),
+            }
+        }
+    }
+
     impl MessageBody for &'static str {
         type Error = Infallible;
 
@@ -379,6 +412,39 @@ mod foreign_impls {
         }
     }
 
+    impl MessageBody for Cow<'static, str> {
+        type Error = Infallible;
+
+        #[inline]
+        fn size(&self) -> BodySize {
+            BodySize::Sized(self.len() as u64)
+        }
+
+        #[inline]
+        fn poll_next(
+            self: Pin<&mut Self>,
+            _cx: &mut Context<'_>,
+        ) -> Poll<Option<Result<Bytes, Self::Error>>> {
+            if self.is_empty() {
+                Poll::Ready(None)
+            } else {
+                let bytes = match mem::take(self.get_mut()) {
+                    Cow::Borrowed(s) => Bytes::from_static(s.as_bytes()),
+                    Cow::Owned(s) => Bytes::from(s.into_bytes()),
+                };
+                Poll::Ready(Some(Ok(bytes)))
+            }
+        }
+
+        #[inline]
+        fn try_into_bytes(self) -> Result<Bytes, Self> {
+            match self {
+                Cow::Borrowed(s) => Ok(Bytes::from_static(s.as_bytes())),
+                Cow::Owned(s) => Ok(Bytes::from(s.into_bytes())),
+            }
+        }
+    }
+
     impl MessageBody for bytestring::ByteString {
         type Error = Infallible;