KokaKiwi/boitalettres
2
1
Fork 2
Code Issues 7 Pull requests 1 Projects 1 Releases Activity

proto: Add support for Stream-based Body

Browse source 
Fixes #2
This commit is contained in:
Jill 2022-06-27 17:57:49 +02:00
parent 01ee8c872b
commit f52d1d44cd
Signed by: KokaKiwi
GPG key ID: 09A5A2688F13FAC1
11 changed files with 185 additions and 31 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

5
Cargo.toml
View file

@ -7,7 +7,7 @@ edition = "2021"
[dependencies] [dependencies]
bytes = "1.1" bytes = "1.1"
miette = "4.7" miette = "5.1"
thiserror = "1.0" thiserror = "1.0"
tracing = "0.1" tracing = "0.1"
@ -18,6 +18,7 @@ imap-codec = "0.5"
# Async # Async
async-compat = "0.2" async-compat = "0.2"
async-stream = "0.3"
futures = "0.3" futures = "0.3"
pin-project = "1.0" pin-project = "1.0"
@ -26,7 +27,7 @@ tokio-tower = "0.6"
tower = { version = "0.4", features = ["full"] } tower = { version = "0.4", features = ["full"] }
[dev-dependencies] [dev-dependencies]
miette = { version = "4.7", features = ["fancy"] } miette = { version = "5.1", features = ["fancy"] }
tracing-subscriber = "0.3" tracing-subscriber = "0.3"
console-subscriber = "0.1" console-subscriber = "0.1"

11
examples/simple.rs
View file

@ -4,13 +4,16 @@ 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::{Capability, Data}; use imap_codec::types::response::{Capability, Data as ImapData};
tracing::debug!("Got request: {:#?}", req); use boitalettres::proto::res::{body::Data, Status};
let capabilities = vec![Capability::Imap4Rev1, Capability::Idle]; let capabilities = vec![Capability::Imap4Rev1, Capability::Idle];
let body = vec![Data::Capability(capabilities)]; let body: Vec<Data> = vec![
Status::ok("Yeah")?.into(),
ImapData::Capability(capabilities).into(),
];
Ok(Response::ok("Done")?.with_body(body)) Ok(Response::ok("Done")?.with_body(body))
} }

1
src/lib.rs
View file

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

1
src/proto/req/mod.rs
View file

@ -1,5 +1,6 @@
use imap_codec::types::command::Command; use imap_codec::types::command::Command;
#[derive(Debug)]
pub struct Request { pub struct Request {
pub command: Command, pub command: Command,
} }

34
src/proto/res/body.rs
View file

@ -1,17 +1,28 @@
use std::io; use std::{fmt, io};
use futures::prelude::*;
use futures::stream::BoxStream;
use imap_codec::codec::Encode; use imap_codec::codec::Encode;
use imap_codec::types::response::{Data as ImapData, Status as ImapStatus}; use imap_codec::types::response::{Data as ImapData, Status as ImapStatus};
#[derive(Debug)] use super::Status;
pub enum Body { pub enum Body {
Once(Vec<Data>), Once(Vec<Data>),
Stream(BoxStream<'static, Data>),
} }
impl Body { impl Body {
pub(crate) fn into_data(self) -> Vec<Data> { pub fn from_stream<St: Stream<Item = Data> + Send + 'static>(stream: St) -> Self {
Body::Stream(stream.boxed())
}
}
impl Body {
pub(crate) fn into_stream(self) -> BoxStream<'static, Data> {
match self { match self {
Body::Once(data) => data, Body::Once(data) => futures::stream::iter(data).boxed(),
Body::Stream(stream) => stream,
} }
} }
} }
@ -46,6 +57,15 @@ impl From<ImapData> for Body {
} }
} }
impl fmt::Debug for Body {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Body::Once(ref data) => f.debug_struct("Body::Once").field("data", data).finish(),
Body::Stream(_) => f.debug_struct("Body::Stream").finish_non_exhaustive(),
}
}
}
#[derive(Debug, Clone)] #[derive(Debug, Clone)]
pub enum Data { pub enum Data {
Data(ImapData), Data(ImapData),
@ -72,3 +92,9 @@ impl From<ImapStatus> for Data {
Data::Status(status) Data::Status(status)
} }
} }
impl From<Status> for Data {
fn from(status: Status) -> Self {
status.into_imap(None).into()
}
}

23
src/proto/res/mod.rs
View file

@ -7,6 +7,7 @@ use self::body::Body;
use crate::errors::{Error, Result}; use crate::errors::{Error, Result};
pub mod body; pub mod body;
pub(crate) mod stream;
#[derive(Debug)] #[derive(Debug)]
pub struct Response { pub struct Response {
@ -51,6 +52,12 @@ impl Response {
} }
} }
impl Response {
pub fn split(self) -> (Option<Body>, Status) {
(self.body, self.status)
}
}
#[derive(Debug, Clone)] #[derive(Debug, Clone)]
pub struct Status { pub struct Status {
pub(crate) code: StatusCode, pub(crate) code: StatusCode,
@ -76,6 +83,22 @@ impl Status {
}) })
} }
pub fn ok(msg: &str) -> Result<Status> {
Self::new(StatusCode::Ok, msg)
}
pub fn no(msg: &str) -> Result<Status> {
Self::new(StatusCode::No, msg)
}
pub fn bad(msg: &str) -> Result<Status> {
Self::new(StatusCode::Bad, msg)
}
pub fn bye(msg: &str) -> Result<Status> {
Self::new(StatusCode::Bye, msg)
}
pub(crate) fn into_imap(self, tag: Option<Tag>) -> ImapStatus { pub(crate) fn into_imap(self, tag: Option<Tag>) -> ImapStatus {
match self.code { match self.code {
StatusCode::Ok => ImapStatus::Ok { StatusCode::Ok => ImapStatus::Ok {

23
src/proto/res/stream.rs Normal file
View file

@ -0,0 +1,23 @@
use futures::prelude::*;
use imap_codec::types::core::Tag;
use super::body::Data;
use super::Response;
pub fn response_stream(res: Response, tag: Option<Tag>) -> impl Stream<Item = Data> {
let (body, status) = res.split();
let body = body.map(|body| body.into_stream());
async_stream::stream! {
if let Some(body) = body {
for await item in body {
tracing::trace!(?item, "response_stream.yield");
yield item;
}
}
let item = status.into_imap(tag);
tracing::trace!(?item, "response_stream.yield");
yield item.into();
}
}

62
src/server/pipeline.rs
View file

@ -2,12 +2,13 @@ use std::pin::Pin;
use std::task::{self, Poll}; use std::task::{self, Poll};
use bytes::BytesMut; use bytes::BytesMut;
use futures::io::{AsyncRead, AsyncWrite}; use futures::prelude::*;
use futures::sink::Sink; use futures::stream::BoxStream;
use futures::stream::Stream;
use imap_codec::types::core::Tag; use imap_codec::types::core::Tag;
use crate::proto::res::body::Data;
use crate::proto::{Request, Response}; use crate::proto::{Request, Response};
use crate::util::stream::ConcatAll;
type Error = tower::BoxError; type Error = tower::BoxError;
type Result<T, E = Error> = std::result::Result<T, E>; type Result<T, E = Error> = std::result::Result<T, E>;
@ -15,8 +16,12 @@ type Result<T, E = Error> = std::result::Result<T, E>;
#[pin_project::pin_project] #[pin_project::pin_project]
pub struct Connection<C> { pub struct Connection<C> {
#[pin] #[pin]
conn: C, pub conn: C,
read_buf: BytesMut, read_buf: BytesMut,
#[pin]
outbox: ConcatAll<BoxStream<'static, Data>>,
write_buf: BytesMut, write_buf: BytesMut,
} }
@ -24,7 +29,10 @@ impl<C> Connection<C> {
pub fn new(conn: C) -> Self { pub fn new(conn: C) -> Self {
Self { Self {
conn, conn,
read_buf: BytesMut::new(),
read_buf: BytesMut::with_capacity(1024),
outbox: ConcatAll::new(),
write_buf: BytesMut::new(), write_buf: BytesMut::new(),
} }
} }
@ -47,7 +55,15 @@ where
Err(e) if e.is_incomplete() => { Err(e) if e.is_incomplete() => {
let mut buf = [0u8; 256]; let mut buf = [0u8; 256];
let read = futures::ready!(this.conn.as_mut().poll_read(cx, &mut buf))?; tracing::trace!("transport.poll_read");
// let read = futures::ready!(this.conn.as_mut().poll_read(cx, &mut buf))?;
let read = match this.conn.as_mut().poll_read(cx, &mut buf) {
Poll::Ready(res) => res?,
Poll::Pending => {
tracing::trace!("transport.pending");
return Poll::Pending;
}
};
tracing::trace!(read = read, "transport.poll_next"); tracing::trace!(read = read, "transport.poll_next");
if read == 0 { if read == 0 {
@ -78,15 +94,29 @@ where
C: AsyncWrite, C: AsyncWrite,
{ {
fn poll_flush_buffer(self: Pin<&mut Self>, cx: &mut task::Context) -> Poll<Result<()>> { fn poll_flush_buffer(self: Pin<&mut Self>, cx: &mut task::Context) -> Poll<Result<()>> {
use bytes::Buf; use bytes::{Buf, BufMut};
use imap_codec::codec::Encode;
let mut this = self.project(); let mut this = self.project();
tracing::debug!(size = this.outbox.len(), "transport.flush_outbox");
let mut writer = this.write_buf.writer();
while let Poll::Ready(Some(data)) = this.outbox.as_mut().poll_next(cx) {
tracing::trace!(?data, "transport.write_buf");
if let Err(err) = data.encode(&mut writer) {
tracing::error!(?err, "transport.encode_error");
return Poll::Ready(Err(Box::new(err)));
}
}
tracing::debug!(size = this.write_buf.len(), "transport.flush_buffer"); tracing::debug!(size = this.write_buf.len(), "transport.flush_buffer");
while !this.write_buf.is_empty() { while !this.write_buf.is_empty() {
let written = futures::ready!(this.conn.as_mut().poll_write(cx, this.write_buf))?; let written = futures::ready!(this.conn.as_mut().poll_write(cx, this.write_buf))?;
this.write_buf.advance(written); this.write_buf.advance(written);
} }
this.write_buf.clear();
Poll::Ready(Ok(())) Poll::Ready(Ok(()))
} }
@ -107,23 +137,13 @@ where
} }
fn start_send( fn start_send(
self: Pin<&mut Self>, mut self: Pin<&mut Self>,
(tag, res): (Option<Tag>, Response), (tag, res): (Option<Tag>, Response),
) -> Result<(), Self::Error> { ) -> Result<(), Self::Error> {
use bytes::BufMut; use crate::proto::res::stream::response_stream;
use imap_codec::codec::Encode;
debug_assert!(self.write_buf.is_empty()); tracing::debug!(?tag, ?res, "transport.start_send");
self.outbox.push(Box::pin(response_stream(res, tag)));
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(())
} }

0
src/util/buf.rs Normal file
View file

2
src/util/mod.rs Normal file
View file

@ -0,0 +1,2 @@
pub mod buf;
pub mod stream;

54
src/util/stream.rs Normal file
View file

@ -0,0 +1,54 @@
use std::fmt;
use std::pin::Pin;
use std::task::{self, Poll};
use futures::prelude::*;
use futures::stream::{FuturesOrdered, Stream, StreamFuture};
/// [`SelectAll`](futures::stream::SelectAll) but ordered
#[pin_project::pin_project]
pub struct ConcatAll<St: Stream + Unpin> {
#[pin]
inner: FuturesOrdered<StreamFuture<St>>,
}
impl<St: Stream + Unpin> ConcatAll<St> {
pub fn new() -> Self {
Self {
inner: FuturesOrdered::new(),
}
}
pub fn len(&self) -> usize {
self.inner.len()
}
pub fn push(&mut self, stream: St) {
use futures::StreamExt;
self.inner.push(stream.into_future());
}
}
impl<St: Stream + Unpin> fmt::Debug for ConcatAll<St> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("ConcatAll").finish()
}
}
impl<St: Stream + Unpin> Stream for ConcatAll<St> {
type Item = St::Item;
fn poll_next(mut self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<Option<Self::Item>> {
loop {
match futures::ready!(self.inner.poll_next_unpin(cx)) {
Some((Some(item), remaining)) => {
self.push(remaining);
return Poll::Ready(Some(item));
}
Some((None, _)) => {}
_ => return Poll::Ready(None),
}
}
}
}