use std::collections::HashMap;
use crate::{
security::xss::{SafeHtml, Sanitizer},
utils::fs,
};
#[derive(Debug, Clone)]
pub enum SameSite {
Strict,
Lax,
None,
}
#[derive(Debug, Clone)]
pub struct Cookie {
pub name: String,
pub value: String,
pub max_age: u64, pub http_only: bool, pub secure: bool, pub same_site: SameSite,
}
impl Cookie {
pub fn new(name: &str, value: &str) -> Self {
Cookie {
name: name.to_string(),
value: value.to_string(),
max_age: 3600, http_only: true, secure: true, same_site: SameSite::Strict, }
}
pub fn set_secure(mut self, secure: bool) -> Self {
self.secure = secure;
self
}
pub fn set_same_site(mut self, same_site: SameSite) -> Self {
self.same_site = same_site;
self
}
}
pub enum ResponseBody {
Html(SafeHtml),
StaticFile(String),
Json(String),
}
pub trait IntoResponseBody {
fn convert(self) -> (ResponseBody, String); }
impl IntoResponseBody for SafeHtml {
fn convert(self) -> (ResponseBody, String) {
(
ResponseBody::Html(self),
"text/html; charset=utf-8".to_string(),
)
}
}
impl IntoResponseBody for String {
fn convert(self) -> (ResponseBody, String) {
(
ResponseBody::Html(Sanitizer::trust(&self)),
"text/html; charset=utf-8".to_string(),
)
}
}
impl IntoResponseBody for &'static str {
fn convert(self) -> (ResponseBody, String) {
(
ResponseBody::Html(Sanitizer::trust(self)),
"text/html; charset=utf-8".to_string(),
)
}
}
pub struct JsonPayload<T>(pub T);
impl<T: serde::Serialize> IntoResponseBody for JsonPayload<T> {
fn convert(self) -> (ResponseBody, String) {
let json_string = serde_json::to_string(&self.0)
.unwrap_or_else(|_| r#"{"error": "Internal Server Serialization Error"}"#.to_string());
(
ResponseBody::Json(json_string),
"application/json; charset=utf-8".to_string(),
)
}
}
impl<K, V> IntoResponseBody for HashMap<K, V>
where
K: serde::Serialize + std::hash::Hash + Eq,
V: serde::Serialize,
{
fn convert(self) -> (ResponseBody, String) {
let json_string = serde_json::to_string(&self)
.unwrap_or_else(|_| r#"{"error": "Internal Server Serialization Error"}"#.to_string());
(
ResponseBody::Json(json_string),
"application/json; charset=utf-8".to_string(),
)
}
}
impl<K, V> IntoResponseBody for &HashMap<K, V>
where
K: serde::Serialize + std::hash::Hash + Eq,
V: serde::Serialize,
{
fn convert(self) -> (ResponseBody, String) {
let json_string = serde_json::to_string(self)
.unwrap_or_else(|_| r#"{"error": "Internal Server Serialization Error"}"#.to_string());
(
ResponseBody::Json(json_string),
"application/json; charset=utf-8".to_string(),
)
}
}
pub struct Response {
pub status: u16,
pub headers: Vec<(String, String)>,
pub cookies: Vec<Cookie>,
pub body: ResponseBody,
}
impl Response {
pub fn new(status: u16, body: SafeHtml) -> Self {
Response {
status,
headers: vec![
(
"Content-Type".to_string(),
"text/html; charset=utf-8".to_string(),
),
("X-Content-Type-Options".to_string(), "nosniff".to_string()),
("X-Frame-Options".to_string(), "DENY".to_string()),
],
cookies: Vec::new(),
body: ResponseBody::Html(body),
}
}
pub fn with_cookie(mut self, cookie: Cookie) -> Self {
self.cookies.push(cookie);
self
}
pub fn static_file(path: &str) -> Self {
Response {
status: 200,
headers: vec![
("X-Content-Type-Options".to_string(), "nosniff".to_string()),
("X-Frame-Options".to_string(), "DENY".to_string()),
],
cookies: Vec::new(),
body: ResponseBody::StaticFile(path.to_string()),
}
}
pub fn to_bytes(&self, body_bytes: &[u8], content_type: &str) -> Vec<u8> {
let mut response = format!("HTTP/1.1 {} OK\r\n", self.status);
response.push_str(&format!("Content-Type: {}\r\n", content_type));
for (key, value) in &self.headers {
response.push_str(&format!("{}: {}\r\n", key, value));
}
for cookie in &self.cookies {
let same_site_str = match cookie.same_site {
SameSite::Strict => "Strict",
SameSite::Lax => "Lax",
SameSite::None => "None",
};
let mut cookie_str = format!(
"Set-Cookie: {}={}; Max-Age={}; SameSite={}; Path=/",
cookie.name, cookie.value, cookie.max_age, same_site_str
);
if cookie.http_only {
cookie_str.push_str("; HttpOnly");
}
if cookie.secure {
cookie_str.push_str("; Secure");
}
println!("{}", cookie_str);
response.push_str(&format!("{}\r\n", cookie_str));
}
response.push_str("\r\n");
let mut raw = response.into_bytes();
raw.extend_from_slice(body_bytes);
raw
}
pub fn json<T: serde::Serialize>(status: u16, data: &T) -> Self {
let json_string = serde_json::to_string(data)
.unwrap_or_else(|_| r#"{"error": "Internal Server Serialization Error"}"#.to_string());
Response {
status,
headers: vec![
(
"Content-Type".to_string(),
"application/json; charset=utf-8".to_string(),
),
("X-Content-Type-Options".to_string(), "nosniff".to_string()),
],
cookies: Vec::new(),
body: ResponseBody::Json(json_string),
}
}
pub fn resolve(&self) -> (Vec<u8>, String) {
match &self.body {
ResponseBody::Html(html) => (html.as_bytes().to_vec(), "text/html".to_string()),
ResponseBody::Json(json_str) => {
(json_str.as_bytes().to_vec(), "application/json".to_string())
}
ResponseBody::StaticFile(path) => fs::serve_static(path).unwrap_or_else(|_| {
(
Sanitizer::trust("<h1>404 File Not Found</h1>")
.as_bytes()
.to_vec(),
"text/html".to_string(),
)
}),
}
}
pub fn redirect(status: u16, location: &str) -> Self {
Response {
status,
headers: vec![
("Location".to_string(), location.to_string()),
(
"Content-Type".to_string(),
"text/html; charset=utf-8".to_string(),
),
("X-Content-Type-Options".to_string(), "nosniff".to_string()),
("X-Frame-Options".to_string(), "DENY".to_string()),
],
cookies: Vec::new(),
body: ResponseBody::Html(Sanitizer::trust(
format!(
"Redirecting to <a href=\"{}\">{}</a>...",
location, location
)
.as_str(),
)),
}
}
pub fn build<B: IntoResponseBody>(status: u16, payload: B) -> Self {
let (body, content_type) = payload.convert();
Response {
status,
headers: vec![
("Content-Type".to_string(), content_type),
("X-Content-Type-Options".to_string(), "nosniff".to_string()),
("X-Frame-Options".to_string(), "DENY".to_string()),
],
cookies: Vec::new(),
body,
}
}
pub fn ok<B: IntoResponseBody>(payload: B) -> Self {
Self::build(200, payload)
}
pub fn created<B: IntoResponseBody>(payload: B) -> Self {
Self::build(201, payload)
}
pub fn bad_request<B: IntoResponseBody>(payload: B) -> Self {
Self::build(400, payload)
}
pub fn unauthorized<B: IntoResponseBody>(payload: B) -> Self {
Self::build(401, payload)
}
pub fn forbidden<B: IntoResponseBody>(payload: B) -> Self {
Self::build(403, payload)
}
pub fn not_found<B: IntoResponseBody>(payload: B) -> Self {
Self::build(404, payload)
}
pub fn internal_error<B: IntoResponseBody>(payload: B) -> Self {
Self::build(500, payload)
}
}