1use 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#[derive(Debug)]
139pub(crate) enum Message {
140 Text,
142 Binary(Bytes),
144 Ping(Bytes),
146 Pong(Bytes),
148 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
186pub 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 NotNegotiated,
255 ExpectedNegotiate {
257 id: u16,
258 },
259 NegotiateHeader {
261 error: format::Error,
262 },
263}
264
265#[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
401pub struct Response {
403 handled: bool,
404}
405
406pub trait IntoResponse
408where
409 Self: 'static + Send,
410{
411 type Error: fmt::Display;
413
414 fn into_response(self) -> Result<Response, Self::Error>;
416}
417
418impl 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
430impl 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
443impl<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
469impl<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
489pub trait Handler
495where
496 Self: 'static + Send + Clone,
497{
498 type Id: Id;
500 type Response: IntoResponse;
502
503 fn open_channel<'this>(
514 &'this self,
515 channel: ChannelId,
516 ) -> impl Future<Output = ()> + Send + 'this {
517 async {
518 _ = channel;
519 }
520 }
521
522 fn close_channel<'this>(
530 &'this self,
531 channel: ChannelId,
532 ) -> impl Future<Output = ()> + Send + 'this {
533 async {
534 _ = channel;
535 }
536 }
537
538 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
569pub trait ChannelAllocator {
576 fn next(&self) -> impl Future<Output = Option<ChannelId>> + Send + '_;
584
585 fn free(&self, channel: ChannelId) -> impl Future<Output = ()> + Send + '_;
587}
588
589pub 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: Option<&'static [Format]>,
746}
747
748impl<S, H> Connect<S, H, Channels>
749where
750 S: ServerImpl,
751 H: Handler,
752{
753 #[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 #[inline]
778 pub fn seed(mut self, seed: u64) -> Self {
779 self.seed = seed;
780 self
781 }
782
783 #[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 #[inline]
801 pub fn handler(&self) -> &H {
802 &self.handler
803 }
804
805 #[inline]
817 pub fn with_formats(mut self, formats: &'static [Format]) -> Self {
818 self.formats = Some(formats);
819 self
820 }
821
822 #[inline]
824 pub fn accepts(&self, format: Format) -> bool {
825 accepts(self.formats, format)
826 }
827
828 #[inline]
838 pub fn max_capacity(mut self, max_capacity: usize) -> Self {
839 self.max_capacity = max_capacity;
840 self
841 }
842
843 #[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 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#[inline]
910fn accepts(formats: Option<&'static [Format]>, format: Format) -> bool {
911 format.is_supported() && formats.is_none_or(|f| f.contains(&format))
912}
913
914pub 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 format: Format,
948 formats: Option<&'static [Format]>,
951}
952
953impl<S, H, C> Server<S, H, C>
954where
955 S: ServerImpl,
956 H: Handler,
957{
958 #[inline]
960 pub fn handler(&self) -> &H {
961 &self.handler
962 }
963
964 #[inline]
970 pub fn format(&self) -> Format {
971 self.format
972 }
973
974 #[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 async fn negotiate(&mut self) -> Result<(), Error> {
998 let mut negotiated = false;
999 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 }
1075 }
1076 }
1077
1078 match failure {
1079 Some(error) => Err(error),
1080 None => Ok(()),
1081 }
1082 }
1083
1084 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 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 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 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 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 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 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 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 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 let format = self.format;
1497
1498 let buf = self.pool.with(|buf| {
1499 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 Close,
1661 Ping,
1663 Recv(Option<Result<Message, E>>),
1665 Send(Result<(), E>),
1667 Flushed(Result<(), E>),
1669 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 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
1747pub 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 pub fn channel(&self) -> ChannelId {
1768 self.channel
1769 }
1770
1771 #[inline]
1776 pub fn format(&self) -> Format {
1777 self.format
1778 }
1779
1780 #[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
1802pub 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 #[inline]
1819 pub fn format(&self) -> Format {
1820 self.format
1821 }
1822
1823 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#[inline]
1866fn scramble_channel(x: u16) -> u16 {
1867 let x = x.wrapping_mul(0x9285);
1868 x ^ (x >> 8)
1869}
1870
1871#[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#[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
1933struct 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}