1use crate::cookies::Cookie;
4use crate::error::Error;
5use bytes::Bytes;
6use http::StatusCode;
7use serde::Serialize;
8
9pub struct Response {
11 pub status: StatusCode,
12 pub headers: http::HeaderMap,
13 pub body: Bytes,
14 cookies: Vec<Cookie>,
15 stream: Option<tokio::sync::mpsc::Receiver<Bytes>>,
16}
17
18pub(crate) type HyperBody = http_body_util::combinators::BoxBody<Bytes, std::convert::Infallible>;
20
21struct ChannelBody {
23 rx: tokio::sync::mpsc::Receiver<Bytes>,
24}
25
26impl hyper::body::Body for ChannelBody {
27 type Data = Bytes;
28 type Error = std::convert::Infallible;
29
30 fn poll_frame(
31 mut self: std::pin::Pin<&mut Self>,
32 cx: &mut std::task::Context<'_>,
33 ) -> std::task::Poll<Option<std::result::Result<hyper::body::Frame<Bytes>, Self::Error>>> {
34 self.rx
35 .poll_recv(cx)
36 .map(|chunk| chunk.map(|b| Ok(hyper::body::Frame::data(b))))
37 }
38}
39
40impl Response {
41 pub fn new(status: StatusCode) -> Self {
42 Self {
43 status,
44 headers: http::HeaderMap::new(),
45 body: Bytes::new(),
46 cookies: Vec::new(),
47 stream: None,
48 }
49 }
50
51 pub fn stream() -> (crate::stream::BodySender, Self) {
75 let (tx, rx) = tokio::sync::mpsc::channel(16);
76 let mut res = Self::ok();
77 res.stream = Some(rx);
78 (crate::stream::BodySender::new(tx), res)
79 }
80
81 pub(crate) fn into_head(mut self) -> Self {
84 let bodiless = self.status.is_informational()
85 || self.status == StatusCode::NO_CONTENT
86 || self.status == StatusCode::NOT_MODIFIED;
87 if self.stream.is_none()
88 && !bodiless
89 && !self.headers.contains_key(http::header::CONTENT_LENGTH)
90 {
91 if let Ok(len) = http::HeaderValue::from_str(&self.body.len().to_string()) {
92 self.headers.insert(http::header::CONTENT_LENGTH, len);
93 }
94 }
95 self.body = Bytes::new();
96 self.stream = None;
97 self
98 }
99
100 pub(crate) fn has_cookies(&self) -> bool {
102 !self.cookies.is_empty()
103 }
104
105 pub fn is_streaming(&self) -> bool {
107 self.stream.is_some()
108 }
109
110 pub fn cookie(mut self, cookie: Cookie) -> Self {
112 self.cookies.push(cookie);
113 self
114 }
115
116 pub fn ok() -> Self {
117 Self::new(StatusCode::OK)
118 }
119
120 pub fn with_status(mut self, status: StatusCode) -> Self {
121 self.status = status;
122 self
123 }
124
125 pub fn with_body(mut self, body: impl Into<Bytes>) -> Self {
126 self.body = body.into();
127 self
128 }
129
130 pub fn with_header(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
131 let key = key.into();
132 let value = value.into();
133
134 if let (Ok(header_name), Ok(header_value)) = (
136 http::HeaderName::from_bytes(key.as_bytes()),
137 http::HeaderValue::from_str(&value),
138 ) {
139 self.headers.insert(header_name, header_value);
140 } else {
141 eprintln!("Warning: Invalid header {}={}", key, value);
142 }
143 self
144 }
145
146 pub fn json<T: Serialize>(data: T) -> Self {
147 match serde_json::to_vec(&data) {
148 Ok(body) => Self::ok()
149 .with_header("content-type", "application/json")
150 .with_body(body),
151 Err(e) => {
152 eprintln!("JSON serialization error: {}", e);
153 Self::new(StatusCode::INTERNAL_SERVER_ERROR)
154 .with_header("content-type", "application/json")
155 .with_body("{\"error\":\"Internal serialization error\"}")
156 }
157 }
158 }
159
160 pub fn text(text: impl Into<String>) -> Self {
161 Self::ok()
162 .with_header("content-type", "text/plain")
163 .with_body(text.into())
164 }
165
166 pub fn html(html: impl Into<String>) -> Self {
168 Self::ok()
169 .with_header("content-type", "text/html; charset=utf-8")
170 .with_body(html.into())
171 }
172
173 pub fn redirect(location: impl Into<String>) -> Self {
175 Self::new(StatusCode::SEE_OTHER).with_header("location", location)
176 }
177
178 pub fn redirect_permanent(location: impl Into<String>) -> Self {
180 Self::new(StatusCode::PERMANENT_REDIRECT).with_header("location", location)
181 }
182
183 pub fn from_error(error: Error) -> Self {
184 let status = error.status();
185 let mut body = serde_json::json!({
186 "error": error.message(),
187 "status": error.status_code(),
188 });
189 if let Error::Validation(fields) = &error {
190 let mut by_field = serde_json::Map::new();
192 for f in fields {
193 if let Some(list) = by_field
194 .entry(f.field.clone())
195 .or_insert_with(|| serde_json::Value::Array(Vec::new()))
196 .as_array_mut()
197 {
198 list.push(serde_json::Value::String(f.message.clone()));
199 }
200 }
201 body["fields"] = serde_json::Value::Object(by_field);
202 }
203 let body_bytes =
204 serde_json::to_vec(&body).unwrap_or_else(|_| b"{\"error\":\"Unknown error\"}".to_vec());
205 Self::new(status)
206 .with_header("content-type", "application/json")
207 .with_body(body_bytes)
208 }
209
210 pub(crate) fn into_hyper(self) -> hyper::Response<HyperBody> {
211 use http_body_util::BodyExt;
212 let mut builder = hyper::Response::builder().status(self.status);
213
214 for (key, value) in self.headers.iter() {
215 builder = builder.header(key, value);
216 }
217
218 for cookie in self.cookies {
220 builder = builder.header("set-cookie", cookie.to_header_value());
221 }
222
223 let body: HyperBody = match self.stream {
224 Some(rx) => ChannelBody { rx }.boxed(),
225 None => http_body_util::Full::new(self.body)
226 .map_err(|never| match never {})
227 .boxed(),
228 };
229 builder.body(body).unwrap_or_else(|e| {
230 eprintln!("Failed to build response: {}", e);
231 hyper::Response::builder()
232 .status(StatusCode::INTERNAL_SERVER_ERROR)
233 .body(
234 http_body_util::Full::new(Bytes::from("Internal Server Error"))
235 .map_err(|never| match never {})
236 .boxed(),
237 )
238 .unwrap()
239 })
240 }
241}
242
243pub struct Json<T>(pub T);
245
246impl<T: Serialize> From<Json<T>> for Response {
247 fn from(json: Json<T>) -> Self {
248 Response::json(json.0)
249 }
250}