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

compio_ws/
lib.rs
