compio_ws/
lib.rs

1//! WebSocket support based on [`tungstenite`].
2//!
3//! This library is an implementation of WebSocket handshakes and streams for
4//! compio. It is based on the tungstenite crate which implements all required
5//! WebSocket protocol logic. This crate brings compio support / compio
6//! integration to it.
7//!
8//! Each WebSocket stream implements message reading and writing.
9//!
10//! [`tungstenite`]: https://docs.rs/tungstenite
11
12#![cfg_attr(docsrs, feature(doc_cfg))]
13#![allow(unused_features)]
14#![warn(missing_docs)]
15#![deny(rustdoc::broken_intra_doc_links)]
16#![doc(
17    html_logo_url = "https://github.com/compio-rs/compio-logo/raw/refs/heads/master/generated/colored-bold.svg"
18)]
19#![doc(
20    html_favicon_url = "https://github.com/compio-rs/compio-logo/raw/refs/heads/master/generated/colored-bold.svg"
21)]
22
23use std::{
24    io,
25    pin::Pin,
26    task::{Context, Poll, ready},
27};
28
29use compio_buf::IntoInner;
30use compio_driver::AsFd;
31use compio_net::{SocketWithState, TcpStream, UnixStream};
32use compio_runtime::fd::PollFd;
33use compio_tls::{MaybeTlsStream, TlsStream};
34use futures_util::{Sink, SinkExt, Stream, StreamExt, stream::FusedStream};
35use pin_project_lite::pin_project;
36use tungstenite::{
37    Error as WsError, Message,
38    client::IntoClientRequest,
39    handshake::server::{Callback, NoCallback},
40    protocol::{CloseFrame, Role, WebSocketConfig},
41};
42
43#[cfg(feature = "connect")]
44mod tls;
45#[cfg(feature = "connect")]
46pub use tls::*;
47pub use tungstenite;
48
49type MaybePollStream<S> = MaybeTlsStream<PollFd<S>>;
50
51/// Configuration for compio-ws.
52///
53/// ## API Interface
54///
55/// `_with_config` functions in this crate accept `impl Into<Config>`, so
56/// following are all valid:
57/// - [`Config`]
58/// - [`WebSocketConfig`] (use custom WebSocket config with default remaining
59///   settings)
60/// - [`None`] (use default value)
61pub struct Config {
62    /// WebSocket configuration from tungstenite.
63    websocket: Option<WebSocketConfig>,
64
65    /// Disable Nagle's algorithm. This only affects
66    /// [`connect_async_with_config()`] and [`connect_async_tls_with_config()`].
67    disable_nagle: bool,
68}
69
70impl Config {
71    /// Creates a new `Config` with default settings.
72    pub fn new() -> Self {
73        Self {
74            websocket: None,
75            disable_nagle: false,
76        }
77    }
78
79    /// Get the WebSocket configuration.
80    pub fn websocket_config(&self) -> Option<&WebSocketConfig> {
81        self.websocket.as_ref()
82    }
83
84    /// Disable Nagle's algorithm, i.e. `set_nodelay(true)`.
85    ///
86    /// Default to `false`. If you don't know what the Nagle's algorithm is,
87    /// better leave it to `false`.
88    pub fn disable_nagle(mut self, disable: bool) -> Self {
89        self.disable_nagle = disable;
90        self
91    }
92}
93
94impl Default for Config {
95    fn default() -> Self {
96        Self::new()
97    }
98}
99
100impl From<WebSocketConfig> for Config {
101    fn from(config: WebSocketConfig) -> Self {
102        Self {
103            websocket: Some(config),
104            ..Default::default()
105        }
106    }
107}
108
109impl From<Option<WebSocketConfig>> for Config {
110    fn from(config: Option<WebSocketConfig>) -> Self {
111        Self {
112            websocket: config,
113            ..Default::default()
114        }
115    }
116}
117
118mod private {
119    use super::*;
120
121    pub trait Sealed<S>
122    where
123        S: AsFd,
124    {
125    }
126
127    impl<S: AsFd> Sealed<S> for PollFd<S> {}
128    impl<S: AsFd> Sealed<S> for MaybePollStream<S> {}
129    impl<S: AsFd> Sealed<S> for TlsStream<PollFd<S>> {}
130    impl Sealed<SocketWithState> for TcpStream {}
131    impl Sealed<SocketWithState> for UnixStream {}
132}
133
134/// Convert a stream into a [`MaybeTlsStream`].
135pub trait IntoMaybeTlsStream<S>: private::Sealed<S>
136where
137    S: AsFd,
138{
139    /// Convert this stream into a [`MaybeTlsStream`].
140    fn into_maybe_tls_stream(self) -> io::Result<MaybePollStream<S>>;
141}
142
143impl<S: AsFd> IntoMaybeTlsStream<S> for PollFd<S> {
144    fn into_maybe_tls_stream(self) -> io::Result<MaybePollStream<S>> {
145        Ok(MaybeTlsStream::new_plain(self))
146    }
147}
148
149impl<S: AsFd> IntoMaybeTlsStream<S> for MaybePollStream<S> {
150    fn into_maybe_tls_stream(self) -> io::Result<MaybePollStream<S>> {
151        Ok(self)
152    }
153}
154
155impl<S: AsFd> IntoMaybeTlsStream<S> for TlsStream<PollFd<S>> {
156    fn into_maybe_tls_stream(self) -> io::Result<MaybePollStream<S>> {
157        Ok(MaybeTlsStream::new_tls(self))
158    }
159}
160
161impl IntoMaybeTlsStream<SocketWithState> for TcpStream {
162    fn into_maybe_tls_stream(self) -> io::Result<MaybePollStream<SocketWithState>> {
163        Ok(MaybeTlsStream::new_plain(self.into_poll_fd()?))
164    }
165}
166
167impl IntoMaybeTlsStream<SocketWithState> for UnixStream {
168    fn into_maybe_tls_stream(self) -> io::Result<MaybePollStream<SocketWithState>> {
169        Ok(MaybeTlsStream::new_plain(self.into_poll_fd()?))
170    }
171}
172
173pin_project! {
174    /// A WebSocket stream that works with compio.
175    #[derive(Debug)]
176    pub struct WebSocketStream<S: AsFd> {
177        #[pin]
178        inner: async_tungstenite::WebSocketStream<MaybePollStream<S>>,
179        next_item: Option<Option<Result<Message, WsError>>>,
180    }
181}
182
183impl<S: AsFd + 'static> WebSocketStream<S> {
184    /// Get a reference to the underlying stream.
185    pub fn get_ref(&self) -> &MaybePollStream<S> {
186        self.inner.get_ref()
187    }
188
189    /// Get a mutable reference to the underlying stream.
190    pub fn get_mut(&mut self) -> &mut MaybePollStream<S> {
191        self.inner.get_mut()
192    }
193
194    /// Convert a raw socket into a [`WebSocketStream`] without performing a
195    /// handshake.
196    ///
197    /// `disable_nagle` will be ignored since the socket is already connected
198    /// and the user can set `nodelay` on the socket directly before calling
199    /// this function if needed.
200    pub async fn from_raw_socket<T: IntoMaybeTlsStream<S>>(
201        stream: T,
202        role: Role,
203        config: impl Into<Config>,
204    ) -> io::Result<Self> {
205        let config = config.into();
206
207        Ok(Self::from_inner(
208            async_tungstenite::WebSocketStream::from_raw_socket(
209                stream.into_maybe_tls_stream()?,
210                role,
211                config.websocket,
212            )
213            .await,
214        ))
215    }
216
217    /// Convert a raw socket into a [`WebSocketStream`] without performing a
218    /// handshake.
219    ///
220    /// `disable_nagle` will be ignored since the socket is already connected
221    /// and the user can set `nodelay` on the socket directly before calling
222    /// this function if needed.
223    pub async fn from_partially_read<T: IntoMaybeTlsStream<S>>(
224        stream: T,
225        part: Vec<u8>,
226        role: Role,
227        config: impl Into<Config>,
228    ) -> io::Result<Self> {
229        let config = config.into();
230
231        Ok(Self::from_inner(
232            async_tungstenite::WebSocketStream::from_partially_read(
233                stream.into_maybe_tls_stream()?,
234                part,
235                role,
236                config.websocket,
237            )
238            .await,
239        ))
240    }
241
242    fn from_inner(inner: async_tungstenite::WebSocketStream<MaybePollStream<S>>) -> Self {
243        WebSocketStream {
244            inner,
245            next_item: None,
246        }
247    }
248
249    /// Send a message on the WebSocket stream.
250    pub async fn send(&mut self, message: Message) -> Result<(), WsError> {
251        SinkExt::send(self, message).await
252    }
253
254    /// Read a message from the WebSocket stream.
255    pub async fn read(&mut self) -> Result<Message, WsError> {
256        self.next()
257            .await
258            .unwrap_or_else(|| Err(WsError::ConnectionClosed))
259    }
260
261    /// Flush the WebSocket stream.
262    pub async fn flush(&mut self) -> Result<(), WsError> {
263        SinkExt::flush(self).await
264    }
265
266    /// Close the WebSocket connection.
267    pub async fn close(&mut self, close_frame: Option<CloseFrame>) -> Result<(), WsError> {
268        self.send(Message::Close(close_frame)).await
269    }
270}
271
272impl<S: AsFd> IntoInner for WebSocketStream<S> {
273    type Inner = MaybePollStream<S>;
274
275    fn into_inner(self) -> Self::Inner {
276        self.inner.into_inner()
277    }
278}
279
280impl<S: AsFd + 'static> Sink<Message> for WebSocketStream<S> {
281    type Error = WsError;
282
283    fn poll_ready(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), WsError>> {
284        self.project().inner.poll_ready(cx)
285    }
286
287    fn start_send(self: Pin<&mut Self>, item: Message) -> Result<(), Self::Error> {
288        self.project().inner.start_send(item)
289    }
290
291    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
292        ready!(self.as_mut().project().inner.poll_flush(cx))?;
293        ready!(futures_util::AsyncWrite::poll_flush(
294            Pin::new(self.project().inner.get_mut().get_mut()),
295            cx
296        ))?;
297        Poll::Ready(Ok(()))
298    }
299
300    fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
301        self.project().inner.poll_close(cx)
302    }
303}
304
305impl<S: AsFd + 'static> Stream for WebSocketStream<S> {
306    type Item = Result<Message, WsError>;
307
308    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
309        let mut this = self.project();
310        loop {
311            if this.next_item.is_some() {
312                ready!(this.inner.as_mut().poll_flush(cx))?;
313                ready!(futures_util::AsyncWrite::poll_flush(
314                    Pin::new(this.inner.get_mut().get_mut()),
315                    cx
316                ))?;
317                break Poll::Ready(this.next_item.take().expect("next_item should be Some"));
318            } else {
319                let item = ready!(this.inner.as_mut().poll_next(cx));
320                *this.next_item = Some(item);
321            }
322        }
323    }
324}
325
326impl<S: AsFd + 'static> FusedStream for WebSocketStream<S> {
327    fn is_terminated(&self) -> bool {
328        self.inner.is_terminated()
329    }
330}
331
332/// Accepts a new WebSocket connection with the provided stream.
333///
334/// This function will internally create a handshake representation and returns
335/// a future representing the resolution of the WebSocket handshake. The
336/// returned future will resolve to either [`WebSocketStream<S>`] or [`WsError`]
337/// depending on if it's successful or not.
338///
339/// This is typically used after a socket has been accepted from a
340/// `TcpListener`. That socket is then passed to this function to perform
341/// the server half of accepting a client's websocket connection.
342pub async fn accept_async<S, T>(stream: T) -> Result<WebSocketStream<S>, WsError>
343where
344    S: AsFd + 'static,
345    T: IntoMaybeTlsStream<S>,
346{
347    accept_hdr_async(stream, NoCallback).await
348}
349
350/// Similar to [`accept_async()`] but user can specify a [`Config`].
351pub async fn accept_async_with_config<S, T>(
352    stream: T,
353    config: impl Into<Config>,
354) -> Result<WebSocketStream<S>, WsError>
355where
356    S: AsFd + 'static,
357    T: IntoMaybeTlsStream<S>,
358{
359    accept_hdr_with_config_async(stream, NoCallback, config).await
360}
361
362/// Accepts a new WebSocket connection with the provided stream.
363///
364/// This function does the same as [`accept_async()`] but accepts an extra
365/// callback for header processing. The callback receives headers of the
366/// incoming requests and is able to add extra headers to the reply.
367pub async fn accept_hdr_async<S, T, C>(
368    stream: T,
369    callback: C,
370) -> Result<WebSocketStream<S>, WsError>
371where
372    S: AsFd + 'static,
373    T: IntoMaybeTlsStream<S>,
374    C: Callback + Unpin,
375{
376    accept_hdr_with_config_async(stream, callback, None).await
377}
378
379/// Similar to [`accept_hdr_async()`] but user can specify a [`Config`].
380pub async fn accept_hdr_with_config_async<S, T, C>(
381    stream: T,
382    callback: C,
383    config: impl Into<Config>,
384) -> Result<WebSocketStream<S>, WsError>
385where
386    S: AsFd + 'static,
387    T: IntoMaybeTlsStream<S>,
388    C: Callback + Unpin,
389{
390    let config = config.into();
391    let inner = async_tungstenite::accept_hdr_async_with_config(
392        stream.into_maybe_tls_stream()?,
393        callback,
394        config.websocket,
395    )
396    .await?;
397    Ok(WebSocketStream::from_inner(inner))
398}
399
400/// Creates a WebSocket handshake from a request and a stream.
401///
402/// For convenience, the user may call this with a url string, a URL,
403/// or a `Request`. Calling with `Request` allows the user to add
404/// a WebSocket protocol or other custom headers.
405///
406/// Internally, this creates a handshake representation and returns
407/// a future representing the resolution of the WebSocket handshake. The
408/// returned future will resolve to either [`WebSocketStream<S>`] or [`WsError`]
409/// depending on whether the handshake is successful.
410///
411/// This is typically used for clients who have already established, for
412/// example, a TCP connection to the remote server.
413pub async fn client_async<R, S, T>(
414    request: R,
415    stream: T,
416) -> Result<(WebSocketStream<S>, tungstenite::handshake::client::Response), WsError>
417where
418    R: IntoClientRequest + Unpin,
419    S: AsFd + 'static,
420    T: IntoMaybeTlsStream<S>,
421{
422    client_async_with_config(request, stream, None).await
423}
424
425/// Similar to [`client_async()`] but user can specify a [`Config`].
426pub async fn client_async_with_config<R, S, T>(
427    request: R,
428    stream: T,
429    config: impl Into<Config>,
430) -> Result<(WebSocketStream<S>, tungstenite::handshake::client::Response), WsError>
431where
432    R: IntoClientRequest + Unpin,
433    S: AsFd + 'static,
434    T: IntoMaybeTlsStream<S>,
435{
436    let config = config.into();
437    let (inner, response) = async_tungstenite::client_async_with_config(
438        request,
439        stream.into_maybe_tls_stream()?,
440        config.websocket,
441    )
442    .await?;
443    Ok((WebSocketStream::from_inner(inner), response))
444}
compio_ws/lib.rs

compio_ws/
lib.rs
