From 534cfe1fda1d65b9cea7d75a0a8c55f46439bb4e Mon Sep 17 00:00:00 2001
From: Sebastian Detert <github@elygor.de>
Date: Fri, 7 Jun 2024 17:22:48 +0200
Subject: [PATCH] feat: add .customize().add_cookie() (#3215)

* feat: add .customize().add_cookie()

* docs: added cookie hint

* fix: added unwrap to test of add_cookie()

* docs: added changelog entry for .customize().add_cookie()

* chore: make append_header infallible

* docs: update changelog

---------

Co-authored-by: Rob Ede <robjtede@icloud.com>
---
 actix-web/CHANGES.md                          |  1 +
 actix-web/src/response/customize_responder.rs | 42 ++++++++++++++++++-
 2 files changed, 42 insertions(+), 1 deletion(-)

diff --git a/actix-web/CHANGES.md b/actix-web/CHANGES.md
index 757fdce68..75f3631c9 100644
--- a/actix-web/CHANGES.md
+++ b/actix-web/CHANGES.md
@@ -4,6 +4,7 @@
 
 ### Added
 
+- Add `CustomizeResponder::add_cookie()` method.
 - Add `guard::GuardContext::app_data()` method.
 - Implement `From<Box<dyn ResponseError>>` for `Error`.
 
diff --git a/actix-web/src/response/customize_responder.rs b/actix-web/src/response/customize_responder.rs
index 4cbd96e20..6a43ac5e6 100644
--- a/actix-web/src/response/customize_responder.rs
+++ b/actix-web/src/response/customize_responder.rs
@@ -7,7 +7,7 @@ use actix_http::{
 
 use crate::{HttpRequest, HttpResponse, Responder};
 
-/// Allows overriding status code and headers for a [`Responder`].
+/// Allows overriding status code and headers (including cookies) for a [`Responder`].
 ///
 /// Created by calling the [`customize`](Responder::customize) method on a [`Responder`] type.
 pub struct CustomizeResponder<R> {
@@ -137,6 +137,29 @@ impl<R: Responder> CustomizeResponder<R> {
             Some(&mut self.inner)
         }
     }
+
+    /// Appends a `cookie` to the final response.
+    ///
+    /// # Errors
+    ///
+    /// Final response will be an error if `cookie` cannot be converted into a valid header value.
+    #[cfg(feature = "cookies")]
+    pub fn add_cookie(mut self, cookie: &crate::cookie::Cookie<'_>) -> Self {
+        use actix_http::header::{TryIntoHeaderValue as _, SET_COOKIE};
+
+        if let Some(inner) = self.inner() {
+            match cookie.to_string().try_into_value() {
+                Ok(val) => {
+                    inner.append_headers.append(SET_COOKIE, val);
+                }
+                Err(err) => {
+                    self.error = Some(err.into());
+                }
+            }
+        }
+
+        self
+    }
 }
 
 impl<T> Responder for CustomizeResponder<T>
@@ -175,6 +198,7 @@ mod tests {
 
     use super::*;
     use crate::{
+        cookie::Cookie,
         http::header::{HeaderValue, CONTENT_TYPE},
         test::TestRequest,
     };
@@ -209,6 +233,22 @@ mod tests {
             to_bytes(res.into_body()).await.unwrap(),
             Bytes::from_static(b"test"),
         );
+
+        let res = "test"
+            .to_string()
+            .customize()
+            .add_cookie(&Cookie::new("name", "value"))
+            .respond_to(&req);
+
+        assert!(res.status().is_success());
+        assert_eq!(
+            res.cookies().collect::<Vec<Cookie<'_>>>(),
+            vec![Cookie::new("name", "value")],
+        );
+        assert_eq!(
+            to_bytes(res.into_body()).await.unwrap(),
+            Bytes::from_static(b"test"),
+        );
     }
 
     #[actix_rt::test]