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

bolt_web/
response.rs
