add resource middleware on actix-web-codegen (#1467)

Co-authored-by: Yuki Okushi <huyuumi.dev@gmail.com>

actix-web-codegen/tests/test_macro.rs

@@ -1,6 +1,11 @@
-use actix_web::{http, test, web::Path, App, HttpResponse, Responder};
+use std::pin::Pin;
+use std::task::{Context, Poll};
+
+use actix_web::{http, test, web::Path, App, HttpResponse, Responder, Error};
+use actix_web::dev::{Service, Transform, ServiceRequest, ServiceResponse};
 use actix_web_codegen::{connect, delete, get, head, options, patch, post, put, trace};
 use futures::{future, Future};
+use actix_web::http::header::{HeaderName, HeaderValue};
 
 // Make sure that we can name function as 'config'
 #[get("/config")]
@@ -73,6 +78,65 @@ async fn get_param_test(_: Path<String>) -> impl Responder {
     HttpResponse::Ok()
 }
 
+pub struct ChangeStatusCode;
+
+impl<S, B> Transform<S> for ChangeStatusCode
+where
+    S: Service<Request = ServiceRequest, Response = ServiceResponse<B>, Error = Error>,
+    S::Future: 'static,
+    B: 'static,
+{
+    type Request = ServiceRequest;
+    type Response = ServiceResponse<B>;
+    type Error = Error;
+    type InitError = ();
+    type Transform = ChangeStatusCodeMiddleware<S>;
+    type Future = future::Ready<Result<Self::Transform, Self::InitError>>;
+
+    fn new_transform(&self, service: S) -> Self::Future {
+        future::ok(ChangeStatusCodeMiddleware { service })
+    }
+}
+
+pub struct ChangeStatusCodeMiddleware<S> {
+    service: S,
+}
+
+impl<S, B> Service for ChangeStatusCodeMiddleware<S>
+where
+    S: Service<Request = ServiceRequest, Response = ServiceResponse<B>, Error = Error>,
+    S::Future: 'static,
+    B: 'static,
+{
+    type Request = ServiceRequest;
+    type Response = ServiceResponse<B>;
+    type Error = Error;
+    type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>>>>;
+
+    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
+        self.service.poll_ready(cx)
+    }
+
+    fn call(&mut self, req: ServiceRequest) -> Self::Future {
+
+        let fut = self.service.call(req);
+
+        Box::pin(async move {
+            let mut res = fut.await?;
+            let headers = res.headers_mut();
+            let header_name = HeaderName::from_lowercase(b"custom-header").unwrap();
+            let header_value = HeaderValue::from_str("hello").unwrap();
+            headers.insert(header_name, header_value);
+            Ok(res)
+        })
+    }
+}
+
+#[get("/test/wrap", wrap = "ChangeStatusCode")]
+async fn get_wrap(_: Path<String>) -> impl Responder {
+    HttpResponse::Ok()
+}
+
 #[actix_rt::test]
 async fn test_params() {
     let srv = test::start(|| {
@@ -155,3 +219,15 @@ async fn test_auto_async() {
     let response = request.send().await.unwrap();
     assert!(response.status().is_success());
 }
+
+#[actix_rt::test]
+async fn test_wrap() {
+    let srv = test::start(|| {
+        App::new()
+            .service(get_wrap)
+    });
+
+    let request = srv.request(http::Method::GET, srv.url("/test/wrap"));
+    let response = request.send().await.unwrap();
+    assert!(response.headers().contains_key("custom-header"));
+}