From 50621bae7168b689d79fe789ce6b26fd202790f6 Mon Sep 17 00:00:00 2001
From: Ali MJ Al-Nasrawy <alimjalnasrawy@gmail.com>
Date: Mon, 13 Sep 2021 13:00:58 +0300
Subject: [PATCH] cors: make middleware generic over body type (#195)

---
 actix-cors/CHANGES.md     |  3 +++
 actix-cors/src/builder.rs | 41 +++++++++++++++++++++++++++++++++++----
 2 files changed, 40 insertions(+), 4 deletions(-)

diff --git a/actix-cors/CHANGES.md b/actix-cors/CHANGES.md
index d75ba99fe..fb96980cd 100644
--- a/actix-cors/CHANGES.md
+++ b/actix-cors/CHANGES.md
@@ -1,8 +1,11 @@
 # Changes
 
 ## Unreleased - 2021-xx-xx
+* Make `Cors` middleware generic over body type [#195]
 * Minimum supported Rust version (MSRV) is now 1.51.
 
+[#195]: https://github.com/actix/actix-extras/pull/195
+
 
 ## 0.6.0-beta.2 - 2021-06-27
 * No notable changes.
diff --git a/actix-cors/src/builder.rs b/actix-cors/src/builder.rs
index e958491dd..8aa914daf 100644
--- a/actix-cors/src/builder.rs
+++ b/actix-cors/src/builder.rs
@@ -1,6 +1,9 @@
-use std::{collections::HashSet, convert::TryInto, iter::FromIterator, rc::Rc};
+use std::{
+    collections::HashSet, convert::TryInto, error::Error as StdError, iter::FromIterator, rc::Rc,
+};
 
 use actix_web::{
+    body::MessageBody,
     dev::{RequestHead, Service, ServiceRequest, ServiceResponse, Transform},
     error::{Error, Result},
     http::{self, header::HeaderName, Error as HttpError, HeaderValue, Method, Uri},
@@ -479,10 +482,12 @@ impl Default for Cors {
     }
 }
 
-impl<S> Transform<S, ServiceRequest> for Cors
+impl<S, B> Transform<S, ServiceRequest> for Cors
 where
-    S: Service<ServiceRequest, Response = ServiceResponse, Error = Error>,
+    S: Service<ServiceRequest, Response = ServiceResponse<B>, Error = Error>,
     S::Future: 'static,
+    B: MessageBody + 'static,
+    B::Error: StdError,
 {
     type Response = ServiceResponse;
     type Error = Error;
@@ -560,11 +565,15 @@ where
 #[cfg(test)]
 mod test {
     use std::convert::{Infallible, TryInto};
+    use std::pin::Pin;
+    use std::task::{Context, Poll};
 
     use actix_web::{
-        dev::Transform,
+        body::{BodySize, MessageBody},
+        dev::{fn_service, Transform},
         http::{HeaderName, StatusCode},
         test::{self, TestRequest},
+        web::{Bytes, HttpResponse},
     };
 
     use super::*;
@@ -616,4 +625,28 @@ mod test {
 
         let _cors = Cors::default().allowed_header(ContentType);
     }
+
+    #[actix_rt::test]
+    async fn middleware_generic_over_body_type() {
+        struct Foo;
+
+        impl MessageBody for Foo {
+            type Error = std::io::Error;
+            fn size(&self) -> BodySize {
+                BodySize::None
+            }
+            fn poll_next(
+                self: Pin<&mut Self>,
+                _: &mut Context<'_>,
+            ) -> Poll<Option<Result<Bytes, Self::Error>>> {
+                Poll::Ready(None)
+            }
+        }
+
+        let srv = fn_service(|req: ServiceRequest| async move {
+            Ok(req.into_response(HttpResponse::Ok().message_body(Foo)?))
+        });
+
+        Cors::default().new_transform(srv).await.unwrap();
+    }
 }