1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142

//! Utility methods for reading and writing data frames.

use std::io::{Read, Write};

use dataframe::{DataFrame, Opcode};
use result::{WebSocketResult, WebSocketError};

use ws::util::header as dfh;
use ws::util::mask;

/// Writes a DataFrame to a Writer.
pub fn write_dataframe<W>(writer: &mut W, mask: bool, dataframe: DataFrame) -> WebSocketResult<()>
	where W: Write {
	
	let mut flags = dfh::DataFrameFlags::empty();
	if dataframe.finished { flags.insert(dfh::FIN); }
	if dataframe.reserved[0] { flags.insert(dfh::RSV1); }
	if dataframe.reserved[1] { flags.insert(dfh::RSV2); }
	if dataframe.reserved[2] { flags.insert(dfh::RSV3); }
	
	let masking_key = if mask { Some(mask::gen_mask()) } else { None };
	
	let header = dfh::DataFrameHeader {
		flags: flags,
		opcode: dataframe.opcode as u8,
		mask: masking_key,
		len: dataframe.data.len() as u64,
	};
	
	try!(dfh::write_header(writer, header));
	
	match masking_key {
		Some(mask) => try!(writer.write_all(&mask::mask_data(mask, &dataframe.data[..])[..])),
		None => try!(writer.write_all(&dataframe.data[..])),
	}
	try!(writer.flush());
	Ok(())
}

/// Reads a DataFrame from a Reader.
pub fn read_dataframe<R>(reader: &mut R, should_be_masked: bool) -> WebSocketResult<DataFrame> 
	where R: Read {

	let header = try!(dfh::read_header(reader));
	
	Ok(DataFrame {
		finished: header.flags.contains(dfh::FIN),
		reserved: [
			header.flags.contains(dfh::RSV1),
			header.flags.contains(dfh::RSV2),
			header.flags.contains(dfh::RSV3)
		],
		opcode: Opcode::new(header.opcode).expect("Invalid header opcode!"),
		data: match header.mask {
			Some(mask) => {
				if !should_be_masked {
					return Err(WebSocketError::DataFrameError(
						"Expected unmasked data frame".to_string()
					));
				}

				let data: Vec<u8> = try!(reader.take(header.len).bytes().collect());
				mask::mask_data(mask, &data)
			}
			None => {
				if should_be_masked {
					return Err(WebSocketError::DataFrameError(
						"Expected masked data frame".to_string()
					));
				}

				try!(reader.take(header.len).bytes().collect())
			}
		}
	})
}

#[cfg(all(feature = "nightly", test))]
mod tests {
	use super::*;
	use dataframe::{DataFrame, Opcode};
	use test;
	#[test]
	fn test_read_dataframe() {
		let data = b"The quick brown fox jumps over the lazy dog";
		let mut dataframe = vec![0x81, 0x2B];
		for i in data.iter() {
			dataframe.push(*i);
		}
		let obtained = read_dataframe(&mut &dataframe[..], false).unwrap();
		let expected = DataFrame {
			finished: true, 
			reserved: [false; 3], 
			opcode: Opcode::Text, 
			data: data.to_vec()
		};
		assert_eq!(obtained, expected);
	}
	#[test]
	fn test_write_dataframe() {
		let data = b"The quick brown fox jumps over the lazy dog";
		let mut expected = vec![0x81, 0x2B];
		for i in data.iter() {
			expected.push(*i);
		}
		let dataframe = DataFrame {
			finished: true, 
			reserved: [false; 3], 
			opcode: Opcode::Text, 
			data: data.to_vec()
		};
		let mut obtained = Vec::new();
		write_dataframe(&mut obtained, false, dataframe).unwrap();
		
		assert_eq!(&obtained[..], &expected[..]);
	}
	#[bench]
	fn bench_read_dataframe(b: &mut test::Bencher) {
		let data = b"The quick brown fox jumps over the lazy dog";
		let mut dataframe = vec![0x81, 0x2B];
		for i in data.iter() {
			dataframe.push(*i);
		}
		b.iter(|| {
			read_dataframe(&mut &dataframe[..], false).unwrap();
		});
	}
	#[bench]
	fn bench_write_dataframe(b: &mut test::Bencher) {
		let data = b"The quick brown fox jumps over the lazy dog";
		let dataframe = DataFrame {
			finished: true, 
			reserved: [false; 3], 
			opcode: Opcode::Text, 
			data: data.to_vec()
		};
		let mut writer = Vec::with_capacity(45);
		b.iter(|| {
			write_dataframe(&mut writer, false, dataframe.clone()).unwrap();
		});
	}
}