Merge pull request #128 from cyderize/fix/dataframe-payload

Dataframe reads all of payload. Fixes #119

Author: Michael Eden

src/client/builder.rs

@@ -55,9 +55,6 @@ mod async_imports {
 #[cfg(feature="async")]
 use self::async_imports::*;
 
-// TODO: add extra funcs for future stuff, like auto ping and auto close
-
-
 /// Build clients with a builder-style API
 /// This makes it easy to create and configure a websocket
 /// connection:
@@ -662,7 +659,6 @@ impl<'u> ClientBuilder<'u> {
 		Box::new(future)
 	}
 
-	// TODO: add timeout option for connecting
 	// TODO: add conveniences like .response_to_pings, .send_close, etc.
 	/// Asynchronously create an insecure (plain TCP) connection to the client.
 	///
@@ -974,6 +970,4 @@ mod tests {
 		assert!(protos.contains(&"electric".to_string()));
 		assert!(!protos.contains(&"rust-websocket".to_string()));
 	}
-
-	// TODO: a few more
 }

src/codec/ws.rs

@@ -24,9 +24,6 @@ use ws::message::Message as MessageTrait;
 use ws::util::header::read_header;
 use result::WebSocketError;
 
-// TODO: IMPORTANT: check if frame_size is correct,
-// do not call .reserve with the entire size
-
 /// Even though a websocket connection may look perfectly symmetrical
 /// in reality there are small differences between clients and servers.
 /// This type is passed to the codecs to inform them of what role they are in
@@ -290,8 +287,6 @@ impl<M> Encoder for MessageCodec<M>
 	}
 }
 
-// TODO: add tests to check boundary cases for reading dataframes
-
 #[cfg(test)]
 mod tests {
 	use super::*;

src/dataframe.rs

@@ -1,5 +1,5 @@
 //! Module containing the default implementation of data frames.
-use std::io::{Read, Write};
+use std::io::{self, Read, Write};
 use result::{WebSocketResult, WebSocketError};
 use ws::dataframe::DataFrame as DataFrameable;
 use ws::util::header::DataFrameHeader;
@@ -61,7 +61,6 @@ impl DataFrame {
 				if !should_be_masked {
 					return Err(WebSocketError::DataFrameError("Expected unmasked data frame"));
 				}
-				// TODO: move data to avoid copy?
 				mask::mask_data(mask, &body)
 			}
 			None => {
@@ -87,7 +86,10 @@ impl DataFrame {
 		let header = try!(dfh::read_header(reader));
 
 		let mut data: Vec<u8> = Vec::with_capacity(header.len as usize);
-		reader.take(header.len).read_to_end(&mut data)?;
+		let read = reader.take(header.len).read_to_end(&mut data)?;
+		if (read as u64) < header.len {
+			return Err(io::Error::new(io::ErrorKind::UnexpectedEof, "incomplete payload").into());
+		}
 
 		DataFrame::read_dataframe_body(header, data, should_be_masked)
 	}
@@ -212,6 +214,25 @@ mod tests {
 		};
 		assert_eq!(obtained, expected);
 	}
+
+	#[test]
+	fn read_incomplete_payloads() {
+		let mut data = vec![0x8au8, 0x08, 0x19, 0xac, 0xab, 0x8a, 0x52, 0x4e, 0x05, 0x00];
+		let payload = vec![25, 172, 171, 138, 82, 78, 5, 0];
+		let short_header = DataFrame::read_dataframe(&mut &data[..1], false);
+		let short_payload = DataFrame::read_dataframe(&mut &data[..6], false);
+		let full_payload = DataFrame::read_dataframe(&mut &data[..], false);
+		data.push(0xff);
+		let more_payload = DataFrame::read_dataframe(&mut &data[..], false);
+
+		match (short_header.unwrap_err(), short_payload.unwrap_err()) {
+			(WebSocketError::NoDataAvailable, WebSocketError::NoDataAvailable) => (),
+			_ => assert!(false),
+		};
+		assert_eq!(full_payload.unwrap().data, payload);
+		assert_eq!(more_payload.unwrap().data, payload);
+	}
+
 	#[bench]
 	fn bench_read_dataframe(b: &mut Bencher) {
 		let data = b"The quick brown fox jumps over the lazy dog";

src/receiver.rs

@@ -132,10 +132,6 @@ impl ws::Receiver for Receiver {
 			}
 		}
 
-		// TODO: this is slow, easy fix
-		let buffer = self.buffer.clone();
-		self.buffer.clear();
-
-		Ok(buffer)
+		Ok(::std::mem::replace(&mut self.buffer, Vec::new()))
 	}
 }