bolt_web/
response.rs

1use base64::{Engine, engine::general_purpose};
2use bytes::Bytes;
3use http_body_util::Full;
4use hyper::{
5    HeaderMap, Response,
6    header::{HeaderName, HeaderValue},
7};
8use serde::Serialize;
9use time::{Duration, OffsetDateTime, format_description::well_known::Rfc2822};
10
11pub struct ResponseWriter {
12    pub body: String,
13    pub headers: HeaderMap,
14    pub status: u16,
15    pub has_error: bool,
16}
17
18#[allow(dead_code)]
19impl ResponseWriter {
20    pub fn new() -> Self {
21        Self {
22            body: "".into(),
23            headers: HeaderMap::new(),
24            status: 200,
25            has_error: false,
26        }
27    }
28
29    pub fn status(&mut self, status: u16) -> &mut Self {
30        self.status = status;
31        self
32    }
33
34    pub fn set_header(&mut self, key: &str, value: &str) -> &mut Self {
35        self.headers.insert(
36            HeaderName::from_bytes(key.as_bytes()).unwrap(),
37            HeaderValue::from_str(value).unwrap(),
38        );
39        self
40    }
41
42    pub fn get_header(&self, key: &str) -> Option<&HeaderValue> {
43        self.headers.get(key)
44    }
45
46    pub fn send(&mut self, body: &str) -> &mut Self {
47        self.body = body.into();
48        self
49    }
50
51    pub fn json<T: Serialize>(&mut self, data: &T) -> &mut Self {
52        match serde_json::to_string(data) {
53            Ok(body) => {
54                self.set_header("Content-Type", "application/json");
55                self.body = body;
56            }
57            Err(_) => {
58                self.set_header("Content-Type", "application/json");
59                self.body = r#"{"error":"Failed to serialize JSON"}"#.to_string();
60                self.status = 500;
61            }
62        }
63        self
64    }
65
66    pub fn html(&mut self, html: &str) -> &mut Self {
67        self.set_header("Content-Type", "text/html; charset=utf-8");
68        self.body = html.to_string();
69        self
70    }
71
72    pub fn bytes(&mut self, bytes: &[u8]) -> &mut Self {
73        let encoded = general_purpose::STANDARD.encode(bytes);
74        self.body = encoded;
75        self.set_header("Content-Type", "application/octet-stream");
76        self
77    }
78
79    pub fn error(&mut self, status: u16, msg: &str) -> &mut Self {
80        self.status = status;
81        self.body = msg.to_string();
82        self.has_error = true;
83        self
84    }
85
86    pub fn has_error(&self) -> bool {
87        self.has_error
88    }
89
90    pub fn cookie(
91        &mut self,
92        name: &str,
93        value: &str,
94        max_age: Option<i64>,
95        path: Option<&str>,
96        domain: Option<&str>,
97        secure: bool,
98        http_only: bool,
99        same_site: Option<&str>,
100    ) -> &mut Self {
101        let mut cookie = format!("{}={}", name, value);
102
103        if let Some(age) = max_age {
104            if let Ok(expires) =
105                (OffsetDateTime::now_utc() + Duration::seconds(age)).format(&Rfc2822)
106            {
107                cookie.push_str(&format!("; Max-Age={}; Expires={}", age, expires));
108            }
109        }
110
111        if let Some(p) = path {
112            cookie.push_str(&format!("; Path={}", p));
113        }
114
115        if let Some(d) = domain {
116            cookie.push_str(&format!("; Domain={}", d));
117        }
118
119        if secure {
120            cookie.push_str("; Secure");
121        }
122
123        if http_only {
124            cookie.push_str("; HttpOnly");
125        }
126
127        if let Some(same_site_val) = same_site {
128            cookie.push_str(&format!("; SameSite={}", same_site_val));
129        }
130
131        self.headers.append(
132            hyper::header::SET_COOKIE,
133            hyper::header::HeaderValue::from_str(&cookie).unwrap(),
134        );
135
136        self
137    }
138
139    pub fn into_response(self) -> Response<Full<Bytes>> {
140        let mut builder = Response::builder().status(self.status);
141
142        for (key, value) in self.headers.iter() {
143            builder = builder.header(key, value);
144        }
145
146        builder.body(Full::new(Bytes::from(self.body))).unwrap()
147    }
148}