add unit tests for auto-mask-write

zephyrchien · Apr 14, 2022 · 47ea8b2 · 47ea8b2
1 parent 387286b
commit 47ea8b2
Show file tree

Hide file tree

Showing 2 changed files with 162 additions and 13 deletions.
diff --git a/src/stream/mod.rs b/src/stream/mod.rs
@@ -219,22 +219,32 @@ mod test {
         fn flush(&mut self) -> Result<()> { Ok(()) }
     }
 
-    pub fn make_frame<R: RoleHelper>(opcode: OpCode, len: usize) -> (Vec<u8>, Vec<u8>) {
-        let data: Vec<u8> = std::iter::repeat(rand::random::<u8>()).take(len).collect();
-        let mut data2 = data.clone();
-
+    pub fn make_head(opcode: OpCode, mask: Mask, len: usize) -> Vec<u8> {
         let mut tmp = vec![0; 14];
-        let head = FrameHead::new(
-            Fin::Y,
-            opcode,
-            R::write_mask_key(&R::new()),
-            PayloadLen::from_num(len as u64),
-        );
+        let head = FrameHead::new(Fin::Y, opcode, mask, PayloadLen::from_num(len as u64));
 
         let head_len = head.encode(&mut tmp).unwrap();
-        let mut frame = Vec::new();
-        let write_n = frame.write(&tmp[..head_len]).unwrap();
+        let mut head = Vec::new();
+        let write_n = head.write(&tmp[..head_len]).unwrap();
         assert_eq!(write_n, head_len);
+        head
+    }
+
+    pub fn make_data(len: usize) -> Vec<u8> {
+        std::iter::repeat(rand::random::<u8>()).take(len).collect()
+    }
+
+    pub fn make_frame<R: RoleHelper>(opcode: OpCode, len: usize) -> (Vec<u8>, Vec<u8>) {
+        make_frame_with_mask(opcode, R::new().write_mask_key(), len)
+    }
+
+    // data is unmasked
+    pub fn make_frame_with_mask(opcode: OpCode, mask: Mask, len: usize) -> (Vec<u8>, Vec<u8>) {
+        let data = make_data(len);
+        let mut data2 = data.clone();
+
+        let mut frame = make_head(opcode, mask, len);
+        let head_len = frame.len();
 
         frame.append(&mut data2);
         assert_eq!(frame.len(), len + head_len);

diff --git a/src/stream/write.rs b/src/stream/write.rs
@@ -56,7 +56,7 @@ impl<IO: Write, Role: RoleHelper> Write for Stream<IO, Role, Guarded> {
 #[cfg(test)]
 mod test {
     use super::*;
-    use super::super::test::{LimitReadWriter, make_frame};
+    use super::super::test::*;
     use crate::frame::*;
     use crate::role::*;
     use std::io::Write;
@@ -120,4 +120,143 @@ mod test {
             }
         }
     }
+
+    #[test]
+    #[cfg(feature = "unsafe_auto_mask_write")]
+    fn write_to_stream_auto_mask_fixed() {
+        fn write<R: RoleHelper>(n: usize) {
+            let key = new_mask_key();
+
+            let (mut frame, data) = make_frame_with_mask(OpCode::Binary, Mask::Key(key), n);
+
+            // manually mask frame data
+            let offset = frame.len() - n;
+            apply_mask4(key, &mut frame[offset..]);
+
+            let io: Vec<u8> = Vec::new();
+            let mut stream = Stream::new(io, R::new());
+            stream.set_write_mask_key(key).unwrap();
+
+            let write_n = stream.write(&data).unwrap();
+
+            assert_eq!(write_n, n);
+
+            assert_eq!(stream.as_ref(), &frame);
+        }
+        for i in 1..=2 {
+            write::<FixedMaskClient>(i);
+        }
+
+        for i in [65536, 65537, 100000] {
+            write::<FixedMaskClient>(i);
+        }
+    }
+
+    #[test]
+    #[cfg(feature = "unsafe_auto_mask_write")]
+    fn write_to_limit_stream_auto_mask_fixed() {
+        fn write<R: RoleHelper>(n: usize, limit: usize) {
+            let key = new_mask_key();
+            let (mut frame, data) = make_frame_with_mask(OpCode::Binary, Mask::Key(key), n);
+
+            // manually mask frame data
+            let offset = frame.len() - n;
+            apply_mask4(key, &mut frame[offset..]);
+
+            let io = LimitReadWriter {
+                buf: Vec::new(),
+                rlimit: 0,
+                wlimit: limit,
+                cursor: 0,
+            };
+
+            let mut stream = Stream::new(io, R::new()).guard();
+            stream.set_write_mask_key(key).unwrap();
+
+            stream.write_all(&data).unwrap();
+
+            assert_eq!(&stream.as_ref().buf, &frame);
+        }
+
+        for i in 1..=256 {
+            for limit in 1..=300 {
+                write::<FixedMaskClient>(i, limit);
+            }
+        }
+
+        for i in [65536, 65537, 100000] {
+            for limit in 1..=1024 {
+                write::<FixedMaskClient>(i, limit);
+            }
+        }
+    }
+
+    #[test]
+    #[cfg(feature = "unsafe_auto_mask_write")]
+    fn write_to_stream_auto_mask_updated() {
+        fn write<R: RoleHelper>(n: usize) {
+            let data = make_data(n);
+            let mut data2 = data.clone();
+
+            let io: Vec<u8> = Vec::new();
+            let mut stream = Stream::new(io, R::new());
+
+            let write_n = stream.write(&data).unwrap();
+            assert_eq!(write_n, n);
+
+            // manually mask frame data
+            let key = stream.write_mask_key().to_key();
+            let head = make_head(OpCode::Binary, Mask::Key(key), n);
+            apply_mask4(key, &mut data2);
+
+            assert_eq!(stream.as_ref()[..head.len()], head);
+            assert_eq!(stream.as_ref()[head.len()..], data2);
+        }
+        for i in 1..=2 {
+            write::<StandardClient>(i);
+        }
+
+        for i in [65536, 65537, 100000] {
+            write::<StandardClient>(i);
+        }
+    }
+
+    #[test]
+    #[cfg(feature = "unsafe_auto_mask_write")]
+    fn write_to_limit_stream_auto_mask_updated() {
+        fn write<R: RoleHelper>(n: usize, limit: usize) {
+            let data = make_data(n);
+            let mut data2 = data.clone();
+
+            let io = LimitReadWriter {
+                buf: Vec::new(),
+                rlimit: 0,
+                wlimit: limit,
+                cursor: 0,
+            };
+
+            let mut stream = Stream::new(io, R::new()).guard();
+            stream.write_all(&data).unwrap();
+
+            // manually mask frame data
+            let key = stream.write_mask_key().to_key();
+            let head = make_head(OpCode::Binary, Mask::Key(key), n);
+            apply_mask4(key, &mut data2);
+
+            assert_eq!(stream.as_ref().buf[..head.len()], head);
+            assert_eq!(stream.as_ref().buf[head.len()..], data2);
+        }
+
+        for i in 1..=256 {
+            for limit in 1..=300 {
+                write::<StandardClient>(i, limit);
+            }
+        }
+
+        for i in [65536, 65537, 100000] {
+            for limit in 1..=1024 {
+                write::<StandardClient>(i, limit);
+            }
+        }
+    }
 }