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
5 use crate::result::WebSocketResult;
6 use crate::ws::util::header as dfh;
7 use crate::ws::util::mask;
8 use crate::ws::util::mask::Masker;
9 use 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)
14 pub trait DataFrame {
15 	/// Is this dataframe the final dataframe of the message?
is_last(&self) -> bool16 	fn is_last(&self) -> bool;
17 	/// What type of data does this dataframe contain?
opcode(&self) -> u818 	fn opcode(&self) -> u8;
19 	/// Reserved bits of this dataframe
reserved(&self) -> &[bool; 3]20 	fn reserved(&self) -> &[bool; 3];
21 
22 	/// How long (in bytes) is this dataframe's payload
size(&self) -> usize23 	fn size(&self) -> usize;
24 
25 	/// Get's the size of the entire dataframe in bytes,
26 	/// i.e. header and payload.
frame_size(&self, masked: bool) -> usize27 	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
write_payload(&self, socket: &mut dyn Write) -> WebSocketResult<()>47 	fn write_payload(&self, socket: &mut dyn Write) -> WebSocketResult<()>;
48 
49 	/// Takes the payload out into a vec
take_payload(self) -> Vec<u8>50 	fn take_payload(self) -> Vec<u8>;
51 
52 	/// Writes a DataFrame to a Writer.
write_to(&self, writer: &mut dyn Write, mask: bool) -> WebSocketResult<()>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 }
94