use crate::{
Conn, Error, HttpContext, Upgrade,
conn::{ConnInner, H1Status, HeadError},
graceful_close::GracefulClose,
};
use futures_lite::io::{AsyncRead, AsyncWrite};
use std::{
fmt::{self, Debug, Formatter},
future::Future,
pin::Pin,
sync::Arc,
task::{Context, Poll},
};
pub struct H1Connection<Transport> {
context: Arc<HttpContext>,
transport: Transport,
initial_bytes: Option<Vec<u8>>,
}
impl<Transport> Debug for H1Connection<Transport> {
fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
f.debug_struct("H1Connection")
.field("initial_bytes", &self.initial_bytes.as_ref().map(Vec::len))
.finish_non_exhaustive()
}
}
impl<Transport> H1Connection<Transport> {
pub fn new(context: Arc<HttpContext>, transport: Transport) -> Self {
Self {
context,
transport,
initial_bytes: None,
}
}
#[must_use]
pub fn with_initial_bytes(mut self, initial_bytes: Vec<u8>) -> Self {
self.initial_bytes = Some(initial_bytes);
self
}
}
impl<Transport> H1Connection<Transport>
where
Transport: AsyncRead + AsyncWrite + Unpin + Send + Sync + 'static,
{
pub async fn run<Handler, Fut>(self, mut handler: Handler) -> H1ConnectionResult<Transport>
where
Handler: FnMut(Conn<Transport>) -> Fut,
Fut: Future<Output = Conn<Transport>>,
{
let Self {
context,
transport,
initial_bytes,
} = self;
let initial_bytes = initial_bytes
.unwrap_or_else(|| Vec::with_capacity(context.config.request_buffer_initial_len));
let _guard = context.swansong.guard();
let mut status = match Conn::new_h1(context, transport, initial_bytes).await {
Ok(conn) => Ok(H1Status::Conn(conn)),
Err(HeadError::BadRequest(bad)) => bad.send().await,
Err(HeadError::Fatal(error)) => Err(error),
};
loop {
status = match status {
Ok(H1Status::Conn(conn)) => handler(conn).await.send().await,
Ok(H1Status::Upgrade(upgrade)) => return H1ConnectionResult::Upgrade(upgrade),
Ok(H1Status::Close(conn)) => return H1ConnectionResult::Close(H1Close::new(conn)),
Ok(H1Status::Ended(conn)) => {
return H1ConnectionResult::Close(H1Close::finished(conn));
}
Err(error) => return H1ConnectionResult::Error(error),
};
}
}
}
#[derive(Debug)]
#[must_use]
#[non_exhaustive]
pub enum H1ConnectionResult<Transport> {
Upgrade(Upgrade<Transport>),
Close(H1Close<Transport>),
Error(Error),
}
impl<Transport> H1ConnectionResult<Transport> {
pub(crate) fn into_upgrade(self) -> crate::Result<Option<Upgrade<Transport>>> {
match self {
Self::Upgrade(upgrade) => Ok(Some(upgrade)),
Self::Close(_) => Ok(None),
Self::Error(error) => Err(error),
}
}
}
pub struct H1Close<Transport>(GracefulClose<Transport>);
impl<Transport> Debug for H1Close<Transport> {
fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
f.debug_tuple("H1Close").finish_non_exhaustive()
}
}
impl<Transport> H1Close<Transport> {
fn new(conn: Conn<Transport>) -> Self {
let Conn { inner, transport } = conn;
let ConnInner {
buffer, context, ..
} = *inner;
Self(GracefulClose::new(
transport,
buffer.into(),
context.config.copy_loops_per_yield,
))
}
fn finished(conn: Conn<Transport>) -> Self {
Self(GracefulClose::finished(conn.transport))
}
}
impl<Transport> Future for H1Close<Transport>
where
Transport: AsyncRead + AsyncWrite + Unpin,
{
type Output = ();
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
Pin::new(&mut self.get_mut().0).poll(cx)
}
}