ahp_ws/
transport.rs

1//! WebSocket transport implementation built on `tokio-tungstenite`.
2
3use ahp::{Transport, TransportError, TransportMessage};
4use futures_util::{SinkExt, StreamExt};
5use thiserror::Error;
6use tokio::net::TcpStream;
7use tokio_tungstenite::tungstenite::Message;
8use tokio_tungstenite::{connect_async, MaybeTlsStream, WebSocketStream};
9use url::Url;
10
11/// Errors that can occur while building a [`WebSocketTransport`].
12#[derive(Debug, Error)]
13pub enum WebSocketTransportError {
14    /// The URL could not be parsed.
15    #[error("invalid url: {0}")]
16    InvalidUrl(#[from] url::ParseError),
17    /// The underlying WebSocket handshake or I/O failed.
18    #[error("websocket error: {0}")]
19    WebSocket(#[from] tokio_tungstenite::tungstenite::Error),
20}
21
22type WsStream = WebSocketStream<MaybeTlsStream<TcpStream>>;
23
24/// A [`Transport`] backed by a `tokio-tungstenite` WebSocket stream.
25///
26/// Use [`WebSocketTransport::connect`] for the common case, or
27/// [`WebSocketTransport::from_stream`] when you need custom connection
28/// options (custom TLS configuration, additional headers, etc.).
29///
30/// # Example
31///
32/// ```no_run
33/// use ahp::{Client, ClientConfig};
34/// use ahp_ws::WebSocketTransport;
35///
36/// # async fn run() -> Result<(), Box<dyn std::error::Error>> {
37/// let transport = WebSocketTransport::connect("ws://localhost:12345").await?;
38/// let client = Client::connect(transport, ClientConfig::default()).await?;
39/// client.initialize("demo".into(), vec!["0.1.0".into()], vec![]).await?;
40/// # Ok(()) }
41/// ```
42pub struct WebSocketTransport {
43    inner: WsStream,
44}
45
46impl WebSocketTransport {
47    /// Open a new WebSocket connection to `url`.
48    ///
49    /// Both `ws://` and `wss://` schemes are accepted; the `wss://` TLS
50    /// backend is chosen at compile time by Cargo feature (see the crate
51    /// docs), defaulting to rustls with OS-trust-store roots. The URL is
52    /// parsed eagerly so malformed URLs surface as
53    /// [`WebSocketTransportError::InvalidUrl`] before any network I/O.
54    pub async fn connect(url: &str) -> Result<Self, WebSocketTransportError> {
55        let _parsed = Url::parse(url)?; // validate early
56        let (stream, _resp) = connect_async(url).await?;
57        Ok(Self { inner: stream })
58    }
59
60    /// Wrap an already-connected WebSocket stream.
61    ///
62    /// Use this when you need to drive the `tokio-tungstenite`
63    /// handshake yourself — for custom TLS configuration, request
64    /// headers, or to reuse an existing socket.
65    pub fn from_stream(inner: WsStream) -> Self {
66        Self { inner }
67    }
68}
69
70impl Transport for WebSocketTransport {
71    async fn send(&mut self, msg: TransportMessage) -> Result<(), TransportError> {
72        let frame = match msg {
73            TransportMessage::Parsed(m) => {
74                let s = serde_json::to_string(&m)
75                    .map_err(|e| TransportError::Protocol(e.to_string()))?;
76                Message::Text(s.into())
77            }
78            TransportMessage::Text(s) => Message::Text(s.into()),
79            TransportMessage::Binary(b) => Message::Binary(b.into()),
80        };
81        self.inner
82            .send(frame)
83            .await
84            .map_err(|e| TransportError::Io(e.to_string()))
85    }
86
87    async fn recv(&mut self) -> Result<Option<TransportMessage>, TransportError> {
88        loop {
89            match self.inner.next().await {
90                None => return Ok(None),
91                Some(Err(e)) => return Err(TransportError::Io(e.to_string())),
92                Some(Ok(Message::Text(s))) => {
93                    return Ok(Some(TransportMessage::Text(s.to_string())))
94                }
95                Some(Ok(Message::Binary(b))) => {
96                    return Ok(Some(TransportMessage::Binary(b.into())))
97                }
98                Some(Ok(Message::Close(_))) => return Ok(None),
99                // Ping, Pong, Frame: keep looping, tungstenite handles protocol-level response.
100                Some(Ok(_)) => continue,
101            }
102        }
103    }
104
105    async fn close(&mut self) -> Result<(), TransportError> {
106        self.inner
107            .close(None)
108            .await
109            .map_err(|e| TransportError::Io(e.to_string()))
110    }
111}
ahp_ws/transport.rs

ahp_ws/
transport.rs
