helix_db/protocol/
response.rs

1use std::collections::HashMap;
2use tokio::io::{AsyncWrite, AsyncWriteExt, Result};
3#[derive(Debug)]
4pub struct Response {
5    pub status: u16,
6    pub headers: HashMap<String, String>,
7    pub body: Vec<u8>,
8}
9
10impl Response {
11    /// Create a new response
12    pub fn new() -> Response {
13        let mut headers = HashMap::new();
14        // TODO: Change to use router config for headers and default routes
15        headers.insert("Content-Type".to_string(), "text/plain".to_string());
16
17        Response {
18            status: 200,
19            headers,
20            body: Vec::new(),
21        }
22    }
23
24    /// Send response back via stream
25    ///
26    /// # Example
27    ///
28    /// ```rust
29    /// use std::io::Cursor;
30    /// use helix_db::protocol::response::Response;
31    ///
32    /// let mut response = Response::new();
33    ///
34    /// response.status = 200;
35    /// response.body = b"Hello World".to_vec();
36    ///
37    /// let mut stream = Cursor::new(Vec::new());
38    /// response.send(&mut stream).unwrap();
39    ///
40    /// let data = stream.into_inner();
41    /// let data = String::from_utf8(data).unwrap();
42    ///
43    /// assert!(data.contains("HTTP/1.1 200 OK"));
44    /// assert!(data.contains("Content-Length: 11"));
45    /// assert!(data.contains("Hello World"));
46
47    pub async fn send<W: AsyncWrite + Unpin>(&mut self, stream: &mut W) -> Result<()> {
48        let status_message = match self.status {
49            200 => "OK",
50            404 => {
51                self.body = b"404 - Route Not Found\n".to_vec();
52                "Not Found"
53            }
54            500 => {
55                // self.body = b"500 - Internal Server Error\n".to_vec();
56                "Internal Server Error"
57            }
58            _ => "Unknown",
59        };
60        let mut writer = tokio::io::BufWriter::new(stream);
61
62        // Write status line
63        writer
64            .write_all(format!("HTTP/1.1 {} {}\r\n", self.status, status_message).as_bytes())
65            .await?;
66
67        // Write headers
68        for (header, value) in &self.headers {
69            writer
70                .write_all(format!("{}: {}\r\n", header, value).as_bytes())
71                .await
72                .map_err(|e| {
73                    std::io::Error::new(
74                        std::io::ErrorKind::Other,
75                        format!("Error writing header: {}", e),
76                    )
77                })?;
78        }
79
80        writer
81            .write_all(format!("Content-Length: {}\r\n\r\n", self.body.len()).as_bytes())
82            .await?;
83
84        // Write body
85        writer.write_all(&self.body).await?;
86        writer.flush().await?;
87        Ok(())
88    }
89}