feat: Add proper Response type

This commit is contained in:
Jill 2022-05-23 03:20:06 +02:00
parent 921ea05f89
commit 667b487427
Signed by untrusted user: KokaKiwi
GPG key ID: 09A5A2688F13FAC1
9 changed files with 236 additions and 51 deletions

View file

@ -1,16 +1,18 @@
use miette::{IntoDiagnostic, Result}; use miette::{IntoDiagnostic, Result};
use boitalettres::proto::{Request, Response}; use boitalettres::proto::{Request, Response};
use boitalettres::server::accept::addr::{AddrIncoming, AddrStream}; use boitalettres::server::accept::addr::{AddrIncoming, AddrStream};
use boitalettres::server::Server; use boitalettres::server::Server;
async fn handle_req(req: Request) -> Result<Response> { async fn handle_req(req: Request) -> Result<Response> {
use imap_codec::types::response::Status; use imap_codec::types::response::{Capability, Data};
tracing::debug!("Got request: {:#?}", req); tracing::debug!("Got request: {:#?}", req);
Ok(Response::Status( let capabilities = vec![Capability::Imap4Rev1, Capability::Idle];
Status::ok(Some(req.tag), None, "Ok").map_err(|e| eyre::eyre!(e))?, let body = vec![Data::Capability(capabilities)];
))
Ok(Response::ok("Done")?.with_body(body))
} }
#[tokio::main] #[tokio::main]
@ -24,9 +26,7 @@ async fn main() -> Result<()> {
let make_service = tower::service_fn(|addr: &AddrStream| { let make_service = tower::service_fn(|addr: &AddrStream| {
tracing::debug!(remote_addr = %addr.remote_addr, local_addr = %addr.local_addr, "accept"); tracing::debug!(remote_addr = %addr.remote_addr, local_addr = %addr.local_addr, "accept");
let service = tower::ServiceBuilder::new() let service = tower::ServiceBuilder::new().service_fn(handle_req);
.buffer(16)
.service_fn(handle_req);
futures::future::ok::<_, std::convert::Infallible>(service) futures::future::ok::<_, std::convert::Infallible>(service)
}); });

33
src/errors.rs Normal file
View file

@ -0,0 +1,33 @@
use miette::Diagnostic;
type BoxedError = tower::BoxError;
#[derive(Debug, thiserror::Error, Diagnostic)]
pub enum Error {
#[error("Error occured when accepting new connections")]
#[diagnostic(code(boitalettres::accept))]
Accept(#[source] BoxedError),
#[error("Error occured on service creation")]
#[diagnostic(code(boitalettres::make_service))]
MakeService(#[source] BoxedError),
#[error("{0}")]
Text(String),
}
impl Error {
pub(crate) fn accept<E: Into<BoxedError>>(err: E) -> Error {
Error::Accept(err.into())
}
pub(crate) fn make_service<E: Into<BoxedError>>(err: E) -> Error {
Error::MakeService(err.into())
}
pub(crate) fn text<E: Into<String>>(err: E) -> Error {
Error::Text(err.into())
}
}
pub type Result<T, E = Error> = std::result::Result<T, E>;

View file

@ -1,2 +1,3 @@
pub mod errors;
pub mod proto; pub mod proto;
pub mod server; pub mod server;

32
src/proto/body.rs Normal file
View file

@ -0,0 +1,32 @@
use imap_codec::types::response::Data as ImapData;
#[derive(Debug)]
pub enum Body {
Once(Vec<ImapData>),
}
impl Body {
pub(crate) fn into_data(self) -> Vec<ImapData> {
match self {
Body::Once(data) => data,
}
}
}
impl FromIterator<ImapData> for Body {
fn from_iter<T: IntoIterator<Item = ImapData>>(iter: T) -> Self {
Body::Once(Vec::from_iter(iter))
}
}
impl From<Vec<ImapData>> for Body {
fn from(data: Vec<ImapData>) -> Self {
Body::from_iter(data)
}
}
impl From<ImapData> for Body {
fn from(data: ImapData) -> Self {
Body::from_iter([data])
}
}

View file

@ -1,5 +1,6 @@
pub use self::req::Request; pub use self::req::Request;
pub use self::res::Response; pub use self::res::Response;
pub mod body;
pub mod req; pub mod req;
pub mod res; pub mod res;

View file

@ -1 +1,104 @@
pub type Response = imap_codec::types::response::Response; use imap_codec::types::{
core::{Tag, Text},
response::{Code as ImapCode, Status as ImapStatus},
};
use super::body::Body;
use crate::errors::{Error, Result};
#[derive(Debug)]
pub struct Response {
pub(crate) status: Status,
pub(crate) body: Option<Body>,
}
impl Response {
pub fn status(code: StatusCode, msg: &str) -> Result<Response> {
Ok(Response {
status: Status::new(code, msg)?,
body: None,
})
}
pub fn ok(msg: &str) -> Result<Response> {
Self::status(StatusCode::Ok, msg)
}
pub fn no(msg: &str) -> Result<Response> {
Self::status(StatusCode::No, msg)
}
pub fn bad(msg: &str) -> Result<Response> {
Self::status(StatusCode::Bad, msg)
}
pub fn bye(msg: &str) -> Result<Response> {
Self::status(StatusCode::Bye, msg)
}
}
impl Response {
pub fn with_extra_code(mut self, extra: ImapCode) -> Self {
self.status.extra = Some(extra);
self
}
pub fn with_body(mut self, body: impl Into<Body>) -> Self {
self.body = Some(body.into());
self
}
}
#[derive(Debug, Clone)]
pub struct Status {
pub(crate) code: StatusCode,
pub(crate) extra: Option<ImapCode>,
pub(crate) text: Text,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum StatusCode {
Ok,
No,
Bad,
PreAuth,
Bye,
}
impl Status {
fn new(code: StatusCode, msg: &str) -> Result<Status> {
Ok(Status {
code,
extra: None,
text: msg.try_into().map_err(Error::text)?,
})
}
pub(crate) fn into_imap(self, tag: Option<Tag>) -> ImapStatus {
match self.code {
StatusCode::Ok => ImapStatus::Ok {
tag,
code: self.extra,
text: self.text,
},
StatusCode::No => ImapStatus::No {
tag,
code: self.extra,
text: self.text,
},
StatusCode::Bad => ImapStatus::Bad {
tag,
code: self.extra,
text: self.text,
},
StatusCode::PreAuth => ImapStatus::PreAuth {
code: self.extra,
text: self.text,
},
StatusCode::Bye => ImapStatus::Bye {
code: self.extra,
text: self.text,
},
}
}
}

View file

@ -3,6 +3,7 @@ use std::pin::Pin;
use std::task::{self, Poll}; use std::task::{self, Poll};
use futures::io::{AsyncRead, AsyncWrite}; use futures::io::{AsyncRead, AsyncWrite};
use imap_codec::types::core::Tag;
use tower::Service; use tower::Service;
use super::pipeline::Connection; use super::pipeline::Connection;
@ -34,6 +35,7 @@ where
F: Future<Output = std::result::Result<S, ME>>, F: Future<Output = std::result::Result<S, ME>>,
ME: std::fmt::Display, ME: std::fmt::Display,
S: Service<Request, Response = Response>, S: Service<Request, Response = Response>,
S::Future: Send + 'static,
S::Error: std::fmt::Display, S::Error: std::fmt::Display,
{ {
type Output = (); type Output = ();
@ -56,27 +58,17 @@ where
// TODO: Properly handle server greeting // TODO: Properly handle server greeting
{ {
use imap_codec::types::response::{Response, Status}; use futures::SinkExt;
let status = match Status::ok(None, None, "Hello") { let greeting = Response::ok("Hello").unwrap(); // "Hello" is a valid
Ok(status) => status, // greeting
Err(err) => { this.conn.start_send_unpin((None, greeting)).unwrap();
tracing::error!("Connection error: {}", err);
return Poll::Ready(());
}
};
let res = Response::Status(status);
if let Err(err) = this.conn.send(res) {
tracing::error!("Connection error: {}", err);
return Poll::Ready(());
};
} }
ConnectingState::Ready { service } ConnectingState::Ready { service }
} }
ConnectingStateProj::Ready { service } => { ConnectingStateProj::Ready { service } => {
let server = PipelineServer::new(this.conn, service); let server = PipelineServer::new(this.conn, PipelineService { inner: service });
futures::pin_mut!(server); futures::pin_mut!(server);
return server.poll(cx).map(|res| { return server.poll(cx).map(|res| {
@ -91,3 +83,29 @@ where
} }
} }
} }
struct PipelineService<S> {
inner: S,
}
impl<S> Service<Request> for PipelineService<S>
where
S: Service<Request>,
S::Future: Send + 'static,
{
type Response = (Option<Tag>, S::Response);
type Error = S::Error;
type Future = futures::future::BoxFuture<'static, Result<Self::Response, Self::Error>>;
fn poll_ready(&mut self, cx: &mut task::Context<'_>) -> Poll<Result<(), Self::Error>> {
self.inner.poll_ready(cx)
}
fn call(&mut self, req: Request) -> Self::Future {
use futures::{FutureExt, TryFutureExt};
let tag = req.tag.clone();
self.inner.call(req).map_ok(|res| (Some(tag), res)).boxed()
}
}

View file

@ -5,6 +5,7 @@ use std::task::{self, Poll};
use futures::io::{AsyncRead, AsyncWrite}; use futures::io::{AsyncRead, AsyncWrite};
use imap_codec::types::response::Capability; use imap_codec::types::response::Capability;
use crate::errors::{Error, Result};
use crate::proto::{Request, Response}; use crate::proto::{Request, Response};
use accept::Accept; use accept::Accept;
use pipeline::Connection; use pipeline::Connection;
@ -15,14 +16,6 @@ mod conn;
mod pipeline; mod pipeline;
mod service; mod service;
#[derive(Debug, thiserror::Error, miette::Diagnostic)]
pub enum Error<A: std::error::Error + 'static> {
#[error("Error occured when accepting new connections")]
Accept(#[source] A),
#[error("Error occured on service creation")]
MakeService(#[source] tower::BoxError),
}
#[derive(Debug, Default, Clone)] #[derive(Debug, Default, Clone)]
pub struct Imap { pub struct Imap {
pub capabilities: Vec<Capability>, pub capabilities: Vec<Capability>,
@ -50,25 +43,25 @@ impl<I, S> Future for Server<I, S>
where where
I: Accept, I: Accept,
I::Conn: AsyncRead + AsyncWrite + Unpin + Send + 'static, I::Conn: AsyncRead + AsyncWrite + Unpin + Send + 'static,
I::Error: 'static, I::Error: Send + Sync + 'static,
S: MakeServiceRef<I::Conn, Request, Response = Response>, S: MakeServiceRef<I::Conn, Request, Response = Response>,
S::MakeError: Into<tower::BoxError> + std::fmt::Display, S::MakeError: Into<tower::BoxError> + std::fmt::Display,
S::Error: std::fmt::Display, S::Error: std::fmt::Display,
S::Future: Send + 'static, S::Future: Send + 'static,
S::Service: Send + 'static, S::Service: Send + 'static,
<S::Service as tower::Service<Request>>::Future: Send + 'static,
{ {
type Output = Result<(), Error<I::Error>>; type Output = Result<()>;
fn poll(mut self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<Self::Output> { fn poll(mut self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<Self::Output> {
loop { loop {
let this = self.as_mut().project(); let this = self.as_mut().project();
if let Some(conn) = futures::ready!(this.incoming.poll_accept(cx)) { if let Some(conn) = futures::ready!(this.incoming.poll_accept(cx)) {
let conn = conn.map_err(Error::Accept)?; let conn = conn.map_err(Error::accept)?;
futures::ready!(this.make_service.poll_ready_ref(cx)) futures::ready!(this.make_service.poll_ready_ref(cx))
.map_err(Into::into) .map_err(Error::make_service)?;
.map_err(Error::MakeService)?;
let service_fut = this.make_service.make_service_ref(&conn); let service_fut = this.make_service.make_service_ref(&conn);

View file

@ -5,6 +5,7 @@ use bytes::BytesMut;
use futures::io::{AsyncRead, AsyncWrite}; use futures::io::{AsyncRead, AsyncWrite};
use futures::sink::Sink; use futures::sink::Sink;
use futures::stream::Stream; use futures::stream::Stream;
use imap_codec::types::core::Tag;
use crate::proto::{Request, Response}; use crate::proto::{Request, Response};
@ -89,21 +90,9 @@ where
Poll::Ready(Ok(())) Poll::Ready(Ok(()))
} }
pub(crate) fn send(&mut self, item: Response) -> Result<()> {
use bytes::BufMut;
use imap_codec::codec::Encode;
let mut writer = BufMut::writer(&mut self.write_buf);
tracing::debug!(item = ?item, "transport.send");
item.encode(&mut writer)?;
Ok(())
}
} }
impl<C> Sink<Response> for Connection<C> impl<C> Sink<(Option<Tag>, Response)> for Connection<C>
where where
C: AsyncWrite + Unpin, C: AsyncWrite + Unpin,
{ {
@ -117,9 +106,24 @@ where
Poll::Ready(Ok(())) Poll::Ready(Ok(()))
} }
fn start_send(self: Pin<&mut Self>, item: Response) -> Result<(), Self::Error> { fn start_send(
self: Pin<&mut Self>,
(tag, res): (Option<Tag>, Response),
) -> Result<(), Self::Error> {
use bytes::BufMut;
use imap_codec::codec::Encode;
debug_assert!(self.write_buf.is_empty()); debug_assert!(self.write_buf.is_empty());
self.get_mut().send(item)?;
let write_buf = &mut self.get_mut().write_buf;
let mut writer = write_buf.writer();
let body = res.body.into_iter().flat_map(|body| body.into_data());
for data in body {
data.encode(&mut writer)?;
}
res.status.into_imap(tag).encode(&mut writer)?;
Ok(()) Ok(())
} }