cataclysm_ws/
frame.rs

1use crate::{FrameParseError, Message};
2
3/// Frame structure from websockets connection
4pub struct Frame {
5    inner_op_code: u8,
6    masking_key: Option<u32>,
7    /// Inner message
8    pub message: Message
9}
10
11impl Frame {
12    /// FIN RSV bytes
13    pub const FIN_RSV: u8 = 0x80;
14
15    // The operation codes are the last 4 bytes of the u8
16
17    /// Operation code for a continuation request
18    pub const OP_CODE_CONTINUATION: u8 = 0x00;
19    /// Operation code for a text message
20    pub const OP_CODE_TEXT: u8 = 0x01;
21    /// Operation code for a binary message
22    pub const OP_CODE_BINARY: u8 = 0x02;
23    /// Operation code for a close request message
24    pub const OP_CODE_CLOSE: u8 = 0x08;
25    /// Operation code for a ping message
26    pub const OP_CODE_PING: u8 = 0x09;
27    /// Operation code for a pong message
28    pub const OP_CODE_PONG: u8 = 0x0A;
29
30    /// Returns the OP CODE of the frame as a u8, where the last 4 bits contain the OP CODE
31    pub fn op_code(&self) -> u8 {
32        self.inner_op_code
33    }
34
35    /// Attempts to parse a frame from a stream of bytes
36    pub fn parse<A: AsRef<[u8]>>(content: A) -> Result<Frame, FrameParseError> {
37        let candidate = content.as_ref();
38
39        if candidate.is_empty() {
40            // Not enough bytes to even read a possible FIN_RSV + OP_CODE, and prevent panics
41            return Err(FrameParseError::NullContent);
42        }
43
44        // The or operation needs to put ones on the first 4 bits, if candidate[0] matches FIN_RSV
45        if ((candidate[0] ^ !Frame::FIN_RSV) >> 4) != 0x0f {
46            return Err(FrameParseError::WrongFinRSV);
47        }
48
49        // We extract the minimum length, removing the masking key
50        let min_length = candidate[1] & (!0x80);
51        let (length, mut offset) = if min_length == 126 {
52            if candidate.len() < 4 {
53                return Err(FrameParseError::Malformed)
54            }
55            (u16::from_be_bytes([candidate[2], candidate[3]]) as usize, 4usize)
56        } else if min_length == 127 {
57            if candidate.len() < 10 {
58                return Err(FrameParseError::Malformed)
59            }
60            (u64::from_be_bytes([candidate[2], candidate[3], candidate[4], candidate[5], candidate[6], candidate[7], candidate[8], candidate[9]]) as usize, 10usize)
61        } else {
62            (min_length as usize, 2usize)
63        };
64
65        // Now, the masking key, if any
66        let masking_key = if 0x80 == (candidate[1] & 0x80) {
67            if candidate.len() < offset + 4  {
68                // Quite likely not a malformed message, just incomplete
69                return Err(FrameParseError::Incomplete{expected: length + offset + 4, obtained: candidate.len()})
70            }
71
72            offset += 4;
73            // Now, we have to add 0, 1, 2 and 3 respectively, without the 4 that we just added.
74            Some([candidate[offset - 4], candidate[offset - 3], candidate[offset - 2], candidate[offset - 1]])
75        } else {
76            None
77        };
78        // Now we read the operation code
79        let inner_op_code = (candidate[0] << 4) >> 4;
80        let mut payload = candidate.get(offset..offset+length).ok_or_else(|| FrameParseError::Incomplete{expected: length + offset, obtained: candidate.len()})?.to_vec();
81        if let Some(masking_key) = &masking_key {
82            // We decode the content in case we have a masking key
83            payload = payload.into_iter().enumerate().map(|(idx, v)| {
84                // According to rfc 6455
85                let j = idx % 4;
86                v ^ masking_key[j]
87            }).collect();
88        }
89        let message = match inner_op_code {
90            Frame::OP_CODE_TEXT => Message::Text(String::from_utf8(payload).map_err(|e| FrameParseError::InvalidUtf8(e))?),
91            Frame::OP_CODE_BINARY => Message::Binary(payload),
92            Frame::OP_CODE_PING => Message::Ping(payload),
93            Frame::OP_CODE_PONG => Message::Pong(payload),
94            Frame::OP_CODE_CLOSE => Message::Close,
95            _ => return Err(FrameParseError::UnsupportedOpCode)
96        };
97
98        Ok(Frame {
99            inner_op_code,
100            masking_key: masking_key.map(u32::from_be_bytes),
101            message
102        })
103    }
104
105    /// Creates a text frame
106    pub fn text<A: Into<String>>(text: A) -> Frame {
107        let payload = text.into();
108        let message = Message::Text(payload);
109        Frame {
110            inner_op_code: Frame::OP_CODE_TEXT,
111            masking_key: None,
112            message
113        }
114    }
115
116    /// Creates a ping message with the given application data
117    pub fn ping<A: Into<Vec<u8>>>(payload: A) -> Frame {
118        let payload = payload.into();
119        let message = Message::Ping(payload);
120        Frame {
121            inner_op_code: Frame::OP_CODE_PING,
122            masking_key: None,
123            message
124        }
125    }
126
127    /// Creates a pong message with the given application data
128    pub fn pong<A: Into<Vec<u8>>>(payload: A) -> Frame {
129        let payload = payload.into();
130        let message = Message::Pong(payload);
131        Frame {
132            inner_op_code: Frame::OP_CODE_PONG,
133            masking_key: None,
134            message
135        }
136    }
137
138    /// Creates a binary frame
139    pub fn binary<A: Into<Vec<u8>>>(binary: A) -> Frame {
140        let payload = binary.into();
141        let message = Message::Binary(payload);
142        Frame {
143            inner_op_code: Frame::OP_CODE_BINARY,
144            masking_key: None,
145            message
146        }
147    }
148
149    /// Creates a close frame
150    pub fn close() -> Frame {
151        let masking_key = None; //Some(rand::random::<u32>());
152        Frame {
153            inner_op_code: Frame::OP_CODE_CLOSE,
154            masking_key,
155            message: Message::Close
156        }
157    }
158
159    /// Takes the frame and returns the contained message, if any
160    pub fn get_message(&self) -> &Message {
161        &self.message
162    }
163
164    /// Indicates if this frame is a closing frame
165    pub fn is_close(&self) -> bool {
166        self.inner_op_code == Frame::OP_CODE_CLOSE
167    }
168}
169
170impl From<Frame> for Message {
171    fn from(source: Frame) -> Message {
172        source.message
173    }
174}
175
176impl From<Frame> for Vec<u8> {
177    fn from(source: Frame) -> Vec<u8> {
178        let mut content = vec![Frame::FIN_RSV ^ source.inner_op_code];
179        let mut payload: Vec<u8> = source.message.into();
180        let payload_length = payload.len();
181        if payload_length < 126 {
182            content.push(payload_length as u8 | if source.masking_key.is_some() {0x80} else {0x00});
183        } else if payload_length <= u16::MAX.into() /*65535*/{
184            content.push(126u8 | if source.masking_key.is_some() {0x80} else {0x00});
185            // And now we push the length as u16
186            content.extend((payload_length as u16).to_be_bytes());
187        } else {
188            content.push(127u8 | if source.masking_key.is_some() {0x80} else {0x00});
189            // And now we push the length as u16
190            content.extend((payload_length as u64).to_be_bytes());
191        }
192
193        if let Some(masking_key) = &source.masking_key {
194            let masking_bytes = masking_key.to_be_bytes();
195            payload = payload.into_iter().enumerate().map(|(idx, v)| {
196                // According to rfc 6455
197                let j = idx % 4;
198                v ^ masking_bytes[j]
199            }).collect();
200            // We add the masking key
201            content.extend(masking_bytes);
202        }
203        if !payload.is_empty() {
204            content.extend(payload);
205        }
206        
207        content
208    }
209}
cataclysm_ws/frame.rs

cataclysm_ws/
frame.rs
