torch_web/
response.rs

1use http::{HeaderMap, HeaderName, HeaderValue, StatusCode};
2use http_body_util::Full;
3use hyper::body::Bytes;
4
5/// HTTP Response builder that provides a fluent API for creating responses
6#[derive(Debug)]
7pub struct Response {
8    status: StatusCode,
9    headers: HeaderMap,
10    body: Vec<u8>,
11}
12
13impl Response {
14    /// Create a new response with 200 OK status
15    pub fn new() -> Self {
16        Self {
17            status: StatusCode::OK,
18            headers: HeaderMap::new(),
19            body: Vec::new(),
20        }
21    }
22
23    /// Create a response with 200 OK status (alias for new)
24    pub fn ok() -> Self {
25        Self::new()
26    }
27
28    /// Create a response with a specific status code
29    pub fn with_status(status: StatusCode) -> Self {
30        Self {
31            status,
32            headers: HeaderMap::new(),
33            body: Vec::new(),
34        }
35    }
36
37    /// Create a 404 Not Found response
38    pub fn not_found() -> Self {
39        Self::with_status(StatusCode::NOT_FOUND)
40            .body("Not Found")
41    }
42
43    /// Create a 500 Internal Server Error response
44    pub fn internal_error() -> Self {
45        Self::with_status(StatusCode::INTERNAL_SERVER_ERROR)
46            .body("Internal Server Error")
47    }
48
49    /// Create a 400 Bad Request response
50    pub fn bad_request() -> Self {
51        Self::with_status(StatusCode::BAD_REQUEST)
52            .body("Bad Request")
53    }
54
55    /// Create a 201 Created response
56    pub fn created() -> Self {
57        Self::with_status(StatusCode::CREATED)
58    }
59
60    /// Create a 204 No Content response
61    pub fn no_content() -> Self {
62        Self::with_status(StatusCode::NO_CONTENT)
63    }
64
65    /// Create a 401 Unauthorized response
66    pub fn unauthorized() -> Self {
67        Self::with_status(StatusCode::UNAUTHORIZED)
68            .body("Unauthorized")
69    }
70
71    /// Create a 403 Forbidden response
72    pub fn forbidden() -> Self {
73        Self::with_status(StatusCode::FORBIDDEN)
74            .body("Forbidden")
75    }
76
77    /// Create a 422 Unprocessable Entity response
78    pub fn unprocessable_entity() -> Self {
79        Self::with_status(StatusCode::UNPROCESSABLE_ENTITY)
80            .body("Unprocessable Entity")
81    }
82
83    /// Create a 429 Too Many Requests response
84    pub fn too_many_requests() -> Self {
85        Self::with_status(StatusCode::TOO_MANY_REQUESTS)
86            .body("Too Many Requests")
87    }
88
89    /// Set the status code
90    pub fn status(mut self, status: StatusCode) -> Self {
91        self.status = status;
92        self
93    }
94
95    /// Set the response body from a string
96    pub fn body<T: Into<Vec<u8>>>(mut self, body: T) -> Self {
97        self.body = body.into();
98        self
99    }
100
101    /// Set the response body from bytes
102    pub fn body_from_bytes(mut self, body: Vec<u8>) -> Self {
103        self.body = body;
104        self
105    }
106
107    /// Set a header
108    pub fn header<K, V>(mut self, key: K, value: V) -> Self
109    where
110        K: TryInto<HeaderName>,
111        V: TryInto<HeaderValue>,
112        K::Error: std::fmt::Debug,
113        V::Error: std::fmt::Debug,
114    {
115        let key = key.try_into().expect("Invalid header name");
116        let value = value.try_into().expect("Invalid header value");
117        self.headers.insert(key, value);
118        self
119    }
120
121    /// Set the Content-Type header
122    pub fn content_type(self, content_type: &str) -> Self {
123        self.header("content-type", content_type)
124    }
125
126    /// Set response as JSON and serialize the value (requires "json" feature)
127    #[cfg(feature = "json")]
128    pub fn json<T: serde::Serialize>(self, value: &T) -> Result<Self, serde_json::Error> {
129        let json_string = serde_json::to_string(value)?;
130        Ok(self
131            .content_type("application/json")
132            .body(json_string))
133    }
134
135    /// Set response as HTML
136    pub fn html<T: Into<Vec<u8>>>(self, html: T) -> Self {
137        self.content_type("text/html; charset=utf-8")
138            .body(html)
139    }
140
141    /// Set response as plain text
142    pub fn text<T: Into<Vec<u8>>>(self, text: T) -> Self {
143        self.content_type("text/plain; charset=utf-8")
144            .body(text)
145    }
146
147    /// Redirect to another URL
148    pub fn redirect(status: StatusCode, location: &str) -> Self {
149        Self::with_status(status)
150            .header("location", location)
151    }
152
153    /// Redirect with 302 Found status
154    pub fn redirect_found(location: &str) -> Self {
155        Self::redirect(StatusCode::FOUND, location)
156    }
157
158    /// Redirect with 301 Moved Permanently status
159    pub fn redirect_permanent(location: &str) -> Self {
160        Self::redirect(StatusCode::MOVED_PERMANENTLY, location)
161    }
162
163    /// Get the status code
164    pub fn status_code(&self) -> StatusCode {
165        self.status
166    }
167
168    /// Get a mutable reference to the status code
169    pub fn status_code_mut(&mut self) -> &mut StatusCode {
170        &mut self.status
171    }
172
173    /// Get the headers
174    pub fn headers(&self) -> &HeaderMap {
175        &self.headers
176    }
177
178    /// Get the body as bytes
179    pub fn body_data(&self) -> &[u8] {
180        &self.body
181    }
182
183    /// Get the body as bytes (alias for body_data)
184    pub fn body_bytes(&self) -> &[u8] {
185        &self.body
186    }
187
188    /// Convert to hyper Response
189    pub fn into_hyper_response(self) -> hyper::Response<Full<Bytes>> {
190        let mut response = hyper::Response::builder()
191            .status(self.status);
192
193        // Add headers
194        for (key, value) in self.headers.iter() {
195            response = response.header(key, value);
196        }
197
198        response
199            .body(Full::new(Bytes::from(self.body)))
200            .expect("Failed to build response")
201    }
202}
203
204impl Default for Response {
205    fn default() -> Self {
206        Self::new()
207    }
208}
209
210impl From<&str> for Response {
211    fn from(body: &str) -> Self {
212        Response::ok().body(body)
213    }
214}
215
216impl From<String> for Response {
217    fn from(body: String) -> Self {
218        Response::ok().body(body)
219    }
220}
221
222impl From<StatusCode> for Response {
223    fn from(status: StatusCode) -> Self {
224        Response::with_status(status)
225    }
226}
227
228#[cfg(test)]
229mod tests {
230    use super::*;
231
232    #[test]
233    fn test_response_creation() {
234        let response = Response::ok().body("Hello, World!");
235        assert_eq!(response.status_code(), StatusCode::OK);
236        assert_eq!(response.body_data(), b"Hello, World!");
237    }
238
239    #[test]
240    fn test_response_with_headers() {
241        let response = Response::ok()
242            .header("x-custom", "value")
243            .content_type("text/plain")
244            .body("test");
245        
246        assert_eq!(response.headers().get("x-custom").unwrap(), "value");
247        assert_eq!(response.headers().get("content-type").unwrap(), "text/plain");
248    }
249
250    #[test]
251    fn test_redirect() {
252        let response = Response::redirect_found("/new-path");
253        assert_eq!(response.status_code(), StatusCode::FOUND);
254        assert_eq!(response.headers().get("location").unwrap(), "/new-path");
255    }
256
257    #[cfg(feature = "json")]
258    #[test]
259    fn test_json_response() {
260        use serde_json::json;
261        
262        let data = json!({"message": "Hello, World!"});
263        let response = Response::ok().json(&data).unwrap();
264        
265        assert_eq!(response.headers().get("content-type").unwrap(), "application/json");
266        assert_eq!(response.body_data(), br#"{"message":"Hello, World!"}"#);
267    }
268}