From 4c9ca7196dab4c745024be19c600630a76388ee4 Mon Sep 17 00:00:00 2001 From: Mohammed Sazid Al Rashid Date: Sun, 5 Dec 2021 04:32:44 +0600 Subject: [PATCH] Add `WsResponseBuilder` to build web socket session response (#1920) Co-authored-by: Rob Ede --- actix-http/CHANGES.md | 4 + actix-http/src/ws/codec.rs | 12 +- actix-web-actors/CHANGES.md | 4 + actix-web-actors/src/ws.rs | 256 ++++++++++++++++++++++++++---- actix-web-actors/tests/test_ws.rs | 201 +++++++++++++++-------- 5 files changed, 379 insertions(+), 98 deletions(-) diff --git a/actix-http/CHANGES.md b/actix-http/CHANGES.md index 23c15296..773f1ff3 100644 --- a/actix-http/CHANGES.md +++ b/actix-http/CHANGES.md @@ -9,6 +9,8 @@ * `body::None` struct. [#2468] * Impl `MessageBody` for `bytestring::ByteString`. [#2468] * `impl Clone for ws::HandshakeError`. [#2468] +* `#[must_use]` for `ws::Codec` to prevent subtle bugs. [#1920] +* `impl Default ` for `ws::Codec`. [#1920] ### Changed * Rename `body::BoxBody::{from_body => new}`. [#2468] @@ -24,9 +26,11 @@ * `impl Future` for `ResponseBuilder`. [#2468] * Remove unnecessary `MessageBody` bound on types passed to `body::AnyBody::new`. [#2468] * Move `body::AnyBody` to `awc`. Replaced with `EitherBody` and `BoxBody`. [#2468] +* `impl Copy` for `ws::Codec`. [#1920] [#2483]: https://github.com/actix/actix-web/pull/2483 [#2468]: https://github.com/actix/actix-web/pull/2468 +[#1920]: https://github.com/actix/actix-web/pull/1920 ## 3.0.0-beta.14 - 2021-11-30 diff --git a/actix-http/src/ws/codec.rs b/actix-http/src/ws/codec.rs index 8655216f..d80613e5 100644 --- a/actix-http/src/ws/codec.rs +++ b/actix-http/src/ws/codec.rs @@ -63,8 +63,8 @@ pub enum Item { Last(Bytes), } -#[derive(Debug, Copy, Clone)] /// WebSocket protocol codec. +#[derive(Debug, Clone)] pub struct Codec { flags: Flags, max_size: usize, @@ -89,7 +89,8 @@ impl Codec { /// Set max frame size. /// - /// By default max size is set to 64kB. + /// By default max size is set to 64KiB. + #[must_use = "This returns the a new Codec, without modifying the original."] pub fn max_size(mut self, size: usize) -> Self { self.max_size = size; self @@ -98,12 +99,19 @@ impl Codec { /// Set decoder to client mode. /// /// By default decoder works in server mode. + #[must_use = "This returns the a new Codec, without modifying the original."] pub fn client_mode(mut self) -> Self { self.flags.remove(Flags::SERVER); self } } +impl Default for Codec { + fn default() -> Self { + Self::new() + } +} + impl Encoder for Codec { type Error = ProtocolError; diff --git a/actix-web-actors/CHANGES.md b/actix-web-actors/CHANGES.md index e3693f0f..898098ed 100644 --- a/actix-web-actors/CHANGES.md +++ b/actix-web-actors/CHANGES.md @@ -1,8 +1,12 @@ # Changes ## Unreleased - 2021-xx-xx +* Add `ws:WsResponseBuilder` for building WebSocket session response. [#1920] +* Deprecate `ws::{start_with_addr, start_with_protocols}`. [#1920] * Minimum supported Rust version (MSRV) is now 1.52. +[#1920]: https://github.com/actix/actix-web/pull/1920 + ## 4.0.0-beta.7 - 2021-09-09 * Minimum supported Rust version (MSRV) is now 1.51. diff --git a/actix-web-actors/src/ws.rs b/actix-web-actors/src/ws.rs index b6323c2e..c41268b0 100644 --- a/actix-web-actors/src/ws.rs +++ b/actix-web-actors/src/ws.rs @@ -1,20 +1,24 @@ //! Websocket integration. -use std::future::Future; -use std::io; -use std::pin::Pin; -use std::task::{Context, Poll}; -use std::{collections::VecDeque, convert::TryFrom}; - -use actix::dev::{ - AsyncContextParts, ContextFut, ContextParts, Envelope, Mailbox, StreamHandler, ToEnvelope, +use std::{ + collections::VecDeque, + convert::TryFrom, + future::Future, + io, mem, + pin::Pin, + task::{Context, Poll}, }; -use actix::fut::ActorFuture; + use actix::{ + dev::{ + AsyncContextParts, ContextFut, ContextParts, Envelope, Mailbox, StreamHandler, + ToEnvelope, + }, + fut::ActorFuture, Actor, ActorContext, ActorState, Addr, AsyncContext, Handler, Message as ActixMessage, SpawnHandle, }; -use actix_codec::{Decoder, Encoder}; +use actix_codec::{Decoder as _, Encoder as _}; pub use actix_http::ws::{ CloseCode, CloseReason, Frame, HandshakeError, Message, ProtocolError, }; @@ -31,9 +35,188 @@ use bytes::{Bytes, BytesMut}; use bytestring::ByteString; use futures_core::Stream; use pin_project_lite::pin_project; -use tokio::sync::oneshot::Sender; +use tokio::sync::oneshot; + +/// Builder for Websocket session response. +/// +/// # Examples +/// +/// Create a Websocket session response with default configuration. +/// ```ignore +/// WsResponseBuilder::new(WsActor, &req, stream).start() +/// ``` +/// +/// Create a Websocket session with a specific max frame size, [`Codec`], and protocols. +/// ```ignore +/// const MAX_FRAME_SIZE: usize = 16_384; // 16KiB +/// +/// ws::WsResponseBuilder::new(WsActor, &req, stream) +/// .codec(Codec::new()) +/// .protocols(&["A", "B"]) +/// .frame_size(MAX_FRAME_SIZE) +/// .start() +/// ``` +pub struct WsResponseBuilder<'a, A, T> +where + A: Actor> + StreamHandler>, + T: Stream> + 'static, +{ + actor: A, + req: &'a HttpRequest, + stream: T, + codec: Option, + protocols: Option<&'a [&'a str]>, + frame_size: Option, +} + +impl<'a, A, T> WsResponseBuilder<'a, A, T> +where + A: Actor> + StreamHandler>, + T: Stream> + 'static, +{ + /// Construct a new `WsResponseBuilder` with actor, request, and payload stream. + /// + /// For usage example, see docs on [`WsResponseBuilder`] struct. + pub fn new(actor: A, req: &'a HttpRequest, stream: T) -> Self { + WsResponseBuilder { + actor, + req, + stream, + codec: None, + protocols: None, + frame_size: None, + } + } + + /// Set the protocols for the session. + pub fn protocols(mut self, protocols: &'a [&'a str]) -> Self { + self.protocols = Some(protocols); + self + } + + /// Set the max frame size for each message (in bytes). + /// + /// **Note**: This will override any given [`Codec`]'s max frame size. + pub fn frame_size(mut self, frame_size: usize) -> Self { + self.frame_size = Some(frame_size); + self + } + + /// Set the [`Codec`] for the session. If [`Self::frame_size`] is also set, the given + /// [`Codec`]'s max frame size will be overridden. + pub fn codec(mut self, codec: Codec) -> Self { + self.codec = Some(codec); + self + } + + fn handshake_resp(&self) -> Result { + match self.protocols { + Some(protocols) => handshake_with_protocols(self.req, protocols), + None => handshake(self.req), + } + } + + fn set_frame_size(&mut self) { + if let Some(frame_size) = self.frame_size { + match &mut self.codec { + Some(codec) => { + // modify existing codec's max frame size + let orig_codec = mem::take(codec); + *codec = orig_codec.max_size(frame_size); + } + + None => { + // create a new codec with the given size + self.codec = Some(Codec::new().max_size(frame_size)); + } + } + } + } + + /// Create a new Websocket context from an actor, request stream, and codec. + /// + /// Returns a pair, where the first item is an addr for the created actor, and the second item + /// is a stream intended to be set as part of the response + /// via [`HttpResponseBuilder::streaming()`]. + fn create_with_codec_addr( + actor: A, + stream: S, + codec: Codec, + ) -> (Addr, impl Stream>) + where + A: StreamHandler>, + S: Stream> + 'static, + { + let mb = Mailbox::default(); + let mut ctx = WebsocketContext { + inner: ContextParts::new(mb.sender_producer()), + messages: VecDeque::new(), + }; + ctx.add_stream(WsStream::new(stream, codec.clone())); + + let addr = ctx.address(); + + (addr, WebsocketContextFut::new(ctx, actor, mb, codec)) + } + + /// Perform WebSocket handshake and start actor. + /// + /// `req` is an [`HttpRequest`] that should be requesting a websocket protocol change. + /// `stream` should be a [`Bytes`] stream (such as `actix_web::web::Payload`) that contains a + /// stream of the body request. + /// + /// If there is a problem with the handshake, an error is returned. + /// + /// If successful, consume the [`WsResponseBuilder`] and return a [`HttpResponse`] wrapped in + /// a [`Result`]. + pub fn start(mut self) -> Result { + let mut res = self.handshake_resp()?; + self.set_frame_size(); + + match self.codec { + Some(codec) => { + let out_stream = WebsocketContext::with_codec(self.actor, self.stream, codec); + Ok(res.streaming(out_stream)) + } + None => { + let out_stream = WebsocketContext::create(self.actor, self.stream); + Ok(res.streaming(out_stream)) + } + } + } + + /// Perform WebSocket handshake and start actor. + /// + /// `req` is an [`HttpRequest`] that should be requesting a websocket protocol change. + /// `stream` should be a [`Bytes`] stream (such as `actix_web::web::Payload`) that contains a + /// stream of the body request. + /// + /// If there is a problem with the handshake, an error is returned. + /// + /// If successful, returns a pair where the first item is an address for the created actor and + /// the second item is the [`HttpResponse`] that should be returned from the websocket request. + pub fn start_with_addr(mut self) -> Result<(Addr, HttpResponse), Error> { + let mut res = self.handshake_resp()?; + self.set_frame_size(); + + match self.codec { + Some(codec) => { + let (addr, out_stream) = + Self::create_with_codec_addr(self.actor, self.stream, codec); + Ok((addr, res.streaming(out_stream))) + } + None => { + let (addr, out_stream) = + WebsocketContext::create_with_addr(self.actor, self.stream); + Ok((addr, res.streaming(out_stream))) + } + } + } +} /// Perform WebSocket handshake and start actor. +/// +/// To customize options, see [`WsResponseBuilder`]. pub fn start(actor: A, req: &HttpRequest, stream: T) -> Result where A: Actor> + StreamHandler>, @@ -45,15 +228,15 @@ where /// Perform WebSocket handshake and start actor. /// -/// `req` is an HTTP Request that should be requesting a websocket protocol -/// change. `stream` should be a `Bytes` stream (such as -/// `actix_web::web::Payload`) that contains a stream of the body request. +/// `req` is an HTTP Request that should be requesting a websocket protocol change. `stream` should +/// be a `Bytes` stream (such as `actix_web::web::Payload`) that contains a stream of the +/// body request. /// /// If there is a problem with the handshake, an error is returned. /// -/// If successful, returns a pair where the first item is an address for the -/// created actor and the second item is the response that should be returned -/// from the WebSocket request. +/// If successful, returns a pair where the first item is an address for the created actor and the +/// second item is the response that should be returned from the WebSocket request. +#[deprecated(since = "4.0.0", note = "Prefer `WsResponseBuilder::start_with_addr`.")] pub fn start_with_addr( actor: A, req: &HttpRequest, @@ -71,6 +254,10 @@ where /// Do WebSocket handshake and start ws actor. /// /// `protocols` is a sequence of known protocols. +#[deprecated( + since = "4.0.0", + note = "Prefer `WsResponseBuilder` for setting protocols." +)] pub fn start_with_protocols( actor: A, protocols: &[&str], @@ -87,20 +274,19 @@ where /// Prepare WebSocket handshake response. /// -/// This function returns handshake `HttpResponse`, ready to send to peer. -/// It does not perform any IO. +/// This function returns handshake `HttpResponse`, ready to send to peer. It does not perform +/// any IO. pub fn handshake(req: &HttpRequest) -> Result { handshake_with_protocols(req, &[]) } /// Prepare WebSocket handshake response. /// -/// This function returns handshake `HttpResponse`, ready to send to peer. -/// It does not perform any IO. +/// This function returns handshake `HttpResponse`, ready to send to peer. It does not perform +/// any IO. /// -/// `protocols` is a sequence of known protocols. On successful handshake, -/// the returned response headers contain the first protocol in this list -/// which the server also knows. +/// `protocols` is a sequence of known protocols. On successful handshake, the returned response +/// headers contain the first protocol in this list which the server also knows. pub fn handshake_with_protocols( req: &HttpRequest, protocols: &[&str], @@ -247,8 +433,8 @@ impl WebsocketContext where A: Actor, { + /// Create a new Websocket context from a request and an actor. #[inline] - /// Create a new Websocket context from a request and an actor pub fn create(actor: A, stream: S) -> impl Stream> where A: StreamHandler>, @@ -258,12 +444,11 @@ where stream } - #[inline] /// Create a new Websocket context from a request and an actor. /// - /// Returns a pair, where the first item is an addr for the created actor, - /// and the second item is a stream intended to be set as part of the - /// response via `HttpResponseBuilder::streaming()`. + /// Returns a pair, where the first item is an addr for the created actor, and the second item + /// is a stream intended to be set as part of the response + /// via [`HttpResponseBuilder::streaming()`]. pub fn create_with_addr( actor: A, stream: S, @@ -284,7 +469,6 @@ where (addr, WebsocketContextFut::new(ctx, actor, mb, Codec::new())) } - #[inline] /// Create a new Websocket context from a request, an actor, and a codec pub fn with_codec( actor: A, @@ -300,7 +484,7 @@ where inner: ContextParts::new(mb.sender_producer()), messages: VecDeque::new(), }; - ctx.add_stream(WsStream::new(stream, codec)); + ctx.add_stream(WsStream::new(stream, codec.clone())); WebsocketContextFut::new(ctx, actor, mb, codec) } @@ -458,12 +642,13 @@ where M: ActixMessage + Send + 'static, M::Result: Send, { - fn pack(msg: M, tx: Option>) -> Envelope { + fn pack(msg: M, tx: Option>) -> Envelope { Envelope::new(msg, tx) } } pin_project! { + #[derive(Debug)] struct WsStream { #[pin] stream: S, @@ -549,9 +734,12 @@ where #[cfg(test)] mod tests { + use actix_web::{ + http::{header, Method}, + test::TestRequest, + }; + use super::*; - use actix_web::http::{header, Method}; - use actix_web::test::TestRequest; #[test] fn test_handshake() { diff --git a/actix-web-actors/tests/test_ws.rs b/actix-web-actors/tests/test_ws.rs index e481b283..a9eb3769 100644 --- a/actix-web-actors/tests/test_ws.rs +++ b/actix-web-actors/tests/test_ws.rs @@ -1,11 +1,9 @@ use actix::prelude::*; -use actix_web::{ - http::{header, StatusCode}, - web, App, HttpRequest, HttpResponse, -}; -use actix_web_actors::*; +use actix_http::ws::Codec; +use actix_web::{web, App, HttpRequest}; +use actix_web_actors::ws; use bytes::Bytes; -use futures_util::{SinkExt as _, StreamExt as _}; +use futures_util::{SinkExt, StreamExt}; struct Ws; @@ -15,37 +13,34 @@ impl Actor for Ws { impl StreamHandler> for Ws { fn handle(&mut self, msg: Result, ctx: &mut Self::Context) { - match msg.unwrap() { - ws::Message::Ping(msg) => ctx.pong(&msg), - ws::Message::Text(text) => ctx.text(text), - ws::Message::Binary(bin) => ctx.binary(bin), - ws::Message::Close(reason) => ctx.close(reason), - _ => {} + match msg { + Ok(ws::Message::Ping(msg)) => ctx.pong(&msg), + Ok(ws::Message::Text(text)) => ctx.text(text), + Ok(ws::Message::Binary(bin)) => ctx.binary(bin), + Ok(ws::Message::Close(reason)) => ctx.close(reason), + _ => ctx.close(Some(ws::CloseCode::Error.into())), } } } -#[actix_rt::test] -async fn test_simple() { - let mut srv = actix_test::start(|| { - App::new().service(web::resource("/").to( - |req: HttpRequest, stream: web::Payload| async move { ws::start(Ws, &req, stream) }, - )) - }); +const MAX_FRAME_SIZE: usize = 10_000; +const DEFAULT_FRAME_SIZE: usize = 10; +async fn common_test_code(mut srv: actix_test::TestServer, frame_size: usize) { // client service let mut framed = srv.ws().await.unwrap(); - framed.send(ws::Message::Text("text".into())).await.unwrap(); + framed.send(ws::Message::Text("text".into())).await.unwrap(); let item = framed.next().await.unwrap().unwrap(); assert_eq!(item, ws::Frame::Text(Bytes::from_static(b"text"))); + let bytes = Bytes::from(vec![0; frame_size]); framed - .send(ws::Message::Binary("text".into())) + .send(ws::Message::Binary(bytes.clone())) .await .unwrap(); let item = framed.next().await.unwrap().unwrap(); - assert_eq!(item, ws::Frame::Binary(Bytes::from_static(b"text"))); + assert_eq!(item, ws::Frame::Binary(bytes)); framed.send(ws::Message::Ping("text".into())).await.unwrap(); let item = framed.next().await.unwrap().unwrap(); @@ -55,55 +50,137 @@ async fn test_simple() { .send(ws::Message::Close(Some(ws::CloseCode::Normal.into()))) .await .unwrap(); - let item = framed.next().await.unwrap().unwrap(); assert_eq!(item, ws::Frame::Close(Some(ws::CloseCode::Normal.into()))); } #[actix_rt::test] -async fn test_with_credentials() { - let mut srv = actix_test::start(|| { +async fn simple_builder() { + let srv = actix_test::start(|| { App::new().service(web::resource("/").to( |req: HttpRequest, stream: web::Payload| async move { - if req.headers().contains_key("Authorization") { - ws::start(Ws, &req, stream) - } else { - Ok(HttpResponse::new(StatusCode::UNAUTHORIZED)) - } + ws::WsResponseBuilder::new(Ws, &req, stream).start() }, )) }); - // client service without credentials - match srv.ws().await { - Ok(_) => panic!("WebSocket client without credentials should panic"), - Err(awc::error::WsClientError::InvalidResponseStatus(status)) => { - assert_eq!(status, StatusCode::UNAUTHORIZED); - } - Err(e) => panic!("Invalid error from WebSocket client: {}", e), - } - - let headers = srv.client_headers().unwrap(); - headers.insert( - header::AUTHORIZATION, - header::HeaderValue::from_static("Bearer Something"), - ); - - // client service with credentials - let client = srv.ws(); - - let mut framed = client.await.unwrap(); - - framed.send(ws::Message::Text("text".into())).await.unwrap(); - - let item = framed.next().await.unwrap().unwrap(); - assert_eq!(item, ws::Frame::Text(Bytes::from_static(b"text"))); - - framed - .send(ws::Message::Close(Some(ws::CloseCode::Normal.into()))) - .await - .unwrap(); - - let item = framed.next().await.unwrap().unwrap(); - assert_eq!(item, ws::Frame::Close(Some(ws::CloseCode::Normal.into()))); + common_test_code(srv, DEFAULT_FRAME_SIZE).await; +} + +#[actix_rt::test] +async fn builder_with_frame_size() { + let srv = actix_test::start(|| { + App::new().service(web::resource("/").to( + |req: HttpRequest, stream: web::Payload| async move { + ws::WsResponseBuilder::new(Ws, &req, stream) + .frame_size(MAX_FRAME_SIZE) + .start() + }, + )) + }); + + common_test_code(srv, MAX_FRAME_SIZE).await; +} + +#[actix_rt::test] +async fn builder_with_frame_size_exceeded() { + const MAX_FRAME_SIZE: usize = 64; + + let mut srv = actix_test::start(|| { + App::new().service(web::resource("/").to( + |req: HttpRequest, stream: web::Payload| async move { + ws::WsResponseBuilder::new(Ws, &req, stream) + .frame_size(MAX_FRAME_SIZE) + .start() + }, + )) + }); + + // client service + let mut framed = srv.ws().await.unwrap(); + + // create a request with a frame size larger than expected + let bytes = Bytes::from(vec![0; MAX_FRAME_SIZE + 1]); + framed.send(ws::Message::Binary(bytes)).await.unwrap(); + + let frame = framed.next().await.unwrap().unwrap(); + let close_reason = match frame { + ws::Frame::Close(Some(reason)) => reason, + _ => panic!("close frame expected"), + }; + assert_eq!(close_reason.code, ws::CloseCode::Error); +} + +#[actix_rt::test] +async fn builder_with_codec() { + let srv = actix_test::start(|| { + App::new().service(web::resource("/").to( + |req: HttpRequest, stream: web::Payload| async move { + ws::WsResponseBuilder::new(Ws, &req, stream) + .codec(Codec::new()) + .start() + }, + )) + }); + + common_test_code(srv, DEFAULT_FRAME_SIZE).await; +} + +#[actix_rt::test] +async fn builder_with_protocols() { + let srv = actix_test::start(|| { + App::new().service(web::resource("/").to( + |req: HttpRequest, stream: web::Payload| async move { + ws::WsResponseBuilder::new(Ws, &req, stream) + .protocols(&["A", "B"]) + .start() + }, + )) + }); + + common_test_code(srv, DEFAULT_FRAME_SIZE).await; +} + +#[actix_rt::test] +async fn builder_with_codec_and_frame_size() { + let srv = actix_test::start(|| { + App::new().service(web::resource("/").to( + |req: HttpRequest, stream: web::Payload| async move { + ws::WsResponseBuilder::new(Ws, &req, stream) + .codec(Codec::new()) + .frame_size(MAX_FRAME_SIZE) + .start() + }, + )) + }); + + common_test_code(srv, DEFAULT_FRAME_SIZE).await; +} + +#[actix_rt::test] +async fn builder_full() { + let srv = actix_test::start(|| { + App::new().service(web::resource("/").to( + |req: HttpRequest, stream: web::Payload| async move { + ws::WsResponseBuilder::new(Ws, &req, stream) + .frame_size(MAX_FRAME_SIZE) + .codec(Codec::new()) + .protocols(&["A", "B"]) + .start() + }, + )) + }); + + common_test_code(srv, MAX_FRAME_SIZE).await; +} + +#[actix_rt::test] +async fn simple_start() { + let srv = actix_test::start(|| { + App::new().service(web::resource("/").to( + |req: HttpRequest, stream: web::Payload| async move { ws::start(Ws, &req, stream) }, + )) + }); + + common_test_code(srv, DEFAULT_FRAME_SIZE).await; }