async_ws/frame/
mod.rs

1mod decode;
2mod frame_head;
3mod frame_payload;
4
5pub use decode::*;
6pub use frame_head::*;
7pub use frame_payload::*;
8
9use crate::message::WsMessageKind;
10use futures::prelude::*;
11use std::error::Error;
12use std::io::Cursor;
13use std::io::Write;
14use strum::Display;
15
16#[derive(Copy, Clone, Debug)]
17pub enum WsFrame {
18    Control(WsControlFrame),
19    Data(WsDataFrame),
20}
21
22impl WsFrame {
23    pub fn decode<T: AsyncRead + Unpin>(transport: T) -> FrameDecoder<T> {
24        FrameDecoderState::new().restore(transport)
25    }
26    // Writes the frame `buffer` and returns the number of bytes written. Panics if
27    // `buffer` is too small or payload size in frame head does not match provided payload.
28    pub fn encode(frame_head: FrameHead, frame_payload: &[u8], buffer: &mut [u8]) -> usize {
29        assert_eq!(frame_head.payload_len, frame_payload.len() as u64);
30        let total = frame_head.len_bytes() + frame_payload.len();
31        frame_head.encode(buffer);
32        let payload_buffer = &mut buffer[frame_head.len_bytes()..total];
33        payload_buffer.copy_from_slice(frame_payload);
34        payload_mask(frame_head.mask, 0, payload_buffer);
35        return total;
36    }
37    pub fn encode_vec(frame_head: FrameHead, frame_payload: &[u8]) -> Vec<u8> {
38        let mut buffer = vec![0u8; frame_head.len_bytes() + frame_payload.len()];
39        WsFrame::encode(frame_head, frame_payload, &mut *buffer);
40        buffer
41    }
42}
43
44#[derive(Copy, Clone, Debug)]
45pub enum WsFrameKind {
46    Control(WsControlFrameKind),
47    Data(WsDataFrameKind),
48}
49
50impl WsFrameKind {
51    pub fn max_payload_len(self) -> u64 {
52        match self {
53            WsFrameKind::Control(_) => 125,
54            WsFrameKind::Data(_) => 1073741824,
55        }
56    }
57    pub fn opcode(self) -> WsOpcode {
58        match self {
59            WsFrameKind::Control(frame) => frame.opcode(),
60            WsFrameKind::Data(frame) => frame.opcode(),
61        }
62    }
63    pub fn is_control(self) -> bool {
64        match self {
65            WsFrameKind::Control(_) => true,
66            WsFrameKind::Data(_) => false,
67        }
68    }
69}
70
71#[derive(Display, Copy, Clone, Debug, Eq, PartialEq)]
72pub enum WsDataFrameKind {
73    Text,
74    Binary,
75    Continuation,
76}
77
78impl WsDataFrameKind {
79    pub fn opcode(self) -> WsOpcode {
80        match self {
81            WsDataFrameKind::Text => WsOpcode::Text,
82            WsDataFrameKind::Binary => WsOpcode::Binary,
83            WsDataFrameKind::Continuation => WsOpcode::Continuation,
84        }
85    }
86    pub fn message_kind(self) -> Option<WsMessageKind> {
87        match self {
88            WsDataFrameKind::Text => Some(WsMessageKind::Text),
89            WsDataFrameKind::Binary => Some(WsMessageKind::Binary),
90            WsDataFrameKind::Continuation => None,
91        }
92    }
93    pub fn frame_kind(self) -> WsFrameKind {
94        WsFrameKind::Data(self)
95    }
96}
97
98#[derive(Copy, Clone, Debug, Eq, PartialEq)]
99pub enum WsControlFrameKind {
100    Ping,
101    Pong,
102    Close,
103}
104
105impl WsControlFrameKind {
106    pub fn opcode(self) -> WsOpcode {
107        match self {
108            WsControlFrameKind::Ping => WsOpcode::Ping,
109            WsControlFrameKind::Pong => WsOpcode::Pong,
110            WsControlFrameKind::Close => WsOpcode::Close,
111        }
112    }
113    pub fn frame_kind(self) -> WsFrameKind {
114        WsFrameKind::Control(self)
115    }
116}
117
118#[derive(Copy, Clone, Debug)]
119pub struct WsDataFrame {
120    pub(crate) kind: WsDataFrameKind,
121    pub(crate) fin: bool,
122    pub(crate) mask: [u8; 4],
123    pub(crate) payload_len: u64,
124}
125
126impl WsDataFrame {
127    pub fn payload_reader(&self) -> FramePayloadReaderState {
128        FramePayloadReaderState::new(self.mask, self.payload_len)
129    }
130    pub fn kind(&self) -> WsDataFrameKind {
131        self.kind
132    }
133    pub fn fin(&self) -> bool {
134        self.fin
135    }
136    pub fn mask(&self) -> [u8; 4] {
137        self.mask
138    }
139    pub fn payload_len(&self) -> u64 {
140        self.payload_len
141    }
142}
143
144#[derive(Copy, Clone, Debug)]
145pub struct WsControlFrame {
146    pub(crate) kind: WsControlFrameKind,
147    pub(crate) payload: WsControlFramePayload,
148}
149
150impl WsControlFrame {
151    pub fn new(kind: WsControlFrameKind, payload: &[u8]) -> Self {
152        let payload = WsControlFramePayload::new(payload);
153        Self { kind, payload }
154    }
155    pub fn payload(&self) -> &[u8] {
156        &self.payload.data()
157    }
158    pub fn kind(&self) -> WsControlFrameKind {
159        self.kind
160    }
161    pub fn head(&self, mask: [u8; 4]) -> FrameHead {
162        FrameHead {
163            fin: true,
164            opcode: self.kind.opcode(),
165            mask,
166            payload_len: self.payload().len() as u64,
167        }
168    }
169}
170
171#[derive(Copy, Clone, Debug)]
172pub(crate) struct WsControlFramePayload {
173    pub(crate) len: u8,
174    pub(crate) buffer: [u8; 125],
175}
176
177impl WsControlFramePayload {
178    pub(crate) fn new(data: &[u8]) -> Self {
179        let len = data.len().min(125);
180        let payload = &data[0..len];
181        let mut buffer = [0u8; 125];
182        buffer[0..len].copy_from_slice(payload);
183        Self {
184            len: len as u8,
185            buffer,
186        }
187    }
188    pub(crate) fn data(&self) -> &[u8] {
189        &self.buffer[0..self.len()]
190    }
191    pub(crate) fn len(&self) -> usize {
192        self.len as usize
193    }
194    pub(crate) fn close_body(&self) -> Result<Option<(u16, &str)>, CloseBodyError> {
195        match self.len() {
196            0 => Ok(None),
197            1 => Err(CloseBodyError::BodyTooShort),
198            _ => {
199                let data = self.data();
200                let code = u16::from_be_bytes([data[0], data[1]]);
201                match code {
202                    0..=999 | 1004..=1006 | 1016..=2999 => Err(CloseBodyError::InvalidCode),
203                    code => match std::str::from_utf8(&data[2..]) {
204                        Ok(reason) => Ok(Some((code, reason))),
205                        Err(_) => Err(CloseBodyError::InvalidUtf8),
206                    },
207                }
208            }
209        }
210    }
211}
212
213#[derive(thiserror::Error, Debug)]
214pub enum CloseBodyError {
215    #[error("close frame body is too short")]
216    BodyTooShort,
217    #[error("invalid utf8 in close body reason")]
218    InvalidUtf8,
219    #[error("invalid close frame body code")]
220    InvalidCode,
221}
222
223impl<E: Error> From<(u16, &E)> for WsControlFramePayload {
224    fn from(err: (u16, &E)) -> Self {
225        let mut buffer = [0u8; 125];
226        buffer[0..2].copy_from_slice(&err.0.to_be_bytes());
227        let mut cursor = Cursor::new(&mut buffer[2..]);
228        write!(cursor, "{}", err.1).ok();
229        let len = 2 + cursor.position() as u8;
230        WsControlFramePayload { len, buffer }
231    }
232}
async_ws/frame/mod.rs

async_ws/frame/
mod.rs
