websocket_base/
dataframe.rs

1//! Module containing the default implementation of data frames.
2use crate::result::{WebSocketError, WebSocketResult};
3use crate::ws::dataframe::DataFrame as DataFrameable;
4use crate::ws::util::header as dfh;
5use crate::ws::util::header::DataFrameHeader;
6use crate::ws::util::mask;
7use std::io::{self, Read, Write};
8
9/// Represents a WebSocket data frame.
10///
11/// The data held in a DataFrame is never masked.
12/// Masking/unmasking is done when sending and receiving the data frame,
13///
14/// This DataFrame, unlike the standard Message implementation (which also
15/// implements the DataFrame trait), owns its entire payload. This means that calls to `payload`
16/// don't allocate extra memory (again unlike the default Message implementation).
17#[derive(Debug, Clone, PartialEq)]
18pub struct DataFrame {
19	/// Whether or no this constitutes the end of a message
20	pub finished: bool,
21	/// The reserved portion of the data frame (RFC6455 5.2)
22	pub reserved: [bool; 3],
23	/// The opcode associated with this data frame
24	pub opcode: Opcode,
25	/// The payload associated with this data frame
26	pub data: Vec<u8>,
27}
28
29impl DataFrame {
30	/// Creates a new DataFrame.
31	pub fn new(finished: bool, opcode: Opcode, data: Vec<u8>) -> DataFrame {
32		DataFrame {
33			finished,
34			reserved: [false; 3],
35			opcode,
36			data,
37		}
38	}
39
40	/// Take the body and header of a dataframe and combine it into a single
41	/// Dataframe struct. A websocket message can be made up of many individual
42	/// dataframes, use the methods from the Message or OwnedMessage structs to
43	/// take many of these and create a websocket message.
44	pub fn read_dataframe_body(
45		header: DataFrameHeader,
46		body: Vec<u8>,
47		should_be_masked: bool,
48	) -> WebSocketResult<Self> {
49		let finished = header.flags.contains(dfh::DataFrameFlags::FIN);
50
51		let reserved = [
52			header.flags.contains(dfh::DataFrameFlags::RSV1),
53			header.flags.contains(dfh::DataFrameFlags::RSV2),
54			header.flags.contains(dfh::DataFrameFlags::RSV3),
55		];
56
57		let opcode = Opcode::new(header.opcode).expect("Invalid header opcode!");
58
59		let data = match header.mask {
60			Some(mask) => {
61				if !should_be_masked {
62					return Err(WebSocketError::DataFrameError(
63						"Expected unmasked data frame",
64					));
65				}
66				mask::mask_data(mask, &body)
67			}
68			None => {
69				if should_be_masked {
70					return Err(WebSocketError::DataFrameError("Expected masked data frame"));
71				}
72				body
73			}
74		};
75
76		Ok(DataFrame {
77			finished,
78			reserved,
79			opcode,
80			data,
81		})
82	}
83
84	/// Reads a DataFrame from a Reader.
85	pub fn read_dataframe<R>(reader: &mut R, should_be_masked: bool) -> WebSocketResult<Self>
86	where
87		R: Read,
88	{
89		let header = dfh::read_header(reader)?;
90
91		let mut data: Vec<u8> = Vec::with_capacity(header.len as usize);
92		let read = reader.take(header.len).read_to_end(&mut data)?;
93		if (read as u64) < header.len {
94			return Err(io::Error::new(io::ErrorKind::UnexpectedEof, "incomplete payload").into());
95		}
96
97		DataFrame::read_dataframe_body(header, data, should_be_masked)
98	}
99
100	/// Reads a DataFrame from a Reader, or error out if header declares exceeding limit you specify
101	pub fn read_dataframe_with_limit<R>(reader: &mut R, should_be_masked: bool, limit: usize) -> WebSocketResult<Self>
102	where
103		R: Read,
104	{
105		let header = dfh::read_header(reader)?;
106
107		if header.len > limit as u64 {
108			return Err(io::Error::new(io::ErrorKind::InvalidData, "exceeded DataFrame length limit").into());
109		}
110		let mut data: Vec<u8> = Vec::with_capacity(header.len as usize);
111		let read = reader.take(header.len).read_to_end(&mut data)?;
112		if (read as u64) < header.len {
113			return Err(io::Error::new(io::ErrorKind::UnexpectedEof, "incomplete payload").into());
114		}
115
116		DataFrame::read_dataframe_body(header, data, should_be_masked)
117	}
118}
119
120impl DataFrameable for DataFrame {
121	#[inline(always)]
122	fn is_last(&self) -> bool {
123		self.finished
124	}
125
126	#[inline(always)]
127	fn opcode(&self) -> u8 {
128		self.opcode as u8
129	}
130
131	#[inline(always)]
132	fn reserved(&self) -> &[bool; 3] {
133		&self.reserved
134	}
135
136	#[inline(always)]
137	fn size(&self) -> usize {
138		self.data.len()
139	}
140
141	#[inline(always)]
142	fn write_payload(&self, socket: &mut dyn Write) -> WebSocketResult<()> {
143		socket.write_all(self.data.as_slice())?;
144		Ok(())
145	}
146
147	#[inline(always)]
148	fn take_payload(self) -> Vec<u8> {
149		self.data
150	}
151}
152
153/// Represents a WebSocket data frame opcode
154#[derive(Clone, Debug, Copy, PartialEq)]
155pub enum Opcode {
156	/// A continuation data frame
157	Continuation,
158	/// A UTF-8 text data frame
159	Text,
160	/// A binary data frame
161	Binary,
162	/// An undefined non-control data frame
163	NonControl1,
164	/// An undefined non-control data frame
165	NonControl2,
166	/// An undefined non-control data frame
167	NonControl3,
168	/// An undefined non-control data frame
169	NonControl4,
170	/// An undefined non-control data frame
171	NonControl5,
172	/// A close data frame
173	Close,
174	/// A ping data frame
175	Ping,
176	/// A pong data frame
177	Pong,
178	/// An undefined control data frame
179	Control1,
180	/// An undefined control data frame
181	Control2,
182	/// An undefined control data frame
183	Control3,
184	/// An undefined control data frame
185	Control4,
186	/// An undefined control data frame
187	Control5,
188}
189
190impl Opcode {
191	/// Attempts to form an Opcode from a nibble.
192	///
193	/// Returns the Opcode, or None if the opcode is out of range.
194	#[warn(clippy::new_ret_no_self)]
195	pub fn new(op: u8) -> Option<Opcode> {
196		Some(match op {
197			0 => Opcode::Continuation,
198			1 => Opcode::Text,
199			2 => Opcode::Binary,
200			3 => Opcode::NonControl1,
201			4 => Opcode::NonControl2,
202			5 => Opcode::NonControl3,
203			6 => Opcode::NonControl4,
204			7 => Opcode::NonControl5,
205			8 => Opcode::Close,
206			9 => Opcode::Ping,
207			10 => Opcode::Pong,
208			11 => Opcode::Control1,
209			12 => Opcode::Control2,
210			13 => Opcode::Control3,
211			14 => Opcode::Control4,
212			15 => Opcode::Control5,
213			_ => return None,
214		})
215	}
216}
217
218#[cfg(all(feature = "nightly", test))]
219mod tests {
220	use super::*;
221	use test::Bencher;
222	use ws::dataframe::DataFrame as DataFrameable;
223
224	#[test]
225	fn test_read_dataframe() {
226		let data = b"The quick brown fox jumps over the lazy dog";
227		let mut dataframe = vec![0x81, 0x2B];
228		for i in data.iter() {
229			dataframe.push(*i);
230		}
231		let obtained = DataFrame::read_dataframe(&mut &dataframe[..], false).unwrap();
232		let expected = DataFrame {
233			finished: true,
234			reserved: [false; 3],
235			opcode: Opcode::Text,
236			data: data.to_vec(),
237		};
238		assert_eq!(obtained, expected);
239	}
240
241	#[test]
242	fn read_incomplete_payloads() {
243		let mut data = vec![0x8au8, 0x08, 0x19, 0xac, 0xab, 0x8a, 0x52, 0x4e, 0x05, 0x00];
244		let payload = vec![25, 172, 171, 138, 82, 78, 5, 0];
245		let short_header = DataFrame::read_dataframe(&mut &data[..1], false);
246		let short_payload = DataFrame::read_dataframe(&mut &data[..6], false);
247		let full_payload = DataFrame::read_dataframe(&mut &data[..], false);
248		data.push(0xff);
249		let more_payload = DataFrame::read_dataframe(&mut &data[..], false);
250
251		match (short_header.unwrap_err(), short_payload.unwrap_err()) {
252			(WebSocketError::NoDataAvailable, WebSocketError::NoDataAvailable) => (),
253			_ => assert!(false),
254		};
255		assert_eq!(full_payload.unwrap().data, payload);
256		assert_eq!(more_payload.unwrap().data, payload);
257	}
258
259	#[bench]
260	fn bench_read_dataframe(b: &mut Bencher) {
261		let data = b"The quick brown fox jumps over the lazy dog";
262		let mut dataframe = vec![0x81, 0x2B];
263		for i in data.iter() {
264			dataframe.push(*i);
265		}
266		b.iter(|| {
267			DataFrame::read_dataframe(&mut &dataframe[..], false).unwrap();
268		});
269	}
270
271	#[test]
272	fn test_write_dataframe() {
273		let data = b"The quick brown fox jumps over the lazy dog";
274		let mut expected = vec![0x81, 0x2B];
275		for i in data.iter() {
276			expected.push(*i);
277		}
278		let dataframe = DataFrame {
279			finished: true,
280			reserved: [false; 3],
281			opcode: Opcode::Text,
282			data: data.to_vec(),
283		};
284		let mut obtained = Vec::new();
285		dataframe.write_to(&mut obtained, false).unwrap();
286
287		assert_eq!(&obtained[..], &expected[..]);
288	}
289
290	#[bench]
291	fn bench_write_dataframe(b: &mut Bencher) {
292		let data = b"The quick brown fox jumps over the lazy dog";
293		let dataframe = DataFrame {
294			finished: true,
295			reserved: [false; 3],
296			opcode: Opcode::Text,
297			data: data.to_vec(),
298		};
299		let mut writer = Vec::with_capacity(45);
300		b.iter(|| {
301			dataframe.write_to(&mut writer, false).unwrap();
302		});
303	}
304}
websocket_base/dataframe.rs

websocket_base/
dataframe.rs
