rivet_envoy_client/
websocket.rs

1use rivet_envoy_protocol as protocol;
2
3use crate::callbacks::BoxFuture;
4
5/// Handler returned by the websocket callback for receiving WebSocket events.
6pub struct WebSocketHandler {
7	pub on_message: Box<dyn Fn(WebSocketMessage) -> BoxFuture<()> + Send + Sync>,
8	pub on_close: Box<dyn Fn(u16, String) -> BoxFuture<()> + Send + Sync>,
9	pub on_open: Option<Box<dyn FnOnce(WebSocketSender) -> BoxFuture<()> + Send>>,
10}
11
12pub struct WebSocketMessage {
13	pub data: Vec<u8>,
14	pub binary: bool,
15	pub gateway_id: protocol::GatewayId,
16	pub request_id: protocol::RequestId,
17	pub message_index: u16,
18	/// Send data back on this WebSocket connection.
19	pub sender: WebSocketSender,
20}
21
22/// Allows sending messages back on a WebSocket connection from within the on_message callback.
23#[derive(Clone)]
24pub struct WebSocketSender {
25	pub(crate) tx: tokio::sync::mpsc::UnboundedSender<WsOutgoing>,
26}
27
28pub(crate) enum WsOutgoing {
29	Message {
30		data: Vec<u8>,
31		binary: bool,
32	},
33	Flush {
34		tx: tokio::sync::oneshot::Sender<()>,
35	},
36	Close {
37		code: Option<u16>,
38		reason: Option<String>,
39	},
40}
41
42impl WebSocketSender {
43	pub fn send(&self, data: Vec<u8>, binary: bool) {
44		let _ = self.tx.send(WsOutgoing::Message { data, binary });
45	}
46
47	pub fn send_text(&self, text: &str) {
48		self.send(text.as_bytes().to_vec(), false);
49	}
50
51	pub async fn flush(&self) {
52		let (tx, rx) = tokio::sync::oneshot::channel();
53		if self.tx.send(WsOutgoing::Flush { tx }).is_ok() {
54			let _ = rx.await;
55		}
56	}
57
58	pub fn close(&self, code: Option<u16>, reason: Option<String>) {
59		let _ = self.tx.send(WsOutgoing::Close { code, reason });
60	}
61}
rivet_envoy_client/websocket.rs

rivet_envoy_client/
websocket.rs
