fix set frame size on existing codec

This commit is contained in:
Rob Ede 2021-12-04 21:16:47 +00:00
parent d3075c0aa4
commit 7b39da96e4
No known key found for this signature in database
GPG Key ID: 97C636207D3EF933
4 changed files with 66 additions and 41 deletions

View File

@ -9,6 +9,7 @@
* `body::None` struct. [#2468] * `body::None` struct. [#2468]
* Impl `MessageBody` for `bytestring::ByteString`. [#2468] * Impl `MessageBody` for `bytestring::ByteString`. [#2468]
* `impl Clone for ws::HandshakeError`. [#2468] * `impl Clone for ws::HandshakeError`. [#2468]
* `#[must_use]` for `ws::Codec` to prevent subtle bugs. [#1920]
### Changed ### Changed
* Rename `body::BoxBody::{from_body => new}`. [#2468] * Rename `body::BoxBody::{from_body => new}`. [#2468]
@ -18,15 +19,18 @@
* `From` implementations on error types now return a `Response<BoxBody>`. [#2468] * `From` implementations on error types now return a `Response<BoxBody>`. [#2468]
* `ResponseBuilder::body(B)` now returns `Response<EitherBody<B>>`. [#2468] * `ResponseBuilder::body(B)` now returns `Response<EitherBody<B>>`. [#2468]
* `ResponseBuilder::finish()` now returns `Response<EitherBody<()>>`. [#2468] * `ResponseBuilder::finish()` now returns `Response<EitherBody<()>>`. [#2468]
* `impl Copy` for `ws::Codec`. [#1920]
### Removed ### Removed
* `ResponseBuilder::streaming`. [#2468] * `ResponseBuilder::streaming`. [#2468]
* `impl Future` for `ResponseBuilder`. [#2468] * `impl Future` for `ResponseBuilder`. [#2468]
* Remove unnecessary `MessageBody` bound on types passed to `body::AnyBody::new`. [#2468] * Remove unnecessary `MessageBody` bound on types passed to `body::AnyBody::new`. [#2468]
* Move `body::AnyBody` to `awc`. Replaced with `EitherBody` and `BoxBody`. [#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 [#2483]: https://github.com/actix/actix-web/pull/2483
[#2468]: https://github.com/actix/actix-web/pull/2468 [#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 ## 3.0.0-beta.14 - 2021-11-30

View File

@ -63,8 +63,8 @@ pub enum Item {
Last(Bytes), Last(Bytes),
} }
#[derive(Debug, Copy, Clone)]
/// WebSocket protocol codec. /// WebSocket protocol codec.
#[derive(Debug, Clone)]
pub struct Codec { pub struct Codec {
flags: Flags, flags: Flags,
max_size: usize, max_size: usize,
@ -89,7 +89,8 @@ impl Codec {
/// Set max frame size. /// 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 { pub fn max_size(mut self, size: usize) -> Self {
self.max_size = size; self.max_size = size;
self self
@ -98,12 +99,19 @@ impl Codec {
/// Set decoder to client mode. /// Set decoder to client mode.
/// ///
/// By default decoder works in server 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 { pub fn client_mode(mut self) -> Self {
self.flags.remove(Flags::SERVER); self.flags.remove(Flags::SERVER);
self self
} }
} }
impl Default for Codec {
fn default() -> Self {
Self::new()
}
}
impl Encoder<Message> for Codec { impl Encoder<Message> for Codec {
type Error = ProtocolError; type Error = ProtocolError;

View File

@ -1,20 +1,24 @@
//! Websocket integration. //! Websocket integration.
use std::future::Future; use std::{
use std::io; collections::VecDeque,
use std::pin::Pin; convert::TryFrom,
use std::task::{Context, Poll}; future::Future,
use std::{collections::VecDeque, convert::TryFrom}; io, mem,
pin::Pin,
use actix::dev::{ task::{Context, Poll},
AsyncContextParts, ContextFut, ContextParts, Envelope, Mailbox, StreamHandler, ToEnvelope,
}; };
use actix::fut::ActorFuture;
use actix::{ use actix::{
dev::{
AsyncContextParts, ContextFut, ContextParts, Envelope, Mailbox, StreamHandler,
ToEnvelope,
},
fut::ActorFuture,
Actor, ActorContext, ActorState, Addr, AsyncContext, Handler, Message as ActixMessage, Actor, ActorContext, ActorState, Addr, AsyncContext, Handler, Message as ActixMessage,
SpawnHandle, SpawnHandle,
}; };
use actix_codec::{Decoder, Encoder}; use actix_codec::{Decoder as _, Encoder as _};
pub use actix_http::ws::{ pub use actix_http::ws::{
CloseCode, CloseReason, Frame, HandshakeError, Message, ProtocolError, CloseCode, CloseReason, Frame, HandshakeError, Message, ProtocolError,
}; };
@ -31,25 +35,25 @@ use bytes::{Bytes, BytesMut};
use bytestring::ByteString; use bytestring::ByteString;
use futures_core::Stream; use futures_core::Stream;
use pin_project_lite::pin_project; use pin_project_lite::pin_project;
use tokio::sync::oneshot::Sender; use tokio::sync::oneshot;
/// Builder for Websocket Session response. /// Builder for Websocket session response.
/// ///
/// # Examples /// # Examples
/// ///
/// Create a Websocket session response with default configs. /// Create a Websocket session response with default configuration.
/// ```ignore /// ```ignore
/// WsResponseBuilder::new(WsActor, &req, stream).start() /// WsResponseBuilder::new(WsActor, &req, stream).start()
/// ``` /// ```
/// ///
/// Create a Websocket session with a specific max frame size, /// Create a Websocket session with a specific max frame size, [`Codec`], and protocols.
/// a [`Codec`] or protocols.
/// ```ignore /// ```ignore
/// const MAX_FRAME_SIZE: usize = 10_000; // in bytes. /// const MAX_FRAME_SIZE: usize = 16_384; // 16KiB
///
/// ws::WsResponseBuilder::new(WsActor, &req, stream) /// ws::WsResponseBuilder::new(WsActor, &req, stream)
/// .codec(Codec::new()) // optional /// .codec(Codec::new())
/// .protocols(&["A", "B"]) // optional /// .protocols(&["A", "B"])
/// .frame_size(MAX_FRAME_SIZE) // optional /// .frame_size(MAX_FRAME_SIZE)
/// .start() /// .start()
/// ``` /// ```
pub struct WsResponseBuilder<'a, A, T> pub struct WsResponseBuilder<'a, A, T>
@ -60,9 +64,9 @@ where
actor: A, actor: A,
req: &'a HttpRequest, req: &'a HttpRequest,
stream: T, stream: T,
codec: Option<Codec>,
protocols: Option<&'a [&'a str]>, protocols: Option<&'a [&'a str]>,
frame_size: Option<usize>, frame_size: Option<usize>,
codec: Option<Codec>,
} }
impl<'a, A, T> WsResponseBuilder<'a, A, T> impl<'a, A, T> WsResponseBuilder<'a, A, T>
@ -78,9 +82,9 @@ where
actor, actor,
req, req,
stream, stream,
codec: None,
protocols: None, protocols: None,
frame_size: None, frame_size: None,
codec: None,
} }
} }
@ -90,7 +94,7 @@ where
self self
} }
/// Set the max frame size for each message. /// Set the max frame size for each message (in bytes).
/// ///
/// **Note**: This will override any given [`Codec`]'s max frame size. /// **Note**: This will override any given [`Codec`]'s max frame size.
pub fn frame_size(mut self, frame_size: usize) -> Self { pub fn frame_size(mut self, frame_size: usize) -> Self {
@ -115,19 +119,21 @@ where
fn set_frame_size(&mut self) { fn set_frame_size(&mut self) {
if let Some(frame_size) = self.frame_size { if let Some(frame_size) = self.frame_size {
match &mut self.codec { match &mut self.codec {
// Modify existing Codec's size.
Some(codec) => { Some(codec) => {
codec.max_size(frame_size); // modify existing codec's max frame size
let orig_codec = mem::take(codec);
*codec = orig_codec.max_size(frame_size);
} }
// Otherwise, create a new codec with the given size.
None => { None => {
// create a new codec with the given size
self.codec = Some(Codec::new().max_size(frame_size)); self.codec = Some(Codec::new().max_size(frame_size));
} }
} }
} }
} }
/// Create a new Websocket context from a request, an actor, and a codec. /// 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 /// 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 /// is a stream intended to be set as part of the response
@ -146,7 +152,7 @@ where
inner: ContextParts::new(mb.sender_producer()), inner: ContextParts::new(mb.sender_producer()),
messages: VecDeque::new(), messages: VecDeque::new(),
}; };
ctx.add_stream(WsStream::new(stream, codec)); ctx.add_stream(WsStream::new(stream, codec.clone()));
let addr = ctx.address(); let addr = ctx.address();
@ -478,7 +484,7 @@ where
inner: ContextParts::new(mb.sender_producer()), inner: ContextParts::new(mb.sender_producer()),
messages: VecDeque::new(), messages: VecDeque::new(),
}; };
ctx.add_stream(WsStream::new(stream, codec)); ctx.add_stream(WsStream::new(stream, codec.clone()));
WebsocketContextFut::new(ctx, actor, mb, codec) WebsocketContextFut::new(ctx, actor, mb, codec)
} }
@ -636,12 +642,13 @@ where
M: ActixMessage + Send + 'static, M: ActixMessage + Send + 'static,
M::Result: Send, 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) Envelope::new(msg, tx)
} }
} }
pin_project! { pin_project! {
#[derive(Debug)]
struct WsStream<S> { struct WsStream<S> {
#[pin] #[pin]
stream: S, stream: S,

View File

@ -1,7 +1,7 @@
use actix::prelude::*; use actix::prelude::*;
use actix_http::ws::Codec; use actix_http::ws::Codec;
use actix_web::{web, App, HttpRequest}; use actix_web::{web, App, HttpRequest};
use actix_web_actors::*; use actix_web_actors::ws;
use bytes::Bytes; use bytes::Bytes;
use futures_util::{SinkExt, StreamExt}; use futures_util::{SinkExt, StreamExt};
@ -13,12 +13,12 @@ impl Actor for Ws {
impl StreamHandler<Result<ws::Message, ws::ProtocolError>> 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) { fn handle(&mut self, msg: Result<ws::Message, ws::ProtocolError>, ctx: &mut Self::Context) {
match msg.unwrap() { match msg {
ws::Message::Ping(msg) => ctx.pong(&msg), Ok(ws::Message::Ping(msg)) => ctx.pong(&msg),
ws::Message::Text(text) => ctx.text(text), Ok(ws::Message::Text(text)) => ctx.text(text),
ws::Message::Binary(bin) => ctx.binary(bin), Ok(ws::Message::Binary(bin)) => ctx.binary(bin),
ws::Message::Close(reason) => ctx.close(reason), Ok(ws::Message::Close(reason)) => ctx.close(reason),
_ => {} _ => ctx.close(Some(ws::CloseCode::Error.into())),
} }
} }
} }
@ -29,8 +29,8 @@ const DEFAULT_FRAME_SIZE: usize = 10;
async fn common_test_code(mut srv: actix_test::TestServer, frame_size: usize) { async fn common_test_code(mut srv: actix_test::TestServer, frame_size: usize) {
// client service // client service
let mut framed = srv.ws().await.unwrap(); 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(); let item = framed.next().await.unwrap().unwrap();
assert_eq!(item, ws::Frame::Text(Bytes::from_static(b"text"))); assert_eq!(item, ws::Frame::Text(Bytes::from_static(b"text")));
@ -50,7 +50,6 @@ async fn common_test_code(mut srv: actix_test::TestServer, frame_size: usize) {
.send(ws::Message::Close(Some(ws::CloseCode::Normal.into()))) .send(ws::Message::Close(Some(ws::CloseCode::Normal.into())))
.await .await
.unwrap(); .unwrap();
let item = framed.next().await.unwrap().unwrap(); let item = framed.next().await.unwrap().unwrap();
assert_eq!(item, ws::Frame::Close(Some(ws::CloseCode::Normal.into()))); assert_eq!(item, ws::Frame::Close(Some(ws::CloseCode::Normal.into())));
} }
@ -85,6 +84,8 @@ async fn test_builder_with_frame_size() {
#[actix_rt::test] #[actix_rt::test]
async fn test_builder_with_frame_size_exceeded() { async fn test_builder_with_frame_size_exceeded() {
const MAX_FRAME_SIZE: usize = 64;
let mut srv = actix_test::start(|| { let mut srv = actix_test::start(|| {
App::new().service(web::resource("/").to( App::new().service(web::resource("/").to(
|req: HttpRequest, stream: web::Payload| async move { |req: HttpRequest, stream: web::Payload| async move {
@ -102,7 +103,12 @@ async fn test_builder_with_frame_size_exceeded() {
let bytes = Bytes::from(vec![0; MAX_FRAME_SIZE + 1]); let bytes = Bytes::from(vec![0; MAX_FRAME_SIZE + 1]);
framed.send(ws::Message::Binary(bytes)).await.unwrap(); framed.send(ws::Message::Binary(bytes)).await.unwrap();
assert!(framed.next().await.is_none()); 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] #[actix_rt::test]