use std::fmt;
#[cfg(feature = "ws-core")]
use std::net::SocketAddr;
use async_trait::async_trait;
#[cfg(feature = "ws-core")]
use eyre::Result;
#[cfg(feature = "ws-core")]
use futures::future::LocalBoxFuture;
use super::WsMessage as Message;
#[cfg(feature = "ws-core")]
use super::WsServerConfig;
#[cfg(feature = "ws-core")]
pub trait RawStream: Send + 'static {
fn read<'a>(
&'a mut self,
buffer: &'a mut [u8],
) -> LocalBoxFuture<'a, Result<usize, StreamError>>;
fn write_all<'a>(&'a mut self, buffer: &'a [u8])
-> LocalBoxFuture<'a, Result<(), StreamError>>;
}
#[cfg(feature = "ws-core")]
impl<T: nagoya::io::Stream + Unpin + Send + 'static> RawStream for T {
fn read<'a>(
&'a mut self,
buffer: &'a mut [u8],
) -> LocalBoxFuture<'a, Result<usize, StreamError>> {
Box::pin(async move {
nagoya::io::Stream::read(self, buffer)
.await
.map_err(|err| StreamError::Other(eyre::eyre!("{err}")))
})
}
fn write_all<'a>(
&'a mut self,
buffer: &'a [u8],
) -> LocalBoxFuture<'a, Result<(), StreamError>> {
Box::pin(async move {
nagoya::io::Stream::write_all(self, buffer)
.await
.map_err(|err| StreamError::Other(eyre::eyre!("{err}")))
})
}
}
#[cfg(feature = "ws-core")]
pub type BoxedStream = Box<dyn RawStream>;
#[cfg(feature = "ws-core")]
impl nagoya::io::Stream for BoxedStream {
async fn read(&mut self, buffer: &mut [u8]) -> Result<usize, nagoya::io::StreamError> {
RawStream::read(self.as_mut(), buffer)
.await
.map_err(|_| nagoya::io::StreamError(libc_eio()))
}
async fn write_all(&mut self, buffer: &[u8]) -> Result<(), nagoya::io::StreamError> {
RawStream::write_all(self.as_mut(), buffer)
.await
.map_err(|_| nagoya::io::StreamError(libc_eio()))
}
}
#[cfg(feature = "ws-core")]
fn libc_eio() -> i32 {
5
}
#[derive(Debug)]
pub enum StreamError {
Closed,
WriteBufferFull,
Protocol(String),
Other(eyre::Error),
}
impl fmt::Display for StreamError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
StreamError::Closed => write!(f, "connection closed"),
StreamError::WriteBufferFull => write!(f, "write buffer full"),
StreamError::Protocol(s) => write!(f, "protocol error: {s}"),
StreamError::Other(e) => write!(f, "{e}"),
}
}
}
impl std::error::Error for StreamError {
fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
match self {
StreamError::Other(e) => Some(e.as_ref()),
_ => None,
}
}
}
#[async_trait(?Send)]
pub trait MessageStream: Unpin + Send {
async fn send(&mut self, msg: Message) -> Result<(), StreamError>;
async fn recv(&mut self) -> Option<Result<Message, StreamError>>;
}
pub use MessageStream as WsStream;
#[cfg(feature = "ws-core")]
pub struct UpgradeEvent {
pub stream: Box<dyn MessageStream>,
pub protocol: String,
}
#[cfg(feature = "ws-core")]
#[async_trait(?Send)]
pub trait WsUpgrader: Send + Sync {
async fn upgrade_stream(
&self,
stream: BoxedStream,
addr: SocketAddr,
config: &WsServerConfig,
) -> Result<Option<UpgradeEvent>>;
}