From 80c4affaa8882e679427a252d02481def7d77db0 Mon Sep 17 00:00:00 2001
From: Valerio Ageno <51341197+Valerioageno@users.noreply.github.com>
Date: Sat, 7 Dec 2024 11:10:47 +0100
Subject: [PATCH] feat: add cookies to the response (#201)

---
 crates/tuono_lib/Cargo.toml      |  1 +
 crates/tuono_lib/src/lib.rs      |  1 +
 crates/tuono_lib/src/response.rs | 71 ++++++++++++++++++++++++++++----
 3 files changed, 66 insertions(+), 7 deletions(-)

diff --git a/crates/tuono_lib/Cargo.toml b/crates/tuono_lib/Cargo.toml
index 7101fddc..d91b1fa1 100644
--- a/crates/tuono_lib/Cargo.toml
+++ b/crates/tuono_lib/Cargo.toml
@@ -19,6 +19,7 @@ include = [
 [dependencies]
 ssr_rs = "0.7.0"
 axum = {version = "0.7.5", features = ["json", "ws"]}
+axum-extra = {version = "0.9.6", features = ["cookie"]}
 tokio = { version = "1.37.0", features = ["full"] }
 serde = { version = "1.0.202", features = ["derive"] }
 erased-serde = "0.4.5"
diff --git a/crates/tuono_lib/src/lib.rs b/crates/tuono_lib/src/lib.rs
index fd844e5d..d94172c6 100644
--- a/crates/tuono_lib/src/lib.rs
+++ b/crates/tuono_lib/src/lib.rs
@@ -19,4 +19,5 @@ pub use tuono_lib_macros::{api, handler};
 
 // Re-exports
 pub use axum;
+pub use axum_extra::extract::cookie;
 pub use tokio;
diff --git a/crates/tuono_lib/src/response.rs b/crates/tuono_lib/src/response.rs
index 5e43c1ac..adfbbad2 100644
--- a/crates/tuono_lib/src/response.rs
+++ b/crates/tuono_lib/src/response.rs
@@ -3,11 +3,13 @@ use crate::{ssr::Js, Payload};
 use axum::http::{HeaderMap, StatusCode};
 use axum::response::{Html, IntoResponse, Redirect};
 use axum::Json;
+use axum_extra::extract::cookie::{Cookie, CookieJar};
 use erased_serde::Serialize;
 
 pub struct Props {
     data: Box<dyn Serialize>,
     http_code: StatusCode,
+    cookies: CookieJar,
 }
 
 pub enum Response {
@@ -57,26 +59,45 @@ impl Props {
         Props {
             data: Box::new(data),
             http_code: StatusCode::OK,
+            cookies: CookieJar::new(),
         }
     }
 
+    pub fn status(&mut self, http_code: StatusCode) {
+        self.http_code = http_code;
+    }
+
     pub fn new_with_status(data: impl Serialize + 'static, http_code: StatusCode) -> Self {
         Props {
             data: Box::new(data),
             http_code,
+            cookies: CookieJar::new(),
         }
     }
+
+    pub fn add_cookie(&mut self, cookie: Cookie) {
+        let jar = self.cookies.clone().add(cookie.into_owned());
+        self.cookies = jar
+    }
 }
 
 impl Response {
     pub fn render_to_string(&self, req: Request) -> impl IntoResponse {
         match self {
-            Self::Props(Props { data, http_code }) => {
-                let payload = Payload::new(&req, data).client_payload().unwrap();
+            Self::Props(Props {
+                data,
+                http_code,
+                cookies,
+            }) => {
+                let payload = Payload::new(&req, data.as_ref()).client_payload().unwrap();
 
                 match Js::render_to_string(Some(&payload)) {
-                    Ok(html) => (*http_code, Html(html)),
-                    Err(_) => (*http_code, Html("500 Internal server error".to_string())),
+                    Ok(html) => (*http_code, cookies.clone(), Html(html)),
+                    Err(_) => (
+                        *http_code,
+                        cookies.clone(),
+                        Html("500 Internal server error".to_string()),
+                    ),
                 }
                 .into_response()
             }
@@ -87,9 +108,16 @@ impl Response {
 
     pub fn json(&self) -> impl IntoResponse {
         match self {
-            Self::Props(Props { data, http_code }) => {
-                (*http_code, Json(JsonResponse::new(data))).into_response()
-            }
+            Self::Props(Props {
+                data,
+                http_code,
+                cookies,
+            }) => (
+                *http_code,
+                cookies.clone(),
+                Json(JsonResponse::new(data.as_ref())),
+            )
+                .into_response(),
             Self::Redirect(destination) => (
                 StatusCode::PERMANENT_REDIRECT,
                 Json(JsonResponse::new_redirect(destination.to_string())),
@@ -101,3 +129,32 @@ impl Response {
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn should_update_the_props_status_and_cookie() {
+        let mut props = Props::new("{}");
+        props.status(StatusCode::NOT_FOUND);
+        props.add_cookie(Cookie::new("test", "cookie"));
+        assert_eq!(props.http_code, StatusCode::NOT_FOUND);
+        assert_eq!(
+            props.cookies.get("test").unwrap(),
+            &Cookie::new("test", "cookie")
+        );
+    }
+
+    #[test]
+    fn should_add_a_cookie_jar() {
+        let mut props = Props::new("{}");
+        props.status(StatusCode::NOT_FOUND);
+        props.add_cookie(Cookie::new("test", "cookie"));
+        assert_eq!(props.http_code, StatusCode::NOT_FOUND);
+        assert_eq!(
+            props.cookies.get("test").unwrap(),
+            &Cookie::new("test", "cookie")
+        );
+    }
+}