use super::errors::HeaderError;
use bytes::{Bytes, BytesMut};
use http::{HeaderMap, HeaderName, HeaderValue, StatusCode};
#[cfg(feature = "json")]
use serde::Serialize;
use std::{fs::File, io::Read, str::FromStr};
#[derive(Debug, Default)]
pub struct Response {
pub status: StatusCode,
pub headers: HeaderMap,
pub body: Option<Bytes>,
pub version: http::Version,
}
impl Response {
const MAX_FILE_SIZE_BYTES: u64 = 4 * 1024 * 1024;
fn set_common_headers(&mut self, content_type: Option<&'static str>, len: usize) {
if let Some(ct) = content_type {
self.headers.insert(HeaderName::from_static("content-type"), HeaderValue::from_static(ct));
}
self.headers.insert(HeaderName::from_static("content-length"), Self::len_to_header_value(len));
}
pub fn set_status(&mut self, status: u16) -> &mut Response {
self.status = StatusCode::from_u16(status).unwrap_or(StatusCode::INTERNAL_SERVER_ERROR);
self
}
pub fn add_header(&mut self, key: &str, value: &str) -> Result<(), HeaderError> {
let val = HeaderValue::from_str(value)?;
let key = HeaderName::from_str(key)?;
self.headers.insert(key, val);
return Ok(());
}
pub fn to_raw(&self) -> Bytes {
let body_len = self.body.as_ref().map_or(0, |b| b.len());
let mut buf = BytesMut::with_capacity(512 + body_len);
buf.extend_from_slice(b"HTTP/1.1 ");
let mut status_buffer = itoa::Buffer::new();
let status_code_str = status_buffer.format(self.status.as_u16());
buf.extend_from_slice(status_code_str.as_bytes());
buf.extend_from_slice(b" ");
buf.extend_from_slice(self.status.canonical_reason().unwrap_or("Unknown").as_bytes());
buf.extend_from_slice(b"\r\n");
for (key, value) in &self.headers {
buf.extend_from_slice(key.as_str().as_bytes());
buf.extend_from_slice(b": ");
buf.extend_from_slice(value.as_bytes());
buf.extend_from_slice(b"\r\n");
}
if !self.headers.contains_key("date") {
let date_str = chrono::Utc::now().to_rfc2822();
buf.extend_from_slice(b"date: ");
buf.extend_from_slice(date_str.as_bytes());
buf.extend_from_slice(b"\r\n");
}
if !self.headers.contains_key("content-length") && body_len > 0 {
buf.extend_from_slice(b"content-length: ");
let mut len_buffer = itoa::Buffer::new();
let len_str = len_buffer.format(body_len);
buf.extend_from_slice(len_str.as_bytes());
buf.extend_from_slice(b"\r\n");
}
buf.extend_from_slice(b"\r\n");
if let Some(ref body) = self.body {
buf.extend_from_slice(body);
}
buf.freeze()
}
pub fn send_text(&mut self, data: impl Into<String>) {
let body = data.into();
self.body = Some(Bytes::from(body));
self.set_common_headers(Some("text/plain;charset=utf-8"), self.body.as_ref().unwrap().len());
}
pub fn send_bytes(&mut self, data: impl Into<Vec<u8>>) {
let body = data.into();
self.body = Some(Bytes::from(body));
self.set_common_headers(None, self.body.as_ref().unwrap().len());
}
pub fn send_html(&mut self, data: impl Into<String>) {
let body = data.into();
self.body = Some(Bytes::from(body));
self.headers.insert(HeaderName::from_static("content-type"), HeaderValue::from_static("text/html"));
let len = self.body.as_ref().unwrap().len();
self.headers.insert(HeaderName::from_static("content-length"), Self::len_to_header_value(len));
}
#[cfg(feature = "json")]
pub fn send_json<T: Serialize>(&mut self, data: &T) {
match serde_json::to_string(&data) {
Ok(json) => {
self.body = Some(Bytes::from(json));
self.headers.insert(HeaderName::from_static("content-type"), HeaderValue::from_static("application/json"));
let len = self.body.as_ref().unwrap().len();
self.headers.insert(HeaderName::from_static("content-length"), Self::len_to_header_value(len));
}
Err(_) => {
self.status = StatusCode::INTERNAL_SERVER_ERROR;
self.body = Some(Bytes::from("Internal Server Error"));
self.headers.insert(HeaderName::from_static("content-type"), HeaderValue::from_static("text/plain"));
let len = self.body.as_ref().unwrap().len();
self.headers.insert(HeaderName::from_static("content-length"), Self::len_to_header_value(len));
}
}
}
pub fn send_file(&mut self, mut file: File) {
let metadata = match file.metadata() {
Ok(m) => m,
Err(_) => {
self.status = StatusCode::INTERNAL_SERVER_ERROR;
self.body = Some(Bytes::from("Failed to read file metadata."));
return;
}
};
if metadata.len() > Self::MAX_FILE_SIZE_BYTES {
self.status = StatusCode::PAYLOAD_TOO_LARGE; self.body = Some(Bytes::from("File size exceeds 4MB limit. Use chunked encoding for larger files."));
return;
}
let mut buffer = Vec::new();
match file.read_to_end(&mut buffer) {
Ok(_) => {
self.body = Some(Bytes::from(buffer));
let len = self.body.as_ref().unwrap().len();
self.headers.insert(HeaderName::from_static("content-length"), Self::len_to_header_value(len));
}
Err(_) => {
self.status = StatusCode::INTERNAL_SERVER_ERROR;
self.body = Some(Bytes::from("Internal Server Error during file read."));
}
}
}
pub fn redirect(&mut self, location: &str, permanent: bool) {
let status = if permanent {
StatusCode::MOVED_PERMANENTLY
} else {
StatusCode::FOUND
};
self.set_status(status.as_u16());
self.headers.insert(HeaderName::from_static("location"), HeaderValue::from_str(location).unwrap());
self.body = Some(Bytes::from(format!("Redirecting to {}", location)));
let len = self.body.as_ref().unwrap().len();
self.headers.insert(HeaderName::from_static("content-length"), Self::len_to_header_value(len));
}
fn len_to_header_value(len: usize) -> HeaderValue {
let mut buffer = itoa::Buffer::new();
let len_str = buffer.format(len);
HeaderValue::from_bytes(len_str.as_bytes()).expect("itoa::Buffer output should be a valid HeaderValue")
}
}