fix: zero padding reader incomplete reads

lib/polka-storage-proofs/src/zero_reader.rs

@@ -28,11 +28,16 @@ impl<R: Read> Read for ZeroPaddingReader<R> {
         // Number of bytes that we read from the inner reader
         let read = self.inner.read(&mut buf[..to_read])?;
 
+        // Incomplete read doesn't mean that we need to pad it yet.
+        // Next call can be complete.
+        if read > 0 {
+            self.remaining -= read as u64;
+            return Ok(read);
+        }
+
         // If we read from the inner reader less then the required bytes, 0-pad
         // the rest of the buffer.
-        if read < to_read {
-            buf[read..to_read].fill(0);
-        }
+        buf[..to_read].fill(0);
 
         // Decrease the number of bytes this 0-padding reader has left to produce.
         self.remaining -= to_read as u64;
@@ -49,6 +54,68 @@ mod tests {
 
     use super::ZeroPaddingReader;
 
+    /// Sole purpose of this reader is to simulate the file reading in the OS.
+    /// When we give it a buf of a certain buf.length(), it might read n < buf.length().
+    struct IncompleteReader {
+        read_count: usize,
+        served_data: Vec<Vec<u8>>,
+    }
+
+    impl Read for IncompleteReader {
+        fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
+            if self.read_count == self.served_data.len() {
+                return Ok(0);
+            }
+
+            let ready = &self.served_data[self.read_count];
+            // this is a simplification for test purposes
+            // assuming buf.len() >= ready.len() always holds
+            buf[..ready.len()].copy_from_slice(&ready);
+
+            self.read_count += 1;
+            Ok(ready.len())
+        }
+    }
+
+    #[test]
+    fn test_zero_padding_reader_with_not_full_reads() {
+        let r = IncompleteReader {
+            read_count: 0,
+            served_data: vec![vec![1, 2, 3, 4], vec![5, 6], vec![7, 8, 9], vec![10]],
+        };
+        let total_size_with_padding = 12;
+        let mut reader = ZeroPaddingReader::new(r, total_size_with_padding);
+        let mut buffer = [0; 4];
+
+        let mut total_read = 0;
+        let read = reader.read(&mut buffer).unwrap();
+        total_read += read;
+        assert_eq!(read, 4);
+        assert_eq!(buffer, [1, 2, 3, 4]);
+
+        let read = reader.read(&mut buffer).unwrap();
+        total_read += read;
+        assert_eq!(read, 2);
+        assert_eq!(buffer, [5, 6, 3, 4]);
+
+        let read = reader.read(&mut buffer).unwrap();
+        total_read += read;
+        assert_eq!(read, 3);
+        assert_eq!(buffer, [7, 8, 9, 4]);
+
+        let read = reader.read(&mut buffer).unwrap();
+        total_read += read;
+        assert_eq!(read, 1);
+        assert_eq!(buffer, [10, 8, 9, 4]);
+
+        let read = reader.read(&mut buffer).unwrap();
+        total_read += read;
+        assert_eq!(read, 2);
+        assert_eq!(buffer, [0, 0, 9, 4]);
+
+        assert_eq!(total_size_with_padding as usize, total_read);
+    }
+
     #[test]
     fn test_zero_padding_reader() {
         let data = vec![1, 2, 3, 4, 5, 6];
@@ -62,11 +129,11 @@ mod tests {
         assert_eq!(buffer, [1, 2, 3, 4]);
         // Second read
         let read = reader.read(&mut buffer).unwrap();
-        assert_eq!(read, 4);
-        assert_eq!(buffer, [5, 6, 0, 0]);
+        assert_eq!(read, 2);
+        assert_eq!(buffer, [5, 6, 3, 4]);
         // Third read
         let read = reader.read(&mut buffer).unwrap();
-        assert_eq!(read, 2);
+        assert_eq!(read, 4);
         assert_eq!(buffer, [0, 0, 0, 0]);
         // Fourth read
         let read = reader.read(&mut buffer).unwrap();