wisp_mux/extensions/
mod.rs

1//! Wisp protocol extensions.
2pub mod password;
3pub mod udp;
4
5use std::ops::{Deref, DerefMut};
6
7use async_trait::async_trait;
8use bytes::{BufMut, Bytes, BytesMut};
9
10use crate::{
11	ws::{LockedWebSocketWrite, WebSocketRead},
12	Role, WispError,
13};
14
15/// Type-erased protocol extension that implements Clone.
16#[derive(Debug)]
17pub struct AnyProtocolExtension(Box<dyn ProtocolExtension + Sync + Send>);
18
19impl AnyProtocolExtension {
20	/// Create a new type-erased protocol extension.
21	pub fn new<T: ProtocolExtension + Sync + Send + 'static>(extension: T) -> Self {
22		Self(Box::new(extension))
23	}
24}
25
26impl Deref for AnyProtocolExtension {
27	type Target = dyn ProtocolExtension;
28	fn deref(&self) -> &Self::Target {
29		self.0.deref()
30	}
31}
32
33impl DerefMut for AnyProtocolExtension {
34	fn deref_mut(&mut self) -> &mut Self::Target {
35		self.0.deref_mut()
36	}
37}
38
39impl Clone for AnyProtocolExtension {
40	fn clone(&self) -> Self {
41		Self(self.0.box_clone())
42	}
43}
44
45impl From<AnyProtocolExtension> for Bytes {
46	fn from(value: AnyProtocolExtension) -> Self {
47		let mut bytes = BytesMut::with_capacity(5);
48		let payload = value.encode();
49		bytes.put_u8(value.get_id());
50		bytes.put_u32_le(payload.len() as u32);
51		bytes.extend(payload);
52		bytes.freeze()
53	}
54}
55
56/// A Wisp protocol extension.
57///
58/// See [the
59/// docs](https://github.com/MercuryWorkshop/wisp-protocol/blob/v2/protocol.md#protocol-extensions).
60#[async_trait]
61pub trait ProtocolExtension: std::fmt::Debug {
62	/// Get the protocol extension ID.
63	fn get_id(&self) -> u8;
64	/// Get the protocol extension's supported packets.
65	///
66	/// Used to decide whether to call the protocol extension's packet handler.
67	fn get_supported_packets(&self) -> &'static [u8];
68
69	/// Encode self into Bytes.
70	fn encode(&self) -> Bytes;
71
72	/// Handle the handshake part of a Wisp connection.
73	///
74	/// This should be used to send or receive data before any streams are created.
75	async fn handle_handshake(
76		&mut self,
77		read: &mut dyn WebSocketRead,
78		write: &LockedWebSocketWrite,
79	) -> Result<(), WispError>;
80
81	/// Handle receiving a packet.
82	async fn handle_packet(
83		&mut self,
84		packet: Bytes,
85		read: &mut dyn WebSocketRead,
86		write: &LockedWebSocketWrite,
87	) -> Result<(), WispError>;
88
89	/// Clone the protocol extension.
90	fn box_clone(&self) -> Box<dyn ProtocolExtension + Sync + Send>;
91}
92
93/// Trait to build a Wisp protocol extension from a payload.
94pub trait ProtocolExtensionBuilder {
95	/// Get the protocol extension ID.
96	///
97	/// Used to decide whether this builder should be used.
98	fn get_id(&self) -> u8;
99
100	/// Build a protocol extension from the extension's metadata.
101	fn build_from_bytes(&self, bytes: Bytes, role: Role)
102		-> Result<AnyProtocolExtension, WispError>;
103
104	/// Build a protocol extension to send to the other side.
105	fn build_to_extension(&self, role: Role) -> AnyProtocolExtension;
106}