use core::convert::Infallible;
use core::fmt::{self, Write};
use core::future::Future;
use core::mem;
use core::num::NonZeroU16;
use core::pin::Pin;
use core::task::{Context, Poll};
use alloc::boxed::Box;
use alloc::collections::{BTreeSet, VecDeque};
use alloc::string::String;
use alloc::sync::Arc;
use alloc::vec::Vec;
use std::sync::{LazyLock, Mutex, MutexGuard, PoisonError};
use bytes::Bytes;
use rand::prelude::*;
use rand::rngs::SmallRng;
use tokio::task::JoinSet;
use tokio::time::{Duration, Instant, Sleep};
use crate::Buf;
use crate::api::{
Broadcast, ChannelId, DecodeBody, EncodeBody, ErrorMessage, Event, Format, Id, IdSet,
MessageId, Mode, RequestHeader, ResponseHeader, VERSION,
};
use crate::buf::{BufPool, InvalidFrame};
use crate::format;
const MAX_CAPACITY: usize = 1048576;
const DEFAULT_MAX_IN_FLIGHT: usize = 128;
const DEFAULT_MAX_OUTBOUND: usize = 8 * 1048576;
const DEFAULT_MAX_CHANNELS: usize = 1024;
const CLOSE_NORMAL: u16 = 1000;
const CLOSE_PROTOCOL_ERROR: u16 = 1002;
const CLOSE_POLICY_VIOLATION: u16 = 1008;
const CLOSE_TIMEOUT: Duration = Duration::from_secs(30);
const PING_TIMEOUT: Duration = Duration::from_secs(10);
const DEFAULT_SEED: u64 = 0xdeadbeef;
#[derive(Debug)]
pub(crate) enum Message {
Text(Bytes),
Binary(Bytes),
Ping(Bytes),
Pong(Bytes),
Close,
}
pub(crate) mod socket_sealed {
pub trait Sealed {}
}
pub(crate) trait SocketImpl
where
Self: self::socket_sealed::Sealed,
{
#[doc(hidden)]
type Message;
#[doc(hidden)]
type Error: fmt::Debug;
#[doc(hidden)]
fn poll_next(
self: Pin<&mut Self>,
ctx: &mut Context<'_>,
) -> Poll<Option<Result<Message, Self::Error>>>;
#[doc(hidden)]
fn poll_ready(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>>;
#[doc(hidden)]
fn start_send(self: Pin<&mut Self>, item: Self::Message) -> Result<(), Self::Error>;
#[doc(hidden)]
fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>>;
}
pub(crate) mod server_sealed {
pub trait Sealed {}
}
pub trait ServerImpl
where
Self: self::server_sealed::Sealed,
{
#[doc(hidden)]
type Error: fmt::Debug;
#[doc(hidden)]
type Message;
#[doc(hidden)]
#[allow(private_bounds)]
type Socket: SocketImpl<Message = Self::Message, Error = Self::Error>;
#[doc(hidden)]
fn ping(data: Bytes) -> Self::Message;
#[doc(hidden)]
fn pong(data: Bytes) -> Self::Message;
#[doc(hidden)]
fn binary(data: &[u8]) -> Self::Message;
#[doc(hidden)]
fn text(data: &[u8]) -> Self::Message;
#[doc(hidden)]
fn close(code: u16, reason: &str) -> Self::Message;
}
#[derive(Debug)]
enum ErrorKind {
#[cfg(feature = "axum-core05")]
AxumCore05 {
error: axum_core05::Error,
},
FormatError,
InvalidFrame {
error: InvalidFrame,
},
Incoming {
error: format::Error,
},
Outgoing {
error: format::Error,
},
EncodeBroadcastHeader {
error: format::Error,
},
EncodeBroadcast {
error: format::Error,
},
EncodeConnectHeader {
error: format::Error,
},
ErrorMessageHeader {
error: format::Error,
},
ErrorMessage {
error: format::Error,
},
EncodeErrorPayload {
error: format::Error,
},
OutOfBounds {
offset: usize,
len: usize,
},
NotNegotiated,
ExpectedNegotiate {
id: u16,
},
NegotiateHeader {
error: format::Error,
},
UnknownHeader {
error: format::Error,
},
UnsupportedVersion {
error: format::Error,
},
Incompatible {
server: u32,
client: u32,
},
}
#[derive(Debug)]
pub struct Error {
kind: ErrorKind,
}
impl Error {
#[inline]
const fn new(kind: ErrorKind) -> Self {
Self { kind }
}
pub(crate) fn incoming(error: format::Error) -> Self {
Self::new(ErrorKind::Incoming { error })
}
pub(crate) fn outgoing(error: format::Error) -> Self {
Self::new(ErrorKind::Outgoing { error })
}
pub(crate) fn encode_broadcast_header(error: format::Error) -> Self {
Self::new(ErrorKind::EncodeBroadcastHeader { error })
}
pub(crate) fn encode_broadcast(error: format::Error) -> Self {
Self::new(ErrorKind::EncodeBroadcast { error })
}
pub(crate) fn encode_connect_header(error: format::Error) -> Self {
Self::new(ErrorKind::EncodeConnectHeader { error })
}
pub(crate) fn encode_error_message_header(error: format::Error) -> Self {
Self::new(ErrorKind::ErrorMessageHeader { error })
}
pub(crate) fn encode_error_message(error: format::Error) -> Self {
Self::new(ErrorKind::ErrorMessage { error })
}
#[inline]
pub fn is_incompatible(&self) -> bool {
matches!(self.kind, ErrorKind::Incompatible { .. })
}
}
impl fmt::Display for Error {
#[inline]
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match &self.kind {
#[cfg(feature = "axum-core05")]
ErrorKind::AxumCore05 { .. } => write!(f, "Error in axum-core"),
ErrorKind::FormatError => write!(f, "Error formatting error response"),
ErrorKind::InvalidFrame { error } => error.fmt(f),
ErrorKind::Incoming { .. } => {
write!(f, "Encoding error when decoding incoming message")
}
ErrorKind::Outgoing { .. } => {
write!(f, "Encoding error when encoding outgoing message")
}
ErrorKind::EncodeBroadcastHeader { .. } => {
write!(f, "Encoding error when encoding broadcast header")
}
ErrorKind::EncodeBroadcast { .. } => {
write!(f, "Encoding error when broadcasting message")
}
ErrorKind::EncodeConnectHeader { .. } => {
write!(f, "Encoding error when encoding connect header")
}
ErrorKind::ErrorMessageHeader { .. } => {
write!(f, "Encoding error when encoding error message header")
}
ErrorKind::ErrorMessage { .. } => {
write!(f, "Encoding error when encoding error message")
}
ErrorKind::EncodeErrorPayload { .. } => {
write!(f, "Encoding error when encoding typed error")
}
ErrorKind::OutOfBounds { offset, len } => {
write!(
f,
"Error when reading message: offset {} is out of bounds for length {}",
offset, len
)
}
ErrorKind::NotNegotiated => {
write!(f, "Connection closed before the format was negotiated")
}
ErrorKind::ExpectedNegotiate { id } => {
write!(
f,
"Expected a negotiation as the first message, but got message id {id}"
)
}
ErrorKind::NegotiateHeader { .. } => {
write!(f, "Encoding error when decoding negotiation header")
}
ErrorKind::UnknownHeader { error } => {
write!(f, "Client sent a header this server does not know: {error}")
}
ErrorKind::UnsupportedVersion { error } => {
write!(f, "Client speaks a different protocol: {error}")
}
ErrorKind::Incompatible { server, client } => {
write!(
f,
"Client states application version {client}, this server states {server}"
)
}
}
}
}
impl core::error::Error for Error {
#[inline]
fn source(&self) -> Option<&(dyn core::error::Error + 'static)> {
match &self.kind {
#[cfg(feature = "axum-core05")]
ErrorKind::AxumCore05 { error } => Some(error),
ErrorKind::Incoming { error } => Some(error),
ErrorKind::Outgoing { error } => Some(error),
ErrorKind::EncodeBroadcastHeader { error } => Some(error),
ErrorKind::EncodeBroadcast { error } => Some(error),
ErrorKind::EncodeConnectHeader { error } => Some(error),
ErrorKind::ErrorMessageHeader { error } => Some(error),
ErrorKind::ErrorMessage { error } => Some(error),
ErrorKind::EncodeErrorPayload { error } => Some(error),
ErrorKind::NegotiateHeader { error } => Some(error),
ErrorKind::UnknownHeader { error } => Some(error),
ErrorKind::UnsupportedVersion { error } => Some(error),
_ => None,
}
}
}
#[cfg(feature = "axum-core05")]
impl From<axum_core05::Error> for Error {
#[inline]
fn from(error: axum_core05::Error) -> Self {
Self::new(ErrorKind::AxumCore05 { error })
}
}
impl From<ErrorKind> for Error {
#[inline]
fn from(kind: ErrorKind) -> Self {
Self::new(kind)
}
}
impl From<InvalidFrame> for Error {
#[inline]
fn from(error: InvalidFrame) -> Self {
Self::new(ErrorKind::InvalidFrame { error })
}
}
type Result<T, E = Error> = core::result::Result<T, E>;
pub struct Response {
handled: bool,
}
pub trait ResponseError
where
Self: 'static + Send,
{
fn fmt_message(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result;
#[inline]
fn encode_payload(&self, payload: ErrorPayload<'_>) -> Result<(), Error> {
_ = payload;
Ok(())
}
}
impl<E> ResponseError for E
where
E: 'static + Send + fmt::Display,
{
#[inline]
fn fmt_message(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
fmt::Display::fmt(self, f)
}
}
impl<E> ResponseError for TypedError<E>
where
E: 'static + Send + fmt::Display + EncodeBody,
{
#[inline]
fn fmt_message(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
fmt::Display::fmt(&self.0, f)
}
#[inline]
fn encode_payload(&self, payload: ErrorPayload<'_>) -> Result<(), Error> {
payload.encode(&self.0)
}
}
pub struct ErrorPayload<'a> {
format: Format,
out: &'a mut Vec<u8>,
encoded: &'a mut bool,
}
impl ErrorPayload<'_> {
#[inline]
pub fn format(&self) -> Format {
self.format
}
pub fn encode<T>(self, value: &T) -> Result<(), Error>
where
T: ?Sized + EncodeBody,
{
self.out.clear();
if let Err(error) = self.format.encode(self.out, value) {
self.out.clear();
return Err(Error::new(ErrorKind::EncodeErrorPayload { error }));
}
*self.encoded = true;
Ok(())
}
}
struct ErrorMessageDisplay<'a, E>(&'a E)
where
E: ?Sized;
impl<E> fmt::Display for ErrorMessageDisplay<'_, E>
where
E: ?Sized + ResponseError,
{
#[inline]
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
self.0.fmt_message(f)
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Default)]
pub struct TypedError<E>(pub E);
impl<E> TypedError<E> {
#[inline]
pub const fn new(error: E) -> Self {
Self(error)
}
#[inline]
pub const fn get(&self) -> &E {
&self.0
}
#[inline]
pub fn into_inner(self) -> E {
self.0
}
}
impl<E> From<E> for TypedError<E> {
#[inline]
fn from(error: E) -> Self {
Self(error)
}
}
pub trait IntoResponse
where
Self: 'static + Send,
{
type Error: ResponseError;
fn into_response(self) -> Result<Response, Self::Error>;
}
impl IntoResponse for () {
type Error = Infallible;
#[inline]
fn into_response(self) -> Result<Response, Self::Error> {
Ok(Response { handled: true })
}
}
impl IntoResponse for bool {
type Error = Infallible;
#[inline]
fn into_response(self) -> Result<Response, Self::Error> {
Ok(Response { handled: self })
}
}
impl<T, E> IntoResponse for Result<T, E>
where
T: IntoResponse<Error = Infallible>,
E: ResponseError,
{
type Error = E;
#[inline]
fn into_response(self) -> Result<Response, E> {
match self {
Ok(into_response) => match IntoResponse::into_response(into_response) {
Ok(response) => Ok(response),
Err(error) => match error {},
},
Err(error) => Err(error),
}
}
}
impl<T> IntoResponse for Option<T>
where
T: IntoResponse,
{
type Error = T::Error;
#[inline]
fn into_response(self) -> Result<Response, Self::Error> {
match self {
Some(value) => value.into_response(),
None => Ok(Response { handled: false }),
}
}
}
pub trait Handler
where
Self: 'static + Send + Clone,
{
type Id: Id;
type Response: IntoResponse;
fn open_channel<'this>(
&'this self,
channel: ChannelId,
) -> impl Future<Output = ()> + Send + 'this {
async {
_ = channel;
}
}
fn close_channel<'this>(
&'this self,
channel: ChannelId,
) -> impl Future<Output = ()> + Send + 'this {
async {
_ = channel;
}
}
fn handle<'this>(
&'this self,
id: Self::Id,
incoming: &'this mut Incoming<'_>,
outgoing: &'this mut Outgoing<'_>,
) -> impl Future<Output = Self::Response> + Send + 'this;
#[inline]
fn or<H>(self, other: H) -> Or<Self, H>
where
H: Handler,
{
Or::new(self, other)
}
}
#[derive(Clone)]
pub struct Or<A, B> {
first: A,
second: B,
}
impl<A, B> Or<A, B>
where
A: Handler,
B: Handler,
{
#[inline]
pub const fn new(first: A, second: B) -> Self {
const {
assert!(
<A::Id as Id>::IDS.is_disjoint(&<B::Id as Id>::IDS),
"The ids of the handlers combined with `ws::Or` overlap"
);
}
Self { first, second }
}
#[inline]
pub const fn first(&self) -> &A {
&self.first
}
#[inline]
pub const fn second(&self) -> &B {
&self.second
}
}
impl<A, B> Handler for Or<A, B>
where
A: Handler,
B: Handler,
{
type Id = Either<A::Id, B::Id>;
type Response = Either<A::Response, B::Response>;
#[inline]
fn open_channel<'this>(
&'this self,
channel: ChannelId,
) -> impl Future<Output = ()> + Send + 'this {
let first = self.first.open_channel(channel);
let second = self.second.open_channel(channel);
async move {
first.await;
second.await;
}
}
#[inline]
fn close_channel<'this>(
&'this self,
channel: ChannelId,
) -> impl Future<Output = ()> + Send + 'this {
let first = self.first.close_channel(channel);
let second = self.second.close_channel(channel);
async move {
first.await;
second.await;
}
}
#[inline]
fn handle<'this>(
&'this self,
id: Self::Id,
incoming: &'this mut Incoming<'_>,
outgoing: &'this mut Outgoing<'_>,
) -> impl Future<Output = Self::Response> + Send + 'this {
let future = match id {
Either::First(id) => Either::First(self.first.handle(id, incoming, outgoing)),
Either::Second(id) => Either::Second(self.second.handle(id, incoming, outgoing)),
};
async move {
match future {
Either::First(future) => Either::First(future.await),
Either::Second(future) => Either::Second(future.await),
}
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub enum Either<A, B> {
First(A),
Second(B),
}
impl<A, B> Id for Either<A, B>
where
A: Id,
B: Id,
{
const IDS: IdSet = IdSet {
__ids: &[],
__includes: &[&A::IDS, &B::IDS],
};
#[inline]
fn id(&self) -> MessageId {
match self {
Either::First(id) => id.id(),
Either::Second(id) => id.id(),
}
}
#[inline]
fn from_id(id: MessageId) -> Self {
match Self::try_from_id(id) {
Some(id) => id,
None => Either::First(A::from_id(id)),
}
}
#[inline]
fn try_from_id(id: MessageId) -> Option<Self> {
if let Some(id) = A::try_from_id(id) {
return Some(Either::First(id));
}
Some(Either::Second(B::try_from_id(id)?))
}
#[doc(hidden)]
fn __do_not_implement_id() {}
}
impl<A, B> IntoResponse for Either<A, B>
where
A: IntoResponse,
B: IntoResponse,
{
type Error = Either<A::Error, B::Error>;
#[inline]
fn into_response(self) -> Result<Response, Self::Error> {
match self {
Either::First(response) => response.into_response().map_err(Either::First),
Either::Second(response) => response.into_response().map_err(Either::Second),
}
}
}
impl<A, B> ResponseError for Either<A, B>
where
A: ResponseError,
B: ResponseError,
{
#[inline]
fn fmt_message(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Either::First(error) => error.fmt_message(f),
Either::Second(error) => error.fmt_message(f),
}
}
#[inline]
fn encode_payload(&self, payload: ErrorPayload<'_>) -> Result<(), Error> {
match self {
Either::First(error) => error.encode_payload(payload),
Either::Second(error) => error.encode_payload(payload),
}
}
}
struct Pinned<S> {
socket: S,
close_sleep: Sleep,
ping_sleep: Sleep,
}
impl<S> Pinned<S> {
#[inline]
fn project(self: Pin<&mut Self>) -> (Pin<&mut Sleep>, Pin<&mut Sleep>, Pin<&mut S>) {
unsafe {
let this = self.get_unchecked_mut();
(
Pin::new_unchecked(&mut this.close_sleep),
Pin::new_unchecked(&mut this.ping_sleep),
Pin::new_unchecked(&mut this.socket),
)
}
}
}
type HandlerOutput<H> = (
Result<<H as Handler>::Response, Error>,
bool,
RequestHeader,
Buf,
);
pub trait ChannelAllocator {
fn next(&self) -> Option<ChannelId>;
fn free(&self, channel: ChannelId);
}
pub struct Connect<S, H, C = Channels>
where
S: ServerImpl,
H: Handler,
{
handler: H,
socket: S::Socket,
channels: C,
seed: u64,
max_capacity: usize,
max_in_flight: usize,
max_outbound: usize,
max_channels: usize,
formats: Option<&'static [Format]>,
app: u32,
}
impl<S, H> Connect<S, H, Channels>
where
S: ServerImpl,
H: Handler,
{
#[inline]
pub(crate) fn new(socket: S::Socket, handler: H) -> Self {
Self {
handler,
socket,
channels: Channels::global(),
seed: DEFAULT_SEED,
max_capacity: MAX_CAPACITY,
max_in_flight: DEFAULT_MAX_IN_FLIGHT,
max_outbound: DEFAULT_MAX_OUTBOUND,
max_channels: DEFAULT_MAX_CHANNELS,
formats: None,
app: 0,
}
}
}
impl<S, H, C> Connect<S, H, C>
where
S: ServerImpl,
H: Handler,
{
#[inline]
pub fn seed(mut self, seed: u64) -> Self {
self.seed = seed;
self
}
#[inline]
pub fn with_channel_allocator<U>(self, channels: U) -> Connect<S, H, U>
where
U: ChannelAllocator,
{
Connect {
handler: self.handler,
socket: self.socket,
channels,
seed: self.seed,
max_capacity: self.max_capacity,
max_in_flight: self.max_in_flight,
max_outbound: self.max_outbound,
max_channels: self.max_channels,
formats: self.formats,
app: self.app,
}
}
#[inline]
pub fn handler(&self) -> &H {
&self.handler
}
#[inline]
pub fn with_formats(mut self, formats: &'static [Format]) -> Self {
self.formats = Some(formats);
self
}
#[inline]
pub fn accepts(&self, format: Format) -> bool {
accepts(self.formats, format)
}
#[inline]
pub fn app_version(mut self, version: u32) -> Self {
self.app = version;
self
}
#[inline]
pub fn max_capacity(mut self, max_capacity: usize) -> Self {
self.max_capacity = max_capacity;
self
}
#[inline]
pub fn with_max_capacity(self, max_capacity: usize) -> Self {
self.max_capacity(max_capacity)
}
#[inline]
pub fn max_in_flight(mut self, max_in_flight: usize) -> Self {
assert!(max_in_flight > 0, "max_in_flight must be non-zero");
self.max_in_flight = max_in_flight;
self
}
#[inline]
pub fn max_outbound(mut self, max_outbound: usize) -> Self {
self.max_outbound = max_outbound;
self
}
#[inline]
pub fn max_channels(mut self, max_channels: usize) -> Self {
self.max_channels = max_channels;
self
}
}
impl<S, H, C> Connect<S, H, C>
where
S: ServerImpl,
Error: From<S::Error>,
H: Handler,
C: ChannelAllocator,
{
pub async fn connect(self) -> Result<Server<S, H, C>, Error> {
let now = Instant::now();
let mut server = Server {
handler: self.handler,
pinned: Box::pin(Pinned {
socket: self.socket,
close_sleep: tokio::time::sleep_until(now + CLOSE_TIMEOUT),
ping_sleep: tokio::time::sleep_until(now + PING_TIMEOUT),
}),
channels: OwnedChannels::new(self.channels, self.max_channels),
closing: false,
pool: BufPool::new(self.max_capacity),
outbound: VecDeque::new(),
outbound_bytes: 0,
max_in_flight: self.max_in_flight,
max_outbound: self.max_outbound,
error: String::new(),
error_payload: Vec::new(),
error_typed: false,
last_ping: None,
rng: SmallRng::seed_from_u64(self.seed),
out: VecDeque::new(),
socket_send: false,
socket_flush: false,
socket_recv: true,
set: JoinSet::new(),
format: Format::DEFAULT,
mode: Mode::DEFAULT,
failure: None,
formats: self.formats,
app: self.app,
};
server.hello()?;
server.negotiate().await?;
Ok(server)
}
}
#[inline]
fn gone(failure: Option<Error>) -> Error {
failure.unwrap_or_else(|| Error::new(ErrorKind::NotNegotiated))
}
#[inline]
fn accepts(formats: Option<&'static [Format]>, format: Format) -> bool {
format.is_supported() && formats.is_none_or(|f| f.contains(&format))
}
pub struct Server<S, H, C = Channels>
where
S: ServerImpl,
H: Handler,
C: ChannelAllocator,
{
handler: H,
pinned: Pin<Box<Pinned<S::Socket>>>,
channels: OwnedChannels<C>,
closing: bool,
pool: BufPool,
outbound: VecDeque<Buf>,
outbound_bytes: usize,
max_in_flight: usize,
max_outbound: usize,
error: String,
error_payload: Vec<u8>,
error_typed: bool,
last_ping: Option<[u8; 4]>,
rng: SmallRng,
out: VecDeque<S::Message>,
socket_send: bool,
socket_flush: bool,
socket_recv: bool,
set: JoinSet<HandlerOutput<H>>,
format: Format,
mode: Mode,
failure: Option<Error>,
formats: Option<&'static [Format]>,
app: u32,
}
impl<S, H, C> Server<S, H, C>
where
S: ServerImpl,
H: Handler,
C: ChannelAllocator,
{
#[inline]
pub fn handler(&self) -> &H {
&self.handler
}
#[inline]
pub fn format(&self) -> Format {
self.format
}
#[inline]
pub fn mode(&self) -> Mode {
self.mode
}
#[inline]
pub fn accepts(&self, format: Format) -> bool {
accepts(self.formats, format)
}
}
impl<S, H, C> Server<S, H, C>
where
S: ServerImpl,
Error: From<S::Error>,
H: Handler,
C: ChannelAllocator,
{
async fn negotiate(&mut self) -> Result<(), Error> {
let mut negotiated = false;
let mut failure = None::<Error>;
loop {
let drained = self.out.is_empty() && !self.socket_flush;
if failure.is_some() && drained {
break;
}
if negotiated && drained && self.outbound.is_empty() {
break;
}
self.handle_send()?;
let result = {
let inner = Select::<S::Socket, H> {
pinned: self.pinned.as_mut(),
wants_socket_recv: self.socket_recv,
wants_socket_send: !self.socket_send,
wants_socket_flush: self.socket_flush,
set: &mut self.set,
};
inner.await
};
match result {
Output::Close => {
return Err(gone(failure));
}
Output::Ping => {
self.handle_ping()?;
}
Output::Recv(message) => {
let Some(message) = message else {
return Err(gone(failure));
};
let message = match message {
Ok(message) => message,
Err(error) => return Err(failure.unwrap_or_else(|| Error::from(error))),
};
match message {
Message::Text(bytes) => match self.handle_negotiate(Mode::Text, bytes) {
Ok(()) => negotiated = true,
Err(error) => failure = Some(error),
},
Message::Binary(bytes) => {
match self.handle_negotiate(Mode::Binary, bytes) {
Ok(()) => negotiated = true,
Err(error) => failure = Some(error),
}
}
Message::Ping(payload) => {
self.out.push_back(S::pong(payload));
}
Message::Pong(data) => {
self.handle_pong(data)?;
}
Message::Close => {
return Err(gone(failure));
}
}
}
Output::Send(result) => {
result?;
self.socket_send = true;
}
Output::Flushed(result) => {
result?;
self.socket_flush = false;
}
Output::Handle(..) => {
}
}
}
match failure {
Some(error) => Err(error),
None => Ok(()),
}
}
fn handle_negotiate(&mut self, mode: Mode, bytes: Bytes) -> Result<(), Error> {
let mut at = 0;
let header: RequestHeader = match format::decode_envelope(mode, &bytes, &mut at) {
Ok(header) => header,
Err(error) if error.unsupported_version().is_some() => {
self.out.push_back(S::close(
CLOSE_PROTOCOL_ERROR,
"Unsupported protocol version",
));
return Err(Error::new(ErrorKind::UnsupportedVersion { error }));
}
Err(error) if error.is_unknown_header() => {
self.out
.push_back(S::close(CLOSE_PROTOCOL_ERROR, "Unknown header"));
return Err(Error::new(ErrorKind::UnknownHeader { error }));
}
Err(error) => {
self.out
.push_back(S::close(CLOSE_PROTOCOL_ERROR, "Invalid request header"));
return Err(Error::new(ErrorKind::NegotiateHeader { error }));
}
};
if MessageId::new(header.id) != Some(MessageId::NEGOTIATE) {
self.out.push_back(S::close(
CLOSE_PROTOCOL_ERROR,
"Expected a negotiation as the first message",
));
return Err(Error::new(ErrorKind::ExpectedNegotiate { id: header.id }));
}
if self.app != 0 && header.app != 0 && header.app != self.app {
tracing::debug!(
server = self.app,
client = header.app,
"Incompatible client"
);
self.out.push_back(S::close(
CLOSE_POLICY_VIOLATION,
"Incompatible application version",
));
return Err(Error::new(ErrorKind::Incompatible {
server: self.app,
client: header.app,
}));
}
let Some(format) = Format::from_u8(header.format) else {
self.format_error_message(format_args!(
"Unknown format id {}, supported: {}",
header.format,
SupportedFormats(self.formats)
))?;
self.format = Format::DEFAULT;
return self.send_error(&header);
};
if !self.accepts(format) {
self.format_error_message(format_args!(
"Unsupported format `{format}`, supported: {}",
SupportedFormats(self.formats)
))?;
tracing::debug!(?format, "Rejected format");
self.format = Format::DEFAULT;
return self.send_error(&header);
}
if !mode.accepts(format) {
self.format_error_message(format_args!(
"Mode `{mode}` cannot carry the `{format}` format"
))?;
tracing::debug!(?mode, ?format, "Rejected mode");
self.format = Format::DEFAULT;
return self.send_error(&header);
}
tracing::debug!(?mode, ?format, "Negotiated format");
self.format = format;
self.mode = mode;
self.send_negotiated(&header, format)
}
fn send_negotiated(&mut self, header: &RequestHeader, format: Format) -> Result<(), Error> {
let mode = self.mode;
let buf = self.pool.with(|buf| {
let mut writer = buf.writer();
let result = writer.envelope(
mode,
&ResponseHeader {
version: VERSION,
serial: header.serial,
broadcast: 0,
error: 0,
format: format.to_u8(),
channel: header.channel,
app: 0,
},
);
result.map_err(Error::encode_connect_header)?;
writer.flush();
Ok::<_, Error>(())
})?;
self.push_outbound(buf);
Ok(())
}
pub async fn run(&mut self) -> Result<(), Error> {
loop {
if self.closing && self.out.is_empty() && self.outbound.is_empty() && !self.socket_flush
{
break;
}
self.handle_send()?;
let wants_socket_recv = self.socket_recv && !self.is_saturated();
let result = {
let inner = Select::<S::Socket, H> {
pinned: self.pinned.as_mut(),
wants_socket_recv,
wants_socket_send: !self.socket_send,
wants_socket_flush: self.socket_flush,
set: &mut self.set,
};
inner.await
};
match result {
Output::Close => {
if self.closing {
break;
}
self.out
.push_back(S::close(CLOSE_NORMAL, "connection timed out"));
self.begin_closing();
}
Output::Ping => {
self.handle_ping()?;
}
Output::Recv(message) => {
let Some(message) = message else {
self.socket_recv = false;
self.begin_closing();
continue;
};
match message? {
Message::Text(bytes) => {
if self.mode.is_text() {
self.handle_message(bytes).await?;
} else {
self.out.push_back(S::close(
CLOSE_PROTOCOL_ERROR,
"Unexpected text message",
));
self.begin_closing();
}
}
Message::Binary(bytes) => {
if self.mode.is_text() {
self.out.push_back(S::close(
CLOSE_PROTOCOL_ERROR,
"Unexpected binary message",
));
self.begin_closing();
} else {
self.handle_message(bytes).await?;
}
}
Message::Ping(payload) => {
self.out.push_back(S::pong(payload));
}
Message::Pong(data) => {
self.handle_pong(data)?;
}
Message::Close => {
self.begin_closing();
}
}
}
Output::Send(result) => {
if let Err(err) = result {
return Err(Error::from(err));
};
self.socket_send = true;
}
Output::Flushed(result) => {
if let Err(err) = result {
return Err(Error::from(err));
};
self.socket_flush = false;
}
Output::Handle(result, written, header, buf) => {
let err = 'err: {
let res = match result {
Ok(res) => res,
Err(error) => {
self.format_error(error)?;
break 'err true;
}
};
let res = match res.into_response() {
Ok(res) => res,
Err(error) => {
self.format_response_error(&error)?;
break 'err true;
}
};
if !res.handled {
self.format_error_message(format_args!(
"No support for request {}",
header.id
))?;
break 'err true;
}
if !written {
self.format_error_message(format_args!(
"Handler for request {} wrote no response",
header.id
))?;
break 'err true;
}
self.push_outbound(buf);
false
};
if err {
self.send_error(&header)?;
}
}
}
}
match self.failure.take() {
Some(error) => Err(error),
None => Ok(()),
}
}
pub fn broadcast<T>(&mut self, message: T) -> Result<(), Error>
where
T: Event,
{
self.broadcast_from(message, ChannelId::NONE)
}
pub fn broadcast_from<T>(&mut self, message: T, origin: ChannelId) -> Result<(), Error>
where
T: Event,
{
tracing::debug!(id = ?<T::Broadcast as Broadcast>::ID, "Broadcast");
let format = self.format;
let mode = self.mode;
let buf = self.pool.with(|buf| {
let mut writer = buf.writer();
writer
.envelope(
mode,
&ResponseHeader {
version: VERSION,
serial: 0,
broadcast: <T::Broadcast as Broadcast>::ID.get(),
error: 0,
format: format.to_u8(),
channel: origin,
app: 0,
},
)
.map_err(Error::encode_broadcast_header)?;
writer
.body(format, &message)
.map_err(Error::encode_broadcast)?;
writer.flush();
Ok::<_, Error>(())
})?;
self.push_outbound(buf);
Ok(())
}
fn hello(&mut self) -> Result<(), Error> {
tracing::debug!("Hello");
let mut buf = self.pool.get();
let result = (|| {
let mut writer = buf.writer();
writer
.envelope(
Mode::DEFAULT,
&ResponseHeader {
version: VERSION,
serial: 0,
broadcast: MessageId::SERVER_HELLO.get(),
error: 0,
format: 0,
channel: ChannelId::NONE,
app: self.app,
},
)
.map_err(Error::encode_broadcast_header)?;
writer.flush();
Ok::<_, Error>(())
})();
if result.is_err() {
self.pool.put(buf);
} else {
self.push_outbound(buf);
}
Ok(())
}
fn format_response_error<E>(&mut self, error: &E) -> Result<(), Error>
where
E: ResponseError,
{
self.format_error_message(ErrorMessageDisplay(error))?;
let mut encoded = false;
let payload = ErrorPayload {
format: self.format,
out: &mut self.error_payload,
encoded: &mut encoded,
};
match error.encode_payload(payload) {
Ok(()) => {
self.error_typed = encoded;
}
Err(error) => {
tracing::warn!(
?error,
"Failed to encode typed error, sending its message only"
);
}
}
Ok(())
}
fn format_error_message(&mut self, error: impl fmt::Display) -> Result<(), Error> {
self.error.clear();
self.error_typed = false;
if write!(self.error, "{error}").is_err() {
self.error.clear();
return Err(Error::new(ErrorKind::FormatError));
}
Ok(())
}
fn format_error(&mut self, error: impl core::error::Error) -> Result<(), Error> {
self.error.clear();
self.error_typed = false;
if write!(self.error, "{error:#}").is_err() {
self.error.clear();
return Err(Error::new(ErrorKind::FormatError));
}
Ok(())
}
#[tracing::instrument(skip(self, bytes))]
async fn handle_message(&mut self, bytes: Bytes) -> Result<(), Error> {
let mut at = 0;
let header: RequestHeader = match format::decode_envelope(self.mode, &bytes, &mut at) {
Ok(header) => header,
Err(error) if error.unsupported_version().is_some() => {
tracing::debug!(?error, "Unsupported protocol version");
self.out.push_back(S::close(
CLOSE_PROTOCOL_ERROR,
"Unsupported protocol version",
));
self.begin_closing();
self.failure
.get_or_insert_with(|| Error::new(ErrorKind::UnsupportedVersion { error }));
return Ok(());
}
Err(error) if error.is_unknown_header() => {
tracing::debug!(?error, "Unknown header");
self.out
.push_back(S::close(CLOSE_PROTOCOL_ERROR, "Unknown header"));
self.begin_closing();
self.failure
.get_or_insert_with(|| Error::new(ErrorKind::UnknownHeader { error }));
return Ok(());
}
Err(error) => {
tracing::debug!(?error, "Invalid request header");
self.out
.push_back(S::close(CLOSE_PROTOCOL_ERROR, "Invalid request header"));
self.begin_closing();
return Ok(());
}
};
let err = 'err: {
let Some(id) = MessageId::new(header.id) else {
self.format_error_message(format_args!("Unsupported message id {}", header.id))?;
break 'err true;
};
match id {
MessageId::CONNECT => {
let channel = match self.channels.open() {
Ok(channel) => channel,
Err(error) => {
self.format_error_message(format_args!("{error}"))?;
break 'err true;
}
};
self.handler.open_channel(channel).await;
let mut buf = self.pool.get();
let mode = self.mode;
let result = (|| {
let mut writer = buf.writer();
let result = writer.envelope(
mode,
&ResponseHeader {
version: VERSION,
serial: header.serial,
broadcast: 0,
error: 0,
format: 0,
channel,
app: 0,
},
);
result.map_err(Error::encode_connect_header)?;
writer.flush();
Ok::<_, Error>(())
})();
if result.is_err() {
self.pool.put(buf);
} else {
self.push_outbound(buf);
}
result?;
break 'err false;
}
MessageId::DISCONNECT => {
if self.channels.close(header.channel) {
self.handler.close_channel(header.channel).await;
} else {
tracing::debug!(channel = ?header.channel, "Ignoring close of a channel which is not open");
}
break 'err false;
}
MessageId::NEGOTIATE => {
let format = self.format;
self.format_error_message(format_args!(
"Format `{format}` has already been negotiated"
))?;
break 'err true;
}
_ => {
let Some(format) = Format::from_u8(header.format) else {
self.format_error_message(format_args!(
"Unknown format id {}",
header.format
))?;
break 'err true;
};
if !self.accepts(format) {
self.format_error_message(format_args!(
"Unsupported format `{format}`, supported: {}",
SupportedFormats(self.formats)
))?;
break 'err true;
}
if !self.mode.accepts(format) {
let mode = self.mode;
self.format_error_message(format_args!(
"Mode `{mode}` cannot carry the `{format}` format"
))?;
break 'err true;
}
let id = <H::Id as Id>::from_id(id);
self.handle_request(bytes, at, header, id, format);
return Ok(());
}
}
};
if err {
self.send_error(&header)?;
}
Ok(())
}
fn send_error(&mut self, header: &RequestHeader) -> Result<(), Error> {
let format = self.format;
let mode = self.mode;
let typed = mem::take(&mut self.error_typed);
let error = if typed {
MessageId::TYPED_ERROR_MESSAGE
} else {
MessageId::ERROR_MESSAGE
};
let buf = self.pool.with(|buf| {
let mut writer = buf.writer();
let result = writer.envelope(
mode,
&ResponseHeader {
version: VERSION,
serial: header.serial,
broadcast: 0,
error: error.get(),
format: format.to_u8(),
channel: header.channel,
app: 0,
},
);
result.map_err(Error::encode_error_message_header)?;
let result = writer.body(
format,
&ErrorMessage {
message: &self.error,
},
);
result.map_err(Error::encode_error_message)?;
if typed {
writer.raw(&self.error_payload);
}
writer.flush();
Ok::<_, Error>(())
})?;
self.push_outbound(buf);
Ok(())
}
fn push_outbound(&mut self, buf: Buf) {
self.outbound_bytes = self.outbound_bytes.saturating_add(buf.remaining());
self.outbound.push_back(buf);
}
fn is_saturated(&self) -> bool {
self.set.len() >= self.max_in_flight || self.outbound_bytes >= self.max_outbound
}
fn begin_closing(&mut self) {
if self.closing {
return;
}
self.closing = true;
let (close_sleep, _, _) = self.pinned.as_mut().project();
close_sleep.reset(Instant::now() + CLOSE_TIMEOUT);
}
#[tracing::instrument(skip(self))]
fn handle_ping(&mut self) -> Result<(), Error> {
let (_, mut ping_sleep, _) = self.pinned.as_mut().project();
let payload = self.rng.random::<u32>();
let payload = payload.to_ne_bytes();
self.last_ping = Some(payload);
tracing::debug!(data = ?&payload[..], "Sending ping");
self.out
.push_back(S::ping(Bytes::from_owner(Vec::from(payload))));
let now = Instant::now();
ping_sleep.as_mut().reset(now + PING_TIMEOUT);
Ok(())
}
#[tracing::instrument(skip(self, payload))]
fn handle_pong(&mut self, payload: Bytes) -> Result<(), Error> {
let (close_sleep, ping_sleep, _) = self.pinned.as_mut().project();
tracing::debug!(payload = ?&payload[..], "Pong");
let Some(expected) = self.last_ping else {
tracing::debug!("No ping sent");
return Ok(());
};
if expected[..] != payload[..] {
tracing::debug!(?expected, ?payload, "Pong doesn't match");
return Ok(());
}
let now = Instant::now();
close_sleep.reset(now + CLOSE_TIMEOUT);
ping_sleep.reset(now + PING_TIMEOUT);
self.last_ping = None;
Ok(())
}
#[tracing::instrument(skip(self))]
fn handle_send(&mut self) -> Result<(), Error> {
let (_, _, mut socket) = self.pinned.as_mut().project();
if self.socket_send {
if let Some(message) = self.out.pop_front() {
socket.as_mut().start_send(message)?;
self.socket_flush = true;
self.socket_send = false;
}
}
while self.socket_send {
let Some(buf) = self.outbound.front_mut() else {
break;
};
let Some(frame) = buf.read()? else {
if let Some(buf) = self.outbound.pop_front() {
self.pool.put(buf);
}
if self.outbound.is_empty() {
self.outbound_bytes = 0;
}
continue;
};
self.outbound_bytes = self
.outbound_bytes
.saturating_sub(frame.len().saturating_add(mem::size_of::<u32>()));
let message = match self.mode {
Mode::Binary => S::binary(frame),
Mode::Text => S::text(frame),
};
socket.as_mut().start_send(message)?;
self.socket_flush = true;
self.socket_send = false;
break;
}
Ok(())
}
fn handle_request(
&mut self,
bytes: Bytes,
offset: usize,
header: RequestHeader,
id: H::Id,
format: Format,
) {
tracing::debug!(header.serial, ?id, ?format, "Got request");
let mut buf = self.pool.get();
let handler = self.handler.clone();
let mode = self.mode;
self.set.spawn(async move {
if offset > bytes.len() {
let kind = ErrorKind::OutOfBounds {
offset,
len: bytes.len(),
};
return (Err(Error::new(kind)), false, header, buf);
}
let mut incoming = Incoming {
error: None,
buf: &bytes,
at: offset,
format,
channel: header.channel,
};
let mut outgoing = Outgoing {
serial: Some(header.serial),
error: None,
buf: &mut buf,
format,
mode,
channel: header.channel,
};
let response = handler.handle(id, &mut incoming, &mut outgoing).await;
let written = outgoing.serial.is_none();
if let Some(error) = incoming.error.take() {
return (Err(Error::incoming(error)), written, header, buf);
}
if let Some(error) = outgoing.error.take() {
return (Err(Error::outgoing(error)), written, header, buf);
}
(Ok(response), written, header, buf)
});
}
}
enum Output<E, R> {
Close,
Ping,
Recv(Option<Result<Message, E>>),
Send(Result<(), E>),
Flushed(Result<(), E>),
Handle(Result<R, Error>, bool, RequestHeader, Buf),
}
struct Select<'a, S, H>
where
H: Handler,
{
pinned: Pin<&'a mut Pinned<S>>,
wants_socket_recv: bool,
wants_socket_send: bool,
wants_socket_flush: bool,
set: &'a mut JoinSet<HandlerOutput<H>>,
}
impl<S, H> Future for Select<'_, S, H>
where
S: SocketImpl,
H: Handler,
{
type Output = Output<S::Error, H::Response>;
#[inline]
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let this = Pin::get_mut(self);
let (close, ping, mut socket) = this.pinned.as_mut().project();
let wants_socket_recv = this.wants_socket_recv;
let wants_socket_send = this.wants_socket_send;
let wants_socket_flush = this.wants_socket_flush;
let set = &mut this.set;
if close.poll(cx).is_ready() {
return Poll::Ready(Output::Close);
}
if ping.poll(cx).is_ready() {
return Poll::Ready(Output::Ping);
}
if wants_socket_recv {
if let Poll::Ready(output) = socket.as_mut().poll_next(cx) {
return Poll::Ready(Output::Recv(output));
}
}
if wants_socket_send {
if let Poll::Ready(result) = socket.as_mut().poll_ready(cx) {
return Poll::Ready(Output::Send(result));
}
}
if wants_socket_flush {
if let Poll::Ready(result) = socket.as_mut().poll_flush(cx) {
return Poll::Ready(Output::Flushed(result));
}
}
if let Poll::Ready(Some(output)) = set.poll_join_next(cx) {
let output = match output {
Ok(output) => output,
Err(error) => {
tracing::debug!(?error, "Join error in handler task");
return Poll::Ready(Output::Close);
}
};
let (result, written, header, buf) = output;
return Poll::Ready(Output::Handle(result, written, header, buf));
}
Poll::Pending
}
}
pub struct Incoming<'de> {
error: Option<format::Error>,
buf: &'de [u8],
at: usize,
format: Format,
channel: ChannelId,
}
impl<'de> Incoming<'de> {
pub fn channel(&self) -> ChannelId {
self.channel
}
#[inline]
pub fn format(&self) -> Format {
self.format
}
#[inline]
pub fn read<T>(&mut self) -> Option<T>
where
T: DecodeBody<'de>,
{
match self.format.decode(self.buf, &mut self.at) {
Ok(value) => Some(value),
Err(error) => {
self.error = Some(error);
None
}
}
}
}
pub struct Outgoing<'a> {
serial: Option<u32>,
error: Option<format::Error>,
buf: &'a mut Buf,
format: Format,
mode: Mode,
channel: ChannelId,
}
impl Outgoing<'_> {
#[inline]
pub fn format(&self) -> Format {
self.format
}
#[inline]
pub fn mode(&self) -> Mode {
self.mode
}
pub fn write<T>(&mut self, value: T)
where
T: EncodeBody,
{
let Some(serial) = self.serial.take() else {
return;
};
let mut writer = self.buf.writer();
let result = writer.envelope(
self.mode,
&ResponseHeader {
version: VERSION,
serial,
broadcast: 0,
error: 0,
format: self.format.to_u8(),
channel: self.channel,
app: 0,
},
);
if let Err(error) = result {
self.error = Some(error);
return;
}
if let Err(error) = writer.body(self.format, &value) {
self.error = Some(error);
}
writer.flush();
}
}
#[inline]
fn scramble_channel(x: u16) -> u16 {
let x = x.wrapping_mul(0x9285);
x ^ (x >> 8)
}
#[inline]
#[cfg(test)]
fn unscramble_channel(x: u16) -> u16 {
let x = x ^ (x >> 8);
x.wrapping_mul(0x964d)
}
#[test]
fn test_scramble() {
assert_eq!(scramble_channel(0), 0);
assert_eq!(unscramble_channel(0), 0);
for i in 1..=u16::MAX {
let scrambled = scramble_channel(i);
let unscrambled = unscramble_channel(scrambled);
assert_eq!(i, unscrambled, "Failed to unscramble channel id");
}
}
#[derive(Default)]
struct ChannelsInner {
last: u16,
free: VecDeque<NonZeroU16>,
}
static GLOBAL_CHANNELS: LazyLock<Channels> = LazyLock::new(Channels::new);
#[derive(Clone)]
pub struct Channels {
inner: Arc<Mutex<ChannelsInner>>,
}
impl Channels {
#[inline]
pub fn new() -> Self {
Self {
inner: Arc::new(Mutex::new(ChannelsInner::default())),
}
}
#[inline]
pub fn global() -> Self {
GLOBAL_CHANNELS.clone()
}
#[inline]
fn lock(&self) -> MutexGuard<'_, ChannelsInner> {
self.inner.lock().unwrap_or_else(PoisonError::into_inner)
}
}
impl Default for Channels {
#[inline]
fn default() -> Self {
Self::new()
}
}
impl ChannelAllocator for Channels {
#[inline]
fn next(&self) -> Option<ChannelId> {
let mut inner = self.lock();
if let Some(id) = inner.free.pop_front() {
return Some(ChannelId::from_u16(id.get()));
}
let id = NonZeroU16::new(inner.last.wrapping_add(1))?;
inner.last = id.get();
tracing::debug!(?id, "Allocated channel id");
Some(ChannelId::from_u16(scramble_channel(id.get())))
}
#[inline]
fn free(&self, id: ChannelId) {
tracing::debug!(?id, "Freeing channel id");
if let Some(id) = NonZeroU16::new(id.raw()) {
self.lock().free.push_back(id);
}
}
}
#[derive(Debug, PartialEq, Eq)]
enum OpenChannelError {
Limit(usize),
Exhausted,
}
impl fmt::Display for OpenChannelError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Self::Limit(limit) => write!(f, "Too many open channels, at most {limit} are allowed"),
Self::Exhausted => f.write_str("No channel id is available"),
}
}
}
struct OwnedChannels<C>
where
C: ChannelAllocator,
{
allocator: C,
open: BTreeSet<ChannelId>,
max: usize,
}
impl<C> OwnedChannels<C>
where
C: ChannelAllocator,
{
#[inline]
fn new(allocator: C, max: usize) -> Self {
Self {
allocator,
open: BTreeSet::new(),
max,
}
}
fn open(&mut self) -> Result<ChannelId, OpenChannelError> {
if self.open.len() >= self.max {
return Err(OpenChannelError::Limit(self.max));
}
let id = self.allocator.next().ok_or(OpenChannelError::Exhausted)?;
if id == ChannelId::NONE || !self.open.insert(id) {
return Err(OpenChannelError::Exhausted);
}
Ok(id)
}
fn close(&mut self, id: ChannelId) -> bool {
if !self.open.remove(&id) {
return false;
}
self.allocator.free(id);
true
}
}
impl<C> Drop for OwnedChannels<C>
where
C: ChannelAllocator,
{
fn drop(&mut self) {
for id in mem::take(&mut self.open) {
self.allocator.free(id);
}
}
}
struct SupportedFormats(Option<&'static [Format]>);
impl fmt::Display for SupportedFormats {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
let mut first = true;
for format in Format::supported() {
if self.0.is_some_and(|formats| !formats.contains(&format)) {
continue;
}
if !first {
f.write_str(", ")?;
}
write!(f, "`{format}`")?;
first = false;
}
if first {
f.write_str("none")?;
}
Ok(())
}
}
#[cfg(test)]
mod tests {
use core::cell::Cell;
use std::thread_local;
use super::*;
const CLOSE_BUDGET: usize = 4;
const RECV_BUDGET: usize = 4;
thread_local! {
static CLOSE_FRAMES: Cell<usize> = const { Cell::new(0) };
}
#[derive(Debug)]
enum TestError {}
impl From<TestError> for Error {
#[inline]
fn from(error: TestError) -> Self {
match error {}
}
}
#[derive(Debug, PartialEq, Eq)]
enum TestMessage {
Ping(Bytes),
Pong(Bytes),
Binary(Vec<u8>),
Text(Vec<u8>),
Close(u16),
}
#[derive(Default)]
struct TestSocket {
sent: Vec<TestMessage>,
ended: bool,
polls_after_end: usize,
write_blocked: bool,
incoming: VecDeque<Message>,
}
impl socket_sealed::Sealed for TestSocket {}
impl SocketImpl for TestSocket {
type Message = TestMessage;
type Error = TestError;
fn poll_next(
self: Pin<&mut Self>,
_: &mut Context<'_>,
) -> Poll<Option<Result<Message, Self::Error>>> {
let this = unsafe { Pin::get_unchecked_mut(self) };
if let Some(message) = this.incoming.pop_front() {
return Poll::Ready(Some(Ok(message)));
}
if !this.ended {
return Poll::Pending;
}
this.polls_after_end += 1;
assert!(
this.polls_after_end <= RECV_BUDGET,
"`Server::run` is spinning: an ended stream was polled {} times without any progress",
this.polls_after_end
);
Poll::Ready(None)
}
fn poll_ready(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
if unsafe { Pin::get_unchecked_mut(self) }.write_blocked {
return Poll::Pending;
}
Poll::Ready(Ok(()))
}
fn start_send(self: Pin<&mut Self>, item: Self::Message) -> Result<(), Self::Error> {
unsafe { Pin::get_unchecked_mut(self).sent.push(item) };
Ok(())
}
fn poll_flush(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
if unsafe { Pin::get_unchecked_mut(self) }.write_blocked {
return Poll::Pending;
}
Poll::Ready(Ok(()))
}
}
struct TestServerImpl;
impl server_sealed::Sealed for TestServerImpl {}
impl ServerImpl for TestServerImpl {
type Error = TestError;
type Message = TestMessage;
type Socket = TestSocket;
fn ping(data: Bytes) -> Self::Message {
TestMessage::Ping(data)
}
fn pong(data: Bytes) -> Self::Message {
TestMessage::Pong(data)
}
fn binary(data: &[u8]) -> Self::Message {
TestMessage::Binary(data.to_vec())
}
fn text(data: &[u8]) -> Self::Message {
TestMessage::Text(data.to_vec())
}
fn close(code: u16, _: &str) -> Self::Message {
let frames = CLOSE_FRAMES.with(|frames| {
let count = frames.get() + 1;
frames.set(count);
count
});
assert!(
frames <= CLOSE_BUDGET,
"`Server::run` is spinning: {frames} close frames without any progress"
);
TestMessage::Close(code)
}
}
#[derive(Debug)]
struct TestId;
impl Id for TestId {
const IDS: IdSet = IdSet::EMPTY;
fn id(&self) -> MessageId {
MessageId::NEGOTIATE
}
fn from_id(_: MessageId) -> Self {
Self
}
fn try_from_id(_: MessageId) -> Option<Self> {
None
}
fn __do_not_implement_id() {}
}
#[derive(Clone, Default)]
struct TestHandler {
block: bool,
}
impl Handler for TestHandler {
type Id = TestId;
type Response = bool;
async fn handle(&self, _: Self::Id, _: &mut Incoming<'_>, _: &mut Outgoing<'_>) -> bool {
if self.block {
core::future::pending::<()>().await;
}
false
}
}
fn timed_out_server(elapsed: Duration) -> Server<TestServerImpl, TestHandler, Channels> {
server_with(Instant::now() - elapsed)
}
fn socket(server: &mut Server<TestServerImpl, TestHandler, Channels>) -> &mut TestSocket {
let (_, _, socket) = server.pinned.as_mut().project();
unsafe { Pin::get_unchecked_mut(socket) }
}
fn response(server: &Server<TestServerImpl, TestHandler, Channels>) -> Buf {
server
.pool
.with(|buf| {
let mut writer = buf.writer();
writer
.envelope(
Mode::DEFAULT,
&ResponseHeader {
version: VERSION,
serial: 1,
broadcast: 0,
error: 0,
format: Format::DEFAULT.to_u8(),
channel: ChannelId::NONE,
app: 0,
},
)
.map_err(Error::encode_connect_header)?;
writer.flush();
Ok::<_, Error>(())
})
.unwrap()
}
fn server_with(close_deadline: Instant) -> Server<TestServerImpl, TestHandler, Channels> {
let now = Instant::now();
Server {
handler: TestHandler::default(),
pinned: Box::pin(Pinned {
socket: TestSocket::default(),
close_sleep: tokio::time::sleep_until(close_deadline),
ping_sleep: tokio::time::sleep_until(now + PING_TIMEOUT),
}),
channels: OwnedChannels::new(Channels::new(), DEFAULT_MAX_CHANNELS),
closing: false,
pool: BufPool::new(MAX_CAPACITY),
outbound: VecDeque::new(),
outbound_bytes: 0,
max_in_flight: DEFAULT_MAX_IN_FLIGHT,
max_outbound: DEFAULT_MAX_OUTBOUND,
error: String::new(),
error_payload: Vec::new(),
error_typed: false,
last_ping: None,
rng: SmallRng::seed_from_u64(DEFAULT_SEED),
out: VecDeque::new(),
socket_send: false,
socket_flush: false,
socket_recv: true,
set: JoinSet::new(),
format: Format::DEFAULT,
mode: Mode::DEFAULT,
failure: None,
formats: None,
app: 0,
}
}
#[tokio::test]
async fn close_deadline_winds_down_once() {
let mut server = timed_out_server(Duration::from_secs(1));
server.run().await.unwrap();
assert_eq!(socket(&mut server).sent, [TestMessage::Close(CLOSE_NORMAL)]);
assert!(server.out.is_empty());
assert!(server.closing);
assert!(!server.socket_flush);
}
#[tokio::test]
async fn close_deadline_while_closing_gives_up() {
let mut server = timed_out_server(Duration::from_secs(1));
server.closing = true;
server.out.push_back(TestMessage::Close(CLOSE_NORMAL));
server.run().await.unwrap();
assert!(socket(&mut server).sent.is_empty());
assert_eq!(server.out.len(), 1);
}
#[tokio::test]
async fn ended_stream_drains_outbound() {
let mut server = server_with(Instant::now() + CLOSE_TIMEOUT);
server.closing = true;
server.socket_flush = true;
server.outbound.push_back(response(&server));
socket(&mut server).ended = true;
server.run().await.unwrap();
let sent = socket(&mut server).sent.as_slice();
assert!(matches!(sent, [TestMessage::Binary(..)]));
assert!(server.outbound.is_empty());
assert!(!server.socket_flush);
assert!(!server.socket_recv);
}
#[tokio::test]
async fn blocked_write_winds_down_on_deadline() {
let mut server = server_with(Instant::now() + Duration::from_millis(50));
server.closing = true;
server.socket_flush = true;
server.outbound.push_back(response(&server));
{
let socket = socket(&mut server);
socket.ended = true;
socket.write_blocked = true;
}
server.run().await.unwrap();
assert!(socket(&mut server).sent.is_empty());
assert_eq!(server.outbound.len(), 1);
assert!(!server.socket_recv);
}
fn request(mode: Mode, serial: u32) -> Message {
let mut out = Vec::new();
let header = RequestHeader {
version: VERSION,
serial,
id: 1,
format: Format::DEFAULT.to_u8(),
channel: ChannelId::NONE,
app: 0,
};
format::encode_envelope(mode, &mut out, &header).unwrap();
match mode {
Mode::Binary => Message::Binary(Bytes::from(out)),
Mode::Text => Message::Text(Bytes::from(out)),
}
}
fn flooded_server(count: u32) -> Server<TestServerImpl, TestHandler, Channels> {
let mut server = server_with(Instant::now() + CLOSE_TIMEOUT);
let mode = server.mode;
socket(&mut server)
.incoming
.extend((1..=count).map(|serial| request(mode, serial)));
server
}
async fn run_until_idle(server: &mut Server<TestServerImpl, TestHandler, Channels>) {
let result = tokio::time::timeout(Duration::from_millis(100), server.run()).await;
assert!(result.is_err(), "The connection should still be up");
}
#[tokio::test]
async fn in_flight_handlers_stop_reading() {
let mut server = flooded_server(10);
server.handler = TestHandler { block: true };
server.max_in_flight = 3;
run_until_idle(&mut server).await;
assert_eq!(server.set.len(), 3);
assert_eq!(socket(&mut server).incoming.len(), 7);
}
#[tokio::test]
async fn queued_output_stops_reading() {
let mut server = flooded_server(10);
server.max_in_flight = 1;
server.max_outbound = 1;
socket(&mut server).write_blocked = true;
run_until_idle(&mut server).await;
assert!(server.set.is_empty());
assert_eq!(server.outbound.len(), 1);
assert!(server.outbound_bytes > 0);
assert_eq!(socket(&mut server).incoming.len(), 9);
assert!(socket(&mut server).sent.is_empty());
}
#[tokio::test]
async fn reading_resumes_once_drained() {
let mut server = flooded_server(10);
server.max_in_flight = 1;
server.max_outbound = 1;
run_until_idle(&mut server).await;
assert!(server.set.is_empty());
assert!(server.outbound.is_empty());
assert_eq!(server.outbound_bytes, 0);
assert!(socket(&mut server).incoming.is_empty());
let sent = &socket(&mut server).sent;
assert_eq!(sent.len(), 10);
assert!(
sent.iter()
.all(|m| matches!(m, TestMessage::Binary(..) | TestMessage::Text(..)))
);
}
#[test]
fn channels_are_distinct_within_a_pool() {
let pool = Channels::new();
let other = pool.clone();
let ids = (0..64)
.map(|n| if n % 2 == 0 { &pool } else { &other })
.map(|pool| pool.next().unwrap())
.collect::<BTreeSet<_>>();
assert_eq!(ids.len(), 64);
assert!(!ids.contains(&ChannelId::NONE));
let a = Channels::global().next().unwrap();
let b = Channels::global().next().unwrap();
assert_ne!(a, b);
Channels::global().free(a);
Channels::global().free(b);
}
#[test]
fn owned_channels_only_close_their_own() {
let pool = Channels::new();
let mut first = OwnedChannels::new(pool.clone(), DEFAULT_MAX_CHANNELS);
let mut second = OwnedChannels::new(pool.clone(), DEFAULT_MAX_CHANNELS);
let a = first.open().unwrap();
assert!(!second.close(a));
assert!(!second.close(ChannelId::NONE));
assert!(!second.close(ChannelId::from_u16(0x1234)));
assert!(pool.lock().free.is_empty());
let b = second.open().unwrap();
assert_ne!(a, b);
assert!(first.close(a));
assert!(!first.close(a));
assert_eq!(pool.lock().free.len(), 1);
}
#[test]
fn owned_channels_free_on_drop() {
let pool = Channels::new();
let mut first = OwnedChannels::new(pool.clone(), DEFAULT_MAX_CHANNELS);
let mut second = OwnedChannels::new(pool.clone(), DEFAULT_MAX_CHANNELS);
let held = [first.open().unwrap(), first.open().unwrap()];
let other = second.open().unwrap();
drop(first);
let freed = pool
.lock()
.free
.iter()
.map(|id| ChannelId::from_u16(id.get()))
.collect::<BTreeSet<_>>();
assert_eq!(freed, BTreeSet::from(held));
let reused = [second.open().unwrap(), second.open().unwrap()];
assert_eq!(BTreeSet::from(reused), BTreeSet::from(held));
assert!(!reused.contains(&other));
}
#[test]
fn owned_channels_limit() {
let mut channels = OwnedChannels::new(Channels::new(), 2);
let a = channels.open().unwrap();
channels.open().unwrap();
assert_eq!(channels.open(), Err(OpenChannelError::Limit(2)));
assert!(channels.close(a));
channels.open().unwrap();
let mut none = OwnedChannels::new(Channels::new(), 0);
assert_eq!(none.open(), Err(OpenChannelError::Limit(0)));
}
#[test]
fn owned_channels_refuse_a_broken_allocator() {
#[derive(Default)]
struct Broken {
next: std::sync::Mutex<Vec<ChannelId>>,
freed: std::sync::Mutex<Vec<ChannelId>>,
}
impl ChannelAllocator for &Broken {
fn next(&self) -> Option<ChannelId> {
self.next.lock().unwrap().pop()
}
fn free(&self, channel: ChannelId) {
self.freed.lock().unwrap().push(channel);
}
}
let one = ChannelId::from_u16(1);
let broken = Broken::default();
broken
.next
.lock()
.unwrap()
.extend([one, ChannelId::NONE, one]);
let mut channels = OwnedChannels::new(&broken, DEFAULT_MAX_CHANNELS);
assert_eq!(channels.open(), Ok(one));
assert_eq!(channels.open(), Err(OpenChannelError::Exhausted));
assert_eq!(channels.open(), Err(OpenChannelError::Exhausted));
assert_eq!(channels.open(), Err(OpenChannelError::Exhausted));
drop(channels);
assert_eq!(*broken.freed.lock().unwrap(), [one]);
}
}