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 mime_guess::from_path;
9use serde::Serialize;
10use std::path::Path;
11use time::{Duration, OffsetDateTime, format_description::well_known::Rfc2822};
12use tokio::fs;
13use tokio::io::AsyncReadExt;
14
15pub struct ResponseWriter {
16    pub body: String,
17    pub headers: HeaderMap,
18    pub status: u16,
19    pub has_error: bool,
20}
21
22#[allow(dead_code)]
23impl ResponseWriter {
24    pub fn new() -> Self {
25        Self {
26            body: "".into(),
27            headers: HeaderMap::new(),
28            status: 200,
29            has_error: false,
30        }
31    }
32
33    pub fn status(&mut self, status: u16) -> &mut Self {
34        self.status = status;
35        self
36    }
37
38    pub fn set_header(&mut self, key: &str, value: &str) -> &mut Self {
39        self.headers.insert(
40            HeaderName::from_bytes(key.as_bytes()).unwrap(),
41            HeaderValue::from_str(value).unwrap(),
42        );
43        self
44    }
45
46    pub fn get_header(&self, key: &str) -> Option<&HeaderValue> {
47        self.headers.get(key)
48    }
49
50    pub fn send(&mut self, body: &str) -> &mut Self {
51        self.body = body.into();
52        self
53    }
54
55    pub fn json<T: Serialize>(&mut self, data: &T) -> &mut Self {
56        match serde_json::to_string(data) {
57            Ok(body) => {
58                self.set_header("Content-Type", "application/json");
59                self.body = body;
60            }
61            Err(_) => {
62                self.set_header("Content-Type", "application/json");
63                self.body = r#"{"error":"Failed to serialize JSON"}"#.to_string();
64                self.status = 500;
65            }
66        }
67        self
68    }
69
70    pub fn html(&mut self, html: &str) -> &mut Self {
71        self.set_header("Content-Type", "text/html; charset=utf-8");
72        self.body = html.to_string();
73        self
74    }
75
76    pub async fn file<P: AsRef<Path>>(&mut self, path: P) {
77        let path_ref = path.as_ref();
78
79        match fs::File::open(path_ref).await {
80            Ok(mut file) => {
81                let mut buf = Vec::new();
82                if let Err(e) = file.read_to_end(&mut buf).await {
83                    self.error(500, &format!("Failed to read file: {}", e));
84                    return;
85                }
86
87                let mime_type = from_path(path_ref).first_or_octet_stream().to_string();
88
89                self.status(200)
90                    .set_header("Content-Type", &mime_type)
91                    .bytes(&buf);
92            }
93            Err(_) => {
94                self.error(404, "File not found");
95            }
96        }
97    }
98
99    pub fn bytes(&mut self, bytes: &[u8]) -> &mut Self {
100        let encoded = general_purpose::STANDARD.encode(bytes);
101        self.body = encoded;
102        self.set_header("Content-Type", "application/octet-stream");
103        self
104    }
105
106    pub fn error(&mut self, status: u16, msg: &str) -> &mut Self {
107        self.status = status;
108        self.body = msg.to_string();
109        self.has_error = true;
110        self
111    }
112
113    pub fn has_error(&self) -> bool {
114        self.has_error
115    }
116
117    pub fn cookie(
118        &mut self,
119        name: &str,
120        value: &str,
121        max_age: Option<i64>,
122        path: Option<&str>,
123        domain: Option<&str>,
124        secure: bool,
125        http_only: bool,
126        same_site: Option<&str>,
127    ) -> &mut Self {
128        let mut cookie = format!("{}={}", name, value);
129
130        if let Some(age) = max_age {
131            if let Ok(expires) =
132                (OffsetDateTime::now_utc() + Duration::seconds(age)).format(&Rfc2822)
133            {
134                cookie.push_str(&format!("; Max-Age={}; Expires={}", age, expires));
135            }
136        }
137
138        if let Some(p) = path {
139            cookie.push_str(&format!("; Path={}", p));
140        }
141
142        if let Some(d) = domain {
143            cookie.push_str(&format!("; Domain={}", d));
144        }
145
146        if secure {
147            cookie.push_str("; Secure");
148        }
149
150        if http_only {
151            cookie.push_str("; HttpOnly");
152        }
153
154        if let Some(same_site_val) = same_site {
155            cookie.push_str(&format!("; SameSite={}", same_site_val));
156        }
157
158        self.headers.append(
159            hyper::header::SET_COOKIE,
160            hyper::header::HeaderValue::from_str(&cookie).unwrap(),
161        );
162
163        self
164    }
165
166    pub fn into_response(self) -> Response<Full<Bytes>> {
167        let mut builder = Response::builder().status(self.status);
168
169        for (key, value) in self.headers.iter() {
170            builder = builder.header(key, value);
171        }
172
173        builder.body(Full::new(Bytes::from(self.body))).unwrap()
174    }
175
176    pub fn strip_header(&mut self, key: &str) {
177        if let Ok(key_name) = hyper::header::HeaderName::from_bytes(key.as_bytes()) {
178            self.headers.remove(key_name);
179        }
180    }
181}
bolt_web/response.rs

bolt_web/
response.rs
