use crate::header::{Append, ByteArrayMap, SetCookie, SetCookieBuilder};
use ohkami_lib::map::TupleMap;
use std::borrow::Cow;
#[derive(Clone)]
pub struct Headers {
standard: ByteArrayMap<N_SERVER_HEADERS, Cow<'static, str>>,
custom: Option<Box<TupleMap<&'static str, Cow<'static, str>>>>,
#[allow(clippy::box_collection)]
pub(super) setcookie: Option<Box<Vec<Cow<'static, str>>>>,
pub(super) size: usize,
}
pub struct SetHeaders<'set>(&'set mut Headers);
impl Headers {
#[inline]
pub fn set(&mut self) -> SetHeaders<'_> {
SetHeaders(self)
}
}
pub trait HeaderAction<'action> {
fn perform(self, set: SetHeaders<'action>, key: Header) -> SetHeaders<'action>;
}
const _: () = {
impl<'a> HeaderAction<'a> for Append {
#[inline]
fn perform(self, set: SetHeaders<'a>, key: Header) -> SetHeaders<'a> {
set.0.append(key, self.0);
set
}
}
impl<'a> HeaderAction<'a> for &'static str {
#[inline(always)]
fn perform(self, set: SetHeaders<'a>, key: Header) -> SetHeaders<'a> {
set.0.insert(key, Cow::Borrowed(self));
set
}
}
impl<'a> HeaderAction<'a> for String {
#[inline(always)]
fn perform(self, set: SetHeaders<'a>, key: Header) -> SetHeaders<'a> {
set.0.insert(key, Cow::Owned(self));
set
}
}
impl<'a> HeaderAction<'a> for std::borrow::Cow<'static, str> {
fn perform(self, set: SetHeaders<'a>, key: Header) -> SetHeaders<'a> {
set.0.insert(key, self);
set
}
}
impl<'a> HeaderAction<'a> for Option<Cow<'static, str>> {
#[inline]
fn perform(self, set: SetHeaders<'a>, key: Header) -> SetHeaders<'a> {
match self {
None => set.0.remove(key),
Some(v) => set.0.insert(key, v),
}
set
}
}
};
pub trait CustomHeadersAction<'action> {
fn perform(self, set: SetHeaders<'action>, key: &'static str) -> SetHeaders<'action>;
}
const _: () = {
impl<'set> CustomHeadersAction<'set> for Append {
fn perform(self, set: SetHeaders<'set>, key: &'static str) -> SetHeaders<'set> {
set.0.append_custom(key, self.0);
set
}
}
impl<'set> CustomHeadersAction<'set> for &'static str {
#[inline(always)]
fn perform(self, set: SetHeaders<'set>, key: &'static str) -> SetHeaders<'set> {
set.0.insert_custom(key, Cow::Borrowed(self));
set
}
}
impl<'set> CustomHeadersAction<'set> for String {
#[inline(always)]
fn perform(self, set: SetHeaders<'set>, key: &'static str) -> SetHeaders<'set> {
set.0.insert_custom(key, Cow::Owned(self));
set
}
}
impl<'set> CustomHeadersAction<'set> for Cow<'static, str> {
fn perform(self, set: SetHeaders<'set>, key: &'static str) -> SetHeaders<'set> {
set.0.insert_custom(key, self);
set
}
}
impl<'set> CustomHeadersAction<'set> for Option<Cow<'static, str>> {
#[inline]
fn perform(self, set: SetHeaders<'set>, key: &'static str) -> SetHeaders<'set> {
match self {
None => set.0.remove_custom(key),
Some(v) => set.0.insert_custom(key, v),
}
set
}
}
};
macro_rules! Header {
($N:literal; $( $method:ident($konst:ident): $name_bytes:literal, )*) => {
pub(crate) const N_SERVER_HEADERS: usize = $N;
const _: [Header; N_SERVER_HEADERS] = [$(Header::$konst),*];
#[derive(Debug, PartialEq, Clone, Copy)]
pub enum Header {
$( $konst, )*
}
impl Header {
#[inline] pub const fn as_bytes(&self) -> &'static [u8] {
match self {
$(
Self::$konst => $name_bytes,
)*
}
}
pub const fn as_str(&self) -> &'static str {
unsafe {std::str::from_utf8_unchecked(self.as_bytes())}
}
#[inline(always)] const fn len(&self) -> usize {
match self {
$(
Self::$konst => $name_bytes.len(),
)*
}
}
pub fn from_bytes(bytes: &[u8]) -> Option<Self> {
(0..N_SERVER_HEADERS)
.map(|i| unsafe {std::mem::transmute::<_, Header>(i as u8)})
.find(|h| h.as_bytes().eq_ignore_ascii_case(bytes))
}
}
impl<T: AsRef<[u8]>> PartialEq<T> for Header {
fn eq(&self, other: &T) -> bool {
self.as_bytes().eq_ignore_ascii_case(other.as_ref())
}
}
#[allow(non_snake_case)]
impl<'set> SetHeaders<'set> {
$(
#[inline]
pub fn $method(self, action: impl HeaderAction<'set>) -> Self {
action.perform(self, Header::$konst)
}
#[deprecated = "Use snake_case method instead"]
pub fn $konst(self, action: impl HeaderAction<'set>) -> Self {
self.$method(action)
}
)*
#[deprecated = "use `.x` instead"]
pub fn custom(self, name: &'static str, action: impl CustomHeadersAction<'set>) -> Self {
self.x(name, action)
}
#[inline]
pub fn x(self, name: &'static str, action: impl CustomHeadersAction<'set>) -> Self {
action.perform(self, name)
}
}
#[allow(non_snake_case)]
impl Headers {
$(
#[inline]
pub fn $method(&self) -> Option<&str> {
self.get_standard(Header::$konst)
}
#[deprecated = "Use snake_case method instead"]
pub fn $konst(&self) -> Option<&str> {
self.$method()
}
)*
#[deprecated = "use `.get` instead"]
pub fn custom(&self, name: &'static str) -> Option<&str> {
self.get(name)
}
#[inline]
pub fn get(&self, name: &'static str) -> Option<&str> {
self.get_custom(name)
.or_else(|| self.get_standard(Header::from_bytes(name.as_bytes())?))
}
}
};
}
Header! {48;
accept_ranges(AcceptRanges): b"Accept-Ranges",
access_control_allow_credentials(AccessControlAllowCredentials): b"Access-Control-Allow-Credentials",
access_control_allow_headers(AccessControlAllowHeaders): b"Access-Control-Allow-Headers",
access_control_allow_methods(AccessControlAllowMethods): b"Access-Control-Allow-Methods",
access_control_allow_origin(AccessControlAllowOrigin): b"Access-Control-Allow-Origin",
access_control_expose_headers(AccessControlExposeHeaders): b"Access-Control-Expose-Headers",
access_control_max_age(AccessControlMaxAge): b"Access-Control-Max-Age",
age(Age): b"Age",
allow(Allow): b"Allow",
alt_svc(AltSvc): b"Alt-Svc",
cache_control(CacheControl): b"Cache-Control",
cache_status(CacheStatus): b"Cache-Status",
cdn_cache_control(CDNCacheControl): b"CDN-Cache-Control",
connection(Connection): b"Connection",
content_dispotision(ContentDisposition): b"Content-Disposition",
content_encoding(ContentEncoding): b"Content-Encoding",
content_language(ContentLanguage): b"Content-Language",
content_length(ContentLength): b"Content-Length",
content_location(ContentLocation): b"Content-Location",
content_range(ContentRange): b"Content-Range",
content_security_policy(ContentSecurityPolicy): b"Content-Security-Policy",
content_security_policy_report_only(ContentSecurityPolicyReportOnly): b"Content-Security-Policy-Report-Only",
content_type(ContentType): b"Content-Type",
cross_origin_embedder_policy(CrossOriginEmbedderPolicy): b"Cross-Origin-Embedder-Policy",
cross_origin_resource_policy(CrossOriginResourcePolicy): b"Cross-Origin-Resource-Policy",
date(Date): b"Date",
etag(ETag): b"ETag",
expires(Expires): b"Expires",
link(Link): b"Link",
location(Location): b"Location",
last_modified(LastModified): b"Last-Modified",
proxy_authenticate(ProxyAuthenticate): b"Proxy-Authenticate",
referrer_policy(ReferrerPolicy): b"Referrer-Policy",
refresh(Refresh): b"Refresh",
retry_after(RetryAfter): b"Retry-After",
sec_websocket_accept(SecWebSocketAccept): b"Sec-WebSocket-Accept",
sec_websocket_protocol(SecWebSocketProtocol): b"Sec-WebSocket-Protocol",
sec_websocket_version(SecWebSocketVersion): b"Sec-WebSocket-Version",
server(Server): b"Server",
strict_transport_security(StrictTransportSecurity): b"Strict-Transport-Security",
trailer(Trailer): b"Trailer",
transfer_encoding(TransferEncoding): b"Transfer-Encoding",
upgrade(Upgrade): b"Upgrade",
vary(Vary): b"Vary",
via(Via): b"Via",
x_content_type_options(XContentTypeOptions): b"X-Content-Type-Options",
x_frame_options(XFrameOptions): b"X-Frame-Options",
www_authenticate(WWWAuthenticate): b"WWW-Authenticate",
}
const _: () = {
#[allow(non_snake_case)]
impl Headers {
pub fn set_cookie(&self) -> impl Iterator<Item = SetCookie<'_>> {
self.setcookie
.as_ref()
.map(|setcookies| {
setcookies
.iter()
.filter_map(|raw| match SetCookie::from_raw(raw) {
Ok(valid) => Some(valid),
Err(_err) => {
#[cfg(debug_assertions)]
crate::WARNING!("Invalid `Set-Cookie`: {_err}");
None
}
})
})
.into_iter()
.flatten()
}
}
#[allow(non_snake_case)]
impl<'s> SetHeaders<'s> {
#[inline]
pub fn set_cookie(
self,
name: &'static str,
value: impl Into<Cow<'static, str>>,
directives: impl FnOnce(SetCookieBuilder) -> SetCookieBuilder,
) -> Self {
let setcookie: Cow<'static, str> = directives(SetCookieBuilder::new(name, value))
.build()
.into();
self.0.size += "Set-Cookie: ".len() + setcookie.len() + "\r\n".len();
match self.0.setcookie.as_mut() {
None => self.0.setcookie = Some(Box::new(vec![setcookie])),
Some(setcookies) => setcookies.push(setcookie),
}
self
}
}
};
impl Headers {
#[inline(always)]
pub(crate) fn insert(&mut self, name: Header, value: Cow<'static, str>) {
let (name_len, value_len) = (name.len(), value.len());
match unsafe { self.standard.get_mut(name as u8) } {
None => {
unsafe { self.standard.insert_new(name as u8, value) };
self.size += name_len + ": ".len() + value_len + "\r\n".len();
}
Some(old) => {
self.size -= old.len();
self.size += value_len;
*old = value
}
}
}
#[inline]
pub(crate) fn insert_custom(&mut self, name: &'static str, value: Cow<'static, str>) {
let self_len = value.len();
match &mut self.custom {
None => {
self.custom = Some(Box::new(TupleMap::from_iter([(name, value)])));
self.size += name.len() + ": ".len() + self_len + "\r\n".len()
}
Some(custom) => {
if let Some(old) = custom.insert(name, value) {
self.size -= old.len();
self.size += self_len;
} else {
self.size += name.len() + ": ".len() + self_len + "\r\n".len()
}
}
}
}
#[inline]
pub(crate) fn remove(&mut self, name: Header) {
let name_len = name.len();
if let Some(v) = unsafe { self.standard.get(name as u8) } {
self.size -= name_len + ": ".len() + v.len() + "\r\n".len();
unsafe { self.standard.delete(name as u8) };
}
}
pub(crate) fn remove_custom(&mut self, name: &'static str) {
if let Some(c) = self.custom.as_mut()
&& let Some(v) = c.remove(name)
{
self.size -= name.len() + ": ".len() + v.len() + "\r\n".len();
}
}
#[inline(always)]
pub(crate) fn get_standard(&self, name: Header) -> Option<&str> {
unsafe { self.standard.get(name as u8) }.map(Cow::as_ref)
}
#[inline]
pub(crate) fn get_custom(&self, name: &'static str) -> Option<&str> {
self.custom.as_ref()?.get(&name).map(Cow::as_ref)
}
pub(crate) fn append(&mut self, name: Header, value: Cow<'static, str>) {
let value_len = value.len();
let target = unsafe { self.standard.get_mut(name as u8) };
self.size += match target {
Some(v) => {
match v {
Cow::Borrowed(slice) => {
let mut appended = String::with_capacity(slice.len() + 2 + value_len);
appended.push_str(slice);
appended.push_str(", ");
appended.push_str(&value);
*v = Cow::Owned(appended);
}
Cow::Owned(string) => {
string.push_str(", ");
string.push_str(&value);
}
}
", ".len() + value_len
}
None => {
unsafe { self.standard.insert_new(name as u8, value) }
name.len() + ": ".len() + value_len + "\r\n".len()
}
};
}
pub(crate) fn append_custom(&mut self, name: &'static str, value: Cow<'static, str>) {
let value_len = value.len();
let custom = {
if self.custom.is_none() {
self.custom = Some(Box::new(TupleMap::new()));
}
unsafe { self.custom.as_mut().unwrap_unchecked() }
};
self.size += match custom.get_mut(&name) {
Some(v) => {
match v {
Cow::Owned(string) => {
string.push_str(", ");
string.push_str(&value);
}
Cow::Borrowed(s) => {
let mut s = s.to_string();
s.push_str(", ");
s.push_str(&value);
*v = Cow::Owned(s);
}
}
", ".len() + value_len
}
None => {
custom.insert(name, value);
name.len() + ": ".len() + value_len + "\r\n".len()
}
};
}
}
impl Headers {
#[inline]
pub(crate) fn new() -> Self {
let mut this = Self {
standard: ByteArrayMap::new(),
custom: None,
setcookie: None,
size: "\r\n".len(),
};
this.set()
.date(ohkami_lib::imf_fixdate(crate::util::unix_timestamp()))
.content_length("0");
this
}
#[cfg(feature = "DEBUG")]
#[doc(hidden)]
pub fn _new() -> Self {
Self::new()
}
pub(crate) fn iter_standard(&self) -> impl Iterator<Item = (&str, &str)> {
self.standard.iter().map(|(i, v)| {
(
unsafe { std::mem::transmute::<u8, Header>(*i) }.as_str(),
&**v,
)
})
}
pub fn iter(&self) -> impl Iterator<Item = (&str, &str)> {
self.iter_standard()
.chain(
self.custom
.as_ref()
.into_iter()
.flat_map(|hm| hm.iter().map(|(k, v)| (*k, &**v))),
)
.chain(
self.setcookie
.as_ref()
.map(|sc| sc.iter().map(Cow::as_ref))
.into_iter()
.flatten()
.map(|sc| ("Set-Cookie", sc)),
)
}
#[cfg(any(feature = "__rt_native__", feature = "DEBUG"))]
pub(crate) unsafe fn write_unchecked_to(&self, buf: &mut Vec<u8>) {
unsafe {
for (i, v) in self.standard.iter() {
let h = std::mem::transmute::<u8, Header>(*i);
{
crate::push_unchecked!(buf <- h.as_bytes());
crate::push_unchecked!(buf <- b": ");
crate::push_unchecked!(buf <- v.as_bytes());
crate::push_unchecked!(buf <- b"\r\n");
}
}
if let Some(custom) = self.custom.as_ref() {
for (k, v) in custom.iter() {
crate::push_unchecked!(buf <- k.as_bytes());
crate::push_unchecked!(buf <- b": ");
crate::push_unchecked!(buf <- v.as_bytes());
crate::push_unchecked!(buf <- b"\r\n");
}
}
if let Some(setcookies) = self.setcookie.as_ref() {
for setcookie in &**setcookies {
crate::push_unchecked!(buf <- b"Set-Cookie: ");
crate::push_unchecked!(buf <- setcookie.as_bytes());
crate::push_unchecked!(buf <- b"\r\n");
}
}
crate::push_unchecked!(buf <- b"\r\n");
}
}
#[cfg(any(feature = "DEBUG", feature = "__rt_native__"))]
pub fn _write_to(&self, buf: &mut Vec<u8>) {
buf.reserve(self.size);
unsafe { self.write_unchecked_to(buf) }
}
}
const _: () = {
impl std::fmt::Debug for Headers {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_map().entries(self.iter()).finish()
}
}
impl PartialEq for Headers {
fn eq(&self, other: &Self) -> bool {
for (k, v) in self.iter_standard() {
if other.get_standard(Header::from_bytes(k.as_bytes()).unwrap()) != Some(v) {
return false;
}
}
if self.custom != other.custom {
return false;
}
true
}
}
impl<A: Into<Cow<'static, str>>> FromIterator<(&'static str, A)> for Headers {
fn from_iter<T: IntoIterator<Item = (&'static str, A)>>(iter: T) -> Self {
let mut this = Headers::new();
for (k, v) in iter {
match Header::from_bytes(k.as_bytes()) {
Some(h) => this.insert(h, v.into()),
None => {
this.set().x(k, v.into());
}
}
}
this
}
}
impl IntoIterator for Headers {
type Item = (&'static str, Cow<'static, str>);
type IntoIter = HeadersIntoIter;
fn into_iter(self) -> Self::IntoIter {
HeadersIntoIter {
standard: self.standard.into_iter(),
custom: self.custom.map(|tm| tm.into_iter()),
setcookie: self.setcookie.map(|sc| sc.into_iter()),
}
}
}
pub struct HeadersIntoIter {
standard: std::vec::IntoIter<(u8, Cow<'static, str>)>,
custom: Option<std::vec::IntoIter<(&'static str, Cow<'static, str>)>>,
setcookie: Option<std::vec::IntoIter<Cow<'static, str>>>,
}
impl Iterator for HeadersIntoIter {
type Item = (&'static str, Cow<'static, str>);
fn next(&mut self) -> Option<Self::Item> {
if let Some((i, v)) = self.standard.next() {
Some((unsafe { std::mem::transmute::<u8, Header>(i) }.as_str(), v))
} else if let Some(custom) = self.custom.as_mut().and_then(|c| c.next()) {
Some(custom)
} else if let Some(setcookie) = self.setcookie.as_mut().and_then(|c| c.next()) {
Some(("Set-Cookie", setcookie))
} else {
None
}
}
}
};
#[cfg(feature = "rt_worker")]
const _: () = {
impl From<Headers> for ::worker::Headers {
#[inline(always)]
fn from(this: Headers) -> ::worker::Headers {
let h = ::worker::Headers::new();
for (k, v) in this.iter() {
if let Err(_e) = h.append(k, v) {
crate::DEBUG!("`worker::Headers::append` failed: {_e:?}");
}
}
h
}
}
};