async_ws/frame/frame_head/
mod.rs

1mod decode;
2
3pub use decode::*;
4
5use crate::frame::{WsControlFrameKind, WsDataFrameKind, WsFrameKind};
6use futures::prelude::*;
7
8#[derive(Copy, Clone, Debug)]
9pub enum WsOpcode {
10    Continuation,
11    Text,
12    Binary,
13    Close,
14    Ping,
15    Pong,
16}
17
18impl WsOpcode {
19    pub fn frame_kind(self) -> WsFrameKind {
20        match self {
21            WsOpcode::Continuation => WsFrameKind::Data(WsDataFrameKind::Continuation),
22            WsOpcode::Text => WsFrameKind::Data(WsDataFrameKind::Text),
23            WsOpcode::Binary => WsFrameKind::Data(WsDataFrameKind::Binary),
24            WsOpcode::Close => WsFrameKind::Control(WsControlFrameKind::Close),
25            WsOpcode::Ping => WsFrameKind::Control(WsControlFrameKind::Ping),
26            WsOpcode::Pong => WsFrameKind::Control(WsControlFrameKind::Pong),
27        }
28    }
29}
30
31#[derive(Copy, Clone, Debug)]
32pub struct FrameHead {
33    pub fin: bool,
34    pub opcode: WsOpcode,
35    pub mask: [u8; 4],
36    pub payload_len: u64,
37}
38
39impl FrameHead {
40    pub fn decode<T: AsyncRead + Unpin>(transport: T) -> FrameHeadDecode<T> {
41        FrameHeadDecodeState::new().restore(transport)
42    }
43    pub fn parse(buffer: &[u8]) -> Result<FrameHead, FrameHeadParseError> {
44        if buffer.len() < 2 {
45            return Err(FrameHeadParseError::Incomplete(2));
46        }
47        let (masked, extra_payload_len_bytes) = match buffer[1] {
48            0..=125 => (false, 0usize),
49            126 => (false, 2usize),
50            127 => (false, 8usize),
51            128..=253 => (true, 0usize),
52            254 => (true, 2usize),
53            255 => (true, 8usize),
54        };
55        let expected_buffer_len = 2 + extra_payload_len_bytes + (masked as usize) * 4;
56        if buffer.len() < expected_buffer_len {
57            return Err(FrameHeadParseError::Incomplete(expected_buffer_len));
58        }
59        let fin = match buffer[0] & 0xF0 {
60            0x00 => false,
61            0x80 => true,
62            _ => return Err(FrameHeadParseError::RsvBit),
63        };
64        let opcode = match buffer[0] & 0x0F {
65            0x0 => WsOpcode::Continuation,
66            0x1 => WsOpcode::Text,
67            0x2 => WsOpcode::Binary,
68            0x8 => WsOpcode::Close,
69            0x9 => WsOpcode::Ping,
70            0xA => WsOpcode::Pong,
71            n => return Err(FrameHeadParseError::InvalidOpcode(n)),
72        };
73        let mut payload_len = [0u8; 8];
74        match extra_payload_len_bytes {
75            0 => payload_len[7] = buffer[1] & 127,
76            2 => payload_len[6..8].copy_from_slice(&buffer[2..4]),
77            8 => payload_len.copy_from_slice(&buffer[2..10]),
78            _ => unreachable!(),
79        };
80        let payload_len = u64::from_be_bytes(payload_len);
81        let mut mask = [0u8; 4];
82        if masked {
83            mask.copy_from_slice(&buffer[2 + extra_payload_len_bytes..6 + extra_payload_len_bytes])
84        }
85        if let WsFrameKind::Control(_) = opcode.frame_kind() {
86            if !fin {
87                return Err(FrameHeadParseError::FragmentedControl);
88            }
89        }
90        if payload_len > opcode.frame_kind().max_payload_len() {
91            return Err(FrameHeadParseError::PayloadLengthTooLong);
92        }
93        Ok(FrameHead {
94            fin,
95            opcode,
96            mask,
97            payload_len,
98        })
99    }
100    // Length of the encoded frame head in bytes ([2..14]).
101    pub fn len_bytes(&self) -> usize {
102        let extra_payload_len_bytes = match self.payload_len {
103            0..=125 => 0usize,
104            126..=65536 => 2usize,
105            _ => 8usize,
106        };
107        2 + extra_payload_len_bytes + self.masked() as usize * 4
108    }
109    pub fn masked(&self) -> bool {
110        self.mask != [0u8, 0u8, 0u8, 0u8]
111    }
112    // Writes the frame head to `buffer`.
113    // Panics if `buffer` is too small (see [len_bytes()][`Self::len_bytes()`]).
114    pub fn encode(&self, buffer: &mut [u8]) {
115        buffer[0] = self.fin as u8 * 0x80;
116        buffer[0] += match self.opcode {
117            WsOpcode::Continuation => 0x0,
118            WsOpcode::Text => 0x1,
119            WsOpcode::Binary => 0x2,
120            WsOpcode::Close => 0x8,
121            WsOpcode::Ping => 0x9,
122            WsOpcode::Pong => 0xA,
123        };
124        buffer[1] = match self.payload_len {
125            0..=125 => self.payload_len as u8,
126            126..=65536 => 126u8,
127            _ => 127u8,
128        };
129        match buffer[1] {
130            126 => buffer[2..4].copy_from_slice(&(self.payload_len as u16).to_be_bytes()),
131            127 => buffer[2..10].copy_from_slice(&self.payload_len.to_be_bytes()),
132            _ => {}
133        }
134        if self.masked() {
135            let mask_buffer = match buffer[1] {
136                126 => &mut buffer[4..8],
137                127 => &mut buffer[10..14],
138                _ => &mut buffer[2..6],
139            };
140            mask_buffer.copy_from_slice(&self.mask);
141            buffer[1] += 128;
142        }
143    }
144}
145
146#[derive(thiserror::Error, Debug)]
147pub enum FrameHeadParseError {
148    #[error("incomplete, need at least {0} bytes")]
149    Incomplete(usize),
150    #[error("one or more RSV bit is set")]
151    RsvBit,
152    #[error("invalid opcode")]
153    InvalidOpcode(u8),
154    #[error("payload length too long")]
155    PayloadLengthTooLong,
156    #[error("fragmented control message")]
157    FragmentedControl,
158}
async_ws/frame/frame_head/mod.rs

async_ws/frame/frame_head/
mod.rs
