wowlab_centrifuge/
transport.rs

1use std::{sync::Once, time::Duration};
2
3use futures_util::{
4    SinkExt, StreamExt,
5    stream::{SplitSink, SplitStream},
6};
7use prost::Message;
8use tokio::net::TcpStream;
9use tokio_tungstenite::{MaybeTlsStream, WebSocketStream, tungstenite::Message as WsMessage};
10
11use crate::{error::Error, proto};
12
13const HANDSHAKE_TIMEOUT: Duration = Duration::from_secs(45);
14
15async fn next_ws_message(read: &mut SplitStream<WsStream>) -> Result<WsMessage, Error> {
16    match read.next().await {
17        Some(result) => Ok(result?),
18        None => Err(Error::connection_closed()),
19    }
20}
21
22fn decode_reply(data: &[u8]) -> Result<proto::Reply, Error> {
23    Ok(proto::Reply::decode_length_delimited(data)?)
24}
25
26static CRYPTO_INIT: Once = Once::new();
27
28fn ensure_crypto_provider() {
29    CRYPTO_INIT.call_once(|| {
30        let _ = rustls::crypto::ring::default_provider().install_default();
31    });
32}
33
34fn log_server_close(code: impl std::fmt::Display, reason: impl std::fmt::Display) {
35    tracing::info!(code = %code, reason = %reason, "WebSocket closed by server");
36}
37
38fn log_server_close_without_frame() {
39    tracing::info!("WebSocket closed by server");
40}
41
42type WsStream = WebSocketStream<MaybeTlsStream<TcpStream>>;
43
44/// WebSocket transport that frames Centrifugo protobuf commands and replies.
45pub(crate) struct Transport {
46    write: SplitSink<WsStream, WsMessage>,
47    read: SplitStream<WsStream>,
48}
49
50impl Transport {
51    pub(crate) async fn connect(url: &str) -> Result<Self, Error> {
52        ensure_crypto_provider();
53
54        let ws_url = http_to_ws(url);
55        let full_url = format!(
56            "{}/connection/websocket?format=protobuf",
57            ws_url.trim_end_matches('/')
58        );
59
60        tracing::info!(url = %ws_url, "Connecting");
61
62        let (ws, _) = tokio_tungstenite::connect_async(&full_url).await?;
63        let (write, read) = ws.split();
64
65        Ok(Self { write, read })
66    }
67
68    pub(crate) async fn send_command(
69        &mut self,
70        cmd: proto::Command,
71    ) -> Result<proto::Reply, Error> {
72        let data = cmd.encode_length_delimited_to_vec();
73
74        self.write.send(WsMessage::Binary(data.into())).await?;
75
76        tokio::time::timeout(HANDSHAKE_TIMEOUT, self.read_message())
77            .await
78            .map_err(|elapsed| {
79                tracing::debug!(?elapsed, "Handshake timed out");
80
81                Error::timeout()
82            })?
83    }
84
85    pub(crate) async fn send_raw(&mut self, data: Vec<u8>) -> Result<(), Error> {
86        self.write.send(WsMessage::Binary(data.into())).await?;
87
88        Ok(())
89    }
90
91    pub(crate) async fn read_message(&mut self) -> Result<proto::Reply, Error> {
92        loop {
93            let msg = next_ws_message(&mut self.read).await?;
94
95            match msg {
96                WsMessage::Binary(data) => return decode_reply(&data),
97                WsMessage::Close(frame) => {
98                    if let Some(frame) = frame {
99                        log_server_close(frame.code, frame.reason);
100                    } else {
101                        log_server_close_without_frame();
102                    }
103
104                    return Err(Error::connection_closed());
105                }
106                WsMessage::Ping(data) => {
107                    self.write.send(WsMessage::Pong(data)).await?;
108                }
109                _ => {}
110            }
111        }
112    }
113
114    pub(crate) async fn close(&mut self) {
115        let _ = self.write.close().await;
116    }
117}
118
119impl std::fmt::Debug for Transport {
120    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
121        f.debug_struct("Transport").finish_non_exhaustive()
122    }
123}
124
125// #t(fn: rust_hardcoded_url) These literals are protocol-prefix transforms, not service endpoints.
126fn http_to_ws(url: &str) -> String {
127    let secure = url
128        .strip_prefix("https://")
129        .map(|rest| format!("wss://{rest}"));
130
131    secure
132        .or_else(|| {
133            url.strip_prefix("http://")
134                .map(|rest| format!("ws://{rest}"))
135        })
136        .unwrap_or_else(|| url.to_string())
137}
wowlab_centrifuge/transport.rs

wowlab_centrifuge/
transport.rs
