WIP fix the seek tests

2 files changed, 103 insertions, 2 deletions

diff --git a/src/lib.rs b/src/lib.rs
index cfc50e1..6b12f17 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -1128,7 +1128,31 @@ impl OutputReader {
     /// reading further, the behavior is unspecified.
     ///
     /// [`Read::read`]: #method.read
-    pub fn fill(&mut self, buf: &mut [u8]) {
+    pub fn fill(&mut self, mut buf: &mut [u8]) {
+        debug_assert!(self.position_within_block < BLOCK_LEN as u8);
+        if self.position_within_block != 0 {
+            let mut partial_block = [0u8; 64];
+            guts::xof(
+                &self.inner.block,
+                self.inner.block_len as u32,
+                &self.inner.input_chaining_value,
+                self.inner.counter,
+                self.inner.flags as u32,
+                &mut partial_block,
+            );
+            let output_bytes = &partial_block[self.position_within_block as usize..];
+            let take = cmp::min(buf.len(), output_bytes.len());
+            buf[..take].copy_from_slice(&output_bytes[..take]);
+            buf = &mut buf[take..];
+            self.position_within_block += take as u8;
+            if self.position_within_block == BLOCK_LEN as u8 {
+                self.position_within_block = 0;
+                self.inner.counter += 1;
+            } else {
+                debug_assert!(buf.is_empty());
+                return;
+            }
+        }
         guts::xof(
             &self.inner.block,
             self.inner.block_len as u32,
@@ -1137,9 +1161,36 @@ impl OutputReader {
             self.inner.flags as u32,
             buf,
         );
+        self.position_within_block = (buf.len() % BLOCK_LEN) as u8;
     }
 
-    pub fn fill_xor(&mut self, buf: &mut [u8]) {
+    pub fn fill_xor(&mut self, mut buf: &mut [u8]) {
+        debug_assert!(self.position_within_block < BLOCK_LEN as u8);
+        if self.position_within_block != 0 {
+            let mut partial_block = [0u8; 64];
+            guts::xof(
+                &self.inner.block,
+                self.inner.block_len as u32,
+                &self.inner.input_chaining_value,
+                self.inner.counter,
+                self.inner.flags as u32,
+                &mut partial_block,
+            );
+            let output_bytes = &partial_block[self.position_within_block as usize..];
+            let take = cmp::min(buf.len(), output_bytes.len());
+            for byte_index in 0..take {
+                buf[byte_index] ^= output_bytes[byte_index];
+            }
+            buf = &mut buf[take..];
+            self.position_within_block += take as u8;
+            if self.position_within_block == BLOCK_LEN as u8 {
+                self.position_within_block = 0;
+                self.inner.counter += 1;
+            } else {
+                debug_assert!(buf.is_empty());
+                return;
+            }
+        }
         guts::xof_xor(
             &self.inner.block,
             self.inner.block_len as u32,
@@ -1148,6 +1199,7 @@ impl OutputReader {
             self.inner.flags as u32,
             buf,
         );
+        self.position_within_block = (buf.len() % BLOCK_LEN) as u8;
     }
 
     /// Return the current read position in the output stream. This is
diff --git a/src/test.rs b/src/test.rs
index 0e2f030..ba3cf46 100644
--- a/src/test.rs
+++ b/src/test.rs
@@ -291,6 +291,55 @@ fn test_xof_seek() {
 }
 
 #[test]
+fn test_xof_xor() {
+    const STEP: usize = 17;
+
+    let mut ref_hasher = reference_impl::Hasher::new();
+    ref_hasher.update(b"foo");
+    let mut ref_output = [0u8; 1000];
+    ref_hasher.finalize(&mut ref_output);
+
+    let mut hasher = crate::Hasher::new();
+    hasher.update(b"foo");
+    let mut reader = hasher.finalize_xof();
+
+    let mut test_output = [0u8; 1000];
+    for chunk in test_output.chunks_mut(STEP) {
+        reader.fill(chunk);
+    }
+    assert_eq!(ref_output, test_output);
+    // Xor'ing the same output should zero the buffer.
+    reader.set_position(0);
+    for chunk in test_output.chunks_mut(STEP) {
+        reader.fill_xor(chunk);
+    }
+    assert_eq!([0u8; 1000], test_output);
+    // Xor'ing the same output again should reproduce the original.
+    reader.set_position(0);
+    for chunk in test_output.chunks_mut(STEP) {
+        reader.fill_xor(chunk);
+    }
+    assert_eq!(ref_output, test_output);
+
+    // Repeat the same test but starting at offset 500.
+    reader.set_position(500);
+    for chunk in test_output[..500].chunks_mut(STEP) {
+        reader.fill(chunk);
+    }
+    assert_eq!(ref_output[500..], test_output[..500]);
+    reader.set_position(500);
+    for chunk in test_output[..500].chunks_mut(STEP) {
+        reader.fill_xor(chunk);
+    }
+    assert_eq!([0u8; 500], test_output[..500]);
+    reader.set_position(500);
+    for chunk in test_output[..500].chunks_mut(STEP) {
+        reader.fill_xor(chunk);
+    }
+    assert_eq!(ref_output[500..], test_output[..500]);
+}
+
+#[test]
 fn test_msg_schedule_permutation() {
     let permutation = [2, 6, 3, 10, 7, 0, 4, 13, 1, 11, 12, 5, 9, 14, 15, 8];