use std::{fmt, str};
use std::borrow::Cow;
use std::collections::HashMap;
use std::pin::Pin;
use tokio::io::{AsyncRead, AsyncSeek};
use crate::http::{Header, HeaderMap, Status, ContentType, Cookie};
use crate::http::uncased::{Uncased, AsUncased};
use crate::data::IoHandler;
use crate::response::Body;
pub struct Builder<'r> {
response: Response<'r>,
}
impl<'r> Builder<'r> {
#[inline(always)]
pub fn new(base: Response<'r>) -> Builder<'r> {
Builder {
response: base,
}
}
#[inline(always)]
pub fn status(&mut self, status: Status) -> &mut Builder<'r> {
self.response.set_status(status);
self
}
#[inline(always)]
pub fn header<'h: 'r, H>(&mut self, header: H) -> &mut Builder<'r>
where H: Into<Header<'h>>
{
self.response.set_header(header);
self
}
#[inline(always)]
pub fn header_adjoin<'h: 'r, H>(&mut self, header: H) -> &mut Builder<'r>
where H: Into<Header<'h>>
{
self.response.adjoin_header(header);
self
}
#[inline(always)]
pub fn raw_header<'a, 'b, N, V>(&mut self, name: N, value: V) -> &mut Builder<'r>
where N: Into<Cow<'a, str>>, V: Into<Cow<'b, str>>, 'a: 'r, 'b: 'r
{
self.response.set_raw_header(name, value);
self
}
#[inline(always)]
pub fn raw_header_adjoin<'a, 'b, N, V>(&mut self, name: N, value: V) -> &mut Builder<'r>
where N: Into<Cow<'a, str>>, V: Into<Cow<'b, str>>, 'a: 'r, 'b: 'r
{
self.response.adjoin_raw_header(name, value);
self
}
pub fn sized_body<B, S>(&mut self, size: S, body: B) -> &mut Builder<'r>
where B: AsyncRead + AsyncSeek + Send + 'r,
S: Into<Option<usize>>
{
self.response.set_sized_body(size, body);
self
}
#[inline(always)]
pub fn streamed_body<B>(&mut self, body: B) -> &mut Builder<'r>
where B: AsyncRead + Send + 'r
{
self.response.set_streamed_body(body);
self
}
#[inline(always)]
pub fn upgrade<P, H>(&mut self, protocol: P, handler: H) -> &mut Builder<'r>
where P: Into<Uncased<'r>>, H: IoHandler + 'r
{
self.response.add_upgrade(protocol.into(), handler);
self
}
#[inline(always)]
pub fn max_chunk_size(&mut self, size: usize) -> &mut Builder<'r> {
self.response.set_max_chunk_size(size);
self
}
#[inline(always)]
pub fn merge(&mut self, other: Response<'r>) -> &mut Builder<'r> {
self.response.merge(other);
self
}
#[inline(always)]
pub fn join(&mut self, other: Response<'r>) -> &mut Builder<'r> {
self.response.join(other);
self
}
pub fn finalize(&mut self) -> Response<'r> {
std::mem::replace(&mut self.response, Response::new())
}
#[inline(always)]
pub fn ok<E>(&mut self) -> Result<Response<'r>, E> {
Ok(self.finalize())
}
}
#[derive(Default)]
pub struct Response<'r> {
status: Option<Status>,
headers: HeaderMap<'r>,
body: Body<'r>,
upgrade: HashMap<Uncased<'r>, Pin<Box<dyn IoHandler + 'r>>>,
}
impl<'r> Response<'r> {
#[inline(always)]
pub fn new() -> Response<'r> {
Response::default()
}
#[inline(always)]
pub fn build() -> Builder<'r> {
Response::build_from(Response::new())
}
#[inline(always)]
pub fn build_from(other: Response<'r>) -> Builder<'r> {
Builder::new(other)
}
#[inline(always)]
pub fn status(&self) -> Status {
self.status.unwrap_or(Status::Ok)
}
#[inline(always)]
pub fn set_status(&mut self, status: Status) {
self.status = Some(status);
}
#[inline(always)]
pub fn content_type(&self) -> Option<ContentType> {
self.headers().get_one("Content-Type").and_then(|v| v.parse().ok())
}
pub fn cookies(&self) -> impl Iterator<Item = Cookie<'_>> {
self.headers()
.get("Set-Cookie")
.filter_map(|header| Cookie::parse_encoded(header).ok())
}
#[inline(always)]
pub fn headers(&self) -> &HeaderMap<'r> {
&self.headers
}
#[inline(always)]
pub fn set_header<'h: 'r, H: Into<Header<'h>>>(&mut self, header: H) -> bool {
self.headers.replace(header)
}
#[inline(always)]
pub fn set_raw_header<'a: 'r, 'b: 'r, N, V>(&mut self, name: N, value: V) -> bool
where N: Into<Cow<'a, str>>, V: Into<Cow<'b, str>>
{
self.set_header(Header::new(name, value))
}
#[inline(always)]
pub fn adjoin_header<'h: 'r, H: Into<Header<'h>>>(&mut self, header: H) {
self.headers.add(header)
}
#[inline(always)]
pub fn adjoin_raw_header<'a: 'r, 'b: 'r, N, V>(&mut self, name: N, value: V)
where N: Into<Cow<'a, str>>, V: Into<Cow<'b, str>>
{
self.adjoin_header(Header::new(name, value));
}
#[inline(always)]
pub fn remove_header(&mut self, name: &str) {
self.headers.remove(name);
}
#[inline(always)]
pub fn body(&self) -> &Body<'r> {
&self.body
}
pub(crate) fn take_upgrade<I: Iterator<Item = &'r str>>(
&mut self,
protocols: I
) -> Result<Option<(Uncased<'r>, Pin<Box<dyn IoHandler + 'r>>)>, ()> {
if self.upgrade.is_empty() {
return Ok(None);
}
let mut protocols = protocols.peekable();
let have_protocols = protocols.peek().is_some();
let found = protocols
.flat_map(|v| v.split(',').map(str::trim))
.find_map(|p| self.upgrade.remove_entry(p.as_uncased()));
match found {
Some(handler) => Ok(Some(handler)),
None if have_protocols => Err(()),
None => Ok(None)
}
}
pub fn upgrade(&mut self, proto: &str) -> Option<Pin<&mut (dyn IoHandler + 'r)>> {
self.upgrade.get_mut(proto.as_uncased()).map(|h| h.as_mut())
}
#[inline(always)]
pub fn body_mut(&mut self) -> &mut Body<'r> {
&mut self.body
}
#[inline(always)]
pub(crate) fn strip_body(&mut self) {
self.body.strip();
}
pub fn set_sized_body<B, S>(&mut self, size: S, body: B)
where B: AsyncRead + AsyncSeek + Send + 'r,
S: Into<Option<usize>>
{
self.body = Body::with_sized(body, size.into());
}
#[inline(always)]
pub fn set_streamed_body<B>(&mut self, body: B)
where B: AsyncRead + Send + 'r
{
self.body = Body::with_unsized(body);
}
pub fn add_upgrade<N, H>(&mut self, protocol: N, handler: H)
where N: Into<Uncased<'r>>, H: IoHandler + 'r
{
self.upgrade.insert(protocol.into(), Box::pin(handler));
}
#[inline(always)]
pub fn set_max_chunk_size(&mut self, size: usize) {
self.body_mut().set_max_chunk_size(size);
}
pub fn merge(&mut self, other: Response<'r>) {
if let Some(status) = other.status {
self.status = Some(status);
}
if other.body().is_some() {
self.body = other.body;
}
for (name, values) in other.headers.into_iter_raw() {
self.headers.replace_all(name.into_cow(), values);
}
}
pub fn join(&mut self, other: Response<'r>) {
if self.status.is_none() {
self.status = other.status;
}
if self.body.is_none() {
self.body = other.body;
}
for (name, mut values) in other.headers.into_iter_raw() {
self.headers.add_all(name.into_cow(), &mut values);
}
}
}
impl fmt::Debug for Response<'_> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
writeln!(f, "{}", self.status())?;
for header in self.headers().iter() {
writeln!(f, "{}", header)?;
}
self.body.fmt(f)
}
}