diff --git a/src/wsframe.rs b/src/wsframe.rs index 094c596d7..efe0fc47c 100644 --- a/src/wsframe.rs +++ b/src/wsframe.rs @@ -124,10 +124,10 @@ impl Frame { let mut length = u64::from(second & 0x7F); if length == 126 { - let mut length_bytes = [0u8; 2]; if size < 2 { return Ok(None) } + let mut length_bytes = [0u8; 2]; length_bytes.copy_from_slice(&buf[idx..idx+2]); size -= 2; idx += 2; @@ -138,13 +138,13 @@ impl Frame { wide}); header_length += 2; } else if length == 127 { - let mut length_bytes = [0u8; 8]; if size < 8 { return Ok(None) } + let mut length_bytes = [0u8; 8]; length_bytes.copy_from_slice(&buf[idx..idx+8]); size -= 8; - idx += 2; + idx += 8; unsafe { length = mem::transmute(length_bytes); } length = u64::from_be(length); @@ -329,3 +329,100 @@ impl fmt::Display for Frame { self.payload.iter().map(|byte| format!("{:x}", byte)).collect::()) } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_parse() { + let mut buf = BytesMut::from(&[0b00000001u8, 0b00000001u8][..]); + assert!(Frame::parse(&mut buf).unwrap().is_none()); + buf.extend(b"1"); + let frame = Frame::parse(&mut buf).unwrap().unwrap(); + println!("FRAME: {:?}", frame); + assert!(!frame.finished); + assert_eq!(frame.opcode, OpCode::Text); + assert_eq!(frame.payload, &b"1"[..]); + } + + #[test] + fn test_parse_length0() { + let mut buf = BytesMut::from(&[0b00000001u8, 0b00000000u8][..]); + let frame = Frame::parse(&mut buf).unwrap().unwrap(); + assert!(!frame.finished); + assert_eq!(frame.opcode, OpCode::Text); + assert!(frame.payload.is_empty()); + } + + #[test] + fn test_parse_length2() { + let mut buf = BytesMut::from(&[0b00000001u8, 126u8][..]); + assert!(Frame::parse(&mut buf).unwrap().is_none()); + buf.extend(&[0u8, 4u8][..]); + buf.extend(b"1234"); + + let frame = Frame::parse(&mut buf).unwrap().unwrap(); + assert!(!frame.finished); + assert_eq!(frame.opcode, OpCode::Text); + assert_eq!(frame.payload, &b"1234"[..]); + } + + #[test] + fn test_parse_length4() { + let mut buf = BytesMut::from(&[0b00000001u8, 127u8][..]); + assert!(Frame::parse(&mut buf).unwrap().is_none()); + buf.extend(&[0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 4u8][..]); + buf.extend(b"1234"); + + let frame = Frame::parse(&mut buf).unwrap().unwrap(); + assert!(!frame.finished); + assert_eq!(frame.opcode, OpCode::Text); + assert_eq!(frame.payload, &b"1234"[..]); + } + + #[test] + fn test_parse_frame_mask() { + let mut buf = BytesMut::from(&[0b00000001u8, 0b10000001u8][..]); + buf.extend(b"0001"); + buf.extend(b"1"); + + let frame = Frame::parse(&mut buf).unwrap().unwrap(); + assert!(!frame.finished); + assert_eq!(frame.opcode, OpCode::Text); + assert_eq!(frame.payload, vec![1u8]); + } + + #[test] + fn test_ping_frame() { + let mut frame = Frame::message(Vec::from("data"), OpCode::Ping, true); + let mut buf = Vec::new(); + frame.format(&mut buf).unwrap(); + + let mut v = vec![137u8, 4u8]; + v.extend(b"data"); + assert_eq!(buf, v); + } + + #[test] + fn test_pong_frame() { + let mut frame = Frame::message(Vec::from("data"), OpCode::Pong, true); + let mut buf = Vec::new(); + frame.format(&mut buf).unwrap(); + + let mut v = vec![138u8, 4u8]; + v.extend(b"data"); + assert_eq!(buf, v); + } + + #[test] + fn test_close_frame() { + let mut frame = Frame::close(CloseCode::Normal, "data"); + let mut buf = Vec::new(); + frame.format(&mut buf).unwrap(); + + let mut v = vec![136u8, 6u8, 3u8, 232u8]; + v.extend(b"data"); + assert_eq!(buf, v); + } +}