Skip to main content

musli_web/
ws.rs

1//! The server side of the websocket protocol.
2//!
3//! See [`server()`] for how to use with [axum].
4//!
5//! A connection starts life as a [`Connect`], which cannot send anything.
6//! [`Connect::connect`] performs the [negotiation protocol] and hands back the
7//! [`Server`], so by the time there is anything able to write a message the
8//! [`Format`] it will be encoded with has been agreed with the client. A client
9//! which does not negotiate never gets a [`Server`] at all.
10//!
11//! Handlers are implemented via the [`Handler`] trait, which allows returning
12//! various forms of responses dictated through the [`IntoResponse`] trait. This
13//! is primarily implemented for `bool`, where returning `false` indicates that
14//! the given request kind is not supported.
15//!
16//! You can also return custom error for a handler by having it return anything
17//! that implements [`fmt::Display`]:
18//!
19//! ```
20//! use musli_web::api::MessageId;
21//! use musli_web::ws;
22//!
23//! mod api {
24//!     use musli::{Decode, Encode};
25//!     use musli_web::api;
26//!
27//!     #[derive(Encode, Decode)]
28//!     pub struct HelloRequest<'de> {
29//!         pub message: &'de str,
30//!     }
31//!
32//!     #[derive(Encode, Decode)]
33//!     pub struct HelloResponse<'de> {
34//!         pub message: &'de str,
35//!     }
36//!
37//!     #[derive(Encode, Decode)]
38//!     pub struct TickEvent<'de> {
39//!         pub message: &'de str,
40//!         pub tick: u32,
41//!     }
42//!
43//!     api::define! {
44//!         pub type Hello;
45//!
46//!         impl Endpoint for Hello {
47//!             impl<'de> Request for HelloRequest<'de>;
48//!             type Response<'de> = HelloResponse<'de>;
49//!         }
50//!
51//!         pub type Tick;
52//!
53//!         impl Broadcast for Tick {
54//!             impl<'de> Event for TickEvent<'de>;
55//!         }
56//!     }
57//! }
58//!
59//! #[derive(Debug, Clone)]
60//! enum Broadcast {
61//!     Tick { tick: u32 },
62//! }
63//!
64//! #[derive(Clone)]
65//! struct MyHandler;
66//!
67//! impl ws::Handler for MyHandler {
68//!     type Id = api::Request;
69//!     type Response = Option<()>;
70//!
71//!     async fn handle(
72//!         &self,
73//!         id: Self::Id,
74//!         incoming: &mut ws::Incoming<'_>,
75//!         outgoing: &mut ws::Outgoing<'_>,
76//!     ) -> Self::Response {
77//!         tracing::info!("Handling: {id:?}");
78//!
79//!         match id {
80//!             api::Request::Hello => {
81//!                 let request = incoming.read::<api::HelloRequest<'_>>()?;
82//!
83//!                 outgoing.write(api::HelloResponse {
84//!                     message: request.message,
85//!                 });
86//!
87//!                 Some(())
88//!             }
89//!             api::Request::Unknown(id) => {
90//!                 tracing::info!("Unknown request id: {}", id.get());
91//!                 None
92//!             }
93//!         }
94//!     }
95//! }
96//! ```
97//!
98//! [`server()`]: crate::axum08::server
99//! [axum]: <https://docs.rs/axum>
100//! [negotiation protocol]: crate::api#negotiating-the-format
101
102use core::convert::Infallible;
103use core::fmt::{self, Write};
104use core::future::Future;
105use core::num::NonZeroU16;
106use core::pin::Pin;
107use core::task::{Context, Poll};
108
109use alloc::boxed::Box;
110use alloc::collections::VecDeque;
111use alloc::string::String;
112use alloc::sync::Arc;
113use alloc::vec::Vec;
114
115use bytes::Bytes;
116use rand::prelude::*;
117use rand::rngs::SmallRng;
118use tokio::sync::Mutex;
119use tokio::task::JoinSet;
120use tokio::time::{Duration, Instant, Sleep};
121
122use crate::Buf;
123use crate::api::{
124    Broadcast, ChannelId, DecodeBody, EncodeBody, ErrorMessage, Event, Format, Id, MessageId,
125    RequestHeader, ResponseHeader,
126};
127use crate::buf::{BufPool, InvalidFrame};
128use crate::format;
129
130const MAX_CAPACITY: usize = 1048576;
131const CLOSE_NORMAL: u16 = 1000;
132const CLOSE_PROTOCOL_ERROR: u16 = 1002;
133const CLOSE_TIMEOUT: Duration = Duration::from_secs(30);
134const PING_TIMEOUT: Duration = Duration::from_secs(10);
135const DEFAULT_SEED: u64 = 0xdeadbeef;
136
137/// A websocket message.
138#[derive(Debug)]
139pub(crate) enum Message {
140    /// A text message.
141    Text,
142    /// A binary message.
143    Binary(Bytes),
144    /// A ping message.
145    Ping(Bytes),
146    /// A pong message.
147    Pong(Bytes),
148    /// A close message.
149    Close,
150}
151
152pub(crate) mod socket_sealed {
153    pub trait Sealed {}
154}
155
156pub(crate) trait SocketImpl
157where
158    Self: self::socket_sealed::Sealed,
159{
160    #[doc(hidden)]
161    type Message;
162
163    #[doc(hidden)]
164    type Error: fmt::Debug;
165
166    #[doc(hidden)]
167    fn poll_next(
168        self: Pin<&mut Self>,
169        ctx: &mut Context<'_>,
170    ) -> Poll<Option<Result<Message, Self::Error>>>;
171
172    #[doc(hidden)]
173    fn poll_ready(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>>;
174
175    #[doc(hidden)]
176    fn start_send(self: Pin<&mut Self>, item: Self::Message) -> Result<(), Self::Error>;
177
178    #[doc(hidden)]
179    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>>;
180}
181
182pub(crate) mod server_sealed {
183    pub trait Sealed {}
184}
185
186/// The details of how a [`Server`] is implemented.
187///
188/// See [`AxumServer`] for an example.
189///
190/// [`AxumServer`]: crate::axum08::AxumServer
191pub trait ServerImpl
192where
193    Self: self::server_sealed::Sealed,
194{
195    #[doc(hidden)]
196    type Error: fmt::Debug;
197
198    #[doc(hidden)]
199    type Message;
200
201    #[doc(hidden)]
202    #[allow(private_bounds)]
203    type Socket: SocketImpl<Message = Self::Message, Error = Self::Error>;
204
205    #[doc(hidden)]
206    fn ping(data: Bytes) -> Self::Message;
207
208    #[doc(hidden)]
209    fn pong(data: Bytes) -> Self::Message;
210
211    #[doc(hidden)]
212    fn binary(data: &[u8]) -> Self::Message;
213
214    #[doc(hidden)]
215    fn close(code: u16, reason: &str) -> Self::Message;
216}
217
218#[derive(Debug)]
219enum ErrorKind {
220    #[cfg(feature = "axum-core05")]
221    AxumCore05 {
222        error: axum_core05::Error,
223    },
224    FormatError,
225    InvalidFrame {
226        error: InvalidFrame,
227    },
228    Incoming {
229        error: format::Error,
230    },
231    Outgoing {
232        error: format::Error,
233    },
234    EncodeBroadcastHeader {
235        error: format::Error,
236    },
237    EncodeBroadcast {
238        error: format::Error,
239    },
240    EncodeConnectHeader {
241        error: format::Error,
242    },
243    ErrorMessageHeader {
244        error: format::Error,
245    },
246    ErrorMessage {
247        error: format::Error,
248    },
249    OutOfBounds {
250        offset: usize,
251        len: usize,
252    },
253    /// The connection went away before the format had been negotiated.
254    NotNegotiated,
255    /// The client sent a message other than a negotiation as its first message.
256    ExpectedNegotiate {
257        id: u16,
258    },
259    /// The client sent a malformed envelope during negotiation.
260    NegotiateHeader {
261        error: format::Error,
262    },
263}
264
265/// The error produced by the server side of the websocket protocol
266#[derive(Debug)]
267pub struct Error {
268    kind: ErrorKind,
269}
270
271impl Error {
272    #[inline]
273    const fn new(kind: ErrorKind) -> Self {
274        Self { kind }
275    }
276
277    pub(crate) fn incoming(error: format::Error) -> Self {
278        Self::new(ErrorKind::Incoming { error })
279    }
280
281    pub(crate) fn outgoing(error: format::Error) -> Self {
282        Self::new(ErrorKind::Outgoing { error })
283    }
284
285    pub(crate) fn encode_broadcast_header(error: format::Error) -> Self {
286        Self::new(ErrorKind::EncodeBroadcastHeader { error })
287    }
288
289    pub(crate) fn encode_broadcast(error: format::Error) -> Self {
290        Self::new(ErrorKind::EncodeBroadcast { error })
291    }
292
293    pub(crate) fn encode_connect_header(error: format::Error) -> Self {
294        Self::new(ErrorKind::EncodeConnectHeader { error })
295    }
296
297    pub(crate) fn encode_error_message_header(error: format::Error) -> Self {
298        Self::new(ErrorKind::ErrorMessageHeader { error })
299    }
300
301    pub(crate) fn encode_error_message(error: format::Error) -> Self {
302        Self::new(ErrorKind::ErrorMessage { error })
303    }
304}
305
306impl fmt::Display for Error {
307    #[inline]
308    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
309        match &self.kind {
310            #[cfg(feature = "axum-core05")]
311            ErrorKind::AxumCore05 { .. } => write!(f, "Error in axum-core"),
312            ErrorKind::FormatError => write!(f, "Error formatting error response"),
313            ErrorKind::InvalidFrame { error } => error.fmt(f),
314            ErrorKind::Incoming { .. } => {
315                write!(f, "Encoding error when decoding incoming message")
316            }
317            ErrorKind::Outgoing { .. } => {
318                write!(f, "Encoding error when encoding outgoing message")
319            }
320            ErrorKind::EncodeBroadcastHeader { .. } => {
321                write!(f, "Encoding error when encoding broadcast header")
322            }
323            ErrorKind::EncodeBroadcast { .. } => {
324                write!(f, "Encoding error when broadcasting message")
325            }
326            ErrorKind::EncodeConnectHeader { .. } => {
327                write!(f, "Encoding error when encoding connect header")
328            }
329            ErrorKind::ErrorMessageHeader { .. } => {
330                write!(f, "Encoding error when encoding error message header")
331            }
332            ErrorKind::ErrorMessage { .. } => {
333                write!(f, "Encoding error when encoding error message")
334            }
335            ErrorKind::OutOfBounds { offset, len } => {
336                write!(
337                    f,
338                    "Error when reading message: offset {} is out of bounds for length {}",
339                    offset, len
340                )
341            }
342            ErrorKind::NotNegotiated => {
343                write!(f, "Connection closed before the format was negotiated")
344            }
345            ErrorKind::ExpectedNegotiate { id } => {
346                write!(
347                    f,
348                    "Expected a negotiation as the first message, but got message id {id}"
349                )
350            }
351            ErrorKind::NegotiateHeader { .. } => {
352                write!(f, "Encoding error when decoding negotiation header")
353            }
354        }
355    }
356}
357
358impl core::error::Error for Error {
359    #[inline]
360    fn source(&self) -> Option<&(dyn core::error::Error + 'static)> {
361        match &self.kind {
362            #[cfg(feature = "axum-core05")]
363            ErrorKind::AxumCore05 { error } => Some(error),
364            ErrorKind::Incoming { error } => Some(error),
365            ErrorKind::Outgoing { error } => Some(error),
366            ErrorKind::EncodeBroadcastHeader { error } => Some(error),
367            ErrorKind::EncodeBroadcast { error } => Some(error),
368            ErrorKind::EncodeConnectHeader { error } => Some(error),
369            ErrorKind::ErrorMessageHeader { error } => Some(error),
370            ErrorKind::ErrorMessage { error } => Some(error),
371            ErrorKind::NegotiateHeader { error } => Some(error),
372            _ => None,
373        }
374    }
375}
376
377#[cfg(feature = "axum-core05")]
378impl From<axum_core05::Error> for Error {
379    #[inline]
380    fn from(error: axum_core05::Error) -> Self {
381        Self::new(ErrorKind::AxumCore05 { error })
382    }
383}
384
385impl From<ErrorKind> for Error {
386    #[inline]
387    fn from(kind: ErrorKind) -> Self {
388        Self::new(kind)
389    }
390}
391
392impl From<InvalidFrame> for Error {
393    #[inline]
394    fn from(error: InvalidFrame) -> Self {
395        Self::new(ErrorKind::InvalidFrame { error })
396    }
397}
398
399type Result<T, E = Error> = core::result::Result<T, E>;
400
401/// The response meta from handling a request.
402pub struct Response {
403    handled: bool,
404}
405
406/// Trait governing how something can be converted into a response.
407pub trait IntoResponse
408where
409    Self: 'static + Send,
410{
411    /// The error variant being produced.
412    type Error: fmt::Display;
413
414    /// Convert self into a response.
415    fn into_response(self) -> Result<Response, Self::Error>;
416}
417
418/// Implement [`IntoResponse`] for unit types `()`.
419///
420/// This indicates that the request has been handled.
421impl IntoResponse for () {
422    type Error = Infallible;
423
424    #[inline]
425    fn into_response(self) -> Result<Response, Self::Error> {
426        Ok(Response { handled: true })
427    }
428}
429
430/// Implement [`IntoResponse`] for `bool`.
431///
432/// On `true`, this means that the request was supported `false` means that it
433/// wasn't.
434impl IntoResponse for bool {
435    type Error = Infallible;
436
437    #[inline]
438    fn into_response(self) -> Result<Response, Self::Error> {
439        Ok(Response { handled: self })
440    }
441}
442
443/// Implement [`IntoResponse`] for [`Result`] types.
444///
445/// Note that this allows anything that implements [`fmt::Display`] to be used
446/// as an [`Err`] variant. The exact message it's being formatted into will be
447/// forwarded as an error to the client.
448///
449/// [`Result`]: core::result::Result
450impl<T, E> IntoResponse for Result<T, E>
451where
452    T: IntoResponse<Error = Infallible>,
453    E: 'static + Send + fmt::Display,
454{
455    type Error = E;
456
457    #[inline]
458    fn into_response(self) -> Result<Response, E> {
459        match self {
460            Ok(into_response) => match IntoResponse::into_response(into_response) {
461                Ok(response) => Ok(response),
462                Err(error) => match error {},
463            },
464            Err(error) => Err(error),
465        }
466    }
467}
468
469/// Implement [`IntoResponse`] for [`Option`] types.
470///
471/// This will propagate any responses for the interior value if present. If the
472/// value is [`None`] this will be treated as unhandled. This can be useful when
473/// used in combination with [`Incoming::read`] since it returns an [`Option`].
474impl<T> IntoResponse for Option<T>
475where
476    T: IntoResponse,
477{
478    type Error = T::Error;
479
480    #[inline]
481    fn into_response(self) -> Result<Response, Self::Error> {
482        match self {
483            Some(value) => value.into_response(),
484            None => Ok(Response { handled: false }),
485        }
486    }
487}
488
489/// A handler for incoming requests.
490///
491/// See [`server()`] for how to use with `axum`.
492///
493/// [`server()`]: crate::axum08::server
494pub trait Handler
495where
496    Self: 'static + Send + Clone,
497{
498    /// The type of message id used.
499    type Id: Id;
500    /// The response type returned by the handler.
501    type Response: IntoResponse;
502
503    /// Indicates that a `channel` has been opened.
504    ///
505    /// This indicates that you are communicating with a client that has opened
506    /// a channel with [`Handle::channel`].
507    ///
508    /// After this has been called, you can expected to receive requests from
509    /// the [`ChannelId`] corresponding to `channel`. The channel id of incoming
510    /// requests can be inspected with [`Incoming::channel`].
511    ///
512    /// [`Handle::channel`]: crate::web::Handle::channel
513    fn open_channel<'this>(
514        &'this self,
515        channel: ChannelId,
516    ) -> impl Future<Output = ()> + Send + 'this {
517        async {
518            _ = channel;
519        }
520    }
521
522    /// Indicates that a `channel` has been cleanly closed.
523    ///
524    /// This indicates that communicating with a client that has opened a
525    /// channel with [`Handle::channel`] has been cleanly closed, which occurs
526    /// when the channel is cleanly closed by dropping the last handle to it.
527    ///
528    /// [`Handle::channel`]: crate::web::Handle::channel
529    fn close_channel<'this>(
530        &'this self,
531        channel: ChannelId,
532    ) -> impl Future<Output = ()> + Send + 'this {
533        async {
534            _ = channel;
535        }
536    }
537
538    /// Handle a request.
539    fn handle<'this>(
540        &'this self,
541        id: Self::Id,
542        incoming: &'this mut Incoming<'_>,
543        outgoing: &'this mut Outgoing<'_>,
544    ) -> impl Future<Output = Self::Response> + Send + 'this;
545}
546
547struct Pinned<S> {
548    socket: S,
549    close_sleep: Sleep,
550    ping_sleep: Sleep,
551}
552
553impl<S> Pinned<S> {
554    #[inline]
555    fn project(self: Pin<&mut Self>) -> (Pin<&mut Sleep>, Pin<&mut Sleep>, Pin<&mut S>) {
556        unsafe {
557            let this = self.get_unchecked_mut();
558            (
559                Pin::new_unchecked(&mut this.close_sleep),
560                Pin::new_unchecked(&mut this.ping_sleep),
561                Pin::new_unchecked(&mut this.socket),
562            )
563        }
564    }
565}
566
567type HandlerOutput<H> = (Result<<H as Handler>::Response, Error>, RequestHeader, Buf);
568
569/// Trait which governs how channel identifiers are allocated with a [`Server`].
570///
571/// By default channel identifiers are scoped to the server which is set up
572/// per-connection. If you want distinct and unique channel identifiers across
573/// multiple websocket connections a custom [`ChannelAllocator`] can be
574/// constructed.
575pub trait ChannelAllocator {
576    /// Allocate the next channel id.
577    ///
578    /// Using `0` is equivalent to [`ChannelId::NONE`] so the allocator must
579    /// avoid constructor identifiers with this value since it is equivalent to
580    /// no channel.
581    ///
582    /// [`ChannelAllocator`]: crate::ws::ChannelAllocator
583    fn next(&self) -> impl Future<Output = Option<ChannelId>> + Send + '_;
584
585    /// Free the given channel id.
586    fn free(&self, channel: ChannelId) -> impl Future<Output = ()> + Send + '_;
587}
588
589/// A connection which has not yet completed the [negotiation protocol].
590///
591/// This is what [`server()`] hands back, and it is the only way to obtain a
592/// [`Server`]. Configuration lives here rather than on [`Server`], since every
593/// setting has to be in place before the first byte goes over the wire.
594///
595/// Crucially this type cannot send messages. A [`Server`] — which can — only
596/// exists once [`Connect::connect`] has resolved, which is precisely the point
597/// at which the [`Format`] for the connection has been settled. Attempting to
598/// broadcast before that is a compile error rather than a message the client
599/// cannot read.
600///
601/// [`server()`]: crate::axum08::server
602/// [negotiation protocol]: crate::api#negotiating-the-format
603///
604/// # Examples
605///
606/// ```
607/// # extern crate axum08 as axum;
608/// # use axum::extract::ws::WebSocket;
609/// use musli_web::api::Format;
610/// use musli_web::{axum08, ws};
611///
612/// mod api {
613///     use musli::{Decode, Encode};
614///     use musli_web::api;
615///
616///     #[derive(Encode, Decode)]
617///     pub struct HelloRequest<'de> {
618///         pub message: &'de str,
619///     }
620///
621///     #[derive(Encode, Decode)]
622///     pub struct HelloResponse<'de> {
623///         pub message: &'de str,
624///     }
625///
626///     api::define! {
627///         pub type Hello;
628///
629///         impl Endpoint for Hello {
630///             impl<'de> Request for HelloRequest<'de>;
631///             type Response<'de> = HelloResponse<'de>;
632///         }
633///     }
634/// }
635///
636/// #[derive(Clone)]
637/// struct MyHandler;
638///
639/// impl ws::Handler for MyHandler {
640///     type Id = api::Request;
641///     type Response = bool;
642///
643///     async fn handle(
644///         &self,
645///         id: Self::Id,
646///         incoming: &mut ws::Incoming<'_>,
647///         outgoing: &mut ws::Outgoing<'_>,
648///     ) -> bool {
649///         false
650///     }
651/// }
652///
653/// # async fn example(socket: WebSocket) -> Result<(), ws::Error> {
654/// let mut server = axum08::server(socket, MyHandler)
655///     .with_formats(&[Format::Wire, Format::Json])
656///     .connect()
657///     .await?;
658///
659/// // Only reachable once the client has negotiated a format.
660/// server.run().await?;
661/// # Ok(())
662/// # }
663/// ```
664///
665/// Skipping the connection step does not compile, since a [`Connect`] has
666/// nothing to broadcast with:
667///
668/// ```compile_fail
669/// # extern crate axum08 as axum;
670/// # use axum::extract::ws::WebSocket;
671/// use musli_web::{axum08, ws};
672///
673/// mod api {
674///     use musli::{Decode, Encode};
675///     use musli_web::api;
676///
677///     #[derive(Encode, Decode)]
678///     pub struct HelloRequest<'de> {
679///         pub message: &'de str,
680///     }
681///
682///     #[derive(Encode, Decode)]
683///     pub struct HelloResponse<'de> {
684///         pub message: &'de str,
685///     }
686///
687///     #[derive(Encode, Decode)]
688///     pub struct TickEvent {
689///         pub tick: u32,
690///     }
691///
692///     api::define! {
693///         pub type Hello;
694///
695///         impl Endpoint for Hello {
696///             impl<'de> Request for HelloRequest<'de>;
697///             type Response<'de> = HelloResponse<'de>;
698///         }
699///
700///         pub type Tick;
701///
702///         impl Broadcast for Tick {
703///             impl Event for TickEvent;
704///         }
705///     }
706/// }
707///
708/// #[derive(Clone)]
709/// struct MyHandler;
710///
711/// impl ws::Handler for MyHandler {
712///     type Id = api::Request;
713///     type Response = bool;
714///
715///     async fn handle(
716///         &self,
717///         id: Self::Id,
718///         incoming: &mut ws::Incoming<'_>,
719///         outgoing: &mut ws::Outgoing<'_>,
720///     ) -> bool {
721///         false
722///     }
723/// }
724///
725/// # async fn example(socket: WebSocket) -> Result<(), ws::Error> {
726/// let mut server = axum08::server(socket, MyHandler);
727/// // `Connect` has no `broadcast`, only the `Server` that `connect()` hands
728/// // back does.
729/// server.broadcast(api::TickEvent { tick: 1 })?;
730/// # Ok(())
731/// # }
732/// ```
733pub struct Connect<S, H, C = Channels>
734where
735    S: ServerImpl,
736    H: Handler,
737{
738    handler: H,
739    socket: S::Socket,
740    channels: C,
741    seed: u64,
742    max_capacity: usize,
743    /// Formats this server is willing to negotiate, or `None` to accept every
744    /// format it was built with support for.
745    formats: Option<&'static [Format]>,
746}
747
748impl<S, H> Connect<S, H, Channels>
749where
750    S: ServerImpl,
751    H: Handler,
752{
753    /// Construct a new pending connection with the specified handler.
754    #[inline]
755    pub(crate) fn new(socket: S::Socket, handler: H) -> Self {
756        Self {
757            handler,
758            socket,
759            channels: Channels::default(),
760            seed: DEFAULT_SEED,
761            max_capacity: MAX_CAPACITY,
762            formats: None,
763        }
764    }
765}
766
767impl<S, H, C> Connect<S, H, C>
768where
769    S: ServerImpl,
770    H: Handler,
771{
772    /// Associate the specified seed with the connection.
773    ///
774    /// This affects the random number generation used for ping messages.
775    ///
776    /// By default the seed is a constant value.
777    #[inline]
778    pub fn seed(mut self, seed: u64) -> Self {
779        self.seed = seed;
780        self
781    }
782
783    /// Associate the specified channel allocator with the connection.
784    #[inline]
785    pub fn with_channel_allocator<U>(self, channels: U) -> Connect<S, H, U>
786    where
787        U: ChannelAllocator,
788    {
789        Connect {
790            handler: self.handler,
791            socket: self.socket,
792            channels,
793            seed: self.seed,
794            max_capacity: self.max_capacity,
795            formats: self.formats,
796        }
797    }
798
799    /// Get a reference to the handler.
800    #[inline]
801    pub fn handler(&self) -> &H {
802        &self.handler
803    }
804
805    /// Restrict the set of [`Format`]s this server is willing to negotiate.
806    ///
807    /// By default every format the crate was built with support for is
808    /// accepted. A client which asks for a format outside of this set is
809    /// rejected during the [negotiation protocol] and the connection settles on
810    /// [`Format::DEFAULT`] instead.
811    ///
812    /// Note that this cannot widen the set, a format which was not compiled in
813    /// is never accepted.
814    ///
815    /// [negotiation protocol]: crate::api#negotiating-the-format
816    #[inline]
817    pub fn with_formats(mut self, formats: &'static [Format]) -> Self {
818        self.formats = Some(formats);
819        self
820    }
821
822    /// Test if this server is willing to negotiate `format`.
823    #[inline]
824    pub fn accepts(&self, format: Format) -> bool {
825        accepts(self.formats, format)
826    }
827
828    /// Modify the max allocated capacity of the buffers used for outgoing
829    /// messages.
830    ///
831    /// This is not a hard limit. A message larger than this is still written,
832    /// but once it has been flushed the allocation is released back down to the
833    /// specified value rather than being kept for the lifetime of the
834    /// connection.
835    ///
836    /// By default, the capacity is 1 MiB.
837    #[inline]
838    pub fn max_capacity(mut self, max_capacity: usize) -> Self {
839        self.max_capacity = max_capacity;
840        self
841    }
842
843    /// Modify the max allocated capacity of the outgoing buffers.
844    ///
845    /// This is an alias for [`Connect::max_capacity`].
846    #[inline]
847    pub fn with_max_capacity(self, max_capacity: usize) -> Self {
848        self.max_capacity(max_capacity)
849    }
850}
851
852impl<S, H, C> Connect<S, H, C>
853where
854    S: ServerImpl,
855    Error: From<S::Error>,
856    H: Handler,
857    C: ChannelAllocator,
858{
859    /// Perform the [negotiation protocol] and hand back the [`Server`] it
860    /// produced.
861    ///
862    /// This sends [`MessageId::SERVER_HELLO`] and then drives the socket —
863    /// including keepalive pings — until the client has answered with a
864    /// [`MessageId::NEGOTIATE`] request and the reply to it has been flushed.
865    ///
866    /// Until that has happened the connection has no [`Server`], so there is no
867    /// way to write a message which the client might not be able to decode.
868    ///
869    /// # Errors
870    ///
871    /// Errors if the connection goes away before the format has been
872    /// negotiated, or if the client sends anything other than a negotiation as
873    /// its first message. Both tear the connection down, since a peer which
874    /// does not negotiate cannot be talked to safely.
875    ///
876    /// [negotiation protocol]: crate::api#negotiating-the-format
877    pub async fn connect(self) -> Result<Server<S, H, C>, Error> {
878        let now = Instant::now();
879
880        let mut server = Server {
881            handler: self.handler,
882            pinned: Box::pin(Pinned {
883                socket: self.socket,
884                close_sleep: tokio::time::sleep_until(now + CLOSE_TIMEOUT),
885                ping_sleep: tokio::time::sleep_until(now + PING_TIMEOUT),
886            }),
887            channels: self.channels,
888            closing: false,
889            pool: BufPool::new(self.max_capacity),
890            outbound: VecDeque::new(),
891            error: String::new(),
892            last_ping: None,
893            rng: SmallRng::seed_from_u64(self.seed),
894            out: VecDeque::new(),
895            socket_send: false,
896            socket_flush: false,
897            set: JoinSet::new(),
898            format: Format::DEFAULT,
899            formats: self.formats,
900        };
901
902        server.hello()?;
903        server.negotiate().await?;
904        Ok(server)
905    }
906}
907
908/// Test if `formats` is willing to negotiate `format`.
909#[inline]
910fn accepts(formats: Option<&'static [Format]>, format: Format) -> bool {
911    format.is_supported() && formats.is_none_or(|f| f.contains(&format))
912}
913
914/// The server side handle of the websocket protocol.
915///
916/// This can only be constructed by completing the [negotiation protocol]
917/// through [`Connect::connect`], so its mere existence means that the
918/// [`Format`] used for everything the server originates has been agreed with
919/// the client.
920///
921/// See [`server()`] for how to use with `axum`.
922///
923/// [`server()`]: crate::axum08::server
924/// [negotiation protocol]: crate::api#negotiating-the-format
925pub struct Server<S, H, C = Channels>
926where
927    S: ServerImpl,
928    H: Handler,
929{
930    handler: H,
931    pinned: Pin<Box<Pinned<S::Socket>>>,
932    channels: C,
933    closing: bool,
934    pool: BufPool,
935    outbound: VecDeque<Buf>,
936    error: String,
937    last_ping: Option<[u8; 4]>,
938    rng: SmallRng,
939    out: VecDeque<S::Message>,
940    socket_send: bool,
941    socket_flush: bool,
942    set: JoinSet<HandlerOutput<H>>,
943    /// The format used for messages the server originates on this connection,
944    /// as agreed by the [negotiation protocol].
945    ///
946    /// [negotiation protocol]: crate::api#negotiating-the-format
947    format: Format,
948    /// Formats this server is willing to negotiate, or `None` to accept every
949    /// format it was built with support for.
950    formats: Option<&'static [Format]>,
951}
952
953impl<S, H, C> Server<S, H, C>
954where
955    S: ServerImpl,
956    H: Handler,
957{
958    /// Get a reference to the handler.
959    #[inline]
960    pub fn handler(&self) -> &H {
961        &self.handler
962    }
963
964    /// The [`Format`] used for messages this server originates, such as
965    /// broadcasts.
966    ///
967    /// This is fixed for the lifetime of the connection and was agreed with the
968    /// client by [`Connect::connect`], which is why it is never in doubt here.
969    #[inline]
970    pub fn format(&self) -> Format {
971        self.format
972    }
973
974    /// Test if this server is willing to negotiate `format`.
975    #[inline]
976    pub fn accepts(&self, format: Format) -> bool {
977        accepts(self.formats, format)
978    }
979}
980
981impl<S, H, C> Server<S, H, C>
982where
983    S: ServerImpl,
984    Error: From<S::Error>,
985    H: Handler,
986    C: ChannelAllocator,
987{
988    /// Drive the socket until the client has negotiated a [`Format`].
989    ///
990    /// This is the connection step every peer is forced through, see
991    /// [`Connect::connect`]. It deliberately understands nothing but
992    /// [`MessageId::NEGOTIATE`] and the keepalive machinery, so no handler is
993    /// ever invoked and nothing the user could write is in flight yet.
994    ///
995    /// It returns once the format is settled *and* the reply confirming it has
996    /// been flushed, so the very next thing on the wire can safely use it.
997    async fn negotiate(&mut self) -> Result<(), Error> {
998        let mut negotiated = false;
999        // NB: Held rather than returned immediately so that the close frame
1000        // explaining the violation makes it onto the wire first.
1001        let mut failure = None::<Error>;
1002
1003        loop {
1004            let drained = self.out.is_empty() && !self.socket_flush;
1005
1006            if failure.is_some() && drained {
1007                break;
1008            }
1009
1010            if negotiated && drained && self.outbound.is_empty() {
1011                break;
1012            }
1013
1014            self.handle_send()?;
1015
1016            let result = {
1017                let inner = Select::<S::Socket, H> {
1018                    pinned: self.pinned.as_mut(),
1019                    wants_socket_send: !self.socket_send,
1020                    wants_socket_flush: self.socket_flush,
1021                    set: &mut self.set,
1022                };
1023
1024                inner.await
1025            };
1026
1027            match result {
1028                Output::Close => {
1029                    return Err(Error::new(ErrorKind::NotNegotiated));
1030                }
1031                Output::Ping => {
1032                    self.handle_ping()?;
1033                }
1034                Output::Recv(message) => {
1035                    let Some(message) = message else {
1036                        return Err(Error::new(ErrorKind::NotNegotiated));
1037                    };
1038
1039                    match message? {
1040                        Message::Text => {
1041                            self.out.push_back(S::close(
1042                                CLOSE_PROTOCOL_ERROR,
1043                                "Unsupported text message",
1044                            ));
1045
1046                            failure = Some(Error::new(ErrorKind::NotNegotiated));
1047                        }
1048                        Message::Binary(bytes) => match self.handle_negotiate(bytes) {
1049                            Ok(()) => negotiated = true,
1050                            Err(error) => failure = Some(error),
1051                        },
1052                        Message::Ping(payload) => {
1053                            self.out.push_back(S::pong(payload));
1054                        }
1055                        Message::Pong(data) => {
1056                            self.handle_pong(data)?;
1057                        }
1058                        Message::Close => {
1059                            return Err(Error::new(ErrorKind::NotNegotiated));
1060                        }
1061                    }
1062                }
1063                Output::Send(result) => {
1064                    result?;
1065                    self.socket_send = true;
1066                }
1067                Output::Flushed(result) => {
1068                    result?;
1069                    self.socket_flush = false;
1070                }
1071                Output::Handle(..) => {
1072                    // NB: No handler can have been spawned yet, since requests
1073                    // are only dispatched after negotiation.
1074                }
1075            }
1076        }
1077
1078        match failure {
1079            Some(error) => Err(error),
1080            None => Ok(()),
1081        }
1082    }
1083
1084    /// Process the single message which is legal before a format has been
1085    /// negotiated.
1086    ///
1087    /// Anything else closes the connection, since a peer which skips
1088    /// negotiation cannot be sent broadcasts safely.
1089    fn handle_negotiate(&mut self, bytes: Bytes) -> Result<(), Error> {
1090        let mut at = 0;
1091
1092        let header: RequestHeader = match format::decode_envelope(&bytes, &mut at) {
1093            Ok(header) => header,
1094            Err(error) => {
1095                self.out
1096                    .push_back(S::close(CLOSE_PROTOCOL_ERROR, "Invalid request header"));
1097                return Err(Error::new(ErrorKind::NegotiateHeader { error }));
1098            }
1099        };
1100
1101        if MessageId::new(header.id) != Some(MessageId::NEGOTIATE) {
1102            self.out.push_back(S::close(
1103                CLOSE_PROTOCOL_ERROR,
1104                "Expected a negotiation as the first message",
1105            ));
1106
1107            return Err(Error::new(ErrorKind::ExpectedNegotiate { id: header.id }));
1108        }
1109
1110        // NB: The connection stays up on a rejected format and settles on the
1111        // default instead, which is what the client falls back to. The error
1112        // tells it which formats it could have asked for.
1113        let Some(format) = Format::from_u8(header.format) else {
1114            self.format_error_message(format_args!(
1115                "Unknown format id {}, supported: {}",
1116                header.format,
1117                SupportedFormats(self.formats)
1118            ))?;
1119
1120            self.format = Format::DEFAULT;
1121            return self.send_error(&header);
1122        };
1123
1124        if !self.accepts(format) {
1125            self.format_error_message(format_args!(
1126                "Unsupported format `{format}`, supported: {}",
1127                SupportedFormats(self.formats)
1128            ))?;
1129
1130            tracing::debug!(?format, "Rejected format");
1131            self.format = Format::DEFAULT;
1132            return self.send_error(&header);
1133        }
1134
1135        tracing::debug!(?format, "Negotiated format");
1136        self.format = format;
1137        self.send_negotiated(&header, format)
1138    }
1139
1140    /// Acknowledge a negotiation by echoing the format that was accepted.
1141    fn send_negotiated(&mut self, header: &RequestHeader, format: Format) -> Result<(), Error> {
1142        let buf = self.pool.with(|buf| {
1143            let mut writer = buf.writer();
1144
1145            let result = writer.envelope(&ResponseHeader {
1146                serial: header.serial,
1147                broadcast: 0,
1148                error: 0,
1149                format: format.to_u8(),
1150                channel: header.channel,
1151            });
1152
1153            result.map_err(Error::encode_connect_header)?;
1154            writer.flush();
1155            Ok::<_, Error>(())
1156        })?;
1157
1158        self.outbound.push_back(buf);
1159        Ok(())
1160    }
1161
1162    /// Run the server.
1163    ///
1164    /// This must be called to handle buffered outgoing and incoming messages.
1165    pub async fn run(&mut self) -> Result<(), Error> {
1166        loop {
1167            if self.closing && self.out.is_empty() && self.outbound.is_empty() {
1168                break;
1169            }
1170
1171            self.handle_send()?;
1172
1173            let result = {
1174                let inner = Select::<S::Socket, H> {
1175                    pinned: self.pinned.as_mut(),
1176                    wants_socket_send: !self.socket_send,
1177                    wants_socket_flush: self.socket_flush,
1178                    set: &mut self.set,
1179                };
1180
1181                inner.await
1182            };
1183
1184            match result {
1185                Output::Close => {
1186                    self.out
1187                        .push_back(S::close(CLOSE_NORMAL, "connection timed out"));
1188                    self.closing = true;
1189                }
1190                Output::Ping => {
1191                    self.handle_ping()?;
1192                }
1193                Output::Recv(message) => {
1194                    let Some(message) = message else {
1195                        self.closing = true;
1196                        continue;
1197                    };
1198
1199                    match message? {
1200                        Message::Text => {
1201                            self.out.push_back(S::close(
1202                                CLOSE_PROTOCOL_ERROR,
1203                                "Unsupported text message",
1204                            ));
1205                            self.closing = true;
1206                        }
1207                        Message::Binary(bytes) => {
1208                            self.handle_message(bytes).await?;
1209                        }
1210                        Message::Ping(payload) => {
1211                            self.out.push_back(S::pong(payload));
1212                        }
1213                        Message::Pong(data) => {
1214                            self.handle_pong(data)?;
1215                        }
1216                        Message::Close => {
1217                            self.closing = true;
1218                        }
1219                    }
1220                }
1221                Output::Send(result) => {
1222                    if let Err(err) = result {
1223                        return Err(Error::from(err));
1224                    };
1225
1226                    self.socket_send = true;
1227                }
1228                Output::Flushed(result) => {
1229                    if let Err(err) = result {
1230                        return Err(Error::from(err));
1231                    };
1232
1233                    self.socket_flush = false;
1234                }
1235                Output::Handle(result, header, buf) => {
1236                    let err = 'err: {
1237                        let res = match result {
1238                            Ok(res) => res,
1239                            Err(error) => {
1240                                self.format_error(error)?;
1241                                break 'err true;
1242                            }
1243                        };
1244
1245                        let res = match res.into_response() {
1246                            Ok(res) => res,
1247                            Err(error) => {
1248                                self.format_error_message(error)?;
1249                                break 'err true;
1250                            }
1251                        };
1252
1253                        if !res.handled {
1254                            self.format_error_message(format_args!(
1255                                "No support for request {}",
1256                                header.id
1257                            ))?;
1258                            break 'err true;
1259                        }
1260
1261                        self.outbound.push_back(buf);
1262                        false
1263                    };
1264
1265                    if err {
1266                        self.send_error(&header)?;
1267                    }
1268                }
1269            }
1270        }
1271
1272        Ok(())
1273    }
1274
1275    /// Write a broadcast message.
1276    ///
1277    /// Note that the written message is buffered, and will be sent when
1278    /// [`Server::run`] is called.
1279    pub fn broadcast<T>(&mut self, message: T) -> Result<(), Error>
1280    where
1281        T: Event,
1282    {
1283        self.broadcast_in(message, ChannelId::NONE)
1284    }
1285
1286    /// Write a broadcast message over a specific `channel`.
1287    ///
1288    /// Note that the written message is buffered, and will be sent when
1289    /// [`Server::run`] is called.
1290    pub fn broadcast_in<T>(&mut self, message: T, channel: ChannelId) -> Result<(), Error>
1291    where
1292        T: Event,
1293    {
1294        tracing::debug!(id = ?<T::Broadcast as Broadcast>::ID, "Broadcast");
1295
1296        let format = self.format;
1297
1298        let buf = self.pool.with(|buf| {
1299            let mut writer = buf.writer();
1300
1301            writer
1302                .envelope(&ResponseHeader {
1303                    serial: 0,
1304                    broadcast: <T::Broadcast as Broadcast>::ID.get(),
1305                    error: 0,
1306                    format: format.to_u8(),
1307                    channel,
1308                })
1309                .map_err(Error::encode_broadcast_header)?;
1310
1311            writer
1312                .body(format, &message)
1313                .map_err(Error::encode_broadcast)?;
1314            writer.flush();
1315            Ok::<_, Error>(())
1316        })?;
1317
1318        self.outbound.push_back(buf);
1319        Ok(())
1320    }
1321
1322    /// Write a broadcast message over a specific connection.
1323    ///
1324    /// Note that the written message is buffered, and will be sent when
1325    /// [`Server::run`] is called.
1326    fn hello(&mut self) -> Result<(), Error> {
1327        tracing::debug!("Hello");
1328
1329        let mut buf = self.pool.get();
1330
1331        let result = (|| {
1332            let mut writer = buf.writer();
1333
1334            writer
1335                .envelope(&ResponseHeader {
1336                    serial: 0,
1337                    broadcast: MessageId::SERVER_HELLO.get(),
1338                    error: 0,
1339                    // NB: Carries no body, so no format applies.
1340                    format: 0,
1341                    channel: ChannelId::NONE,
1342                })
1343                .map_err(Error::encode_broadcast_header)?;
1344
1345            writer.flush();
1346            Ok::<_, Error>(())
1347        })();
1348
1349        if result.is_err() {
1350            self.pool.put(buf);
1351        } else {
1352            self.outbound.push_back(buf);
1353        }
1354
1355        Ok(())
1356    }
1357
1358    fn format_error_message(&mut self, error: impl fmt::Display) -> Result<(), Error> {
1359        self.error.clear();
1360
1361        if write!(self.error, "{error}").is_err() {
1362            self.error.clear();
1363            return Err(Error::new(ErrorKind::FormatError));
1364        }
1365
1366        Ok(())
1367    }
1368
1369    fn format_error(&mut self, error: impl core::error::Error) -> Result<(), Error> {
1370        self.error.clear();
1371
1372        if write!(self.error, "{error:#}").is_err() {
1373            self.error.clear();
1374            return Err(Error::new(ErrorKind::FormatError));
1375        }
1376
1377        Ok(())
1378    }
1379
1380    #[tracing::instrument(skip(self, bytes))]
1381    async fn handle_message(&mut self, bytes: Bytes) -> Result<(), Error> {
1382        let mut at = 0;
1383
1384        let header: RequestHeader = match format::decode_envelope(&bytes, &mut at) {
1385            Ok(header) => header,
1386            Err(error) => {
1387                tracing::debug!(?error, "Invalid request header");
1388                self.out
1389                    .push_back(S::close(CLOSE_PROTOCOL_ERROR, "Invalid request header"));
1390                self.closing = true;
1391                return Ok(());
1392            }
1393        };
1394
1395        let err = 'err: {
1396            let Some(id) = MessageId::new(header.id) else {
1397                self.format_error_message(format_args!("Unsupported message id {}", header.id))?;
1398                break 'err true;
1399            };
1400
1401            match id {
1402                MessageId::CONNECT => {
1403                    let Some(channel) = self.channels.next().await else {
1404                        self.format_error_message(format_args!(
1405                            "Failed to allocate connection ID"
1406                        ))?;
1407
1408                        break 'err true;
1409                    };
1410
1411                    self.handler.open_channel(channel).await;
1412
1413                    let mut buf = self.pool.get();
1414
1415                    let result = (|| {
1416                        let mut writer = buf.writer();
1417
1418                        let result = writer.envelope(&ResponseHeader {
1419                            serial: header.serial,
1420                            broadcast: 0,
1421                            error: 0,
1422                            format: 0,
1423                            channel,
1424                        });
1425
1426                        result.map_err(Error::encode_connect_header)?;
1427                        writer.flush();
1428                        Ok::<_, Error>(())
1429                    })();
1430
1431                    if result.is_err() {
1432                        self.pool.put(buf);
1433                    } else {
1434                        self.outbound.push_back(buf);
1435                    }
1436
1437                    result?;
1438                    break 'err false;
1439                }
1440                MessageId::DISCONNECT => {
1441                    self.channels.free(header.channel).await;
1442                    self.handler.close_channel(header.channel).await;
1443                    break 'err false;
1444                }
1445                MessageId::NEGOTIATE => {
1446                    // NB: The format is settled once and for all by the
1447                    // connection step, see `Connect::connect`. Letting it move
1448                    // afterwards would mean messages already queued for the old
1449                    // format go out under the new one.
1450                    let format = self.format;
1451
1452                    self.format_error_message(format_args!(
1453                        "Format `{format}` has already been negotiated"
1454                    ))?;
1455
1456                    break 'err true;
1457                }
1458                _ => {
1459                    let Some(format) = Format::from_u8(header.format) else {
1460                        self.format_error_message(format_args!(
1461                            "Unknown format id {}",
1462                            header.format
1463                        ))?;
1464
1465                        break 'err true;
1466                    };
1467
1468                    if !self.accepts(format) {
1469                        self.format_error_message(format_args!(
1470                            "Unsupported format `{format}`, supported: {}",
1471                            SupportedFormats(self.formats)
1472                        ))?;
1473
1474                        break 'err true;
1475                    }
1476
1477                    let id = <H::Id as Id>::from_id(id);
1478                    self.handle_request(bytes, at, header, id, format);
1479                    return Ok(());
1480                }
1481            }
1482        };
1483
1484        if err {
1485            self.send_error(&header)?;
1486        }
1487
1488        Ok(())
1489    }
1490
1491    fn send_error(&mut self, header: &RequestHeader) -> Result<(), Error> {
1492        // NB: Errors are encoded with the connection format rather than the
1493        // format the request asked for, since the request might have failed
1494        // precisely because that format is not supported. The client reads the
1495        // format back out of the response envelope either way.
1496        let format = self.format;
1497
1498        let buf = self.pool.with(|buf| {
1499            // Reset the buffer to the previous start point.
1500            let mut writer = buf.writer();
1501
1502            let result = writer.envelope(&ResponseHeader {
1503                serial: header.serial,
1504                broadcast: 0,
1505                error: MessageId::ERROR_MESSAGE.get(),
1506                format: format.to_u8(),
1507                channel: header.channel,
1508            });
1509
1510            result.map_err(Error::encode_error_message_header)?;
1511
1512            let result = writer.body(
1513                format,
1514                &ErrorMessage {
1515                    message: &self.error,
1516                },
1517            );
1518
1519            result.map_err(Error::encode_error_message)?;
1520            writer.flush();
1521            Ok::<_, Error>(())
1522        })?;
1523
1524        self.outbound.push_back(buf);
1525        Ok(())
1526    }
1527
1528    #[tracing::instrument(skip(self))]
1529    fn handle_ping(&mut self) -> Result<(), Error> {
1530        let (_, mut ping_sleep, _) = self.pinned.as_mut().project();
1531
1532        let payload = self.rng.random::<u32>();
1533        let payload = payload.to_ne_bytes();
1534
1535        self.last_ping = Some(payload);
1536
1537        tracing::debug!(data = ?&payload[..], "Sending ping");
1538
1539        self.out
1540            .push_back(S::ping(Bytes::from_owner(Vec::from(payload))));
1541
1542        let now = Instant::now();
1543        ping_sleep.as_mut().reset(now + PING_TIMEOUT);
1544        Ok(())
1545    }
1546
1547    #[tracing::instrument(skip(self, payload))]
1548    fn handle_pong(&mut self, payload: Bytes) -> Result<(), Error> {
1549        let (close_sleep, ping_sleep, _) = self.pinned.as_mut().project();
1550
1551        tracing::debug!(payload = ?&payload[..], "Pong");
1552
1553        let Some(expected) = self.last_ping else {
1554            tracing::debug!("No ping sent");
1555            return Ok(());
1556        };
1557
1558        if expected[..] != payload[..] {
1559            tracing::debug!(?expected, ?payload, "Pong doesn't match");
1560            return Ok(());
1561        }
1562
1563        let now = Instant::now();
1564
1565        close_sleep.reset(now + CLOSE_TIMEOUT);
1566        ping_sleep.reset(now + PING_TIMEOUT);
1567        self.last_ping = None;
1568        Ok(())
1569    }
1570
1571    #[tracing::instrument(skip(self))]
1572    fn handle_send(&mut self) -> Result<(), Error> {
1573        let (_, _, mut socket) = self.pinned.as_mut().project();
1574
1575        if self.socket_send
1576            && let Some(message) = self.out.pop_front()
1577        {
1578            socket.as_mut().start_send(message)?;
1579            self.socket_flush = true;
1580            self.socket_send = false;
1581        }
1582
1583        while self.socket_send
1584            && let Some(buf) = self.outbound.front_mut()
1585        {
1586            let Some(frame) = buf.read()? else {
1587                if let Some(buf) = self.outbound.pop_front() {
1588                    self.pool.put(buf);
1589                }
1590
1591                continue;
1592            };
1593
1594            socket.as_mut().start_send(S::binary(frame))?;
1595
1596            self.socket_flush = true;
1597            self.socket_send = false;
1598            break;
1599        }
1600
1601        Ok(())
1602    }
1603
1604    fn handle_request(
1605        &mut self,
1606        bytes: Bytes,
1607        offset: usize,
1608        header: RequestHeader,
1609        id: H::Id,
1610        format: Format,
1611    ) {
1612        tracing::debug!(header.serial, ?id, ?format, "Got request");
1613
1614        let mut buf = self.pool.get();
1615        let handler = self.handler.clone();
1616
1617        self.set.spawn(async move {
1618            if offset > bytes.len() {
1619                let kind = ErrorKind::OutOfBounds {
1620                    offset,
1621                    len: bytes.len(),
1622                };
1623
1624                return (Err(Error::new(kind)), header, buf);
1625            }
1626
1627            let mut incoming = Incoming {
1628                error: None,
1629                buf: &bytes,
1630                at: offset,
1631                format,
1632                channel: header.channel,
1633            };
1634
1635            let mut outgoing = Outgoing {
1636                serial: Some(header.serial),
1637                error: None,
1638                buf: &mut buf,
1639                format,
1640                channel: header.channel,
1641            };
1642
1643            let response = handler.handle(id, &mut incoming, &mut outgoing).await;
1644
1645            if let Some(error) = incoming.error.take() {
1646                return (Err(Error::incoming(error)), header, buf);
1647            }
1648
1649            if let Some(error) = outgoing.error.take() {
1650                return (Err(Error::outgoing(error)), header, buf);
1651            }
1652
1653            (Ok(response), header, buf)
1654        });
1655    }
1656}
1657
1658enum Output<E, R> {
1659    /// The connection should be closed.
1660    Close,
1661    /// A ping message was received.
1662    Ping,
1663    /// A message was received.
1664    Recv(Option<Result<Message, E>>),
1665    /// A message is ready to be sent.
1666    Send(Result<(), E>),
1667    /// Outgoing messages have been successfully flushed.
1668    Flushed(Result<(), E>),
1669    /// Handle output.
1670    Handle(Result<R, Error>, RequestHeader, Buf),
1671}
1672
1673struct Select<'a, S, H>
1674where
1675    H: Handler,
1676{
1677    pinned: Pin<&'a mut Pinned<S>>,
1678    wants_socket_send: bool,
1679    wants_socket_flush: bool,
1680    set: &'a mut JoinSet<HandlerOutput<H>>,
1681}
1682
1683impl<S, H> Future for Select<'_, S, H>
1684where
1685    S: SocketImpl,
1686    H: Handler,
1687{
1688    type Output = Output<S::Error, H::Response>;
1689
1690    #[inline]
1691    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
1692        let close;
1693        let ping;
1694        let mut socket;
1695        let wants_socket_send;
1696        let wants_socket_flush;
1697        let set;
1698
1699        // SAFETY: This type is not Unpin.
1700        unsafe {
1701            let this = Pin::get_unchecked_mut(self);
1702            (close, ping, socket) = this.pinned.as_mut().project();
1703            wants_socket_send = this.wants_socket_send;
1704            wants_socket_flush = this.wants_socket_flush;
1705            set = &mut this.set;
1706        };
1707
1708        if close.poll(cx).is_ready() {
1709            return Poll::Ready(Output::Close);
1710        }
1711
1712        if ping.poll(cx).is_ready() {
1713            return Poll::Ready(Output::Ping);
1714        }
1715
1716        if let Poll::Ready(output) = socket.as_mut().poll_next(cx) {
1717            return Poll::Ready(Output::Recv(output));
1718        }
1719
1720        if wants_socket_send && let Poll::Ready(result) = socket.as_mut().poll_ready(cx) {
1721            return Poll::Ready(Output::Send(result));
1722        }
1723
1724        if wants_socket_flush && let Poll::Ready(result) = socket.as_mut().poll_flush(cx) {
1725            return Poll::Ready(Output::Flushed(result));
1726        }
1727
1728        if let Poll::Ready(output) = set.poll_join_next(cx)
1729            && let Some(output) = output
1730        {
1731            let output = match output {
1732                Ok(output) => output,
1733                Err(error) => {
1734                    tracing::debug!(?error, "Join error in handler task");
1735                    return Poll::Ready(Output::Close);
1736                }
1737            };
1738
1739            let (result, header, buf) = output;
1740            return Poll::Ready(Output::Handle(result, header, buf));
1741        }
1742
1743        Poll::Pending
1744    }
1745}
1746
1747/// The buffer for incoming requests.
1748///
1749/// See [`server()`] for how to use with `axum`.
1750///
1751/// [`server()`]: crate::axum08::server
1752pub struct Incoming<'de> {
1753    error: Option<format::Error>,
1754    buf: &'de [u8],
1755    at: usize,
1756    format: Format,
1757    channel: ChannelId,
1758}
1759
1760impl<'de> Incoming<'de> {
1761    /// The channel over which the incoming request was received.
1762    ///
1763    /// This is [`ChannelId::NONE`] unless the packet belongs to a response to a
1764    /// handle constructed with [`Handle::channel`].
1765    ///
1766    /// [`Handle::channel`]: crate::web::Handle::channel
1767    pub fn channel(&self) -> ChannelId {
1768        self.channel
1769    }
1770
1771    /// The [`Format`] the incoming request body is encoded with.
1772    ///
1773    /// This is the format the client declared for this particular request, and
1774    /// is also the format the response will be written with.
1775    #[inline]
1776    pub fn format(&self) -> Format {
1777        self.format
1778    }
1779
1780    /// Read a request and return `Some(T)` if the request was successfully
1781    /// decoded.
1782    ///
1783    /// Note that any failure to decode will be propagated as an error
1784    /// automatically, the user does not have to deal with it themselves.
1785    /// Instead, failure to decode should be treated as if the request was
1786    /// unhandled by returning for example `false` or `Option::None`.
1787    #[inline]
1788    pub fn read<T>(&mut self) -> Option<T>
1789    where
1790        T: DecodeBody<'de>,
1791    {
1792        match self.format.decode(self.buf, &mut self.at) {
1793            Ok(value) => Some(value),
1794            Err(error) => {
1795                self.error = Some(error);
1796                None
1797            }
1798        }
1799    }
1800}
1801
1802/// The buffer for outgoing responses.
1803///
1804/// See [`server()`] for how to use with `axum`.
1805///
1806/// [`server()`]: crate::axum08::server
1807pub struct Outgoing<'a> {
1808    serial: Option<u32>,
1809    error: Option<format::Error>,
1810    buf: &'a mut Buf,
1811    format: Format,
1812    channel: ChannelId,
1813}
1814
1815impl Outgoing<'_> {
1816    /// The [`Format`] the response will be encoded with, which is the format
1817    /// the corresponding request declared.
1818    #[inline]
1819    pub fn format(&self) -> Format {
1820        self.format
1821    }
1822
1823    /// Write a response.
1824    ///
1825    /// This can only be called once. Calling this multiple times has no effect.
1826    ///
1827    /// See [`server()`] for how to use with `axum`.
1828    ///
1829    /// [`server()`]: crate::axum08::server
1830    pub fn write<T>(&mut self, value: T)
1831    where
1832        T: EncodeBody,
1833    {
1834        let Some(serial) = self.serial.take() else {
1835            return;
1836        };
1837
1838        let mut writer = self.buf.writer();
1839
1840        let result = writer.envelope(&ResponseHeader {
1841            serial,
1842            broadcast: 0,
1843            error: 0,
1844            format: self.format.to_u8(),
1845            channel: self.channel,
1846        });
1847
1848        if let Err(error) = result {
1849            self.error = Some(error);
1850            return;
1851        }
1852
1853        if let Err(error) = writer.body(self.format, &value) {
1854            self.error = Some(error);
1855        }
1856
1857        writer.flush();
1858    }
1859}
1860
1861/// Scramble a sequential channel id into a value that looks random to clients.
1862///
1863/// Uses an odd multiply followed by an XOR-shift (bijective over u16, preserves
1864/// 0, self-inverse).
1865#[inline]
1866fn scramble_channel(x: u16) -> u16 {
1867    let x = x.wrapping_mul(0x9285);
1868    x ^ (x >> 8)
1869}
1870
1871/// Inverse of [`scramble_channel`].
1872#[inline]
1873#[cfg(test)]
1874fn unscramble_channel(x: u16) -> u16 {
1875    let x = x ^ (x >> 8);
1876    x.wrapping_mul(0x964d)
1877}
1878
1879#[test]
1880fn test_scramble() {
1881    assert_eq!(scramble_channel(0), 0);
1882    assert_eq!(unscramble_channel(0), 0);
1883
1884    for i in 1..=u16::MAX {
1885        let scrambled = scramble_channel(i);
1886        let unscrambled = unscramble_channel(scrambled);
1887        assert_eq!(i, unscrambled, "Failed to unscramble channel id");
1888    }
1889}
1890
1891#[derive(Default)]
1892struct ChannelsInner {
1893    last: u16,
1894    free: VecDeque<NonZeroU16>,
1895}
1896
1897/// A global channel allocator which can be cloned and re-used across multiple
1898/// servers allowing channels across servers to have distinct channel
1899/// identifiers.
1900#[derive(Default, Clone)]
1901pub struct Channels {
1902    inner: Arc<Mutex<ChannelsInner>>,
1903}
1904
1905impl ChannelAllocator for Channels {
1906    #[inline]
1907    async fn next(&self) -> Option<ChannelId> {
1908        let mut inner = self.inner.lock().await;
1909
1910        if let Some(id) = inner.free.pop_front() {
1911            return Some(ChannelId::from_u16(id.get()));
1912        }
1913
1914        let id = NonZeroU16::new(inner.last.wrapping_add(1))?;
1915        inner.last = id.get();
1916
1917        tracing::debug!(?id, "Allocated channel id");
1918        Some(ChannelId::from_u16(scramble_channel(id.get())))
1919    }
1920
1921    #[inline]
1922    async fn free(&self, id: ChannelId) {
1923        tracing::debug!(?id, "Freeing channel id");
1924
1925        let mut inner = self.inner.lock().await;
1926
1927        if let Some(id) = NonZeroU16::new(id.raw()) {
1928            inner.free.push_back(id);
1929        }
1930    }
1931}
1932
1933/// Renders the formats a server is willing to negotiate, for error messages.
1934struct SupportedFormats(Option<&'static [Format]>);
1935
1936impl fmt::Display for SupportedFormats {
1937    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1938        let mut first = true;
1939
1940        for format in Format::supported() {
1941            if let Some(formats) = self.0
1942                && !formats.contains(&format)
1943            {
1944                continue;
1945            }
1946
1947            if !first {
1948                f.write_str(", ")?;
1949            }
1950
1951            write!(f, "`{format}`")?;
1952            first = false;
1953        }
1954
1955        if first {
1956            f.write_str("none")?;
1957        }
1958
1959        Ok(())
1960    }
1961}