actix/actix-web
1
0
Fork 0
mirror of https://github.com/fafhrd91/actix-web synced 2025-06-25 14:49:20 +02:00
Code Issues Releases Wiki Activity

Add websockets continuation frame support

Browse Source
This commit is contained in:
Nikolay Kim
2019-12-12 14:06:54 +06:00
parent 4a1695f719
commit b4b3350b3e
7 changed files with 269 additions and 68 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@ -12,6 +12,8 @@ pub enum Message {
Text(String),
/// Binary message
Binary(Bytes),
/// Continuation
Continuation(Item),
/// Ping message
Ping(Bytes),
/// Pong message
@ -26,9 +28,11 @@ pub enum Message {
#[derive(Debug, PartialEq)]
pub enum Frame {
/// Text frame, codec does not verify utf8 encoding
Text(Option<BytesMut>),
Text(Bytes),
/// Binary frame
Binary(Option<BytesMut>),
Binary(Bytes),
/// Continuation
Continuation(Item),
/// Ping message
Ping(Bytes),
/// Pong message
@ -37,11 +41,28 @@ pub enum Frame {
Close(Option<CloseReason>),
}
/// `WebSocket` continuation item
#[derive(Debug, PartialEq)]
pub enum Item {
FirstText(Bytes),
FirstBinary(Bytes),
Continue(Bytes),
Last(Bytes),
}
#[derive(Debug, Copy, Clone)]
/// WebSockets protocol codec
pub struct Codec {
flags: Flags,
max_size: usize,
server: bool,
}
bitflags::bitflags! {
struct Flags: u8 {
const SERVER = 0b0000_0001;
const CONTINUATION = 0b0000_0010;
const W_CONTINUATION = 0b0000_0100;
}
}
impl Codec {
@ -49,7 +70,7 @@ impl Codec {
pub fn new() -> Codec {
Codec {
max_size: 65_536,
server: true,
flags: Flags::SERVER,
}
}
@ -65,7 +86,7 @@ impl Codec {
///
/// By default decoder works in server mode.
pub fn client_mode(mut self) -> Self {
self.server = false;
self.flags.remove(Flags::SERVER);
self
}
}
@ -76,19 +97,94 @@ impl Encoder for Codec {
fn encode(&mut self, item: Message, dst: &mut BytesMut) -> Result<(), Self::Error> {
match item {
Message::Text(txt) => {
Parser::write_message(dst, txt, OpCode::Text, true, !self.server)
Message::Text(txt) => Parser::write_message(
dst,
txt,
OpCode::Text,
true,
!self.flags.contains(Flags::SERVER),
),
Message::Binary(bin) => Parser::write_message(
dst,
bin,
OpCode::Binary,
true,
!self.flags.contains(Flags::SERVER),
),
Message::Ping(txt) => Parser::write_message(
dst,
txt,
OpCode::Ping,
true,
!self.flags.contains(Flags::SERVER),
),
Message::Pong(txt) => Parser::write_message(
dst,
txt,
OpCode::Pong,
true,
!self.flags.contains(Flags::SERVER),
),
Message::Close(reason) => {
Parser::write_close(dst, reason, !self.flags.contains(Flags::SERVER))
}
Message::Binary(bin) => {
Parser::write_message(dst, bin, OpCode::Binary, true, !self.server)
}
Message::Ping(txt) => {
Parser::write_message(dst, txt, OpCode::Ping, true, !self.server)
}
Message::Pong(txt) => {
Parser::write_message(dst, txt, OpCode::Pong, true, !self.server)
}
Message::Close(reason) => Parser::write_close(dst, reason, !self.server),
Message::Continuation(cont) => match cont {
Item::FirstText(data) => {
if self.flags.contains(Flags::W_CONTINUATION) {
return Err(ProtocolError::ContinuationStarted);
} else {
self.flags.insert(Flags::W_CONTINUATION);
Parser::write_message(
dst,
&data[..],
OpCode::Binary,
false,
!self.flags.contains(Flags::SERVER),
)
}
}
Item::FirstBinary(data) => {
if self.flags.contains(Flags::W_CONTINUATION) {
return Err(ProtocolError::ContinuationStarted);
} else {
self.flags.insert(Flags::W_CONTINUATION);
Parser::write_message(
dst,
&data[..],
OpCode::Text,
false,
!self.flags.contains(Flags::SERVER),
)
}
}
Item::Continue(data) => {
if self.flags.contains(Flags::W_CONTINUATION) {
Parser::write_message(
dst,
&data[..],
OpCode::Continue,
false,
!self.flags.contains(Flags::SERVER),
)
} else {
return Err(ProtocolError::ContinuationNotStarted);
}
}
Item::Last(data) => {
if self.flags.contains(Flags::W_CONTINUATION) {
self.flags.remove(Flags::W_CONTINUATION);
Parser::write_message(
dst,
&data[..],
OpCode::Continue,
true,
!self.flags.contains(Flags::SERVER),
)
} else {
return Err(ProtocolError::ContinuationNotStarted);
}
}
},
Message::Nop => (),
}
Ok(())
@ -100,15 +196,64 @@ impl Decoder for Codec {
type Error = ProtocolError;
fn decode(&mut self, src: &mut BytesMut) -> Result<Option<Self::Item>, Self::Error> {
match Parser::parse(src, self.server, self.max_size) {
match Parser::parse(src, self.flags.contains(Flags::SERVER), self.max_size) {
Ok(Some((finished, opcode, payload))) => {
// continuation is not supported
if !finished {
return Err(ProtocolError::NoContinuation);
return match opcode {
OpCode::Continue => {
if self.flags.contains(Flags::CONTINUATION) {
Ok(Some(Frame::Continuation(Item::Continue(
payload
.map(|pl| pl.freeze())
.unwrap_or_else(Bytes::new),
))))
} else {
Err(ProtocolError::ContinuationNotStarted)
}
}
OpCode::Binary => {
if !self.flags.contains(Flags::CONTINUATION) {
self.flags.insert(Flags::CONTINUATION);
Ok(Some(Frame::Continuation(Item::FirstBinary(
payload
.map(|pl| pl.freeze())
.unwrap_or_else(Bytes::new),
))))
} else {
Err(ProtocolError::ContinuationStarted)
}
}
OpCode::Text => {
if !self.flags.contains(Flags::CONTINUATION) {
self.flags.insert(Flags::CONTINUATION);
Ok(Some(Frame::Continuation(Item::FirstText(
payload
.map(|pl| pl.freeze())
.unwrap_or_else(Bytes::new),
))))
} else {
Err(ProtocolError::ContinuationStarted)
}
}
_ => {
error!("Unfinished fragment {:?}", opcode);
Err(ProtocolError::ContinuationFragment(opcode))
}
};
}
match opcode {
OpCode::Continue => Err(ProtocolError::NoContinuation),
OpCode::Continue => {
if self.flags.contains(Flags::CONTINUATION) {
self.flags.remove(Flags::CONTINUATION);
Ok(Some(Frame::Continuation(Item::Last(
payload.map(|pl| pl.freeze()).unwrap_or_else(Bytes::new),
))))
} else {
Err(ProtocolError::ContinuationNotStarted)
}
}
OpCode::Bad => Err(ProtocolError::BadOpCode),
OpCode::Close => {
if let Some(ref pl) = payload {
@ -118,29 +263,18 @@ impl Decoder for Codec {
Ok(Some(Frame::Close(None)))
}
}
OpCode::Ping => {
if let Some(pl) = payload {
Ok(Some(Frame::Ping(pl.freeze())))
} else {
Ok(Some(Frame::Ping(Bytes::new())))
}
}
OpCode::Pong => {
if let Some(pl) = payload {
Ok(Some(Frame::Pong(pl.freeze())))
} else {
Ok(Some(Frame::Pong(Bytes::new())))
}
}
OpCode::Binary => Ok(Some(Frame::Binary(payload))),
OpCode::Text => {
Ok(Some(Frame::Text(payload)))
//let tmp = Vec::from(payload.as_ref());
//match String::from_utf8(tmp) {
// Ok(s) => Ok(Some(Message::Text(s))),
// Err(_) => Err(ProtocolError::BadEncoding),
//}
}
OpCode::Ping => Ok(Some(Frame::Ping(
payload.map(|pl| pl.freeze()).unwrap_or_else(Bytes::new),
))),
OpCode::Pong => Ok(Some(Frame::Pong(
payload.map(|pl| pl.freeze()).unwrap_or_else(Bytes::new),
))),
OpCode::Binary => Ok(Some(Frame::Binary(
payload.map(|pl| pl.freeze()).unwrap_or_else(Bytes::new),
))),
OpCode::Text => Ok(Some(Frame::Text(
payload.map(|pl| pl.freeze()).unwrap_or_else(Bytes::new),
))),
}
}
Ok(None) => Ok(None),

6
actix-http/src/ws/frame.rs
View File

@ -1,6 +1,6 @@
use std::convert::TryFrom;
use bytes::{BufMut, Bytes, BytesMut};
use bytes::{BufMut, BytesMut};
use log::debug;
use rand;
@ -154,14 +154,14 @@ impl Parser {
}
/// Generate binary representation
pub fn write_message<B: Into<Bytes>>(
pub fn write_message<B: AsRef<[u8]>>(
dst: &mut BytesMut,
pl: B,
op: OpCode,
fin: bool,
mask: bool,
) {
let payload = pl.into();
let payload = pl.as_ref();
let one: u8 = if fin {
0x80 | Into::<u8>::into(op)
} else {

17
actix-http/src/ws/mod.rs
View File

@ -18,7 +18,7 @@ mod frame;
mod mask;
mod proto;
pub use self::codec::{Codec, Frame, Message};
pub use self::codec::{Codec, Frame, Item, Message};
pub use self::dispatcher::Dispatcher;
pub use self::frame::Parser;
pub use self::proto::{hash_key, CloseCode, CloseReason, OpCode};
@ -44,12 +44,15 @@ pub enum ProtocolError {
/// A payload reached size limit.
#[display(fmt = "A payload reached size limit.")]
Overflow,
/// Continuation is not supported
#[display(fmt = "Continuation is not supported.")]
NoContinuation,
/// Bad utf-8 encoding
#[display(fmt = "Bad utf-8 encoding.")]
BadEncoding,
/// Continuation is not started
#[display(fmt = "Continuation is not started.")]
ContinuationNotStarted,
/// Received new continuation but it is already started
#[display(fmt = "Received new continuation but it is already started")]
ContinuationStarted,
/// Unknown continuation fragment
#[display(fmt = "Unknown continuation fragment.")]
ContinuationFragment(OpCode),
/// Io error
#[display(fmt = "io error: {}", _0)]
Io(io::Error),