use std::{ cell::{Ref, RefCell, RefMut}, net, rc::Rc, }; use bitflags::bitflags; use crate::{ header::{self, HeaderMap}, Extensions, Method, StatusCode, Uri, Version, }; /// Represents various types of connection #[derive(Copy, Clone, PartialEq, Debug)] pub enum ConnectionType { /// Close connection after response Close, /// Keep connection alive after response KeepAlive, /// Connection is upgraded to different type Upgrade, } bitflags! { pub(crate) struct Flags: u8 { const CLOSE = 0b0000_0001; const KEEP_ALIVE = 0b0000_0010; const UPGRADE = 0b0000_0100; const EXPECT = 0b0000_1000; const NO_CHUNKING = 0b0001_0000; const CAMEL_CASE = 0b0010_0000; } } #[doc(hidden)] pub trait Head: Default + 'static { fn clear(&mut self); fn with_pool(f: F) -> R where F: FnOnce(&MessagePool) -> R; } #[derive(Debug, Clone)] pub struct RequestHead { pub method: Method, pub uri: Uri, pub version: Version, pub headers: HeaderMap, pub peer_addr: Option, flags: Flags, } impl Default for RequestHead { fn default() -> RequestHead { RequestHead { method: Method::default(), uri: Uri::default(), version: Version::HTTP_11, headers: HeaderMap::with_capacity(16), peer_addr: None, flags: Flags::empty(), } } } impl Head for RequestHead { fn clear(&mut self) { self.flags = Flags::empty(); self.headers.clear(); } fn with_pool(f: F) -> R where F: FnOnce(&MessagePool) -> R, { REQUEST_POOL.with(|p| f(p)) } } impl RequestHead { /// Read the message headers. pub fn headers(&self) -> &HeaderMap { &self.headers } /// Mutable reference to the message headers. pub fn headers_mut(&mut self) -> &mut HeaderMap { &mut self.headers } /// Is to uppercase headers with Camel-Case. /// Default is `false` #[inline] pub fn camel_case_headers(&self) -> bool { self.flags.contains(Flags::CAMEL_CASE) } /// Set `true` to send headers which are formatted as Camel-Case. #[inline] pub fn set_camel_case_headers(&mut self, val: bool) { if val { self.flags.insert(Flags::CAMEL_CASE); } else { self.flags.remove(Flags::CAMEL_CASE); } } #[inline] /// Set connection type of the message pub fn set_connection_type(&mut self, ctype: ConnectionType) { match ctype { ConnectionType::Close => self.flags.insert(Flags::CLOSE), ConnectionType::KeepAlive => self.flags.insert(Flags::KEEP_ALIVE), ConnectionType::Upgrade => self.flags.insert(Flags::UPGRADE), } } #[inline] /// Connection type pub fn connection_type(&self) -> ConnectionType { if self.flags.contains(Flags::CLOSE) { ConnectionType::Close } else if self.flags.contains(Flags::KEEP_ALIVE) { ConnectionType::KeepAlive } else if self.flags.contains(Flags::UPGRADE) { ConnectionType::Upgrade } else if self.version < Version::HTTP_11 { ConnectionType::Close } else { ConnectionType::KeepAlive } } /// Connection upgrade status pub fn upgrade(&self) -> bool { self.headers() .get(header::CONNECTION) .map(|hdr| { if let Ok(s) = hdr.to_str() { s.to_ascii_lowercase().contains("upgrade") } else { false } }) .unwrap_or(false) } #[inline] /// Get response body chunking state pub fn chunked(&self) -> bool { !self.flags.contains(Flags::NO_CHUNKING) } #[inline] pub fn no_chunking(&mut self, val: bool) { if val { self.flags.insert(Flags::NO_CHUNKING); } else { self.flags.remove(Flags::NO_CHUNKING); } } #[inline] /// Request contains `EXPECT` header pub fn expect(&self) -> bool { self.flags.contains(Flags::EXPECT) } #[inline] pub(crate) fn set_expect(&mut self) { self.flags.insert(Flags::EXPECT); } } #[derive(Debug)] #[allow(clippy::large_enum_variant)] pub enum RequestHeadType { Owned(RequestHead), Rc(Rc, Option), } impl RequestHeadType { pub fn extra_headers(&self) -> Option<&HeaderMap> { match self { RequestHeadType::Owned(_) => None, RequestHeadType::Rc(_, headers) => headers.as_ref(), } } } impl AsRef for RequestHeadType { fn as_ref(&self) -> &RequestHead { match self { RequestHeadType::Owned(head) => head, RequestHeadType::Rc(head, _) => head.as_ref(), } } } impl From for RequestHeadType { fn from(head: RequestHead) -> Self { RequestHeadType::Owned(head) } } #[derive(Debug)] pub struct ResponseHead { pub version: Version, pub status: StatusCode, pub headers: HeaderMap, pub reason: Option<&'static str>, pub(crate) extensions: RefCell, flags: Flags, } impl ResponseHead { /// Create new instance of `ResponseHead` type #[inline] pub fn new(status: StatusCode) -> ResponseHead { ResponseHead { status, version: Version::default(), headers: HeaderMap::with_capacity(12), reason: None, flags: Flags::empty(), extensions: RefCell::new(Extensions::new()), } } /// Message extensions #[inline] pub fn extensions(&self) -> Ref<'_, Extensions> { self.extensions.borrow() } /// Mutable reference to a the message's extensions #[inline] pub fn extensions_mut(&self) -> RefMut<'_, Extensions> { self.extensions.borrow_mut() } #[inline] /// Read the message headers. pub fn headers(&self) -> &HeaderMap { &self.headers } #[inline] /// Mutable reference to the message headers. pub fn headers_mut(&mut self) -> &mut HeaderMap { &mut self.headers } #[inline] /// Set connection type of the message pub fn set_connection_type(&mut self, ctype: ConnectionType) { match ctype { ConnectionType::Close => self.flags.insert(Flags::CLOSE), ConnectionType::KeepAlive => self.flags.insert(Flags::KEEP_ALIVE), ConnectionType::Upgrade => self.flags.insert(Flags::UPGRADE), } } #[inline] pub fn connection_type(&self) -> ConnectionType { if self.flags.contains(Flags::CLOSE) { ConnectionType::Close } else if self.flags.contains(Flags::KEEP_ALIVE) { ConnectionType::KeepAlive } else if self.flags.contains(Flags::UPGRADE) { ConnectionType::Upgrade } else if self.version < Version::HTTP_11 { ConnectionType::Close } else { ConnectionType::KeepAlive } } /// Check if keep-alive is enabled #[inline] pub fn keep_alive(&self) -> bool { self.connection_type() == ConnectionType::KeepAlive } /// Check upgrade status of this message #[inline] pub fn upgrade(&self) -> bool { self.connection_type() == ConnectionType::Upgrade } /// Get custom reason for the response #[inline] pub fn reason(&self) -> &str { self.reason.unwrap_or_else(|| { self.status .canonical_reason() .unwrap_or("") }) } #[inline] pub(crate) fn conn_type(&self) -> Option { if self.flags.contains(Flags::CLOSE) { Some(ConnectionType::Close) } else if self.flags.contains(Flags::KEEP_ALIVE) { Some(ConnectionType::KeepAlive) } else if self.flags.contains(Flags::UPGRADE) { Some(ConnectionType::Upgrade) } else { None } } #[inline] /// Get response body chunking state pub fn chunked(&self) -> bool { !self.flags.contains(Flags::NO_CHUNKING) } #[inline] /// Set no chunking for payload pub fn no_chunking(&mut self, val: bool) { if val { self.flags.insert(Flags::NO_CHUNKING); } else { self.flags.remove(Flags::NO_CHUNKING); } } } pub struct Message { /// Rc here should not be cloned by anyone. /// It's used to reuse allocation of T and no shared ownership is allowed. head: Rc, } impl Message { /// Get new message from the pool of objects pub fn new() -> Self { T::with_pool(MessagePool::get_message) } } impl std::ops::Deref for Message { type Target = T; fn deref(&self) -> &Self::Target { self.head.as_ref() } } impl std::ops::DerefMut for Message { fn deref_mut(&mut self) -> &mut Self::Target { Rc::get_mut(&mut self.head).expect("Multiple copies exist") } } impl Drop for Message { fn drop(&mut self) { T::with_pool(|p| p.release(self.head.clone())) } } pub(crate) struct BoxedResponseHead { head: Option>, } impl BoxedResponseHead { /// Get new message from the pool of objects pub fn new(status: StatusCode) -> Self { RESPONSE_POOL.with(|p| p.get_message(status)) } } impl std::ops::Deref for BoxedResponseHead { type Target = ResponseHead; fn deref(&self) -> &Self::Target { self.head.as_ref().unwrap() } } impl std::ops::DerefMut for BoxedResponseHead { fn deref_mut(&mut self) -> &mut Self::Target { self.head.as_mut().unwrap() } } impl Drop for BoxedResponseHead { fn drop(&mut self) { if let Some(head) = self.head.take() { RESPONSE_POOL.with(move |p| p.release(head)) } } } #[doc(hidden)] /// Request's objects pool pub struct MessagePool(RefCell>>); #[doc(hidden)] #[allow(clippy::vec_box)] /// Request's objects pool pub struct BoxedResponsePool(RefCell>>); thread_local!(static REQUEST_POOL: MessagePool = MessagePool::::create()); thread_local!(static RESPONSE_POOL: BoxedResponsePool = BoxedResponsePool::create()); impl MessagePool { fn create() -> MessagePool { MessagePool(RefCell::new(Vec::with_capacity(128))) } /// Get message from the pool #[inline] fn get_message(&self) -> Message { if let Some(mut msg) = self.0.borrow_mut().pop() { // Message is put in pool only when it's the last copy. // which means it's guaranteed to be unique when popped out. Rc::get_mut(&mut msg) .expect("Multiple copies exist") .clear(); Message { head: msg } } else { Message { head: Rc::new(T::default()), } } } #[inline] /// Release request instance fn release(&self, msg: Rc) { let v = &mut self.0.borrow_mut(); if v.len() < 128 { v.push(msg); } } } impl BoxedResponsePool { fn create() -> BoxedResponsePool { BoxedResponsePool(RefCell::new(Vec::with_capacity(128))) } /// Get message from the pool #[inline] fn get_message(&self, status: StatusCode) -> BoxedResponseHead { if let Some(mut head) = self.0.borrow_mut().pop() { head.reason = None; head.status = status; head.headers.clear(); head.flags = Flags::empty(); BoxedResponseHead { head: Some(head) } } else { BoxedResponseHead { head: Some(Box::new(ResponseHead::new(status))), } } } #[inline] /// Release request instance fn release(&self, mut msg: Box) { let v = &mut self.0.borrow_mut(); if v.len() < 128 { msg.extensions.get_mut().clear(); v.push(msg); } } }