coa_website/

response.rs

use std::collections::HashMap;
use std::io::Write;

use httpdate::HttpDate;

use crate::http::StatusCode;

/// Represents an HTTP Response send to a client
pub struct HttpResponse {
    /// HTTP status code
    pub status: StatusCode,
    /// HTTP headers
    pub headers: HashMap<String, String>,
    /// Body
    pub body: Option<Vec<u8>>,
}

impl HttpResponse {

    /// Converts the response to bytes to be sent in the TCPStream
    pub fn to_bytes(&self) -> Vec<u8> {
        let mut buffer = Vec::new();
        let _ = write!(
            &mut buffer,
            "HTTP/1.1 {} {}\r\n",
            self.status.code(),
            self.status.reason()
        );

        for (k, v) in &self.headers {
            let _ = write!(&mut buffer, "{}: {}\r\n", k, v);
        }
        buffer.extend_from_slice(b"\r\n");

        match(self.headers.get("Transfer-Encoding").map(|v| v.eq_ignore_ascii_case("chunked")), &self.body) {
            (Some(true), Some(body)) => {
                if !body.is_empty() {
                    let _ = write!(&mut buffer, "{:X}\r\n", body.len());
                    buffer.extend_from_slice(body);
                    buffer.extend_from_slice(b"\r\n");
                }
                buffer.extend_from_slice(b"0\r\n\r\n");
            },
            (_, Some(body)) => buffer.extend_from_slice(body),
            _ => {}
        }

        buffer
    }

    /// Create a new builder to construct a HttpResponse
    pub fn builder() -> HttpResponseBuilder {
        HttpResponseBuilder::default()
    }

    /// Create a preconfigured builder for 200 OK response
    pub fn ok(content_type: &str) -> HttpResponseBuilder {
        HttpResponse::builder().with_status(StatusCode::Ok).with_content_type(content_type)
    }

    /// Create a preconfigured builder for 404 Not Found response
    pub fn not_found() -> HttpResponseBuilder {
        HttpResponse::builder().with_status(StatusCode::NotFound).with_body(vec![])
    }

    /// Create a preconfigured builder for 405 Method Not Allowed response
    pub fn method_not_allowed() -> HttpResponseBuilder {
        HttpResponse::builder().with_status(StatusCode::MethodNotAllowed).with_body(vec![])
    }

    /// Create a preconfigured builder for 500 Internal Server Error
    pub fn internal_server_error() -> HttpResponseBuilder {
        HttpResponse::builder().with_status(StatusCode::InternalServerError).with_body(vec![])
    }
}

/// Builder pattern to build a `HttpResponse`
#[derive(Debug, Default)]
pub struct HttpResponseBuilder {
    pub status: Option<StatusCode>,
    pub headers: HashMap<String, String>,
    pub body: Option<Vec<u8>>,
}


impl HttpResponseBuilder {

    /// Sets the status code of the response
    pub fn with_status(mut self, status: StatusCode) -> Self {
        self.status = Some(status);
        self
    }

    /// Adds an header to the response
    pub fn with_header(mut self, key: &str, value: &str) -> Self {
        self.headers.insert(key.to_string(), value.to_string());
        self
    }

    /// Sets the Content-Type header
    pub fn with_content_type(self, content_type: &str) -> Self {
        self.with_header("Content-Type", content_type)
    }

    /// Sets the response body
    pub fn with_body(mut self, body: Vec<u8>) -> Self {
        self.body = Some(body);
        self
    }

    /// Build the response into `HttpResponse`, adding default headers if missing
    /// 
    /// # Errors
    /// Returns an error if the status is not specified
    pub fn build(mut self) -> Result<HttpResponse, &'static str> {
        if !self.headers.contains_key("Server") {
            self.headers.insert("Server".to_string(), "CoaWebserver/1.0".to_string());
        }
        if !self.headers.contains_key("Date") {
            self.headers.insert("Date".to_string(), format!("{}", HttpDate::from(std::time::SystemTime::now())));
        }

        if !self.headers.contains_key("Transfer-Encoding") && !self.headers.contains_key("Content-Length") {
            if let Some(ref b) = self.body {
                self.headers.insert("Content-Length".to_string(), b.len().to_string());
            } else {
                self.headers.insert("Transfer-Encoding".to_string(), "chunked".to_string());
            }
        }

        match self.status.take() {
            Some(status) => Ok(HttpResponse { status, headers: self.headers, body: self.body}),
            None => Err("Status is not specified.")
        }
    }

}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_response_to_bytes() {
        let res = HttpResponse::ok("text/plain").with_body(b"Hello".to_vec()).build().unwrap();
        let res_str = String::from_utf8(res.to_bytes()).unwrap();
        assert!(res_str.contains("HTTP/1.1 200 OK"));
        assert!(res_str.contains("Content-Type: text/plain"));
        assert!(res_str.contains("Content-Length: 5"));
        assert!(res_str.contains("Hello"));
    }

    #[test]
    fn test_chunked_transfer_encoding() {
        let response = HttpResponse::builder()
            .with_status(StatusCode::Ok)
            .with_header("Transfer-Encoding", "chunked")
            .with_body(b"HelloChunk".to_vec())
            .build()
            .unwrap();

        let bytes = response.to_bytes();
        let output = String::from_utf8_lossy(&bytes);

        assert!(output.contains("Transfer-Encoding: chunked"));
        assert!(!output.contains("Content-Length"));
        assert!(output.contains("A\r\nHelloChunk\r\n0\r\n\r\n"));
    }

    #[test]
    fn test_chunked_empty_body() {
        let res = HttpResponse::builder()
            .with_status(StatusCode::Ok)
            .with_header("Transfer-Encoding", "chunked")
            .with_body(vec![])
            .build()
            .unwrap();

        let bytes = res.to_bytes();
        let out = String::from_utf8_lossy(&bytes);
        assert!(out.contains("0\r\n\r\n"));
    }

    #[test]
    fn test_auto_headers_inserted() {
        let response = HttpResponse::builder()
            .with_status(StatusCode::Ok)
            .with_body(b"Ping".to_vec())
            .build()
            .unwrap();

        let bytes = response.to_bytes();
        let s = String::from_utf8_lossy(&bytes);
        assert!(s.contains("Server: CoaWebserver/1.0"));
        assert!(s.contains("Date:"));
    }

    #[test]
    fn test_missing_status_should_fail() {
        let result = HttpResponse::builder().with_body(b"Missing".to_vec()).build();
        assert!(result.is_err());
    }

}