Skip to main content

ullun/
response.rs

1//! Response types and builders
2
3use crate::cookies::Cookie;
4use crate::error::Error;
5use bytes::Bytes;
6use http::StatusCode;
7use serde::Serialize;
8
9/// HTTP Response wrapper
10pub 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
18/// Body type of everything sent to hyper: buffered or streamed
19pub(crate) type HyperBody = http_body_util::combinators::BoxBody<Bytes, std::convert::Infallible>;
20
21/// Chunks arriving through a channel, ended when every sender is dropped
22struct 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    /// A response whose body is produced while it is being sent
52    ///
53    /// Returns the sender for the body chunks and the response; set status and
54    /// headers on the response as usual. The body ends when the last
55    /// [`BodySender`](crate::stream::BodySender) is dropped, and `send` fails
56    /// with [`Disconnected`](crate::stream::Disconnected) once the client is gone.
57    /// Middleware that rewrites `body` (compression, ETags) leaves streams alone.
58    ///
59    /// ```
60    /// use ullun::prelude::*;
61    ///
62    /// async fn countdown(_req: Request) -> Result<Response> {
63    ///     let (tx, res) = Response::stream();
64    ///     tokio::spawn(async move {
65    ///         for n in (1..=3).rev() {
66    ///             if tx.send(format!("{n}\n")).await.is_err() {
67    ///                 return; // client went away
68    ///             }
69    ///         }
70    ///     });
71    ///     Ok(res.with_header("content-type", "text/plain"))
72    /// }
73    /// ```
74    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    /// The response a `HEAD` request gets: same status and headers, no body.
82    /// `content-length` is set from the buffered body unless already present.
83    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    /// Whether `Response::cookie` was used
101    pub(crate) fn has_cookies(&self) -> bool {
102        !self.cookies.is_empty()
103    }
104
105    /// Whether the body is streamed instead of held in `body`
106    pub fn is_streaming(&self) -> bool {
107        self.stream.is_some()
108    }
109
110    /// Set a cookie
111    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        // Validate and insert header, skip invalid ones
135        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    /// `200 OK` with an HTML body
167    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    /// `303 See Other`, the right answer after a form POST (Post/Redirect/Get)
174    pub fn redirect(location: impl Into<String>) -> Self {
175        Self::new(StatusCode::SEE_OTHER).with_header("location", location)
176    }
177
178    /// `308 Permanent Redirect`, keeps method and body
179    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            // field -> list of messages, in the order they were found
191            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
213        let body: HyperBody = match self.stream {
214            Some(rx) => ChannelBody { rx }.boxed(),
215            None => http_body_util::Full::new(self.body)
216                .map_err(|never| match never {})
217                .boxed(),
218        };
219        let mut res = hyper::Response::new(body);
220        *res.status_mut() = self.status;
221        // headers were validated when they were set: move the map as it is
222        *res.headers_mut() = self.headers;
223        for cookie in self.cookies {
224            match http::HeaderValue::from_str(&cookie.to_header_value()) {
225                Ok(v) => {
226                    res.headers_mut().append(http::header::SET_COOKIE, v);
227                }
228                Err(e) => eprintln!("Failed to build response: {}", e),
229            }
230        }
231        res
232    }
233}
234
235/// JSON response helper
236pub struct Json<T>(pub T);
237
238impl<T: Serialize> From<Json<T>> for Response {
239    fn from(json: Json<T>) -> Self {
240        Response::json(json.0)
241    }
242}