websocket_base/ws/
dataframe.rs

1//! Describes the generic DataFrame, defining a trait
2//! that all dataframes should share. This is so one can
3//! optimize the memory footprint of a dataframe for their
4//! own needs, and be able to use custom dataframes quickly
5use crate::result::WebSocketResult;
6use crate::ws::util::header as dfh;
7use crate::ws::util::mask;
8use crate::ws::util::mask::Masker;
9use std::io::Write;
10
11/// A generic DataFrame. Every dataframe should be able to
12/// provide these methods. (If the payload is not known in advance then
13/// rewrite the write_payload method)
14pub trait DataFrame {
15	/// Is this dataframe the final dataframe of the message?
16	fn is_last(&self) -> bool;
17	/// What type of data does this dataframe contain?
18	fn opcode(&self) -> u8;
19	/// Reserved bits of this dataframe
20	fn reserved(&self) -> &[bool; 3];
21
22	/// How long (in bytes) is this dataframe's payload
23	fn size(&self) -> usize;
24
25	/// Get's the size of the entire dataframe in bytes,
26	/// i.e. header and payload.
27	fn frame_size(&self, masked: bool) -> usize {
28		// one byte for the opcode & reserved & fin
29		1
30        // depending on the size of the payload, add the right payload len bytes
31        + match self.size() {
32            s if s <= 125 => 1,
33            s if s <= 65535 => 3,
34            _ => 9,
35        }
36        // add the mask size if there is one
37        + if masked {
38            4
39        } else {
40            0
41        }
42        // finally add the payload len
43        + self.size()
44	}
45
46	/// Write the payload to a writer
47	fn write_payload(&self, socket: &mut dyn Write) -> WebSocketResult<()>;
48
49	/// Takes the payload out into a vec
50	fn take_payload(self) -> Vec<u8>;
51
52	/// Writes a DataFrame to a Writer.
53	fn write_to(&self, writer: &mut dyn Write, mask: bool) -> WebSocketResult<()> {
54		let mut flags = dfh::DataFrameFlags::empty();
55		if self.is_last() {
56			flags.insert(dfh::DataFrameFlags::FIN);
57		}
58		{
59			let reserved = self.reserved();
60			if reserved[0] {
61				flags.insert(dfh::DataFrameFlags::RSV1);
62			}
63			if reserved[1] {
64				flags.insert(dfh::DataFrameFlags::RSV2);
65			}
66			if reserved[2] {
67				flags.insert(dfh::DataFrameFlags::RSV3);
68			}
69		}
70
71		let masking_key = if mask { Some(mask::gen_mask()) } else { None };
72
73		let header = dfh::DataFrameHeader {
74			flags,
75			opcode: self.opcode() as u8,
76			mask: masking_key,
77			len: self.size() as u64,
78		};
79
80		let mut data = Vec::<u8>::new();
81		dfh::write_header(&mut data, header)?;
82
83		match masking_key {
84			Some(mask) => {
85				let mut masker = Masker::new(mask, &mut data);
86				self.write_payload(&mut masker)?
87			}
88			None => self.write_payload(&mut data)?,
89		};
90		writer.write_all(data.as_slice())?;
91		Ok(())
92	}
93}