actix/actix-web
1
0
Fork 0
mirror of https://github.com/fafhrd91/actix-web synced 2024-11-30 10:42:55 +01:00
Code Issues Releases Wiki Activity

Add WsResponseBuilder to build web socket session response (#1920)

Browse Source 
Co-authored-by: Rob Ede <robjtede@icloud.com>
This commit is contained in:
Mohammed Sazid Al Rashid 2021-12-05 04:32:44 +06:00 committed by GitHub
parent fa7f3e6908
commit 4c9ca7196d
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
5 changed files with 379 additions and 98 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
actix-http/CHANGES.md
View File

@ -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

12
actix-http/src/ws/codec.rs
View File

@ -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<Message> for Codec {
type Error = ProtocolError;

4
actix-web-actors/CHANGES.md
View File

@ -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.

256
actix-web-actors/src/ws.rs
View File

@ -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<Context = WebsocketContext<A>> + StreamHandler<Result<Message, ProtocolError>>,
T: Stream<Item = Result<Bytes, PayloadError>> + 'static,
{
actor: A,
req: &'a HttpRequest,
stream: T,
codec: Option<Codec>,
protocols: Option<&'a [&'a str]>,
frame_size: Option<usize>,
}
impl<'a, A, T> WsResponseBuilder<'a, A, T>
where
A: Actor<Context = WebsocketContext<A>> + StreamHandler<Result<Message, ProtocolError>>,
T: Stream<Item = Result<Bytes, PayloadError>> + '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<HttpResponseBuilder, HandshakeError> {
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<S>(
actor: A,
stream: S,
codec: Codec,
) -> (Addr<A>, impl Stream<Item = Result<Bytes, Error>>)
where
A: StreamHandler<Result<Message, ProtocolError>>,
S: Stream<Item = Result<Bytes, PayloadError>> + '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<HttpResponse, Error> {
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<A>, 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<A, T>(actor: A, req: &HttpRequest, stream: T) -> Result<HttpResponse, Error>
where
A: Actor<Context = WebsocketContext<A>> + StreamHandler<Result<Message, ProtocolError>>,
@ -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<A, T>(
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<A, T>(
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<HttpResponseBuilder, HandshakeError> {
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<A> WebsocketContext<A>
where
A: Actor<Context = Self>,
{
/// 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<S>(actor: A, stream: S) -> impl Stream<Item = Result<Bytes, Error>>
where
A: StreamHandler<Result<Message, ProtocolError>>,
@ -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<S>(
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<S>(
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<Sender<M::Result>>) -> Envelope<A> {
fn pack(msg: M, tx: Option<oneshot::Sender<M::Result>>) -> Envelope<A> {
Envelope::new(msg, tx)
}
}
pin_project! {
#[derive(Debug)]
struct WsStream<S> {
#[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() {

195
actix-web-actors/tests/test_ws.rs
View File

@ -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<Result<ws::Message, ws::ProtocolError>> for Ws {
fn handle(&mut self, msg: Result<ws::Message, ws::ProtocolError>, 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),
common_test_code(srv, DEFAULT_FRAME_SIZE).await;
}
let headers = srv.client_headers().unwrap();
headers.insert(
header::AUTHORIZATION,
header::HeaderValue::from_static("Bearer Something"),
);
#[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()
},
))
});
// 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, 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;
}