pub mod poll {
use std::task::{Context, Poll, ready};
use bytes::{Buf, BufMut, Bytes};
use web_transport_trait::{MaybeSend, MaybeSync, poll};
pub trait Session: poll::Session<SendStream: SendStream, RecvStream: RecvStream> + Clone + 'static {
fn accept_uni(&mut self) -> AcceptUni<'_, Self> {
AcceptUni(self)
}
fn accept_bi(&mut self) -> AcceptBi<'_, Self> {
AcceptBi(self)
}
fn open_uni(&mut self) -> OpenUni<'_, Self> {
OpenUni(self)
}
fn open_bi(&mut self) -> OpenBi<'_, Self> {
OpenBi(self)
}
fn recv_datagram(&mut self) -> RecvDatagram<'_, Self> {
RecvDatagram(self)
}
fn send_datagram(&mut self, payload: &[u8]) -> Result<(), Self::Error> {
let mut cx = Context::from_waker(std::task::Waker::noop());
match self.poll_send_datagram(&mut cx, payload) {
Poll::Ready(res) => res,
Poll::Pending => Ok(()),
}
}
fn closed(&mut self) -> SessionClosed<'_, Self> {
SessionClosed(self)
}
}
macro_rules! poll_future {
($(#[$doc:meta])* $name:ident, $bound:path, $poll:ident, $out:ty) => {
$(#[$doc])*
pub struct $name<'a, S: ?Sized>(&'a mut S);
impl<S: $bound> Future for $name<'_, S> {
type Output = $out;
fn poll(mut self: std::pin::Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
self.0.$poll(cx)
}
}
};
}
poll_future!(
AcceptUni, poll::Session, poll_accept_uni, Result<S::RecvStream, S::Error>);
poll_future!(
AcceptBi, poll::Session, poll_accept_bi, Result<poll::BiStreams<S>, S::Error>);
poll_future!(
OpenUni, poll::Session, poll_open_uni, Result<S::SendStream, S::Error>);
poll_future!(
OpenBi, poll::Session, poll_open_bi, Result<poll::BiStreams<S>, S::Error>);
poll_future!(
RecvDatagram, poll::Session, poll_recv_datagram, Result<Bytes, S::Error>);
poll_future!(
SessionClosed, poll::Session, poll_closed, S::Error);
poll_future!(
SendClosed, poll::SendStream, poll_closed, Result<(), S::Error>);
poll_future!(
RecvClosed, poll::RecvStream, poll_closed, Result<(), S::Error>);
impl<S> Session for S where S: poll::Session<SendStream: SendStream, RecvStream: RecvStream> + Clone + 'static {}
pub trait Boxable:
Session<SendStream: MaybeSend, RecvStream: MaybeSend, Error: MaybeSend> + MaybeSend + MaybeSync
{
}
impl<S> Boxable for S where
S: Session<SendStream: MaybeSend, RecvStream: MaybeSend, Error: MaybeSend> + MaybeSend + MaybeSync
{
}
pub trait SendStream: poll::SendStream + 'static {
fn write<'a>(&'a mut self, buf: &'a [u8]) -> Write<'a, Self> {
Write { stream: self, buf }
}
fn write_buf<'a, B: Buf>(&'a mut self, buf: &'a mut B) -> WriteBuf<'a, Self, B> {
WriteBuf { stream: self, buf }
}
fn write_chunk(&mut self, chunk: Bytes) -> WriteChunk<'_, Self> {
WriteChunk { stream: self, chunk }
}
fn closed(&mut self) -> SendClosed<'_, Self> {
SendClosed(self)
}
}
pub struct Write<'a, S: ?Sized> {
stream: &'a mut S,
buf: &'a [u8],
}
impl<S: poll::SendStream> Future for Write<'_, S> {
type Output = Result<usize, S::Error>;
fn poll(mut self: std::pin::Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let this = &mut *self;
this.stream.poll_write(cx, this.buf)
}
}
pub struct WriteBuf<'a, S: ?Sized, B> {
stream: &'a mut S,
buf: &'a mut B,
}
impl<S: poll::SendStream, B: Buf> Future for WriteBuf<'_, S, B> {
type Output = Result<usize, S::Error>;
fn poll(mut self: std::pin::Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let this = &mut *self;
this.stream.poll_write_buf(cx, this.buf)
}
}
pub struct WriteChunk<'a, S: ?Sized> {
stream: &'a mut S,
chunk: Bytes,
}
impl<S: poll::SendStream> Future for WriteChunk<'_, S> {
type Output = Result<(), S::Error>;
fn poll(mut self: std::pin::Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let this = &mut *self;
while !this.chunk.is_empty() {
ready!(this.stream.poll_write_buf(cx, &mut this.chunk))?;
}
Poll::Ready(Ok(()))
}
}
impl<S> SendStream for S where S: poll::SendStream + 'static {}
pub trait RecvStream: poll::RecvStream + 'static {
fn read<'a>(&'a mut self, dst: &'a mut [u8]) -> Read<'a, Self> {
Read { stream: self, dst }
}
fn read_buf<'a, B: BufMut>(&'a mut self, buf: &'a mut B) -> ReadBuf<'a, Self, B> {
ReadBuf { stream: self, buf }
}
fn read_chunk(&mut self, max: usize) -> ReadChunk<'_, Self> {
ReadChunk { stream: self, max }
}
fn closed(&mut self) -> RecvClosed<'_, Self> {
RecvClosed(self)
}
}
pub struct Read<'a, S: ?Sized> {
stream: &'a mut S,
dst: &'a mut [u8],
}
impl<S: poll::RecvStream> Future for Read<'_, S> {
type Output = Result<Option<usize>, S::Error>;
fn poll(mut self: std::pin::Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let this = &mut *self;
this.stream.poll_read(cx, this.dst)
}
}
pub struct ReadBuf<'a, S: ?Sized, B> {
stream: &'a mut S,
buf: &'a mut B,
}
impl<S: poll::RecvStream, B: BufMut> Future for ReadBuf<'_, S, B> {
type Output = Result<Option<usize>, S::Error>;
fn poll(mut self: std::pin::Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let this = &mut *self;
this.stream.poll_read_buf(cx, this.buf)
}
}
pub struct ReadChunk<'a, S: ?Sized> {
stream: &'a mut S,
max: usize,
}
impl<S: poll::RecvStream> Future for ReadChunk<'_, S> {
type Output = Result<Option<Bytes>, S::Error>;
fn poll(mut self: std::pin::Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let this = &mut *self;
this.stream.poll_read_chunk(cx, this.max)
}
}
impl<S> RecvStream for S where S: poll::RecvStream + 'static {}
}