ws-bridge 0.2.1

Strongly-typed WebSocket endpoints for Rust — define once, use everywhere (axum server, yew/gloo browser client, tokio-tungstenite native client)
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143

use std::fmt;
use std::marker::PhantomData;

use crate::codec::{DecodeError, EncodeError, WsCodec, WsMessage};

/// Error when sending a message.
#[derive(Debug)]
pub enum SendError {
    Encode(EncodeError),
    Closed,
}

impl fmt::Display for SendError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            SendError::Encode(e) => write!(f, "send error: {e}"),
            SendError::Closed => write!(f, "connection closed"),
        }
    }
}

impl std::error::Error for SendError {}

impl From<EncodeError> for SendError {
    fn from(e: EncodeError) -> Self {
        SendError::Encode(e)
    }
}

/// Error when receiving a message.
#[derive(Debug)]
pub enum RecvError {
    Decode(DecodeError),
    Closed,
}

impl fmt::Display for RecvError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            RecvError::Decode(e) => write!(f, "recv error: {e}"),
            RecvError::Closed => write!(f, "connection closed"),
        }
    }
}

impl std::error::Error for RecvError {}

impl From<DecodeError> for RecvError {
    fn from(e: DecodeError) -> Self {
        RecvError::Decode(e)
    }
}

/// The sending half of a typed WebSocket connection.
///
/// Obtained by calling [`WsConnection::split`].
pub struct WsSender<S: WsCodec> {
    pub(crate) sink: Box<dyn ErasedSink>,
    pub(crate) _send: PhantomData<S>,
}

impl<S: WsCodec> WsSender<S> {
    pub async fn send(&mut self, msg: S) -> Result<(), SendError> {
        let ws_msg = msg.encode()?;
        self.sink.send(ws_msg).await.map_err(|_| SendError::Closed)
    }

    pub async fn close(&mut self) -> Result<(), SendError> {
        self.sink.close().await.map_err(|_| SendError::Closed)
    }
}

/// The receiving half of a typed WebSocket connection.
///
/// Obtained by calling [`WsConnection::split`].
pub struct WsReceiver<R: WsCodec> {
    pub(crate) stream: Box<dyn ErasedStream>,
    pub(crate) _recv: PhantomData<R>,
}

impl<R: WsCodec> WsReceiver<R> {
    pub async fn recv(&mut self) -> Option<Result<R, RecvError>> {
        match self.stream.next().await {
            None => None,
            Some(Err(_)) => Some(Err(RecvError::Closed)),
            Some(Ok(ws_msg)) => Some(R::decode(ws_msg).map_err(RecvError::Decode)),
        }
    }
}

/// A typed WebSocket connection parameterized by send and receive types.
///
/// On the server side: `Send = E::ServerMsg`, `Recv = E::ClientMsg`.
/// On the client side: `Send = E::ClientMsg`, `Recv = E::ServerMsg`.
pub struct WsConnection<S: WsCodec, R: WsCodec> {
    pub(crate) sink: Box<dyn ErasedSink>,
    pub(crate) stream: Box<dyn ErasedStream>,
    pub(crate) _types: PhantomData<(S, R)>,
}

impl<S: WsCodec, R: WsCodec> WsConnection<S, R> {
    pub async fn send(&mut self, msg: S) -> Result<(), SendError> {
        let ws_msg = msg.encode()?;
        self.sink.send(ws_msg).await.map_err(|_| SendError::Closed)
    }

    pub async fn recv(&mut self) -> Option<Result<R, RecvError>> {
        match self.stream.next().await {
            None => None,
            Some(Err(_)) => Some(Err(RecvError::Closed)),
            Some(Ok(ws_msg)) => Some(R::decode(ws_msg).map_err(RecvError::Decode)),
        }
    }

    pub fn split(self) -> (WsSender<S>, WsReceiver<R>) {
        (
            WsSender {
                sink: self.sink,
                _send: PhantomData,
            },
            WsReceiver {
                stream: self.stream,
                _recv: PhantomData,
            },
        )
    }
}

// -- Erased trait objects for transport independence --

pub(crate) type BoxFuture<'a, T> =
    std::pin::Pin<Box<dyn std::future::Future<Output = T> + Send + 'a>>;

/// Type-erased sink that can send `WsMessage`s.
pub(crate) trait ErasedSink: Send {
    fn send(&mut self, msg: WsMessage) -> BoxFuture<'_, Result<(), ()>>;
    fn close(&mut self) -> BoxFuture<'_, Result<(), ()>>;
}

/// Type-erased stream that yields `WsMessage`s.
pub(crate) trait ErasedStream: Send {
    fn next(&mut self) -> BoxFuture<'_, Option<Result<WsMessage, ()>>>;
}