use crate::errors::{ProtocolError, WsError};
use crate::frame::{get_bit, parse_opcode, parse_payload_len, Frame, OpCode};
use crate::protocol::{cal_accept_key, standard_handshake_req_check};
use bytes::{Buf, BytesMut};
use std::fmt::Debug;
#[derive(Debug, Clone)]
pub struct FrameConfig {
pub check_rsv: bool,
pub mask: bool,
pub max_frame_payload_size: usize,
pub auto_fragment_size: usize,
pub merge_frame: bool,
}
impl Default for FrameConfig {
fn default() -> Self {
Self {
check_rsv: true,
mask: true,
max_frame_payload_size: 0,
auto_fragment_size: 0,
merge_frame: true,
}
}
}
type IOResult<T> = std::io::Result<T>;
struct FrameReadState {
fragmented: bool,
config: FrameConfig,
read_buf: [u8; 1024 * 4],
read_data: BytesMut,
fragmented_data: BytesMut,
fragmented_type: OpCode,
}
impl FrameReadState {
pub fn new() -> Self {
Self {
config: FrameConfig::default(),
fragmented: false,
read_buf: [0; 1024 * 4],
read_data: BytesMut::with_capacity(1024 * 4),
fragmented_data: BytesMut::new(),
fragmented_type: OpCode::default(),
}
}
pub fn with_config(config: FrameConfig) -> Self {
Self {
config,
..Self::new()
}
}
pub fn leading_bits_ok(&self) -> bool {
self.read_data.len() >= 2
}
pub fn get_leading_bits(&self) -> u8 {
self.read_data[0] >> 4
}
pub fn body_ok(&self, expect_len: usize) -> bool {
self.read_data.len() >= expect_len
}
fn parse_frame_header(&mut self) -> Result<usize, WsError> {
let leading_bits = self.get_leading_bits();
if self.config.check_rsv && !(leading_bits == 0b00001000 || leading_bits == 0b00000000) {
return Err(WsError::ProtocolError {
close_code: 1008,
error: ProtocolError::InvalidLeadingBits(leading_bits),
});
}
parse_opcode(self.read_data[0]).map_err(|e| WsError::ProtocolError {
error: ProtocolError::InvalidOpcode(e),
close_code: 1008,
})?;
let (payload_len, len_occ_bytes) =
parse_payload_len(self.read_data.as_ref()).map_err(|e| WsError::ProtocolError {
close_code: 1008,
error: e,
})?;
let max_payload_size = self.config.max_frame_payload_size;
if max_payload_size > 0 && payload_len > max_payload_size {
return Err(WsError::ProtocolError {
close_code: 1008,
error: ProtocolError::PayloadTooLarge(max_payload_size),
});
}
let mut expected_len = 1 + len_occ_bytes + payload_len;
let mask = get_bit(&self.read_data, 1, 0);
if mask {
expected_len += 4;
};
Ok(expected_len)
}
fn consume_frame(&mut self, len: usize) -> Frame {
let mut data = BytesMut::with_capacity(len);
data.extend_from_slice(&self.read_data[..len]);
self.read_data.advance(len);
Frame(data)
}
fn check_frame(&mut self, frame: Frame) -> Result<Option<Frame>, WsError> {
let opcode = frame.opcode();
match opcode {
OpCode::Continue => {
if !self.fragmented {
return Err(WsError::ProtocolError {
close_code: 1002,
error: ProtocolError::MissInitialFragmentedFrame,
});
}
self.fragmented_data
.extend_from_slice(&frame.payload_data_unmask());
if frame.fin() {
let completed_frame = Frame::new_with_payload(
self.fragmented_type.clone(),
&self.fragmented_data,
);
Ok(Some(completed_frame))
} else {
Ok(None)
}
}
OpCode::Text | OpCode::Binary => {
if self.fragmented {
return Err(WsError::ProtocolError {
close_code: 1002,
error: ProtocolError::NotContinueFrameAfterFragmented,
});
}
if !frame.fin() {
self.fragmented = true;
self.fragmented_type = opcode.clone();
let payload = frame.payload_data_unmask();
self.fragmented_data.extend_from_slice(&payload);
Ok(None)
} else {
Ok(Some(frame))
}
}
OpCode::Close | OpCode::Ping | OpCode::Pong => {
if !frame.fin() {
return Err(WsError::ProtocolError {
close_code: 1002,
error: ProtocolError::FragmentedControlFrame,
});
}
let payload_len = frame.payload_len();
if payload_len > 125 {
let error = ProtocolError::ControlFrameTooBig(payload_len as usize);
return Err(WsError::ProtocolError {
close_code: 1002,
error,
});
}
if opcode == OpCode::Close {
if payload_len == 1 {
let error = ProtocolError::InvalidCloseFramePayload;
return Err(WsError::ProtocolError {
close_code: 1002,
error,
});
}
if payload_len >= 2 {
let payload = frame.payload_data_unmask();
let mut code_byte = [0u8; 2];
code_byte.copy_from_slice(&payload[..2]);
let code = u16::from_be_bytes(code_byte);
if code < 1000
|| (1004..=1006).contains(&code)
|| (1015..=2999).contains(&code)
|| code >= 5000
{
let error = ProtocolError::InvalidCloseCode(code);
return Err(WsError::ProtocolError {
close_code: 1002,
error,
});
}
if String::from_utf8(payload[2..].to_vec()).is_err() {
let error = ProtocolError::InvalidUtf8;
return Err(WsError::ProtocolError {
close_code: 1007,
error,
});
}
}
}
if opcode == OpCode::Close || !self.fragmented {
Ok(Some(frame))
} else {
tracing::debug!("{:?} frame between self.fragmented data", opcode);
Ok(Some(frame))
}
}
OpCode::ReservedNonControl | OpCode::ReservedControl => {
return Err(WsError::UnsupportedFrame(opcode));
}
}
}
}
struct FrameWriteState {
config: FrameConfig,
}
impl FrameWriteState {
pub fn new() -> Self {
Self {
config: FrameConfig::default(),
}
}
pub fn with_config(config: FrameConfig) -> Self {
Self { config }
}
}
#[cfg(feature = "blocking")]
mod blocking {
use super::{FrameConfig, FrameReadState, FrameWriteState, IOResult};
use crate::{
errors::WsError,
frame::{Frame, OpCode},
protocol::standard_handshake_resp_check,
};
use std::io::{Read, Write};
impl FrameReadState {
fn poll<S: Read>(&mut self, stream: &mut S) -> IOResult<usize> {
let count = stream.read(&mut self.read_buf)?;
if count == 0 {
return Err(std::io::Error::new(
std::io::ErrorKind::ConnectionAborted,
"read eof",
));
}
self.read_data.extend_from_slice(&self.read_buf[..count]);
Ok(count)
}
fn read_one_frame<S: Read>(&mut self, stream: &mut S) -> Result<Frame, WsError> {
while !self.leading_bits_ok() {
self.poll(stream)?;
}
let len = self.parse_frame_header()?;
while !self.body_ok(len) {
self.poll(stream)?;
}
Ok(self.consume_frame(len))
}
pub fn receive<S: Read>(&mut self, stream: &mut S) -> Result<Frame, WsError> {
loop {
let frame = self.read_one_frame(stream)?;
if let Some(frame) = self.check_frame(frame)? {
break Ok(frame);
}
}
}
}
impl FrameWriteState {
fn send_one<S: Write>(
&mut self,
payload: &[u8],
fin: bool,
mask: bool,
code: OpCode,
stream: &mut S,
) -> IOResult<usize> {
let mut frame = Frame::new_with_opcode(code);
frame.set_fin(fin);
frame.set_mask(mask);
frame.set_payload(payload);
stream.write(&frame.0)
}
pub fn send<S: Write>(
&mut self,
payload: &[u8],
code: OpCode,
stream: &mut S,
) -> IOResult<usize> {
let split_size = self.config.auto_fragment_size;
let data = payload;
if split_size > 0 {
let mut idx = 0;
let mut fin = idx + split_size >= data.len();
let mut total = 0;
while (idx + split_size) <= data.len() {
let end = data.len().min(idx + split_size);
let send = self.send_one(
&data[idx..end],
fin,
self.config.mask,
code.clone(),
stream,
)?;
total += send;
idx = end;
fin = idx + split_size >= data.len();
}
Ok(total)
} else {
self.send_one(&data, true, self.config.mask, code, stream)
}
}
}
pub struct WsFrameCodec<S: Read + Write> {
stream: S,
read_state: FrameReadState,
write_state: FrameWriteState,
}
impl<S: Read + Write> WsFrameCodec<S> {
pub fn new(stream: S) -> Self {
Self {
stream,
read_state: FrameReadState::new(),
write_state: FrameWriteState::new(),
}
}
pub fn new_with(stream: S, config: FrameConfig) -> Self {
Self {
stream,
read_state: FrameReadState::with_config(config.clone()),
write_state: FrameWriteState::with_config(config),
}
}
pub fn stream_mut(&mut self) -> &mut S {
&mut self.stream
}
pub fn factory(_req: http::Request<()>, stream: S) -> Result<Self, WsError> {
let mut config = FrameConfig::default();
config.mask = false;
Ok(Self::new_with(stream, config))
}
pub fn check_fn(key: String, resp: http::Response<()>, stream: S) -> Result<Self, WsError> {
standard_handshake_resp_check(key.as_bytes(), &resp)?;
Ok(Self::new_with(stream, FrameConfig::default()))
}
pub fn receive(&mut self) -> Result<Frame, WsError> {
self.read_state.receive(&mut self.stream)
}
pub fn send(&mut self, code: OpCode, payload: &[u8]) -> Result<usize, WsError> {
self.write_state
.send(payload, code, &mut self.stream)
.map_err(|e| WsError::IOError(Box::new(e)))
}
}
}
#[cfg(feature = "blocking")]
pub use blocking::WsFrameCodec;
#[cfg(feature = "async")]
mod non_block {
use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
use super::{FrameConfig, FrameReadState, FrameWriteState, IOResult};
use crate::{
errors::WsError,
frame::{Frame, OpCode},
protocol::standard_handshake_resp_check,
};
impl FrameReadState {
async fn async_poll<S: AsyncRead + Unpin>(&mut self, stream: &mut S) -> IOResult<usize> {
let count = stream.read(&mut self.read_buf).await?;
if count == 0 {
return Err(std::io::Error::new(
std::io::ErrorKind::ConnectionAborted,
"read eof",
));
}
self.read_data.extend_from_slice(&self.read_buf[..count]);
Ok(count)
}
async fn async_read_one_frame<S: AsyncRead + Unpin>(
&mut self,
stream: &mut S,
) -> Result<Frame, WsError> {
while !self.leading_bits_ok() {
self.async_poll(stream).await?;
}
let len = self.parse_frame_header()?;
while !self.body_ok(len) {
self.async_poll(stream).await?;
}
Ok(self.consume_frame(len))
}
pub async fn async_receive<S: AsyncRead + Unpin>(
&mut self,
stream: &mut S,
) -> Result<Frame, WsError> {
loop {
let frame = self.async_read_one_frame(stream).await?;
if let Some(frame) = self.check_frame(frame)? {
break Ok(frame);
}
}
}
}
impl FrameWriteState {
async fn async_send_one<S: AsyncWrite + Unpin>(
&mut self,
payload: &[u8],
fin: bool,
mask: bool,
code: OpCode,
stream: &mut S,
) -> IOResult<usize> {
let mut frame = Frame::new_with_opcode(code);
frame.set_fin(fin);
frame.set_mask(mask);
frame.set_payload(payload);
stream.write(&frame.0).await
}
pub async fn async_send<S: AsyncWrite + Unpin>(
&mut self,
payload: &[u8],
code: OpCode,
stream: &mut S,
) -> IOResult<usize> {
let split_size = self.config.auto_fragment_size;
let data = payload;
if split_size > 0 {
let mut idx = 0;
let mut fin = idx + split_size >= data.len();
let mut total = 0;
while (idx + split_size) <= data.len() {
let end = data.len().min(idx + split_size);
let send = self
.async_send_one(
&data[idx..end],
fin,
self.config.mask,
code.clone(),
stream,
)
.await?;
total += send;
idx = end;
fin = idx + split_size >= data.len();
}
Ok(total)
} else {
self.async_send_one(&data, true, self.config.mask, code, stream)
.await
}
}
}
pub struct AsyncWsFrameCodec<S: AsyncRead + AsyncWrite> {
stream: S,
read_state: FrameReadState,
write_state: FrameWriteState,
}
impl<S: AsyncRead + AsyncWrite + Unpin> AsyncWsFrameCodec<S> {
pub fn new(stream: S) -> Self {
Self {
stream,
read_state: FrameReadState::new(),
write_state: FrameWriteState::new(),
}
}
pub fn new_with(stream: S, config: FrameConfig) -> Self {
Self {
stream,
read_state: FrameReadState::with_config(config.clone()),
write_state: FrameWriteState::with_config(config),
}
}
pub fn factory(_req: http::Request<()>, stream: S) -> Result<Self, WsError> {
let mut config = FrameConfig::default();
config.mask = false;
Ok(Self::new_with(stream, config))
}
pub fn check_fn(key: String, resp: http::Response<()>, stream: S) -> Result<Self, WsError> {
standard_handshake_resp_check(key.as_bytes(), &resp)?;
Ok(Self::new_with(stream, FrameConfig::default()))
}
pub fn stream_mut(&mut self) -> &mut S {
&mut self.stream
}
pub async fn receive(&mut self) -> Result<Frame, WsError> {
self.read_state.async_receive(&mut self.stream).await
}
pub async fn send(&mut self, code: OpCode, payload: &[u8]) -> Result<usize, WsError> {
self.write_state
.async_send(payload, code, &mut self.stream)
.await
.map_err(|e| WsError::IOError(Box::new(e)))
}
}
}
#[cfg(feature = "async")]
pub use non_block::AsyncWsFrameCodec;
pub fn default_handshake_handler(
req: http::Request<()>,
) -> Result<(http::Request<()>, http::Response<String>), WsError> {
match standard_handshake_req_check(&req) {
Ok(_) => {
let key = req.headers().get("sec-websocket-key").unwrap();
let resp = http::Response::builder()
.version(http::Version::HTTP_11)
.status(http::StatusCode::SWITCHING_PROTOCOLS)
.header("Upgrade", "WebSocket")
.header("Connection", "Upgrade")
.header("Sec-WebSocket-Accept", cal_accept_key(key.as_bytes()))
.body(String::new())
.unwrap();
Ok((req, resp))
}
Err(e) => {
let resp = http::Response::builder()
.version(http::Version::HTTP_11)
.status(http::StatusCode::BAD_REQUEST)
.header("Content-Type", "text/html")
.body(e.to_string())
.unwrap();
Ok((req, resp))
}
}
}