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}