cataclysm_ws/
frame.rs

1use crate::{Error, Message};
2
3/// Frame structure from websockets connection
4pub struct Frame {
5    inner_op_code: u8,
6    masking_key: Option<u32>,
7    /// Payload
8    pub(crate) message: Option<Message>
9}
10
11impl Frame {
12    pub const FIN_RSV: u8 = 0x80;
13    pub const OP_CODE_CONTINUATION: u8 = 0x00;
14    pub const OP_CODE_TEXT: u8 = 0x01;
15    pub const OP_CODE_BINARY: u8 = 0x02;
16    pub const OP_CODE_CLOSE: u8 = 0x08;
17    pub const OP_CODE_PING: u8 = 0x09;
18    pub const OP_CODE_PONG: u8 = 0x0A;
19
20    /// Returns the OP CODE of the frame
21    pub fn op_code(&self) -> u8 {
22        self.inner_op_code
23    }
24
25    /// Attempts to parse a frame from a stream of bytes
26    pub fn parse<A: AsRef<[u8]>>(content: A) -> Result<Frame, Error> {
27        let candidate = content.as_ref();
28        if candidate.len() < 2 {
29            return Err(Error::Incomplete)
30        }
31        // We check if the first two bytes are correct
32        if ((candidate[0] ^ !Frame::FIN_RSV) >> 4) == 0x08 {
33            return Err(Error::Parse(format!("Malformed message")))
34        }
35        // We extract the minimum length, removing the masking key
36        let min_length = candidate[1] & (!0x80);
37        let (length, mut offset) = if min_length == 126 {
38            if candidate.len() < 4 {
39                return Err(Error::Incomplete)
40            }
41            (u16::from_be_bytes([candidate[2], candidate[3]]) as usize, 4usize)
42        } else if min_length == 127 {
43            if candidate.len() < 10 {
44                return Err(Error::Incomplete)
45            }
46            (u64::from_be_bytes([candidate[2], candidate[3], candidate[4], candidate[5], candidate[6], candidate[7], candidate[8], candidate[9]]) as usize, 10usize)
47        } else {
48            (min_length as usize, 2usize)
49        };
50        // Now, the masking key, if any
51        let masking_key = if 0x80 == (candidate[1] & 0x80) {
52            offset += 4;
53            // Now, we have to add 0, 1, 2 and 3 respectively, without the 4 that we just added.
54            Some([candidate[offset - 4], candidate[offset - 3], candidate[offset - 2], candidate[offset - 1]])
55        } else {
56            None
57        };
58        // Now we read the operation code
59        let inner_op_code = (candidate[0] << 4) >> 4;
60        let mut payload = candidate.get(offset..offset+length).ok_or_else(|| Error::Incomplete)?.to_vec();
61        if let Some(masking_key) = &masking_key {
62            // We decode the content in case we have a masking key
63            payload = payload.into_iter().enumerate().map(|(idx, v)| {
64                // According to rfc 6455
65                let j = idx % 4;
66                v ^ masking_key[j]
67            }).collect();
68        }
69        let message = match inner_op_code {
70            Frame::OP_CODE_TEXT => Some(Message::Text(String::from_utf8(payload).map_err(|e| Error::Parse(format!("{}", e)))?)),
71            Frame::OP_CODE_BINARY => Some(Message::Binary(payload)),
72            Frame::OP_CODE_PING => Some(Message::Ping),
73            Frame::OP_CODE_PONG => Some(Message::Pong),
74            Frame::OP_CODE_CLOSE => None,
75            _ => return Err(Error::UnsupportedOpCode)
76        };
77
78        Ok(Frame {
79            inner_op_code,
80            masking_key: masking_key.map(u32::from_be_bytes),
81            message
82        })
83    }
84
85    /// Creates a text frame
86    pub fn text<A: Into<String>>(text: A) -> Frame {
87        let payload = text.into();
88        let message = Some(Message::Text(payload));
89        Frame {
90            inner_op_code: Frame::OP_CODE_TEXT,
91            masking_key: None,
92            message
93        }
94    }
95
96    /// Creates a binary frame
97    pub fn binary<A: Into<Vec<u8>>>(binary: A) -> Frame {
98        let payload = binary.into();
99        let message = Some(Message::Binary(payload));
100        Frame {
101            inner_op_code: Frame::OP_CODE_BINARY,
102            masking_key: None,
103            message
104        }
105    }
106
107    /// Creates a close frame
108    pub fn close() -> Frame {
109        let masking_key = Some(rand::random::<u32>());
110        Frame {
111            inner_op_code: Frame::OP_CODE_CLOSE,
112            masking_key,
113            message: None
114        }
115    }
116
117    /// Takes the frame and returns the contained message, if any
118    pub fn get_message(&self) -> Option<&Message> {
119        self.message.as_ref()
120    }
121
122    /// Indicates if this frame is a closing frame
123    pub fn is_close(&self) -> bool {
124        self.inner_op_code == Frame::OP_CODE_CLOSE
125    }
126}
127
128impl From<Frame> for Option<Message> {
129    fn from(source: Frame) -> Option<Message> {
130        source.message
131    }
132}
133
134impl From<Frame> for Vec<u8> {
135    fn from(source: Frame) -> Vec<u8> {
136        let mut content = vec![Frame::FIN_RSV ^ source.inner_op_code];
137        let mut payload: Option<Vec<u8>> = source.message.map(|m| m.into());
138        let payload_length = payload.iter().map(|m| m.len()).next().unwrap_or(0);
139        if payload_length < 126 {
140            content.push(payload_length as u8 | if source.masking_key.is_some() {0x80} else {0x00});
141        } else if payload_length <= u16::MAX.into() /*65535*/{
142            content.push(126u8 | if source.masking_key.is_some() {0x80} else {0x00});
143            // And now we push the length as u16
144            content.extend((payload_length as u16).to_be_bytes());
145        } else {
146            content.push(127u8 | if source.masking_key.is_some() {0x80} else {0x00});
147            // And now we push the length as u16
148            content.extend((payload_length as u64).to_be_bytes());
149        }
150
151        if let Some(masking_key) = &source.masking_key {
152            let masking_bytes = masking_key.to_be_bytes();
153            payload = payload.map(|b| b.into_iter().enumerate().map(|(idx, v)| {
154                // According to rfc 6455
155                let j = idx % 4;
156                v ^ masking_bytes[j]
157            }).collect());
158            content.extend(masking_bytes);
159        }
160        if let Some(payload) = payload {
161            content.extend(payload);
162        }
163        
164        content
165    }
166}
cataclysm_ws/frame.rs

cataclysm_ws/
frame.rs
