use super::Response;
use crate::error::Result;
use ahash::AHashMap;
use bytes::Bytes;
use http::{HeaderMap, HeaderName, HeaderValue, StatusCode};
use once_cell::sync::Lazy;
use serde::Serialize;
use std::borrow::Cow;
static COMMON_RESPONSES: Lazy<AHashMap<&'static str, Bytes>> = Lazy::new(|| {
let mut map = AHashMap::new();
map.insert("health_ok", Bytes::from_static(b"{\"status\":\"healthy\"}"));
map.insert(
"not_found",
Bytes::from_static(b"{\"error\":\"Not Found\"}"),
);
map.insert(
"server_error",
Bytes::from_static(b"{\"error\":\"Internal Server Error\"}"),
);
map.insert(
"unauthorized",
Bytes::from_static(b"{\"error\":\"Unauthorized\"}"),
);
map.insert(
"forbidden",
Bytes::from_static(b"{\"error\":\"Forbidden\"}"),
);
map.insert(
"bad_request",
Bytes::from_static(b"{\"error\":\"Bad Request\"}"),
);
map.insert(
"method_not_allowed",
Bytes::from_static(b"{\"error\":\"Method Not Allowed\"}"),
);
map.insert("empty_json", Bytes::from_static(b"{}"));
map.insert("empty_array", Bytes::from_static(b"[]"));
map.insert("ok_message", Bytes::from_static(b"{\"message\":\"OK\"}"));
map.insert("success", Bytes::from_static(b"{\"success\":true}"));
map.insert("pong", Bytes::from_static(b"{\"message\":\"pong\"}"));
map
});
static COMMON_HEADERS: Lazy<AHashMap<&'static str, HeaderValue>> = Lazy::new(|| {
let mut map = AHashMap::new();
map.insert("json", HeaderValue::from_static("application/json"));
map.insert(
"text",
HeaderValue::from_static("text/plain; charset=utf-8"),
);
map.insert("html", HeaderValue::from_static("text/html; charset=utf-8"));
map.insert("xml", HeaderValue::from_static("application/xml"));
map.insert("css", HeaderValue::from_static("text/css"));
map.insert("js", HeaderValue::from_static("application/javascript"));
map.insert("png", HeaderValue::from_static("image/png"));
map.insert("jpg", HeaderValue::from_static("image/jpeg"));
map.insert("gif", HeaderValue::from_static("image/gif"));
map.insert("svg", HeaderValue::from_static("image/svg+xml"));
map.insert("pdf", HeaderValue::from_static("application/pdf"));
map.insert(
"octet",
HeaderValue::from_static("application/octet-stream"),
);
map.insert("cors_any", HeaderValue::from_static("*"));
map.insert(
"cors_methods",
HeaderValue::from_static("GET, POST, PUT, DELETE, OPTIONS"),
);
map.insert(
"cors_headers",
HeaderValue::from_static("Content-Type, Authorization"),
);
map
});
static CONTENT_TYPE: Lazy<HeaderName> = Lazy::new(|| HeaderName::from_static("content-type"));
static CONTENT_LENGTH: Lazy<HeaderName> = Lazy::new(|| HeaderName::from_static("content-length"));
static CACHE_CONTROL: Lazy<HeaderName> = Lazy::new(|| HeaderName::from_static("cache-control"));
static ACCESS_CONTROL_ALLOW_ORIGIN: Lazy<HeaderName> =
Lazy::new(|| HeaderName::from_static("access-control-allow-origin"));
static ACCESS_CONTROL_ALLOW_METHODS: Lazy<HeaderName> =
Lazy::new(|| HeaderName::from_static("access-control-allow-methods"));
static ACCESS_CONTROL_ALLOW_HEADERS: Lazy<HeaderName> =
Lazy::new(|| HeaderName::from_static("access-control-allow-headers"));
#[derive(Debug, Clone)]
pub struct ResponseBuilder {
status: StatusCode,
headers: HeaderMap,
body: Option<ResponseBody>,
}
#[derive(Debug, Clone)]
enum ResponseBody {
Static(&'static [u8]),
Dynamic(Bytes),
Cow(Cow<'static, str>),
}
impl ResponseBuilder {
#[inline]
pub fn new() -> Self {
Self {
status: StatusCode::OK,
headers: HeaderMap::with_capacity(8), body: None,
}
}
#[inline]
pub fn with_status(status: StatusCode) -> Self {
Self {
status,
headers: HeaderMap::with_capacity(8),
body: None,
}
}
#[inline]
pub fn status(mut self, status: StatusCode) -> Self {
self.status = status;
self
}
#[inline]
pub fn status_code(mut self, status_code: u16) -> Self {
if let Ok(status) = StatusCode::from_u16(status_code) {
self.status = status;
}
self
}
#[inline]
pub fn body_static(mut self, body: &'static [u8]) -> Self {
self.body = Some(ResponseBody::Static(body));
self
}
#[inline]
pub fn body_static_str(mut self, body: &'static str) -> Self {
self.body = Some(ResponseBody::Static(body.as_bytes()));
self
}
#[inline]
pub fn body_bytes(mut self, body: Bytes) -> Self {
self.body = Some(ResponseBody::Dynamic(body));
self
}
#[inline]
pub fn body_cow(mut self, body: Cow<'static, str>) -> Self {
self.body = Some(ResponseBody::Cow(body));
self
}
#[inline]
pub fn body<T: Into<Bytes>>(mut self, body: T) -> Self {
self.body = Some(ResponseBody::Dynamic(body.into()));
self
}
pub fn header<K, V>(mut self, key: K, value: V) -> Self
where
K: TryInto<HeaderName>,
V: TryInto<HeaderValue>,
K::Error: std::fmt::Debug,
V::Error: std::fmt::Debug,
{
if let (Ok(name), Ok(val)) = (key.try_into(), value.try_into()) {
self.headers.insert(name, val);
}
self
}
#[inline]
pub fn content_type_json(mut self) -> Self {
self.headers
.insert(CONTENT_TYPE.clone(), COMMON_HEADERS["json"].clone());
self
}
#[inline]
pub fn content_type_text(mut self) -> Self {
self.headers
.insert(CONTENT_TYPE.clone(), COMMON_HEADERS["text"].clone());
self
}
#[inline]
pub fn content_type_html(mut self) -> Self {
self.headers
.insert(CONTENT_TYPE.clone(), COMMON_HEADERS["html"].clone());
self
}
#[inline]
pub fn json_cow<T: Into<Cow<'static, str>>>(mut self, text: T) -> Self {
self.headers
.insert(CONTENT_TYPE.clone(), COMMON_HEADERS["json"].clone());
self.body = Some(ResponseBody::Cow(text.into()));
self
}
#[inline]
pub fn text<T: Into<Cow<'static, str>>>(mut self, text: T) -> Self {
self.headers
.insert(CONTENT_TYPE.clone(), COMMON_HEADERS["text"].clone());
self.body = Some(ResponseBody::Cow(text.into()));
self
}
#[inline]
pub fn html<T: Into<Cow<'static, str>>>(mut self, html: T) -> Self {
self.headers
.insert(CONTENT_TYPE.clone(), COMMON_HEADERS["html"].clone());
self.body = Some(ResponseBody::Cow(html.into()));
self
}
pub fn json_static(mut self, json_key: &'static str) -> Self {
if let Some(body) = COMMON_RESPONSES.get(json_key) {
self.headers
.insert(CONTENT_TYPE.clone(), COMMON_HEADERS["json"].clone());
self.body = Some(ResponseBody::Dynamic(body.clone()));
} else {
self.headers
.insert(CONTENT_TYPE.clone(), COMMON_HEADERS["json"].clone());
self.body = Some(ResponseBody::Static(b"{}"));
}
self
}
pub fn json<T: Serialize>(mut self, data: &T) -> Result<Self> {
let mut buf = Vec::with_capacity(1024); serde_json::to_writer(&mut buf, data)?;
self.headers
.insert(CONTENT_TYPE.clone(), COMMON_HEADERS["json"].clone());
if let Ok(len_str) = buf.len().to_string().parse::<HeaderValue>() {
self.headers.insert(CONTENT_LENGTH.clone(), len_str);
}
self.body = Some(ResponseBody::Dynamic(Bytes::from(buf)));
Ok(self)
}
pub fn json_with_capacity<T: Serialize>(mut self, data: &T, capacity: usize) -> Result<Self> {
let mut buf = Vec::with_capacity(capacity);
serde_json::to_writer(&mut buf, data)?;
self.headers
.insert(CONTENT_TYPE.clone(), COMMON_HEADERS["json"].clone());
if let Ok(len_str) = buf.len().to_string().parse::<HeaderValue>() {
self.headers.insert(CONTENT_LENGTH.clone(), len_str);
}
self.body = Some(ResponseBody::Dynamic(Bytes::from(buf)));
Ok(self)
}
#[inline]
pub fn cache_1_hour(mut self) -> Self {
self.headers.insert(
CACHE_CONTROL.clone(),
HeaderValue::from_static("public, max-age=3600"),
);
self
}
#[inline]
pub fn cache_1_day(mut self) -> Self {
self.headers.insert(
CACHE_CONTROL.clone(),
HeaderValue::from_static("public, max-age=86400"),
);
self
}
#[inline]
pub fn cache_1_week(mut self) -> Self {
self.headers.insert(
CACHE_CONTROL.clone(),
HeaderValue::from_static("public, max-age=604800"),
);
self
}
#[inline]
pub fn cache_no_store(mut self) -> Self {
self.headers.insert(
CACHE_CONTROL.clone(),
HeaderValue::from_static("no-store, no-cache, must-revalidate"),
);
self
}
#[inline]
pub fn cors_any(mut self) -> Self {
self.headers.insert(
ACCESS_CONTROL_ALLOW_ORIGIN.clone(),
COMMON_HEADERS["cors_any"].clone(),
);
self.headers.insert(
ACCESS_CONTROL_ALLOW_METHODS.clone(),
COMMON_HEADERS["cors_methods"].clone(),
);
self.headers.insert(
ACCESS_CONTROL_ALLOW_HEADERS.clone(),
COMMON_HEADERS["cors_headers"].clone(),
);
self
}
#[inline]
pub fn cors_origin(mut self, origin: &str) -> Self {
if let Ok(origin_val) = HeaderValue::from_str(origin) {
self.headers
.insert(ACCESS_CONTROL_ALLOW_ORIGIN.clone(), origin_val);
}
self.headers.insert(
ACCESS_CONTROL_ALLOW_METHODS.clone(),
COMMON_HEADERS["cors_methods"].clone(),
);
self.headers.insert(
ACCESS_CONTROL_ALLOW_HEADERS.clone(),
COMMON_HEADERS["cors_headers"].clone(),
);
self
}
pub fn build(self) -> Response {
let body_bytes = match self.body {
Some(ResponseBody::Static(bytes)) => Bytes::from_static(bytes),
Some(ResponseBody::Dynamic(bytes)) => bytes,
Some(ResponseBody::Cow(cow)) => match cow {
Cow::Borrowed(s) => Bytes::from_static(s.as_bytes()),
Cow::Owned(s) => Bytes::from(s),
},
None => Bytes::new(),
};
Response {
status: self.status,
headers: self.headers,
body: body_bytes,
cache_control: None, }
}
#[inline]
pub fn health() -> Response {
ResponseBuilder::new().json_static("health_ok").build()
}
#[inline]
pub fn not_found() -> Response {
ResponseBuilder::with_status(StatusCode::NOT_FOUND)
.json_static("not_found")
.build()
}
#[inline]
pub fn server_error() -> Response {
ResponseBuilder::with_status(StatusCode::INTERNAL_SERVER_ERROR)
.json_static("server_error")
.build()
}
#[inline]
pub fn unauthorized() -> Response {
ResponseBuilder::with_status(StatusCode::UNAUTHORIZED)
.json_static("unauthorized")
.build()
}
#[inline]
pub fn forbidden() -> Response {
ResponseBuilder::with_status(StatusCode::FORBIDDEN)
.json_static("forbidden")
.build()
}
#[inline]
pub fn bad_request(message: &'static str) -> Response {
ResponseBuilder::with_status(StatusCode::BAD_REQUEST)
.json_static(message)
.build()
}
#[inline]
pub fn method_not_allowed() -> Response {
ResponseBuilder::with_status(StatusCode::METHOD_NOT_ALLOWED)
.json_static("method_not_allowed")
.build()
}
#[inline]
pub fn ok() -> Response {
ResponseBuilder::new().json_static("ok_message").build()
}
#[inline]
pub fn success() -> Response {
ResponseBuilder::new().json_static("success").build()
}
#[inline]
pub fn pong() -> Response {
ResponseBuilder::new().json_static("pong").build()
}
#[inline]
pub fn empty_json() -> Response {
ResponseBuilder::new().json_static("empty_json").build()
}
#[inline]
pub fn empty_array() -> Response {
ResponseBuilder::new().json_static("empty_array").build()
}
}
impl Default for ResponseBuilder {
fn default() -> Self {
Self::new()
}
}
impl Response {
pub fn static_json(key: &'static str) -> Self {
if let Some(body) = COMMON_RESPONSES.get(key) {
Self {
status: StatusCode::OK,
headers: {
let mut headers = HeaderMap::with_capacity(1);
headers.insert(CONTENT_TYPE.clone(), COMMON_HEADERS["json"].clone());
headers
},
body: body.clone(),
cache_control: None,
}
} else {
ResponseBuilder::empty_json()
}
}
pub fn json_static(json_str: &'static str) -> Self {
Self {
status: StatusCode::OK,
headers: {
let mut headers = HeaderMap::with_capacity(1);
headers.insert(CONTENT_TYPE.clone(), COMMON_HEADERS["json"].clone());
headers
},
body: Bytes::from_static(json_str.as_bytes()),
cache_control: None,
}
}
pub fn text_static(text: &'static str) -> Self {
Self {
status: StatusCode::OK,
headers: {
let mut headers = HeaderMap::with_capacity(1);
headers.insert(CONTENT_TYPE.clone(), COMMON_HEADERS["text"].clone());
headers
},
body: Bytes::from_static(text.as_bytes()),
cache_control: None,
}
}
pub fn html_static(html: &'static str) -> Self {
Self {
status: StatusCode::OK,
headers: {
let mut headers = HeaderMap::with_capacity(1);
headers.insert(CONTENT_TYPE.clone(), COMMON_HEADERS["html"].clone());
headers
},
body: Bytes::from_static(html.as_bytes()),
cache_control: None,
}
}
pub fn clone_body(&self) -> Bytes {
self.body.clone()
}
pub fn empty_json() -> Self {
ResponseBuilder::empty_json()
}
pub fn health_check() -> Self {
ResponseBuilder::health()
}
pub fn server_error() -> Self {
ResponseBuilder::server_error()
}
pub fn unauthorized() -> Self {
ResponseBuilder::unauthorized()
}
pub fn forbidden() -> Self {
ResponseBuilder::forbidden()
}
pub fn bad_request(message: &'static str) -> Self {
ResponseBuilder::bad_request(message)
}
pub fn method_not_allowed() -> Self {
ResponseBuilder::method_not_allowed()
}
pub fn success() -> Self {
ResponseBuilder::success()
}
pub fn pong() -> Self {
ResponseBuilder::pong()
}
pub fn empty_array() -> Self {
ResponseBuilder::empty_array()
}
pub fn has_cache_control(&self) -> bool {
self.headers.contains_key(CACHE_CONTROL.as_str())
}
pub fn get_cache_control(&self) -> Option<&HeaderValue> {
self.headers.get(CACHE_CONTROL.as_str())
}
pub fn with_cors_any(mut self) -> Self {
self.headers.insert(
ACCESS_CONTROL_ALLOW_ORIGIN.clone(),
COMMON_HEADERS["cors_any"].clone(),
);
self.headers.insert(
ACCESS_CONTROL_ALLOW_METHODS.clone(),
COMMON_HEADERS["cors_methods"].clone(),
);
self.headers.insert(
ACCESS_CONTROL_ALLOW_HEADERS.clone(),
COMMON_HEADERS["cors_headers"].clone(),
);
self
}
pub fn with_header<K, V>(mut self, key: K, value: V) -> Self
where
K: TryInto<HeaderName>,
V: TryInto<HeaderValue>,
K::Error: std::fmt::Debug,
V::Error: std::fmt::Debug,
{
if let (Ok(name), Ok(val)) = (key.try_into(), value.try_into()) {
self.headers.insert(name, val);
}
self
}
}