Skip to main content

http_ws/
lib.rs

1//! WebSocket protocol using high level API that operate over `futures_core::Stream` trait.
2//!
3//! # HTTP type
4//! - `http` crate types are used for input and output
5//! - support `http/1.1` and `http/2`
6//! ## Examples
7//! ```rust
8//! use http::{header, Request, StatusCode};
9//! use http_ws::handshake;
10//!
11//! // an incoming http request.
12//! let request = Request::get("/")
13//!     .header(header::UPGRADE, header::HeaderValue::from_static("websocket"))
14//!     .header(header::CONNECTION, header::HeaderValue::from_static("upgrade"))
15//!     .header(header::SEC_WEBSOCKET_VERSION, header::HeaderValue::from_static("13"))
16//!     .header(header::SEC_WEBSOCKET_KEY, header::HeaderValue::from_static("some_key"))
17//!     .body(())
18//!     .unwrap();
19//!
20//! let method = request.method();
21//! let headers = request.headers();
22//!
23//! // handshake with request and return a response builder on success.
24//! let response_builder = handshake(method, headers).unwrap();
25//!
26//! // add body to builder and finalized it.
27//! let response = response_builder.body(()).unwrap();
28//!
29//! // response is valid response to websocket request.
30//! assert_eq!(response.status(), StatusCode::SWITCHING_PROTOCOLS);
31//! ```
32//!
33//! # async HTTP body
34//! Please reference [ws] function
35
36extern crate alloc;
37
38use core::ops::Deref;
39
40use http::{
41    Method, StatusCode, Version,
42    header::{
43        ALLOW, CONNECTION, HeaderMap, HeaderValue, SEC_WEBSOCKET_ACCEPT, SEC_WEBSOCKET_KEY, SEC_WEBSOCKET_VERSION,
44        UPGRADE,
45    },
46    request::Request,
47    response::{Builder, Response},
48    uri::Uri,
49};
50
51mod codec;
52mod crypto;
53mod error;
54mod frame;
55mod mask;
56mod proto;
57
58pub use self::{
59    codec::{Codec, Item, Message},
60    error::{HandshakeError, ProtocolError},
61    proto::{CloseCode, CloseReason, OpCode, hash_key},
62};
63
64#[allow(clippy::declare_interior_mutable_const)]
65mod const_header {
66    use super::HeaderValue;
67
68    pub(super) const WEBSOCKET: HeaderValue = HeaderValue::from_static("websocket");
69    pub(super) const UPGRADE_VALUE: HeaderValue = HeaderValue::from_static("upgrade");
70    pub(super) const SEC_WEBSOCKET_VERSION_VALUE: HeaderValue = HeaderValue::from_static("13");
71}
72
73use const_header::*;
74
75impl From<HandshakeError> for Builder {
76    fn from(e: HandshakeError) -> Self {
77        match e {
78            HandshakeError::GetMethodRequired => Response::builder()
79                .status(StatusCode::METHOD_NOT_ALLOWED)
80                .header(ALLOW, "GET"),
81
82            _ => Response::builder().status(StatusCode::BAD_REQUEST),
83        }
84    }
85}
86
87/// Prepare a [Request] with given [Uri] and [Version]  for websocket connection.
88///
89/// Only [Version::HTTP_11] and [Version::HTTP_2] are supported.
90/// After process the request would be ready to be sent to server.
91pub fn client_request_from_uri(uri: Uri, version: Version) -> Request<()> {
92    let mut req = Request::new(());
93    *req.uri_mut() = uri;
94    *req.version_mut() = version;
95
96    client_request_extend(&mut req);
97
98    req
99}
100
101/// Extend a [Request] with websocket associated headers and methods.
102/// After extension the request would be ready to be sent to server.
103///
104/// # HTTP/2 specific behavior
105///
106/// For HTTP/2 websocket a [`Http2WsProtocol`] type is injected into [`Extensions`]
107/// It can be used for extending protocol pseudo header for :protocol
108///
109/// [`Extensions`]: http::Extensions
110pub fn client_request_extend<B>(req: &mut Request<B>) {
111    match req.version() {
112        Version::HTTP_11 => {
113            req.headers_mut().insert(UPGRADE, WEBSOCKET);
114            req.headers_mut().insert(CONNECTION, UPGRADE_VALUE);
115
116            // generate 24 bytes base64 encoded random key.
117            let output = crypto::base64::<16, 24>(&crypto::random());
118
119            req.headers_mut()
120                .insert(SEC_WEBSOCKET_KEY, HeaderValue::from_bytes(&output).unwrap());
121        }
122        Version::HTTP_2 => {
123            *req.method_mut() = Method::CONNECT;
124            req.extensions_mut().insert(Http2WsProtocol::new());
125        }
126        _ => {}
127    }
128
129    req.headers_mut()
130        .insert(SEC_WEBSOCKET_VERSION, SEC_WEBSOCKET_VERSION_VALUE);
131}
132
133#[derive(Clone)]
134pub struct Http2WsProtocol(&'static str);
135
136impl AsRef<str> for Http2WsProtocol {
137    fn as_ref(&self) -> &str {
138        self.0
139    }
140}
141
142impl Deref for Http2WsProtocol {
143    type Target = str;
144
145    fn deref(&self) -> &Self::Target {
146        self.0
147    }
148}
149
150impl Http2WsProtocol {
151    const fn new() -> Self {
152        Self("websocket")
153    }
154}
155
156/// Verify HTTP/1.1 WebSocket handshake request and create handshake response.
157pub fn handshake(method: &Method, headers: &HeaderMap) -> Result<Builder, HandshakeError> {
158    let key = verify_handshake(method, headers)?;
159    let builder = handshake_response(key);
160    Ok(builder)
161}
162
163/// Verify HTTP/2 WebSocket handshake request and create handshake response.
164///
165/// # Protocol validation
166/// This function does **not** verify the `:protocol` pseudo-header. Per [RFC 8441], the caller
167/// must ensure the request's `:protocol` is `"websocket"` before calling this function.
168/// Typically the HTTP/2 transport layer exposes the parsed pseudo-header; the caller should
169/// check it and only proceed to this handshake when the value matches.
170///
171/// [RFC 8441]: https://www.rfc-editor.org/rfc/rfc8441
172pub fn handshake_h2(method: &Method, headers: &HeaderMap) -> Result<Builder, HandshakeError> {
173    // Check for method
174    if method != Method::CONNECT {
175        return Err(HandshakeError::ConnectMethodRequired);
176    }
177
178    ws_version_check(headers)?;
179
180    Ok(Response::builder().status(StatusCode::OK))
181}
182
183/// Verify WebSocket handshake request and return `SEC_WEBSOCKET_KEY` header value as `&[u8]`
184fn verify_handshake<'a>(method: &'a Method, headers: &'a HeaderMap) -> Result<&'a [u8], HandshakeError> {
185    // Check for method
186    if method != Method::GET {
187        return Err(HandshakeError::GetMethodRequired);
188    }
189
190    // Check for "Upgrade" header
191    let has_upgrade_hd = headers
192        .get(UPGRADE)
193        .and_then(|hdr| hdr.to_str().ok())
194        .filter(|s| s.to_ascii_lowercase().contains("websocket"))
195        .is_some();
196
197    if !has_upgrade_hd {
198        return Err(HandshakeError::NoWebsocketUpgrade);
199    }
200
201    // Check for "Connection" header
202    let has_connection_hd = headers
203        .get(CONNECTION)
204        .and_then(|hdr| hdr.to_str().ok())
205        .filter(|s| s.to_ascii_lowercase().contains("upgrade"))
206        .is_some();
207
208    if !has_connection_hd {
209        return Err(HandshakeError::NoConnectionUpgrade);
210    }
211
212    ws_version_check(headers)?;
213
214    // check client handshake for validity
215    let value = headers.get(SEC_WEBSOCKET_KEY).ok_or(HandshakeError::BadWebsocketKey)?;
216
217    Ok(value.as_bytes())
218}
219
220/// Create WebSocket handshake response.
221///
222/// This function returns handshake `http::response::Builder`, ready to send to peer.
223fn handshake_response(key: &[u8]) -> Builder {
224    let key = hash_key(key);
225
226    Response::builder()
227        .status(StatusCode::SWITCHING_PROTOCOLS)
228        .header(UPGRADE, WEBSOCKET)
229        .header(CONNECTION, UPGRADE_VALUE)
230        .header(
231            SEC_WEBSOCKET_ACCEPT,
232            // key is known to be header value safe ascii
233            HeaderValue::from_bytes(&key).unwrap(),
234        )
235}
236
237// check supported version
238fn ws_version_check(headers: &HeaderMap) -> Result<(), HandshakeError> {
239    let value = headers
240        .get(SEC_WEBSOCKET_VERSION)
241        .ok_or(HandshakeError::NoVersionHeader)?;
242
243    if value != "13" && value != "8" && value != "7" {
244        Err(HandshakeError::UnsupportedVersion)
245    } else {
246        Ok(())
247    }
248}
249
250#[cfg(feature = "stream")]
251pub mod stream;
252
253#[cfg(feature = "stream")]
254pub use self::stream::{RequestStream, ResponseSender, ResponseStream, ResponseWeakSender, WsError};
255
256#[cfg(feature = "stream")]
257pub type WsOutput<B> = (RequestStream<B>, Response<ResponseStream>, ResponseSender);
258
259#[cfg(feature = "stream")]
260/// A shortcut for generating a set of response types with given [Request] and `<Body>` type.
261///
262/// `<Body>` must be a type impl [futures_core::Stream] trait with `Result<T: AsRef<[u8]>, E>`
263/// as `Stream::Item` associated type.
264///
265/// # HTTP/2
266/// For HTTP/2 requests, the caller must verify the `:protocol` pseudo-header is `"websocket"`
267/// before calling this function. See [`handshake_h2`] for details.
268///
269/// # Examples:
270/// ```rust
271/// # use std::pin::Pin;
272/// # use std::task::{Context, Poll};
273/// # use http::{header, Request};
274/// # use futures_core::Stream;
275/// # #[derive(Default)]
276/// # struct DummyRequestBody;
277/// #
278/// # impl Stream for DummyRequestBody {
279/// #   type Item = Result<Vec<u8>, ()>;
280/// #   fn poll_next(self:Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
281/// #        Poll::Ready(Some(Ok(vec![1, 2, 3])))
282/// #    }
283/// # }
284/// # async fn ws() {
285/// use http_ws::{ws, Message};
286///
287/// // an incoming http request.
288/// let mut req = Request::get("/")
289///     .header(header::UPGRADE, header::HeaderValue::from_static("websocket"))
290///     .header(header::CONNECTION, header::HeaderValue::from_static("upgrade"))
291///     .header(header::SEC_WEBSOCKET_VERSION, header::HeaderValue::from_static("13"))
292///     .header(header::SEC_WEBSOCKET_KEY, header::HeaderValue::from_static("some_key"))
293///     .body(())
294///     .unwrap();
295///
296/// // http request body associated with http request.
297/// let body = DummyRequestBody;
298///
299/// // generate response from request and it's body.
300/// let (mut req_stream, response, res_stream) = ws(&mut req, DummyRequestBody).unwrap();
301///
302/// // req_stream must be polled with Stream interface to receive websocket message
303/// use futures_util::stream::StreamExt;
304/// if let Some(Ok(msg)) = req_stream.next().await {
305///     if let Message::Text(text) = msg {
306///         res_stream.text(text).await.unwrap();
307///     }
308/// }
309///
310/// # }
311/// ```
312pub fn ws<ReqB, B, T, E>(req: &Request<ReqB>, body: B) -> Result<WsOutput<B>, HandshakeError>
313where
314    B: futures_core::Stream<Item = Result<T, E>>,
315    T: AsRef<[u8]>,
316{
317    let builder = match req.version() {
318        Version::HTTP_2 => handshake_h2(req.method(), req.headers())?,
319        _ => handshake(req.method(), req.headers())?,
320    };
321
322    let decode = RequestStream::new(body);
323    let (res, tx) = decode.response_stream();
324
325    let res = builder
326        .body(res)
327        .expect("handshake function failed to generate correct Response Builder");
328
329    Ok((decode, res, tx))
330}
331
332#[cfg(test)]
333mod tests {
334    use super::*;
335
336    #[test]
337    fn test_handshake() {
338        let req = Request::builder().method(Method::POST).body(()).unwrap();
339        assert_eq!(
340            HandshakeError::GetMethodRequired,
341            verify_handshake(req.method(), req.headers()).unwrap_err(),
342        );
343
344        let req = Request::builder().body(()).unwrap();
345        assert_eq!(
346            HandshakeError::NoWebsocketUpgrade,
347            verify_handshake(req.method(), req.headers()).unwrap_err(),
348        );
349
350        let req = Request::builder()
351            .header(UPGRADE, HeaderValue::from_static("test"))
352            .body(())
353            .unwrap();
354        assert_eq!(
355            HandshakeError::NoWebsocketUpgrade,
356            verify_handshake(req.method(), req.headers()).unwrap_err(),
357        );
358
359        let req = Request::builder().header(UPGRADE, WEBSOCKET).body(()).unwrap();
360        assert_eq!(
361            HandshakeError::NoConnectionUpgrade,
362            verify_handshake(req.method(), req.headers()).unwrap_err(),
363        );
364
365        let req = Request::builder()
366            .header(UPGRADE, WEBSOCKET)
367            .header(CONNECTION, UPGRADE_VALUE)
368            .body(())
369            .unwrap();
370        assert_eq!(
371            HandshakeError::NoVersionHeader,
372            verify_handshake(req.method(), req.headers()).unwrap_err(),
373        );
374
375        let req = Request::builder()
376            .header(UPGRADE, WEBSOCKET)
377            .header(CONNECTION, UPGRADE_VALUE)
378            .header(SEC_WEBSOCKET_VERSION, HeaderValue::from_static("5"))
379            .body(())
380            .unwrap();
381        assert_eq!(
382            HandshakeError::UnsupportedVersion,
383            verify_handshake(req.method(), req.headers()).unwrap_err(),
384        );
385
386        let builder = || {
387            Request::builder()
388                .header(UPGRADE, WEBSOCKET)
389                .header(CONNECTION, UPGRADE_VALUE)
390                .header(SEC_WEBSOCKET_VERSION, SEC_WEBSOCKET_VERSION_VALUE)
391        };
392
393        let req = builder().body(()).unwrap();
394        assert_eq!(
395            HandshakeError::BadWebsocketKey,
396            verify_handshake(req.method(), req.headers()).unwrap_err(),
397        );
398
399        let req = builder()
400            .header(SEC_WEBSOCKET_KEY, SEC_WEBSOCKET_VERSION_VALUE)
401            .body(())
402            .unwrap();
403        let key = verify_handshake(req.method(), req.headers()).unwrap();
404        assert_eq!(
405            StatusCode::SWITCHING_PROTOCOLS,
406            handshake_response(key).body(()).unwrap().status()
407        );
408    }
409
410    #[test]
411    fn test_ws_error_http_response() {
412        let res = Builder::from(HandshakeError::GetMethodRequired).body(()).unwrap();
413        assert_eq!(res.status(), StatusCode::METHOD_NOT_ALLOWED);
414        let res = Builder::from(HandshakeError::NoWebsocketUpgrade).body(()).unwrap();
415        assert_eq!(res.status(), StatusCode::BAD_REQUEST);
416        let res = Builder::from(HandshakeError::NoConnectionUpgrade).body(()).unwrap();
417        assert_eq!(res.status(), StatusCode::BAD_REQUEST);
418        let res = Builder::from(HandshakeError::NoVersionHeader).body(()).unwrap();
419        assert_eq!(res.status(), StatusCode::BAD_REQUEST);
420        let res = Builder::from(HandshakeError::UnsupportedVersion).body(()).unwrap();
421        assert_eq!(res.status(), StatusCode::BAD_REQUEST);
422        let res = Builder::from(HandshakeError::BadWebsocketKey).body(()).unwrap();
423        assert_eq!(res.status(), StatusCode::BAD_REQUEST);
424    }
425}