use std::{
io::{self, Read, Write},
pin::Pin,
};
use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
use bytes::{Buf, Bytes};
enum ResponseBody {
Async(Pin<Box<dyn AsyncRead + Send + Sync>>),
Sync(Box<dyn Read + Send + Sync>),
}
pub struct Response {
pub code: u32,
pub meta: String,
body: Option<ResponseBody>,
}
impl Response {
pub fn new(code: u32, meta: impl Into<String>) -> Self {
Self {
code,
meta: meta.into(),
body: None,
}
}
pub fn body(self, body: impl Into<Bytes>) -> Self {
self.body_sync(Bytes::from(body.into()).reader())
}
pub fn body_sync<R>(mut self, body: R) -> Self
where
R: Read + Send + Sync + 'static,
{
self.body = Some(ResponseBody::Sync(Box::new(body)));
self
}
pub fn body_async<R>(mut self, body: R) -> Self
where
R: AsyncRead + Send + Sync + 'static,
{
self.body = Some(ResponseBody::Async(Box::pin(body)));
self
}
#[inline]
pub fn input(request: impl Into<String>) -> Self {
Self::new(10, request)
}
#[inline]
pub fn input_sensitive(request: impl Into<String>) -> Self {
Self::new(11, request)
}
#[inline]
pub fn success(mime: impl Into<String>, body: impl Into<Bytes>) -> Self {
Self::new(20, mime).body(body)
}
#[inline]
pub fn success_sync<M, R>(mime: M, body: R) -> Self
where
M: Into<String>,
R: Read + Send + Sync + 'static,
{
Self::new(20, mime).body_sync(body)
}
#[inline]
pub fn success_async<M, R>(mime: M, body: R) -> Self
where
M: Into<String>,
R: AsyncReadExt + Send + Sync + 'static,
{
Self::new(20, mime).body_async(body)
}
#[inline]
pub fn redirect(redirect: impl Into<String>) -> Self {
Self::new(30, redirect)
}
#[inline]
pub fn redirect_perm(redirect: impl Into<String>) -> Self {
Self::new(31, redirect)
}
#[inline]
pub fn error_temp(message: impl Into<String>) -> Self {
Self::new(40, message)
}
#[inline]
pub fn unavailable(message: impl Into<String>) -> Self {
Self::new(41, message)
}
#[inline]
pub fn error_cgi(message: impl Into<String>) -> Self {
Self::new(42, message)
}
#[inline]
pub fn error_proxy(message: impl Into<String>) -> Self {
Self::new(43, message)
}
#[inline]
pub fn slow_down(seconds: u32) -> Self {
Self::new(44, seconds.to_string())
}
#[inline]
pub fn error_perm(message: impl Into<String>) -> Self {
Self::new(50, message)
}
#[inline]
pub fn not_found(message: impl Into<String>) -> Self {
Self::new(51, message)
}
#[inline]
pub fn gone(message: impl Into<String>) -> Self {
Self::new(52, message)
}
#[inline]
pub fn proxy_refused(message: impl Into<String>) -> Self {
Self::new(53, message)
}
#[inline]
pub fn bad_request(message: impl Into<String>) -> Self {
Self::new(59, message)
}
#[inline]
pub fn cert_required(message: impl Into<String>) -> Self {
Self::new(60, message)
}
#[inline]
pub fn cert_not_authorised(message: impl Into<String>) -> Self {
Self::new(61, message)
}
#[inline]
pub fn cert_not_valid(message: impl Into<String>) -> Self {
Self::new(62, message)
}
pub fn header(&self) -> String {
let meta = self.meta.lines().next().unwrap();
format!("{} {}\r\n", self.code, meta)
}
pub async fn send_async<W>(self, writer: &mut W) -> Result<(), io::Error>
where
W: AsyncWrite + Unpin + ?Sized,
{
let header = self.header();
writer.write_all(header.as_bytes()).await?;
match self.body {
Some(ResponseBody::Async(mut reader)) => {
tokio::io::copy(&mut reader, writer).await?;
}
Some(ResponseBody::Sync(mut reader)) => {
let mut buf = [0; 1024];
loop {
let read = reader.read(&mut buf)?;
if read == 0 {
break;
}
writer.write_all(&buf[..read]).await?;
}
}
None => {}
}
Ok(())
}
pub async fn send_sync<W>(self, writer: &mut W) -> Result<(), io::Error>
where
W: Write + ?Sized,
{
let header = self.header();
writer.write_all(header.as_bytes())?;
match self.body {
Some(ResponseBody::Async(mut reader)) => {
let mut buf = [0; 1024];
loop {
let read = reader.read(&mut buf).await?;
if read == 0 {
break;
}
writer.write_all(&buf[..read])?;
}
}
Some(ResponseBody::Sync(mut reader)) => {
io::copy(&mut reader, writer)?;
}
None => {}
};
Ok(())
}
}