Skip to main content

rustdtp/
server.rs

1//! Protocol server implementation.
2
3use super::command_channel::*;
4use super::timeout::*;
5use crate::crypto::*;
6use crate::error::{Error, Result};
7use crate::util::*;
8use serde::de::DeserializeOwned;
9use serde::ser::Serialize;
10use std::collections::HashMap;
11use std::future::Future;
12use std::marker::PhantomData;
13use std::net::SocketAddr;
14use std::pin::Pin;
15use std::sync::Arc;
16use tokio::io::{AsyncReadExt, AsyncWriteExt};
17use tokio::net::{TcpListener, TcpStream, ToSocketAddrs};
18use tokio::sync::mpsc::{channel, Receiver, Sender};
19use tokio::task::JoinHandle;
20
21/// Configuration for a server's event callbacks.
22///
23/// # Events
24///
25/// There are four events for which callbacks can be registered:
26///
27///  - `connect`
28///  - `disconnect`
29///  - `receive`
30///  - `stop`
31///
32/// All callbacks are optional, and can be registered for any combination of
33/// these events. Note that each callback must be provided as a function or
34/// closure returning a thread-safe future. The future will be awaited by the
35/// runtime.
36///
37/// # Example
38///
39/// ```no_run
40/// # use rustdtp::prelude::*;
41///
42/// # #[tokio::main]
43/// # async fn main() {
44/// let server = Server::builder()
45///     .sending::<usize>()
46///     .receiving::<String>()
47///     .with_event_callbacks(
48///         ServerEventCallbacks::new()
49///             .on_connect(move |client_id| async move {
50///                 // some async operation...
51///                 println!("Client with ID {} connected", client_id);
52///             })
53///             .on_disconnect(move |client_id| async move {
54///                 // some async operation...
55///                 println!("Client with ID {} disconnected", client_id);
56///             })
57///             .on_receive(move |client_id, data| async move {
58///                 // some async operation...
59///                 println!("Received data from client with ID {}: {}", client_id, data);
60///             })
61///             .on_stop(move || async move {
62///                 // some async operation...
63///                 println!("Server closed");
64///             })
65///     )
66///     .start(("127.0.0.1", 29275))
67///     .await
68///     .unwrap();
69/// # }
70/// ```
71#[allow(clippy::type_complexity)]
72#[must_use = "event callbacks do nothing unless you configure them for a server"]
73pub struct ServerEventCallbacks<R>
74where
75    R: DeserializeOwned + 'static,
76{
77    /// The `connect` event callback.
78    connect: Option<Arc<dyn Fn(usize) -> Pin<Box<dyn Future<Output = ()> + Send>> + Send + Sync>>,
79    /// The `disconnect` event callback.
80    disconnect:
81        Option<Arc<dyn Fn(usize) -> Pin<Box<dyn Future<Output = ()> + Send>> + Send + Sync>>,
82    /// The `receive` event callback.
83    receive:
84        Option<Arc<dyn Fn(usize, R) -> Pin<Box<dyn Future<Output = ()> + Send>> + Send + Sync>>,
85    /// The `stop` event callback.
86    stop: Option<Arc<dyn Fn() -> Pin<Box<dyn Future<Output = ()> + Send>> + Send + Sync>>,
87}
88
89impl<R> ServerEventCallbacks<R>
90where
91    R: DeserializeOwned + 'static,
92{
93    /// Creates a new server event callbacks configuration with all callbacks
94    /// empty.
95    pub const fn new() -> Self {
96        Self {
97            connect: None,
98            disconnect: None,
99            receive: None,
100            stop: None,
101        }
102    }
103
104    /// Registers a callback on the `connect` event.
105    pub fn on_connect<C, F>(mut self, callback: C) -> Self
106    where
107        C: Fn(usize) -> F + Send + Sync + 'static,
108        F: Future<Output = ()> + Send + 'static,
109    {
110        self.connect = Some(Arc::new(move |client_id| Box::pin((callback)(client_id))));
111        self
112    }
113
114    /// Registers a callback on the `disconnect` event.
115    pub fn on_disconnect<C, F>(mut self, callback: C) -> Self
116    where
117        C: Fn(usize) -> F + Send + Sync + 'static,
118        F: Future<Output = ()> + Send + 'static,
119    {
120        self.disconnect = Some(Arc::new(move |client_id| Box::pin((callback)(client_id))));
121        self
122    }
123
124    /// Registers a callback on the `receive` event.
125    pub fn on_receive<C, F>(mut self, callback: C) -> Self
126    where
127        C: Fn(usize, R) -> F + Send + Sync + 'static,
128        F: Future<Output = ()> + Send + 'static,
129    {
130        self.receive = Some(Arc::new(move |client_id, data| {
131            Box::pin((callback)(client_id, data))
132        }));
133        self
134    }
135
136    /// Registers a callback on the `stop` event.
137    pub fn on_stop<C, F>(mut self, callback: C) -> Self
138    where
139        C: Fn() -> F + Send + Sync + 'static,
140        F: Future<Output = ()> + Send + 'static,
141    {
142        self.stop = Some(Arc::new(move || Box::pin((callback)())));
143        self
144    }
145}
146
147impl<R> Default for ServerEventCallbacks<R>
148where
149    R: DeserializeOwned + 'static,
150{
151    fn default() -> Self {
152        Self::new()
153    }
154}
155
156/// An event handling trait for the server.
157///
158/// # Events
159///
160/// There are four events for which methods can be implemented:
161///
162///  - `connect`
163///  - `disconnect`
164///  - `receive`
165///  - `stop`
166///
167/// All method implementations are optional, and can be registered for any
168/// combination of these events. Note that the type that implements the trait
169/// must be `Send + Sync`, and that all event method futures must be `Send`.
170///
171/// # Example
172///
173/// ```no_run
174/// # use rustdtp::prelude::*;
175///
176/// # #[tokio::main]
177/// # async fn main() {
178/// struct MyServerHandler;
179///
180/// impl ServerEventHandler<String> for MyServerHandler {
181///     async fn on_connect(&self, client_id: usize) {
182///         // some async operation...
183///         println!("Client with ID {} connected", client_id);
184///     }
185///
186///     async fn on_disconnect(&self, client_id: usize) {
187///         // some async operation...
188///         println!("Client with ID {} disconnected", client_id);
189///     }
190///
191///     async fn on_receive(&self, client_id: usize, data: String) {
192///         // some async operation...
193///         println!("Received data from client with ID {}: {}", client_id, data);
194///     }
195///
196///     async fn on_stop(&self) {
197///         // some async operation...
198///         println!("Server closed");
199///     }
200/// }
201///
202/// let server = Server::builder()
203///     .sending::<usize>()
204///     .receiving::<String>()
205///     .with_event_handler(MyServerHandler)
206///     .start(("127.0.0.1", 29275))
207///     .await
208///     .unwrap();
209/// # }
210/// ```
211pub trait ServerEventHandler<R>
212where
213    Self: Send + Sync,
214    R: DeserializeOwned + 'static,
215{
216    /// Handles the `connect` event.
217    #[allow(unused_variables)]
218    fn on_connect(&self, client_id: usize) -> impl Future<Output = ()> + Send {
219        async {}
220    }
221
222    /// Handles the `disconnect` event.
223    #[allow(unused_variables)]
224    fn on_disconnect(&self, client_id: usize) -> impl Future<Output = ()> + Send {
225        async {}
226    }
227
228    /// Handles the `receive` event.
229    #[allow(unused_variables)]
230    fn on_receive(&self, client_id: usize, data: R) -> impl Future<Output = ()> + Send {
231        async {}
232    }
233
234    /// Handles the `stop` event.
235    fn on_stop(&self) -> impl Future<Output = ()> + Send {
236        async {}
237    }
238}
239
240/// Unknown server sending type.
241pub struct ServerSendingUnknown;
242
243/// Known server sending type, stored as the type parameter `S`.
244pub struct ServerSending<S>(PhantomData<fn() -> S>)
245where
246    S: Serialize + 'static;
247
248/// A server sending marker trait.
249trait ServerSendingConfig {}
250
251impl ServerSendingConfig for ServerSendingUnknown {}
252
253impl<S> ServerSendingConfig for ServerSending<S> where S: Serialize + 'static {}
254
255/// Unknown server receiving type.
256pub struct ServerReceivingUnknown;
257
258/// Known server receiving type, stored as the type parameter `R`.
259pub struct ServerReceiving<R>(PhantomData<fn() -> R>)
260where
261    R: DeserializeOwned + 'static;
262
263/// A server receiving marker trait.
264trait ServerReceivingConfig {}
265
266impl ServerReceivingConfig for ServerReceivingUnknown {}
267
268impl<R> ServerReceivingConfig for ServerReceiving<R> where R: DeserializeOwned + 'static {}
269
270/// Unknown server event reporting type.
271pub struct ServerEventReportingUnknown;
272
273/// Known server event reporting type, stored as the type parameter `E`.
274pub struct ServerEventReporting<E>(E);
275
276/// Server event reporting via callbacks.
277pub struct ServerEventReportingCallbacks<R>(ServerEventCallbacks<R>)
278where
279    R: DeserializeOwned + 'static;
280
281/// Server event reporting via an event handler.
282pub struct ServerEventReportingHandler<R, H>
283where
284    R: DeserializeOwned + 'static,
285    H: ServerEventHandler<R>,
286{
287    /// The event handler instance.
288    handler: H,
289    /// Phantom `R` owner.
290    phantom_receive: PhantomData<fn() -> R>,
291}
292
293/// Server event reporting via a channel.
294pub struct ServerEventReportingChannel;
295
296/// A server event reporting marker trait.
297trait ServerEventReportingConfig {}
298
299impl ServerEventReportingConfig for ServerEventReportingUnknown {}
300
301impl<R> ServerEventReportingConfig for ServerEventReporting<ServerEventReportingCallbacks<R>> where
302    R: DeserializeOwned + 'static
303{
304}
305
306impl<R, H> ServerEventReportingConfig for ServerEventReporting<ServerEventReportingHandler<R, H>>
307where
308    R: DeserializeOwned + 'static,
309    H: ServerEventHandler<R>,
310{
311}
312
313impl ServerEventReportingConfig for ServerEventReporting<ServerEventReportingChannel> {}
314
315/// A builder for the [`Server`].
316///
317/// An instance of this can be constructed using `ServerBuilder::new()` or
318/// `Server::builder()`. The configuration information exists primarily at the
319/// type-level, so it is impossible to misconfigure this.
320///
321/// This method of configuration is technically not necessary, but it is far
322/// clearer and more explicit than simply configuring the `Server` type. Plus,
323/// it provides additional ways of detecting events.
324///
325/// # Configuration
326///
327/// To configure the server, first provide the types that will be sent and
328/// received through the server using the `.sending::<...>()` and
329/// `.receiving::<...>()` methods. Then specify the way in which events will
330/// be detected. There are three methods of receiving events:
331///
332/// - via callback functions (`.with_event_callbacks(...)`)
333/// - via implementation of a handler trait (`.with_event_handler(...)`)
334/// - via a channel (`.with_event_channel()`)
335///
336/// The channel method is the most versatile, hence why it's the `Server`'s
337/// default implementation. The other methods are provided to support a
338/// greater variety of program architectures.
339///
340/// Once configured, the `.start(...)` method, which is effectively identical
341/// to the `Server::start(...)` method, can be called to start the server.
342///
343/// # Example
344///
345/// ```no_run
346/// # use rustdtp::prelude::*;
347///
348/// # #[tokio::main]
349/// # async fn main() {
350/// let (server, server_events) = Server::builder()
351///     .sending::<usize>()
352///     .receiving::<String>()
353///     .with_event_channel()
354///     .start(("127.0.0.1", 29275))
355///     .await
356///     .unwrap();
357/// # }
358/// ```
359#[allow(private_bounds)]
360#[must_use = "server builders do nothing unless `start` is called"]
361pub struct ServerBuilder<SC, RC, EC>
362where
363    SC: ServerSendingConfig,
364    RC: ServerReceivingConfig,
365    EC: ServerEventReportingConfig,
366{
367    /// Phantom marker for `SC` and `RC`.
368    marker: PhantomData<fn() -> (SC, RC)>,
369    /// The event reporting configuration.
370    event_reporting: EC,
371}
372
373impl ServerBuilder<ServerSendingUnknown, ServerReceivingUnknown, ServerEventReportingUnknown> {
374    /// Creates a new server builder.
375    pub const fn new() -> Self {
376        Self {
377            marker: PhantomData,
378            event_reporting: ServerEventReportingUnknown,
379        }
380    }
381}
382
383impl Default
384    for ServerBuilder<ServerSendingUnknown, ServerReceivingUnknown, ServerEventReportingUnknown>
385{
386    fn default() -> Self {
387        Self::new()
388    }
389}
390
391#[allow(private_bounds)]
392impl<RC, EC> ServerBuilder<ServerSendingUnknown, RC, EC>
393where
394    RC: ServerReceivingConfig,
395    EC: ServerEventReportingConfig,
396{
397    /// Configures the type of data the server intends to send to clients.
398    pub fn sending<S>(self) -> ServerBuilder<ServerSending<S>, RC, EC>
399    where
400        S: Serialize + 'static,
401    {
402        ServerBuilder {
403            marker: PhantomData,
404            event_reporting: self.event_reporting,
405        }
406    }
407}
408
409#[allow(private_bounds)]
410impl<SC, EC> ServerBuilder<SC, ServerReceivingUnknown, EC>
411where
412    SC: ServerSendingConfig,
413    EC: ServerEventReportingConfig,
414{
415    /// Configures the type of data the server intends to receive from
416    /// clients.
417    pub fn receiving<R>(self) -> ServerBuilder<SC, ServerReceiving<R>, EC>
418    where
419        R: DeserializeOwned + 'static,
420    {
421        ServerBuilder {
422            marker: PhantomData,
423            event_reporting: self.event_reporting,
424        }
425    }
426}
427
428impl<S, R> ServerBuilder<ServerSending<S>, ServerReceiving<R>, ServerEventReportingUnknown>
429where
430    S: Serialize + 'static,
431    R: DeserializeOwned + 'static,
432{
433    /// Configures the server to receive events via callbacks.
434    ///
435    /// Using callbacks is typically considered an anti-pattern in Rust, so
436    /// this should only be used if it makes sense in the context of the
437    /// design of the code utilizing this API.
438    ///
439    /// See [`ServerEventCallbacks`] for more information and examples.
440    pub fn with_event_callbacks(
441        self,
442        callbacks: ServerEventCallbacks<R>,
443    ) -> ServerBuilder<
444        ServerSending<S>,
445        ServerReceiving<R>,
446        ServerEventReporting<ServerEventReportingCallbacks<R>>,
447    >
448    where
449        R: DeserializeOwned + 'static,
450    {
451        ServerBuilder {
452            marker: PhantomData,
453            event_reporting: ServerEventReporting(ServerEventReportingCallbacks(callbacks)),
454        }
455    }
456
457    /// Configures the server to receive events via a trait implementation.
458    ///
459    /// This provides an approach to event handling that closely aligns with
460    /// object-oriented practices.
461    ///
462    /// See [`ServerEventHandler`] for more information and examples.
463    pub fn with_event_handler<H>(
464        self,
465        handler: H,
466    ) -> ServerBuilder<
467        ServerSending<S>,
468        ServerReceiving<R>,
469        ServerEventReporting<ServerEventReportingHandler<R, H>>,
470    >
471    where
472        H: ServerEventHandler<R>,
473    {
474        ServerBuilder {
475            marker: PhantomData,
476            event_reporting: ServerEventReporting(ServerEventReportingHandler {
477                handler,
478                phantom_receive: PhantomData,
479            }),
480        }
481    }
482
483    /// Configures the server to receive events via a channel.
484    ///
485    /// This is the most versatile event handling strategy. In fact, all other
486    /// event handling options use this implementation under the hood.
487    /// Because of its flexibility, this will typically be the desired
488    /// approach.
489    pub fn with_event_channel(
490        self,
491    ) -> ServerBuilder<
492        ServerSending<S>,
493        ServerReceiving<R>,
494        ServerEventReporting<ServerEventReportingChannel>,
495    > {
496        ServerBuilder {
497            marker: PhantomData,
498            event_reporting: ServerEventReporting(ServerEventReportingChannel),
499        }
500    }
501}
502
503impl<S, R>
504    ServerBuilder<
505        ServerSending<S>,
506        ServerReceiving<R>,
507        ServerEventReporting<ServerEventReportingCallbacks<R>>,
508    >
509where
510    S: Serialize + 'static,
511    R: DeserializeOwned + 'static,
512{
513    /// Starts the server. This is effectively identical to [`Server::start`].
514    ///
515    /// # Errors
516    ///
517    /// The set of errors that can occur are identical to that of
518    /// [`Server::start`].
519    #[allow(clippy::future_not_send)]
520    pub async fn start<A>(self, addr: A) -> Result<ServerHandle<S>>
521    where
522        A: ToSocketAddrs,
523    {
524        let (server, mut server_events) = Server::<S, R>::start(addr).await?;
525        let callbacks = self.event_reporting.0 .0;
526
527        tokio::spawn(async move {
528            while let Ok(event) = server_events.next_raw().await {
529                match event {
530                    ServerEventRawSafe::Connect { client_id } => {
531                        if let Some(ref connect) = callbacks.connect {
532                            let connect = Arc::clone(connect);
533                            tokio::spawn(async move {
534                                (*connect)(client_id).await;
535                            });
536                        }
537                    }
538                    ServerEventRawSafe::Disconnect { client_id } => {
539                        if let Some(ref disconnect) = callbacks.disconnect {
540                            let disconnect = Arc::clone(disconnect);
541                            tokio::spawn(async move {
542                                (*disconnect)(client_id).await;
543                            });
544                        }
545                    }
546                    ServerEventRawSafe::Receive { client_id, data } => {
547                        if let Some(ref receive) = callbacks.receive {
548                            let receive = Arc::clone(receive);
549                            tokio::spawn(async move {
550                                let data = data.deserialize();
551                                (*receive)(client_id, data).await;
552                            });
553                        }
554                    }
555                    ServerEventRawSafe::Stop => {
556                        if let Some(ref stop) = callbacks.stop {
557                            let stop = Arc::clone(stop);
558                            tokio::spawn(async move {
559                                (*stop)().await;
560                            });
561                        }
562                    }
563                }
564            }
565        });
566
567        Ok(server)
568    }
569}
570
571impl<S, R, H>
572    ServerBuilder<
573        ServerSending<S>,
574        ServerReceiving<R>,
575        ServerEventReporting<ServerEventReportingHandler<R, H>>,
576    >
577where
578    S: Serialize + 'static,
579    R: DeserializeOwned + 'static,
580    H: ServerEventHandler<R> + 'static,
581{
582    /// Starts the server. This is effectively identical to [`Server::start`].
583    ///
584    /// # Errors
585    ///
586    /// The set of errors that can occur are identical to that of
587    /// [`Server::start`].
588    #[allow(clippy::future_not_send)]
589    pub async fn start<A>(self, addr: A) -> Result<ServerHandle<S>>
590    where
591        A: ToSocketAddrs,
592    {
593        let (server, mut server_events) = Server::<S, R>::start(addr).await?;
594        let handler = Arc::new(self.event_reporting.0.handler);
595
596        tokio::spawn(async move {
597            while let Ok(event) = server_events.next_raw().await {
598                match event {
599                    ServerEventRawSafe::Connect { client_id } => {
600                        let handler = Arc::clone(&handler);
601                        tokio::spawn(async move {
602                            handler.on_connect(client_id).await;
603                        });
604                    }
605                    ServerEventRawSafe::Disconnect { client_id } => {
606                        let handler = Arc::clone(&handler);
607                        tokio::spawn(async move {
608                            handler.on_disconnect(client_id).await;
609                        });
610                    }
611                    ServerEventRawSafe::Receive { client_id, data } => {
612                        let handler = Arc::clone(&handler);
613                        tokio::spawn(async move {
614                            let data = data.deserialize();
615                            handler.on_receive(client_id, data).await;
616                        });
617                    }
618                    ServerEventRawSafe::Stop => {
619                        let handler = Arc::clone(&handler);
620                        tokio::spawn(async move {
621                            handler.on_stop().await;
622                        });
623                    }
624                }
625            }
626        });
627
628        Ok(server)
629    }
630}
631
632impl<S, R>
633    ServerBuilder<
634        ServerSending<S>,
635        ServerReceiving<R>,
636        ServerEventReporting<ServerEventReportingChannel>,
637    >
638where
639    S: Serialize + 'static,
640    R: DeserializeOwned + 'static,
641{
642    /// Starts the server. This is effectively identical to [`Server::start`].
643    ///
644    /// # Errors
645    ///
646    /// The set of errors that can occur are identical to that of
647    /// [`Server::start`].
648    #[allow(clippy::future_not_send)]
649    pub async fn start<A>(self, addr: A) -> Result<(ServerHandle<S>, ServerEventStream<R>)>
650    where
651        A: ToSocketAddrs,
652    {
653        Server::<S, R>::start(addr).await
654    }
655}
656
657/// A command sent from the server handle to the background server task.
658pub enum ServerCommand {
659    /// Stop the server.
660    Stop,
661    /// Send data to a client.
662    Send {
663        /// The ID of the client to send the data to.
664        client_id: usize,
665        /// The data to send.
666        data: Vec<u8>,
667    },
668    /// Send data to all clients.
669    SendAll {
670        /// The data to send.
671        data: Vec<u8>,
672    },
673    /// Get the local server address.
674    GetAddr,
675    /// Get the address of a client.
676    GetClientAddr {
677        /// The ID of the client.
678        client_id: usize,
679    },
680    /// Disconnect a client from the server.
681    RemoveClient {
682        /// The ID of the client.
683        client_id: usize,
684    },
685}
686
687/// The return value of a command executed on the background server task.
688pub enum ServerCommandReturn {
689    /// Stop return value.
690    Stop(Result<()>),
691    /// Sent data return value.
692    Send(Result<()>),
693    /// Sent data to all return value.
694    SendAll(Result<()>),
695    /// Local server address return value.
696    GetAddr(Result<SocketAddr>),
697    /// Client address return value.
698    GetClientAddr(Result<SocketAddr>),
699    /// Disconnect client return value.
700    RemoveClient(Result<()>),
701}
702
703/// A command sent from the server background task to a client background task.
704pub enum ServerClientCommand {
705    /// Send data to the client.
706    Send {
707        /// The serialized data to send.
708        data: Arc<[u8]>,
709    },
710    /// Get the address of the client.
711    GetAddr,
712    /// Disconnect the client.
713    Remove,
714}
715
716/// The return value of a command executed on a client background task.
717pub enum ServerClientCommandReturn {
718    /// Send data return value.
719    Send(Result<()>),
720    /// Client address return value.
721    GetAddr(Result<SocketAddr>),
722    /// Disconnect client return value.
723    Remove(Result<()>),
724}
725
726/// An event from the server.
727///
728/// ```no_run
729/// use rustdtp::prelude::*;
730///
731/// #[tokio::main]
732/// async fn main() {
733///     // Create the server
734///     let (mut server, mut server_events) = Server::builder()
735///         .sending::<()>()
736///         .receiving::<String>()
737///         .with_event_channel()
738///         .start(("127.0.0.1", 29275))
739///         .await
740///         .unwrap();
741///
742///     // Iterate over events
743///     while let Ok(event) = server_events.next().await {
744///         match event {
745///             ServerEvent::Connect { client_id } => {
746///                 println!("Client with ID {} connected", client_id);
747///             }
748///             ServerEvent::Disconnect { client_id } => {
749///                 println!("Client with ID {} disconnected", client_id);
750///             }
751///             ServerEvent::Receive { client_id, data } => {
752///                 println!("Client with ID {} sent: {}", client_id, data);
753///             }
754///             ServerEvent::Stop => {
755///                 // No more events will be sent, and the loop will end
756///                 println!("Server closed");
757///             }
758///         }
759///     }
760/// }
761/// ```
762#[derive(Debug, Clone)]
763pub enum ServerEvent<R>
764where
765    R: DeserializeOwned + 'static,
766{
767    /// A client connected.
768    Connect {
769        /// The ID of the client that connected.
770        client_id: usize,
771    },
772    /// A client disconnected.
773    Disconnect {
774        /// The ID of the client that disconnected.
775        client_id: usize,
776    },
777    /// Data received from a client.
778    Receive {
779        /// The ID of the client that sent the data.
780        client_id: usize,
781        /// The data itself.
782        data: R,
783    },
784    /// Server stopped.
785    Stop,
786}
787
788/// Identical to `ServerEvent`, but with the received data in serialized form.
789#[derive(Debug, Clone)]
790enum ServerEventRaw {
791    /// A client connected.
792    Connect {
793        /// The ID of the client that connected.
794        client_id: usize,
795    },
796    /// A client disconnected.
797    Disconnect {
798        /// The ID of the client that disconnected.
799        client_id: usize,
800    },
801    /// Data received from a client.
802    Receive {
803        /// The ID of the client that sent the data.
804        client_id: usize,
805        /// The data itself.
806        data: Vec<u8>,
807    },
808    /// Server stopped.
809    Stop,
810}
811
812impl ServerEventRaw {
813    /// Deserializes this instance into a `ServerEvent`.
814    fn deserialize<R>(&self) -> Result<ServerEvent<R>>
815    where
816        R: DeserializeOwned + 'static,
817    {
818        match self {
819            Self::Connect { client_id } => Ok(ServerEvent::Connect {
820                client_id: *client_id,
821            }),
822            Self::Disconnect { client_id } => Ok(ServerEvent::Disconnect {
823                client_id: *client_id,
824            }),
825            Self::Receive { client_id, data } => {
826                Ok(
827                    serde_json::from_slice(data).map(|data| ServerEvent::Receive {
828                        client_id: *client_id,
829                        data,
830                    })?,
831                )
832            }
833            Self::Stop => Ok(ServerEvent::Stop),
834        }
835    }
836}
837
838/// The serialized data component of a server receive event. The data is
839/// guaranteed to be deserializable into an instance of `R`.
840#[derive(Debug, Clone)]
841struct ServerEventRawSafeData<R>
842where
843    R: DeserializeOwned + 'static,
844{
845    /// The raw data.
846    data: Vec<u8>,
847    /// Phantom marker for `R`.
848    marker: PhantomData<fn() -> R>,
849}
850
851/// Identical to `ServerEventRaw`, but with the guarantee that the data can be
852/// deserialized into an instance of `R`.
853#[derive(Debug, Clone)]
854enum ServerEventRawSafe<R>
855where
856    R: DeserializeOwned + 'static,
857{
858    /// A client connected.
859    Connect {
860        /// The ID of the client that connected.
861        client_id: usize,
862    },
863    /// A client disconnected.
864    Disconnect {
865        /// The ID of the client that disconnected.
866        client_id: usize,
867    },
868    /// Data received from a client.
869    Receive {
870        /// The ID of the client that sent the data.
871        client_id: usize,
872        /// The data itself.
873        data: ServerEventRawSafeData<R>,
874    },
875    /// Server stopped.
876    Stop,
877}
878
879impl<R> TryFrom<ServerEventRaw> for ServerEventRawSafe<R>
880where
881    R: DeserializeOwned + 'static,
882{
883    type Error = Error;
884
885    fn try_from(value: ServerEventRaw) -> std::result::Result<Self, Self::Error> {
886        value.deserialize::<R>()?;
887
888        Ok(match value {
889            ServerEventRaw::Connect { client_id } => Self::Connect { client_id },
890            ServerEventRaw::Disconnect { client_id } => Self::Disconnect { client_id },
891            ServerEventRaw::Receive { client_id, data } => Self::Receive {
892                client_id,
893                data: ServerEventRawSafeData {
894                    data,
895                    marker: PhantomData,
896                },
897            },
898            ServerEventRaw::Stop => Self::Stop,
899        })
900    }
901}
902
903impl<R> ServerEventRawSafeData<R>
904where
905    R: DeserializeOwned + 'static,
906{
907    /// Deserialize the raw data into an instance of `R`. This is guaranteed to
908    /// succeed.
909    fn deserialize(&self) -> R {
910        serde_json::from_slice(&self.data).unwrap()
911    }
912}
913
914impl<R> ServerEventRawSafe<R>
915where
916    R: DeserializeOwned + 'static,
917{
918    /// Deserializes this instance into a `ServerEvent`.
919    #[allow(dead_code)]
920    fn deserialize(&self) -> ServerEvent<R> {
921        match self {
922            Self::Connect { client_id } => ServerEvent::Connect {
923                client_id: *client_id,
924            },
925            Self::Disconnect { client_id } => ServerEvent::Disconnect {
926                client_id: *client_id,
927            },
928            Self::Receive { client_id, data } => ServerEvent::Receive {
929                client_id: *client_id,
930                data: data.deserialize(),
931            },
932            Self::Stop => ServerEvent::Stop,
933        }
934    }
935}
936
937/// An asynchronous stream of server events.
938pub struct ServerEventStream<R>
939where
940    R: DeserializeOwned + 'static,
941{
942    /// The event receiver channel.
943    event_receiver: Receiver<ServerEventRaw>,
944    /// Phantom marker for `R`.
945    marker: PhantomData<fn() -> R>,
946}
947
948impl<R> ServerEventStream<R>
949where
950    R: DeserializeOwned + 'static,
951{
952    /// Consumes and returns the next value in the stream.
953    ///
954    /// # Errors
955    ///
956    /// This will return an error if the stream is closed, or if there was an
957    /// error while deserializing data received.
958    pub async fn next(&mut self) -> Result<ServerEvent<R>> {
959        match self.event_receiver.recv().await {
960            Some(serialized_event) => serialized_event.deserialize(),
961            None => Err(Error::ConnectionClosed),
962        }
963    }
964
965    /// Identical to `next`, but doesn't deserialize the event. It does,
966    /// however, validate that the event can be deserialized without error.
967    async fn next_raw(&mut self) -> Result<ServerEventRawSafe<R>> {
968        match self.event_receiver.recv().await {
969            Some(serialized_event) => serialized_event.try_into(),
970            None => Err(Error::ConnectionClosed),
971        }
972    }
973}
974
975/// A handle to the server.
976pub struct ServerHandle<S>
977where
978    S: Serialize + 'static,
979{
980    /// The channel through which commands can be sent to the background task.
981    server_command_sender: CommandChannelSender<ServerCommand, ServerCommandReturn>,
982    /// The handle to the background task.
983    server_task_handle: JoinHandle<Result<()>>,
984    /// Phantom marker for `S`.
985    marker: PhantomData<fn() -> S>,
986}
987
988impl<S> ServerHandle<S>
989where
990    S: Serialize + 'static,
991{
992    /// Stop the server, disconnect all clients, and shut down all network
993    /// interfaces.
994    ///
995    /// Returns a result of the error variant if an error occurred while
996    /// disconnecting clients.
997    ///
998    /// ```no_run
999    /// use rustdtp::prelude::*;
1000    ///
1001    /// #[tokio::main]
1002    /// async fn main() {
1003    ///     // Create the server
1004    ///     let (mut server, mut server_events) = Server::builder()
1005    ///         .sending::<()>()
1006    ///         .receiving::<String>()
1007    ///         .with_event_channel()
1008    ///         .start(("127.0.0.1", 29275))
1009    ///         .await
1010    ///         .unwrap();
1011    ///
1012    ///     // Wait for events until a client requests the server be stopped
1013    ///     while let Ok(event) = server_events.next().await {
1014    ///         match event {
1015    ///             // Stop the server when a client requests it be stopped
1016    ///             ServerEvent::Receive { client_id, data } => {
1017    ///                 if data.as_str() == "Stop the server!" {
1018    ///                     println!("Server stop requested");
1019    ///                     server.stop().await.unwrap();
1020    ///                     break;
1021    ///                 }
1022    ///             }
1023    ///             _ => {}  // Do nothing for other events
1024    ///         }
1025    ///     }
1026    ///
1027    ///     // The last event should be a stop event
1028    ///     assert!(matches!(server_events.next().await.unwrap(), ServerEvent::Stop));
1029    /// }
1030    /// ```
1031    ///
1032    /// # Errors
1033    ///
1034    /// This will return an error if the server socket has already closed, or if
1035    /// the underlying server loop returned an error.
1036    #[allow(clippy::missing_panics_doc)]
1037    pub async fn stop(mut self) -> Result<()> {
1038        let value = self
1039            .server_command_sender
1040            .send_command(ServerCommand::Stop)
1041            .await?;
1042        // `unwrap` is allowed, as an error is returned only when the underlying
1043        // task panics, which it never should
1044        self.server_task_handle.await.unwrap()?;
1045        unwrap_enum!(value, ServerCommandReturn::Stop)
1046    }
1047
1048    /// Send data to a client.
1049    ///
1050    /// - `client_id`: the ID of the client to send the data to.
1051    /// - `data`: the data to send.
1052    ///
1053    /// Returns a result of the error variant if an error occurred while
1054    /// sending.
1055    ///
1056    /// ```no_run
1057    /// use rustdtp::prelude::*;
1058    ///
1059    /// #[tokio::main]
1060    /// async fn main() {
1061    ///     // Create the server
1062    ///     let (mut server, mut server_events) = Server::builder()
1063    ///         .sending::<String>()
1064    ///         .receiving::<()>()
1065    ///         .with_event_channel()
1066    ///         .start(("127.0.0.1", 29275))
1067    ///         .await
1068    ///         .unwrap();
1069    ///
1070    ///     // Iterate over events
1071    ///     while let Ok(event) = server_events.next().await {
1072    ///         match event {
1073    ///             // When a client connects, send a greeting
1074    ///             ServerEvent::Connect { client_id } => {
1075    ///                 server.send(client_id, format!("Hello, client {}!", client_id)).await.unwrap();
1076    ///             }
1077    ///             _ => {}  // Do nothing for other events
1078    ///         }
1079    ///     }
1080    /// }
1081    /// ```
1082    ///
1083    /// # Errors
1084    ///
1085    /// This will return an error if the server socket has closed, or if data
1086    /// serialization fails.
1087    #[allow(clippy::future_not_send)]
1088    pub async fn send(&mut self, client_id: usize, data: S) -> Result<()> {
1089        let data_serialized = serde_json::to_vec(&data)?;
1090        let value = self
1091            .server_command_sender
1092            .send_command(ServerCommand::Send {
1093                client_id,
1094                data: data_serialized,
1095            })
1096            .await?;
1097        unwrap_enum!(value, ServerCommandReturn::Send)
1098    }
1099
1100    /// Send data to all clients.
1101    ///
1102    /// - `data`: the data to send.
1103    ///
1104    /// Returns a result of the error variant if an error occurred while
1105    /// sending.
1106    ///
1107    /// ```no_run
1108    /// use rustdtp::prelude::*;
1109    ///
1110    /// #[tokio::main]
1111    /// async fn main() {
1112    ///     // Create the server
1113    ///     let (mut server, mut server_events) = Server::builder()
1114    ///         .sending::<String>()
1115    ///         .receiving::<()>()
1116    ///         .with_event_channel()
1117    ///         .start(("127.0.0.1", 29275))
1118    ///         .await
1119    ///         .unwrap();
1120    ///
1121    ///     // Iterate over events
1122    ///     while let Ok(event) = server_events.next().await {
1123    ///         match event {
1124    ///             // When a client connects, notify all clients
1125    ///             ServerEvent::Connect { client_id } => {
1126    ///                 server.send_all(format!("A new client with ID {} has joined!", client_id)).await.unwrap();
1127    ///             }
1128    ///             _ => {}  // Do nothing for other events
1129    ///         }
1130    ///     }
1131    /// }
1132    /// ```
1133    ///
1134    /// # Errors
1135    ///
1136    /// This will return an error if the server socket has closed, or if data
1137    /// serialization fails.
1138    #[allow(clippy::future_not_send)]
1139    pub async fn send_all(&mut self, data: S) -> Result<()> {
1140        let data_serialized = serde_json::to_vec(&data)?;
1141        let value = self
1142            .server_command_sender
1143            .send_command(ServerCommand::SendAll {
1144                data: data_serialized,
1145            })
1146            .await?;
1147        unwrap_enum!(value, ServerCommandReturn::SendAll)
1148    }
1149
1150    /// Get the address the server is listening on.
1151    ///
1152    /// Returns a result containing the address the server is listening on, or
1153    /// the error variant if an error occurred.
1154    ///
1155    /// ```no_run
1156    /// use rustdtp::prelude::*;
1157    ///
1158    /// #[tokio::main]
1159    /// async fn main() {
1160    ///     // Create the server
1161    ///     let (mut server, mut server_events) = Server::builder()
1162    ///         .sending::<()>()
1163    ///         .receiving::<()>()
1164    ///         .with_event_channel()
1165    ///         .start(("127.0.0.1", 29275))
1166    ///         .await
1167    ///         .unwrap();
1168    ///
1169    ///     // Get the server address
1170    ///     let addr = server.get_addr().await.unwrap();
1171    ///     println!("Server listening on {}", addr);
1172    /// }
1173    /// ```
1174    ///
1175    /// # Errors
1176    ///
1177    /// This will return an error if the server socket has closed.
1178    pub async fn get_addr(&mut self) -> Result<SocketAddr> {
1179        let value = self
1180            .server_command_sender
1181            .send_command(ServerCommand::GetAddr)
1182            .await?;
1183        unwrap_enum!(value, ServerCommandReturn::GetAddr)
1184    }
1185
1186    /// Get the address of a connected client.
1187    ///
1188    /// - `client_id`: the ID of the client.
1189    ///
1190    /// Returns a result containing the address of the client, or the error
1191    /// variant if the client ID is invalid.
1192    ///
1193    /// ```no_run
1194    /// use rustdtp::prelude::*;
1195    ///
1196    /// #[tokio::main]
1197    /// async fn main() {
1198    ///     // Create the server
1199    ///     let (mut server, mut server_events) = Server::builder()
1200    ///         .sending::<()>()
1201    ///         .receiving::<()>()
1202    ///         .with_event_channel()
1203    ///         .start(("127.0.0.1", 29275))
1204    ///         .await
1205    ///         .unwrap();
1206    ///
1207    ///     // Iterate over events
1208    ///     while let Ok(event) = server_events.next().await {
1209    ///         match event {
1210    ///             // When a client connects, get their address
1211    ///             ServerEvent::Connect { client_id } => {
1212    ///                 let addr = server.get_client_addr(client_id).await.unwrap();
1213    ///                 println!("Client with ID {} connected from {}", client_id, addr);
1214    ///             }
1215    ///             _ => {}  // Do nothing for other events
1216    ///         }
1217    ///     }
1218    /// }
1219    /// ```
1220    ///
1221    /// # Errors
1222    ///
1223    /// This will return an error if the server socket has closed, or if the
1224    /// client ID is invalid.
1225    pub async fn get_client_addr(&mut self, client_id: usize) -> Result<SocketAddr> {
1226        let value = self
1227            .server_command_sender
1228            .send_command(ServerCommand::GetClientAddr { client_id })
1229            .await?;
1230        unwrap_enum!(value, ServerCommandReturn::GetClientAddr)
1231    }
1232
1233    /// Disconnect a client from the server.
1234    ///
1235    /// - `client_id`: the ID of the client.
1236    ///
1237    /// Returns a result of the error variant if an error occurred while
1238    /// disconnecting the client, or if the client ID is invalid.
1239    ///
1240    /// ```no_run
1241    /// use rustdtp::prelude::*;
1242    ///
1243    /// #[tokio::main]
1244    /// async fn main() {
1245    ///     // Create the server
1246    ///     let (mut server, mut server_events) = Server::builder()
1247    ///         .sending::<String>()
1248    ///         .receiving::<i32>()
1249    ///         .with_event_channel()
1250    ///         .start(("127.0.0.1", 29275))
1251    ///         .await
1252    ///         .unwrap();
1253    ///
1254    ///     // Iterate over events
1255    ///     while let Ok(event) = server_events.next().await {
1256    ///         match event {
1257    ///             // Disconnect a client if they send an even number
1258    ///             ServerEvent::Receive { client_id, data } => {
1259    ///                 if data % 2 == 0 {
1260    ///                     println!("Disconnecting client with ID {}", client_id);
1261    ///                     server.send(client_id, "Even numbers are not allowed".to_owned()).await.unwrap();
1262    ///                     server.remove_client(client_id).await.unwrap();
1263    ///                 }
1264    ///             }
1265    ///             _ => {}  // Do nothing for other events
1266    ///         }
1267    ///     }
1268    ///
1269    ///     // The last event should be a stop event
1270    ///     assert!(matches!(server_events.next().await.unwrap(), ServerEvent::Stop));
1271    /// }
1272    /// ```
1273    ///
1274    /// # Errors
1275    ///
1276    /// This will return an error if the server socket has closed, or if the
1277    /// client ID is invalid.
1278    pub async fn remove_client(&mut self, client_id: usize) -> Result<()> {
1279        let value = self
1280            .server_command_sender
1281            .send_command(ServerCommand::RemoveClient { client_id })
1282            .await?;
1283        unwrap_enum!(value, ServerCommandReturn::RemoveClient)
1284    }
1285}
1286
1287/// A socket server.
1288///
1289/// The server takes two generic parameters:
1290///
1291/// - `S`: the type of data that will be **sent** to clients.
1292/// - `R`: the type of data that will be **received** from clients.
1293///
1294/// Both types must be serializable in order to be sent through the socket. When
1295/// creating clients, the types should be swapped, since the server's send type will be the client's receive type and vice versa.
1296///
1297/// ```no_run
1298/// use rustdtp::prelude::*;
1299///
1300/// #[tokio::main]
1301/// async fn main() {
1302///     // Create a server that receives strings and returns the length of each string
1303///     let (mut server, mut server_events) = Server::builder()
1304///         .sending::<usize>()
1305///         .receiving::<String>()
1306///         .with_event_channel()
1307///         .start(("127.0.0.1", 29275))
1308///         .await
1309///         .unwrap();
1310///
1311///     // Iterate over events
1312///     while let Ok(event) = server_events.next().await {
1313///         match event {
1314///             ServerEvent::Connect { client_id } => {
1315///                 println!("Client with ID {} connected", client_id);
1316///             }
1317///             ServerEvent::Disconnect { client_id } => {
1318///                 println!("Client with ID {} disconnected", client_id);
1319///             }
1320///             ServerEvent::Receive { client_id, data } => {
1321///                 // Send back the length of the string
1322///                 server.send(client_id, data.len()).await.unwrap();
1323///             }
1324///             ServerEvent::Stop => {
1325///                 // No more events will be sent, and the loop will end
1326///                 println!("Server closed");
1327///             }
1328///         }
1329///     }
1330/// }
1331/// ```
1332pub struct Server<S, R>
1333where
1334    S: Serialize + 'static,
1335    R: DeserializeOwned + 'static,
1336{
1337    /// Phantom marker for `S` and `R`.
1338    marker: PhantomData<fn() -> (S, R)>,
1339}
1340
1341impl Server<(), ()> {
1342    /// Constructs a server builder. Use this for a clearer, more explicit,
1343    /// and more featureful server configuration. See [`ServerBuilder`] for
1344    /// more information.
1345    pub const fn builder(
1346    ) -> ServerBuilder<ServerSendingUnknown, ServerReceivingUnknown, ServerEventReportingUnknown>
1347    {
1348        ServerBuilder::new()
1349    }
1350}
1351
1352impl<S, R> Server<S, R>
1353where
1354    S: Serialize + 'static,
1355    R: DeserializeOwned + 'static,
1356{
1357    /// Start a socket server.
1358    ///
1359    /// - `addr`: the address for the server to listen on.
1360    ///
1361    /// Returns a result containing a handle to the server and a channel from
1362    /// which to receive server events, or the error variant if an error
1363    /// occurred while starting the server.
1364    ///
1365    /// ```no_run
1366    /// use rustdtp::prelude::*;
1367    ///
1368    /// #[tokio::main]
1369    /// async fn main() {
1370    ///     let (mut server, mut server_events) = Server::builder()
1371    ///         .sending::<()>()
1372    ///         .receiving::<()>()
1373    ///         .with_event_channel()
1374    ///         .start(("127.0.0.1", 29275))
1375    ///         .await
1376    ///         .unwrap();
1377    /// }
1378    /// ```
1379    ///
1380    /// Neither the server handle nor the event receiver should be dropped until
1381    /// the server has been stopped. Prematurely dropping either one can cause
1382    /// unintended behavior.
1383    ///
1384    /// # Errors
1385    ///
1386    /// This will return an error if a TCP listener cannot be bound to the
1387    /// provided address.
1388    #[allow(clippy::future_not_send)]
1389    pub async fn start<A>(addr: A) -> Result<(ServerHandle<S>, ServerEventStream<R>)>
1390    where
1391        A: ToSocketAddrs,
1392    {
1393        // Server TCP listener
1394        let listener = TcpListener::bind(addr).await?;
1395        // Channels for sending commands from the server handle to the background server task
1396        let (server_command_sender, server_command_receiver) = command_channel();
1397        // Channels for sending event notifications from the background server task
1398        let (server_event_sender, server_event_receiver) = channel(CHANNEL_BUFFER_SIZE);
1399
1400        // Start the background server task, saving the join handle for when the server is stopped
1401        let server_task_handle = tokio::spawn(server_handler(
1402            listener,
1403            server_event_sender,
1404            server_command_receiver,
1405        ));
1406
1407        // Create a handle for the server
1408        let server_handle = ServerHandle {
1409            server_command_sender,
1410            server_task_handle,
1411            marker: PhantomData,
1412        };
1413
1414        // Create an event stream for the server
1415        let server_event_stream = ServerEventStream {
1416            event_receiver: server_event_receiver,
1417            marker: PhantomData,
1418        };
1419
1420        Ok((server_handle, server_event_stream))
1421    }
1422}
1423
1424/// The server client loop. Handles received data and commands.
1425#[allow(clippy::too_many_lines)]
1426async fn server_client_loop(
1427    client_id: usize,
1428    mut socket: TcpStream,
1429    server_client_event_sender: Sender<ServerEventRaw>,
1430    mut client_command_receiver: CommandChannelReceiver<
1431        ServerClientCommand,
1432        ServerClientCommandReturn,
1433    >,
1434) -> Result<()> {
1435    // Generate X25519 keys
1436    let (public_key, secret_key) = dh_key_pair().await;
1437    // Send the public key to the client
1438    socket.write_all(public_key.as_bytes()).await?;
1439    // Flush the stream
1440    socket.flush().await?;
1441
1442    // Buffer in which to receive the client's public key
1443    let mut other_public_key = [0; PUBLIC_KEY_SIZE];
1444    // Read the public key from the client
1445    handshake_timeout! {
1446        socket.read_exact(&mut other_public_key)
1447    }??;
1448    // Establish the shared AES key
1449    let aes_key = dh_shared_key(secret_key, other_public_key).await;
1450
1451    // Buffer in which to receive the size portion of a message
1452    let mut size_buffer = [0; LEN_SIZE];
1453
1454    // Client loop
1455    loop {
1456        // Await messages from the client
1457        // and commands from the background server task
1458        tokio::select! {
1459            // Read the size portion from the client socket
1460            read_value = socket.read(&mut size_buffer[..]) => {
1461                // Return an error if the socket could not be read
1462                let n_size = read_value?;
1463
1464                // If there were no bytes read, or if there were fewer bytes
1465                // read than there should have been, close the socket
1466                if n_size != LEN_SIZE {
1467                    socket.shutdown().await?;
1468                    break;
1469                }
1470
1471                // Decode the size portion of the message
1472                let encrypted_data_size = decode_message_size(&size_buffer);
1473                // Initialize the buffer for the data portion of the message
1474                let mut encrypted_data_buffer = vec![0; encrypted_data_size];
1475
1476                // Read the data portion from the client socket, returning an
1477                // error if the socket could not be read
1478                let n_data = data_read_timeout! {
1479                    socket.read_exact(&mut encrypted_data_buffer[..])
1480                }??;
1481
1482                // If there were no bytes read, or if there were fewer bytes
1483                // read than there should have been, close the socket
1484                if n_data != encrypted_data_size {
1485                    socket.shutdown().await?;
1486                    break;
1487                }
1488
1489                // Decrypt the data
1490                let data_serialized = aes_decrypt(aes_key, encrypted_data_buffer.into()).await?;
1491
1492                // Send an event to note that a piece of data has been received from
1493                // a client
1494                if let Err(_e) = server_client_event_sender.send(ServerEventRaw::Receive { client_id, data: data_serialized }).await {
1495                    // Sending failed, disconnect the client
1496                    socket.shutdown().await?;
1497                    break;
1498                }
1499            }
1500            // Process a command sent to the client
1501            client_command_value = client_command_receiver.recv_command() => {
1502                // Handle the command, or lack thereof if the channel is closed
1503                match client_command_value {
1504                    Ok(client_command) => {
1505                        // Process the command
1506                        match client_command {
1507                            ServerClientCommand::Send { data } => {
1508                                let value = 'val: {
1509                                    // Encrypt the serialized data
1510                                    let encrypted_data_buffer = break_on_err!(aes_encrypt(aes_key, data).await, 'val);
1511                                    // Encode the message size to a buffer
1512                                    let size_buffer = encode_message_size(encrypted_data_buffer.len());
1513
1514                                    // Initialize the message buffer
1515                                    let mut buffer = vec![];
1516                                    // Extend the buffer to contain the payload
1517                                    // size
1518                                    buffer.extend_from_slice(&size_buffer);
1519                                    // Extend the buffer to contain the payload
1520                                    // data
1521                                    buffer.extend(&encrypted_data_buffer);
1522
1523                                    // Write the data to the client socket
1524                                    break_on_err!(socket.write_all(&buffer).await, 'val);
1525                                    // Flush the stream
1526                                    break_on_err!(socket.flush().await, 'val);
1527
1528                                    Ok(())
1529                                };
1530
1531                                let error_occurred = value.is_err();
1532
1533                                // Return the status of the send operation
1534                                if let Err(_e) = client_command_receiver.command_return(ServerClientCommandReturn::Send(value)).await {
1535                                    // Channel is closed, disconnect the client
1536                                    socket.shutdown().await?;
1537                                    break;
1538                                }
1539
1540                                // If the send failed, disconnect the client
1541                                if error_occurred {
1542                                    socket.shutdown().await?;
1543                                    break;
1544                                }
1545                            },
1546                            ServerClientCommand::GetAddr => {
1547                                // Get the client socket's address
1548                                let addr = socket.peer_addr();
1549
1550                                // Return the address
1551                                if let Err(_e) = client_command_receiver.command_return(ServerClientCommandReturn::GetAddr(addr.map_err(Into::into))).await {
1552                                    // Channel is closed, disconnect the client
1553                                    socket.shutdown().await?;
1554                                    break;
1555                                }
1556                            },
1557                            ServerClientCommand::Remove => {
1558                                // Disconnect the client
1559                                let value = socket.shutdown().await;
1560
1561                                // Return the status of the remove operation,
1562                                // ignoring failures, since a failure indicates
1563                                // that the client has probably already
1564                                // disconnected
1565                                _ = client_command_receiver.command_return(ServerClientCommandReturn::Remove(value.map_err(Into::into))).await;
1566
1567                                // Break the client loop
1568                                break;
1569                            },
1570                        }
1571                    },
1572                    Err(_e) => {
1573                        // Channel is closed, disconnect the client
1574                        socket.shutdown().await?;
1575                        break;
1576                    },
1577                }
1578            }
1579        }
1580    }
1581
1582    Ok(())
1583}
1584
1585/// Starts a server client loop in the background.
1586fn server_client_handler(
1587    client_id: usize,
1588    socket: TcpStream,
1589    server_client_event_sender: Sender<ServerEventRaw>,
1590    client_cleanup_sender: Sender<usize>,
1591) -> (
1592    CommandChannelSender<ServerClientCommand, ServerClientCommandReturn>,
1593    JoinHandle<Result<()>>,
1594) {
1595    // Channels for sending commands from the background server task to a background client task
1596    let (client_command_sender, client_command_receiver) = command_channel();
1597
1598    // Start a background client task, saving the join handle for when the
1599    // server is stopped
1600    let client_task_handle = tokio::spawn(async move {
1601        let res = server_client_loop(
1602            client_id,
1603            socket,
1604            server_client_event_sender,
1605            client_command_receiver,
1606        )
1607        .await;
1608
1609        // Tell the server to clean up after the client, ignoring failures,
1610        // since a failure indicates that the server has probably closed
1611        _ = client_cleanup_sender.send(client_id).await;
1612
1613        res
1614    });
1615
1616    (client_command_sender, client_task_handle)
1617}
1618
1619/// The server loop. Handles incoming connections and commands.
1620#[allow(clippy::too_many_lines)]
1621async fn server_loop(
1622    listener: TcpListener,
1623    server_event_sender: Sender<ServerEventRaw>,
1624    mut server_command_receiver: CommandChannelReceiver<ServerCommand, ServerCommandReturn>,
1625    client_command_senders: &mut HashMap<
1626        usize,
1627        CommandChannelSender<ServerClientCommand, ServerClientCommandReturn>,
1628    >,
1629    client_join_handles: &mut HashMap<usize, JoinHandle<Result<()>>>,
1630) -> Result<()> {
1631    // ID assigned to the next client
1632    let mut next_client_id = 0usize;
1633    // Channel for indicating that a client needs to be cleaned up after
1634    let (server_client_cleanup_sender, mut server_client_cleanup_receiver) =
1635        channel::<usize>(CHANNEL_BUFFER_SIZE);
1636
1637    // Server loop
1638    loop {
1639        // Await new clients connecting,
1640        // commands from the server handle,
1641        // and notifications of clients disconnecting
1642        tokio::select! {
1643            // Accept a connecting client
1644            accept_value = listener.accept() => {
1645                // Get the client socket, exiting if an error occurs
1646                let (socket, _) = accept_value?;
1647                // New client ID
1648                let client_id = next_client_id;
1649                // Increment next client ID
1650                next_client_id += 1;
1651                // Clone the event sender so the background client tasks can
1652                // send events
1653                let server_client_event_sender = server_event_sender.clone();
1654                // Clone the client cleanup sender to the background client
1655                // tasks can be cleaned up properly
1656                let client_cleanup_sender = server_client_cleanup_sender.clone();
1657
1658                // Handle the new connection
1659                let (client_command_sender, client_task_handle) = server_client_handler(client_id, socket, server_client_event_sender, client_cleanup_sender);
1660                // Keep track of client command senders
1661                client_command_senders.insert(client_id, client_command_sender);
1662                // Keep track of client task handles
1663                client_join_handles.insert(client_id, client_task_handle);
1664
1665                // Send an event to note that a client has connected
1666                // successfully
1667                if let Err(_e) = server_event_sender
1668                    .send(ServerEventRaw::Connect { client_id })
1669                    .await
1670                {
1671                    // Server is probably closed
1672                    break;
1673                }
1674            },
1675            // Process a command from the server handle
1676            command_value = server_command_receiver.recv_command() => {
1677                // Handle the command, or lack thereof if the channel is closed
1678                match command_value {
1679                    Ok(command) => {
1680                        match command {
1681                            ServerCommand::Stop => {
1682                                // If a command fails to send, the server has
1683                                // already closed, and the error can be ignored.
1684                                // It should be noted that this is not where the
1685                                // stop method actually returns its `Result`.
1686                                // This immediately returns with an `Ok` status.
1687                                // The real return value is the `Result`
1688                                // returned from the server task join handle.
1689                                _ = server_command_receiver.command_return(ServerCommandReturn::Stop(Ok(()))).await;
1690
1691                                // Break the server loop, the clients will be
1692                                // disconnected before the task ends
1693                                break;
1694                            },
1695                            ServerCommand::Send { client_id, data } => {
1696                                let value = match client_command_senders.get_mut(&client_id) {
1697                                    Some(client_command_sender) => {
1698                                        // Turn `Vec<u8>` into `Arc<[u8]>`,
1699                                        // making it more easily shareable
1700                                        let shareable_data = Arc::<[u8]>::from(data);
1701
1702                                        match client_command_sender.send_command(ServerClientCommand::Send { data: shareable_data }).await {
1703                                            Ok(return_value) => unwrap_enum!(return_value, ServerClientCommandReturn::Send),
1704                                            Err(_e) => {
1705                                                // The channel is closed, and
1706                                                // the client has probably been
1707                                                // disconnected, so the error
1708                                                // can be ignored
1709                                                Ok(())
1710                                            },
1711                                        }
1712                                    },
1713                                    None => Err(Error::InvalidClientId(client_id)),
1714                                };
1715
1716                                // If a command fails to send, the client has probably disconnected,
1717                                // and the error can be ignored
1718                                _ = server_command_receiver.command_return(ServerCommandReturn::Send(value)).await;
1719                            },
1720                            ServerCommand::SendAll { data } => {
1721                                let value = {
1722                                    // Turn `Vec<u8>` into `Arc<[u8]>`, making
1723                                    // it more easily shareable
1724                                    let shareable_data = Arc::<[u8]>::from(data);
1725
1726                                    let send_futures = client_command_senders.iter_mut().map(|(_client_id, client_command_sender)| async {
1727                                        match client_command_sender.send_command(ServerClientCommand::Send { data: Arc::clone(&shareable_data) }).await {
1728                                            Ok(return_value) => unwrap_enum!(return_value, ServerClientCommandReturn::Send),
1729                                            Err(_e) => {
1730                                                // The channel is closed, and
1731                                                // the client has probably been
1732                                                // disconnected, so the error
1733                                                // can be ignored
1734                                                Ok(())
1735                                            }
1736                                        }
1737                                    });
1738
1739                                    let resolved = futures::future::join_all(send_futures).await;
1740                                    resolved.into_iter().collect::<Result<Vec<_>>>().map(|_| ())
1741                                };
1742
1743                                // If a command fails to send, the client has
1744                                // probably disconnected, and the error can be
1745                                // ignored
1746                                _ = server_command_receiver.command_return(ServerCommandReturn::SendAll(value)).await;
1747                            },
1748                            ServerCommand::GetAddr => {
1749                                // Get the server listener's address
1750                                let addr = listener.local_addr();
1751
1752                                // If a command fails to send, the client has
1753                                // probably disconnected, and the error can be
1754                                // ignored
1755                                _ = server_command_receiver.command_return(ServerCommandReturn::GetAddr(addr.map_err(Into::into))).await;
1756                            },
1757                            ServerCommand::GetClientAddr { client_id } => {
1758                                let value = match client_command_senders.get_mut(&client_id) {
1759                                    Some(client_command_sender) => match client_command_sender.send_command(ServerClientCommand::GetAddr).await {
1760                                        Ok(return_value) => unwrap_enum!(return_value, ServerClientCommandReturn::GetAddr),
1761                                        Err(_e) => {
1762                                            // The channel is closed, and the
1763                                            // client has probably been
1764                                            // disconnected, so the error can be
1765                                            // treated as an invalid client
1766                                            // error
1767                                            Err(Error::InvalidClientId(client_id))
1768                                        },
1769                                    },
1770                                    None => Err(Error::InvalidClientId(client_id)),
1771                                };
1772
1773                                // If a command fails to send, the client has
1774                                // probably disconnected, and the error can be
1775                                // ignored
1776                                _ = server_command_receiver.command_return(ServerCommandReturn::GetClientAddr(value)).await;
1777                            },
1778                            ServerCommand::RemoveClient { client_id } => {
1779                                let value = match client_command_senders.get_mut(&client_id) {
1780                                    Some(client_command_sender) => match client_command_sender.send_command(ServerClientCommand::Remove).await {
1781                                        Ok(return_value) => unwrap_enum!(return_value, ServerClientCommandReturn::Remove),
1782                                        Err(_e) => {
1783                                            // The channel is closed, and the
1784                                            // client has probably been
1785                                            // disconnected, so the error can be
1786                                            // ignored
1787                                            Ok(())
1788                                        },
1789                                    },
1790                                    None => Err(Error::InvalidClientId(client_id)),
1791                                };
1792
1793                                // If a command fails to send, the client has
1794                                // probably disconnected already, and the error
1795                                // can be ignored
1796                                _ = server_command_receiver.command_return(ServerCommandReturn::RemoveClient(value)).await;
1797                            },
1798                        }
1799                    },
1800                    Err(_e) => {
1801                        // Server is probably closed, exit
1802                        break;
1803                    },
1804                }
1805            }
1806            // Clean up after a disconnecting client
1807            disconnecting_client_id = server_client_cleanup_receiver.recv() => {
1808                match disconnecting_client_id {
1809                    Some(client_id) => {
1810                        // Remove the client's command sender, which will be
1811                        // dropped after this block ends
1812                        client_command_senders.remove(&client_id);
1813
1814                        // Remove the client's join handle
1815                        if let Some(handle) = client_join_handles.remove(&client_id) {
1816                            // Join the client's handle
1817                            if let Err(e) = handle.await.unwrap() {
1818                                if cfg!(test) {
1819                                    // If testing, fail
1820                                    Err(e)?;
1821                                } else {
1822                                    // If not testing, ignore client handler
1823                                    // errors
1824                                }
1825                            }
1826                        }
1827
1828                        // Send an event to note that a client has disconnected
1829                        if let Err(_e) = server_event_sender.send(ServerEventRaw::Disconnect { client_id }).await {
1830                            // Server is probably closed, exit
1831                            break;
1832                        }
1833                    },
1834                    None => {
1835                        // Server is probably closed, exit
1836                        break;
1837                    },
1838                }
1839            }
1840        }
1841    }
1842
1843    Ok(())
1844}
1845
1846/// Starts the server loop task in the background.
1847async fn server_handler(
1848    listener: TcpListener,
1849    server_event_sender: Sender<ServerEventRaw>,
1850    server_command_receiver: CommandChannelReceiver<ServerCommand, ServerCommandReturn>,
1851) -> Result<()> {
1852    // Collection of channels for sending commands from the background server
1853    // task to a background client task
1854    let mut client_command_senders: HashMap<
1855        usize,
1856        CommandChannelSender<ServerClientCommand, ServerClientCommandReturn>,
1857    > = HashMap::new();
1858    // Background client task join handles
1859    let mut client_join_handles: HashMap<usize, JoinHandle<Result<()>>> = HashMap::new();
1860
1861    // Wrap server loop in a block to catch all exit scenarios
1862    let server_exit = server_loop(
1863        listener,
1864        server_event_sender.clone(),
1865        server_command_receiver,
1866        &mut client_command_senders,
1867        &mut client_join_handles,
1868    )
1869    .await;
1870
1871    // Send a remove command to all clients
1872    futures::future::join_all(client_command_senders.into_values().map(
1873        |mut client_command_sender| async move {
1874            // If a command fails to send, the client has probably disconnected
1875            // already, and the error can be ignored
1876            _ = client_command_sender
1877                .send_command(ServerClientCommand::Remove)
1878                .await;
1879        },
1880    ))
1881    .await;
1882
1883    // Join all background client tasks before exiting
1884    futures::future::join_all(client_join_handles.into_values().map(|handle| async move {
1885        if let Err(e) = handle.await.unwrap() {
1886            if cfg!(test) {
1887                // If testing, fail
1888                Err(e)?;
1889            } else {
1890                // If not testing, ignore client handler errors
1891            }
1892        }
1893
1894        Ok(())
1895    }))
1896    .await
1897    .into_iter()
1898    .collect::<Result<Vec<_>>>()?;
1899
1900    // Send a stop event, ignoring send errors
1901    _ = server_event_sender.send(ServerEventRaw::Stop).await;
1902
1903    // Return server loop result
1904    server_exit
1905}