use std::fmt;
use std::pin::Pin;
use std::task::{Context, Poll, ready};
use futures_core::Stream;
use futures_sink::Sink;
use crate::core::Dir;
use crate::error::{ConnectionLost, DatagramError, MessageError};
use crate::identity::Identity;
use crate::packet::Handshake;
use crate::shell::{BiStream, Connection, Endpoint, Intro, Notification, RecvStream, WakerSlot};
impl<S: Handshake + 'static> Connection<S> {
pub fn messages(&self) -> Messages<'_, S> {
Messages {
conn: self,
slot: self.message_reader_slot_boxed(),
}
}
pub fn datagrams(&self) -> Datagrams<'_, S> {
Datagrams {
conn: self,
slot: self.datagram_reader_slot_boxed(),
}
}
pub fn incoming_bi(&self) -> IncomingBi<'_, S> {
IncomingBi {
conn: self,
slot: self.acceptor_slot_boxed(Dir::Bi),
}
}
pub fn incoming_uni(&self) -> IncomingUni<'_, S> {
IncomingUni {
conn: self,
slot: self.acceptor_slot_boxed(Dir::Uni),
}
}
pub fn notifications(&self) -> Notifications<'_, S> {
Notifications {
conn: self,
slot: self.notification_slot_boxed(),
}
}
pub fn message_sink(&self) -> MessageSink<'_, S> {
MessageSink {
conn: self,
slot: self.message_sender_slot_boxed(),
pending: None,
}
}
pub fn datagram_sink(&self) -> DatagramSink<'_, S> {
DatagramSink { conn: self }
}
}
impl<I: Identity> Endpoint<I> {
pub fn incoming(&self) -> Incoming<'_, I> {
Incoming {
endpoint: self,
pending: None,
}
}
}
pub struct Messages<'a, S: Handshake> {
conn: &'a Connection<S>,
slot: WakerSlot<Box<dyn FnMut(u64)>>,
}
impl<S: Handshake> Stream for Messages<'_, S> {
type Item = Result<Vec<u8>, ConnectionLost>;
fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
let this = self.get_mut();
this.conn.poll_recv_message(cx, this.slot.key()).map(Some)
}
}
pub struct Datagrams<'a, S: Handshake> {
conn: &'a Connection<S>,
slot: WakerSlot<Box<dyn FnMut(u64)>>,
}
impl<S: Handshake> Stream for Datagrams<'_, S> {
type Item = Result<Vec<u8>, ConnectionLost>;
fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
let this = self.get_mut();
this.conn.poll_recv_datagram(cx, this.slot.key()).map(Some)
}
}
pub struct IncomingBi<'a, S: Handshake> {
conn: &'a Connection<S>,
slot: WakerSlot<Box<dyn FnMut(u64)>>,
}
impl<S: Handshake> Stream for IncomingBi<'_, S> {
type Item = Result<BiStream<S>, ConnectionLost>;
fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
let this = self.get_mut();
this.conn.poll_accept_bi(cx, this.slot.key()).map(Some)
}
}
pub struct IncomingUni<'a, S: Handshake> {
conn: &'a Connection<S>,
slot: WakerSlot<Box<dyn FnMut(u64)>>,
}
impl<S: Handshake> Stream for IncomingUni<'_, S> {
type Item = Result<RecvStream<S>, ConnectionLost>;
fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
let this = self.get_mut();
this.conn.poll_accept_uni(cx, this.slot.key()).map(Some)
}
}
pub struct Notifications<'a, S: Handshake> {
conn: &'a Connection<S>,
slot: WakerSlot<Box<dyn FnMut(u64)>>,
}
impl<S: Handshake> Stream for Notifications<'_, S> {
type Item = Result<Notification, ConnectionLost>;
fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
let this = self.get_mut();
this.conn.poll_notified(cx, this.slot.key()).map(Some)
}
}
pub struct Incoming<'a, I: Identity> {
endpoint: &'a Endpoint<I>,
pending: Option<tokio::sync::oneshot::Receiver<Intro<I>>>,
}
impl<I: Identity> Stream for Incoming<'_, I> {
type Item = Intro<I>;
fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
let this = self.get_mut();
this.endpoint.poll_accept(cx, &mut this.pending)
}
}
pub struct MessageSink<'a, S: Handshake> {
conn: &'a Connection<S>,
slot: WakerSlot<Box<dyn FnMut(u64)>>,
pending: Option<Vec<u8>>,
}
impl<S: Handshake> MessageSink<'_, S> {
fn poll_drain(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), MessageError>> {
let Some(msg) = self.pending.as_deref() else {
return Poll::Ready(Ok(()));
};
let outcome = ready!(self.conn.poll_send_message(cx, msg, self.slot.key()));
self.pending = None;
Poll::Ready(outcome)
}
}
impl<S: Handshake> Sink<Vec<u8>> for MessageSink<'_, S> {
type Error = MessageError;
fn poll_ready(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
self.get_mut().poll_drain(cx)
}
fn start_send(self: Pin<&mut Self>, item: Vec<u8>) -> Result<(), Self::Error> {
let this = self.get_mut();
debug_assert!(
this.pending.is_none(),
"Sink::start_send without a preceding Ready poll_ready overwrites the slot"
);
this.pending = Some(item);
Ok(())
}
fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
self.get_mut().poll_drain(cx)
}
fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
self.get_mut().poll_drain(cx)
}
}
pub struct DatagramSink<'a, S: Handshake> {
conn: &'a Connection<S>,
}
impl<S: Handshake> Sink<Vec<u8>> for DatagramSink<'_, S> {
type Error = DatagramError;
fn poll_ready(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
Poll::Ready(Ok(()))
}
fn start_send(self: Pin<&mut Self>, item: Vec<u8>) -> Result<(), Self::Error> {
self.get_mut().conn.send_datagram(&item)
}
fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
Poll::Ready(Ok(()))
}
fn poll_close(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
Poll::Ready(Ok(()))
}
}
impl<S: Handshake> fmt::Debug for Messages<'_, S> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("Messages").finish_non_exhaustive()
}
}
impl<S: Handshake> fmt::Debug for Datagrams<'_, S> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("Datagrams").finish_non_exhaustive()
}
}
impl<S: Handshake> fmt::Debug for IncomingBi<'_, S> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("IncomingBi").finish_non_exhaustive()
}
}
impl<S: Handshake> fmt::Debug for IncomingUni<'_, S> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("IncomingUni").finish_non_exhaustive()
}
}
impl<S: Handshake> fmt::Debug for Notifications<'_, S> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("Notifications").finish_non_exhaustive()
}
}
impl<I: Identity> fmt::Debug for Incoming<'_, I> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("Incoming")
.field("request_outstanding", &self.pending.is_some())
.finish_non_exhaustive()
}
}
impl<S: Handshake> fmt::Debug for MessageSink<'_, S> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("MessageSink")
.field("pending", &self.pending.is_some())
.finish_non_exhaustive()
}
}
impl<S: Handshake> fmt::Debug for DatagramSink<'_, S> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("DatagramSink").finish_non_exhaustive()
}
}