swadhinbiswas 📦 v2.0.8: security fixes, documentation, testing improvements · 7/28/2026
3e270a0
85 lines · 3.0 KB
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
//! WebSocket handler types and default echo implementation.

#[cfg(feature = "ws")]
pub type WsRead = Box<
    dyn futures::Stream<
            Item = Result<
                tokio_tungstenite::tungstenite::Message,
                tokio_tungstenite::tungstenite::Error,
            >,
        > + Unpin
        + Send,
>;

#[cfg(feature = "ws")]
pub type WsWrite = Box<
    dyn futures::Sink<
            tokio_tungstenite::tungstenite::Message,
            Error = tokio_tungstenite::tungstenite::Error,
        > + Unpin
        + Send,
>;

/// Connection metadata handed to a WebSocket handler on upgrade. Mirrors the
/// subset of the HTTP request a WebSocket scope needs (`path`, decoded query
/// string, raw headers) so handler frameworks can build a Starlette-style
/// `scope` without re-reading the (already upgraded) socket.
#[cfg(feature = "ws")]
#[derive(Debug, Clone)]
pub struct WsConnInfo {
    /// Request path the upgrade arrived on.
    pub path: String,
    /// Raw query string (percent-encoded, without the leading `?`).
    pub query_string: Vec<u8>,
    /// Raw request headers observed on the upgrade request.
    pub headers: Vec<(Vec<u8>, Vec<u8>)>,
    /// Remote peer address as `(host, port)`, if known.
    pub client: Option<(String, u16)>,
}

#[cfg(feature = "ws")]
pub type WsHandler = std::sync::Arc<
    dyn Fn(
            WsConnInfo,
            WsRead,
            WsWrite,
        )
            -> std::pin::Pin<Box<dyn std::future::Future<Output = anyhow::Result<()>> + Send>>
        + Send
        + Sync,
>;

/// Default WebSocket handler used when no application handler is registered.
/// Echoes text/binary frames (mirroring the legacy raw-TCP behavior) so that
/// standalone servers remain WebSocket-compatible out of the box. `with_ws`
/// replaces this with an application-provided handler.
#[cfg(feature = "ws")]
pub fn default_ws_echo() -> WsHandler {
    std::sync::Arc::new(|_info, mut read, mut write| {
        Box::pin(async move {
            use futures::{SinkExt, StreamExt};
            while let Some(msg) = read.next().await {
                match msg {
                    Ok(m @ tokio_tungstenite::tungstenite::Message::Text(_))
                    | Ok(m @ tokio_tungstenite::tungstenite::Message::Binary(_)) => {
                        if write.send(m).await.is_err() {
                            break;
                        }
                    }
                    Ok(tokio_tungstenite::tungstenite::Message::Ping(p)) => {
                        let _ = write.send(tokio_tungstenite::tungstenite::Message::Pong(p)).await;
                    }
                    Ok(tokio_tungstenite::tungstenite::Message::Close(_)) | Err(_) => {
                        let _ =
                            write.send(tokio_tungstenite::tungstenite::Message::Close(None)).await;
                        break;
                    }
                    _ => {}
                }
            }
            let _ = write.close().await;
            Ok(())
        })
    })
}