age: impl BufRead for ArmoredReader

age/CHANGELOG.md

@@ -9,6 +9,8 @@ and this project adheres to Rust's notion of
 to 1.0.0 are beta releases.
 
 ## [Unreleased]
+### Added
+- `impl std::io::BufRead for age::armor::ArmoredReader`
 
 ## [0.9.1] - 2022-03-24
 ### Added

age/src/primitives/armor.rs

@@ -726,6 +726,7 @@ impl<R> ArmoredReader<R> {
     ///
     /// Returns the number of bytes read into the buffer, or None if there was no cached
     /// data.
+    #[cfg(feature = "async")]
     fn read_cached_data(&mut self, buf: &mut [u8]) -> Option<usize> {
         if self.byte_start >= self.byte_end {
             None
@@ -820,8 +821,8 @@ impl<R> ArmoredReader<R> {
     }
 }
 
-impl<R: BufRead> Read for ArmoredReader<R> {
+impl<R: BufRead> BufRead for ArmoredReader<R> {
-    fn read(&mut self, mut buf: &mut [u8]) -> io::Result<usize> {
+    fn fill_buf(&mut self) -> io::Result<&[u8]> {
         loop {
             match self.is_armored {
                 None => {
@@ -829,60 +830,97 @@ impl<R: BufRead> Read for ArmoredReader<R> {
                     self.detect_armor()?
                 }
                 Some(false) => {
-                    // Return any leftover data from armor detection
+                    break if self.byte_start >= self.byte_end {
-                    return if let Some(read) = self.read_cached_data(buf) {
+                        self.inner.read(&mut self.byte_buf[..]).map(|read| {
-                        Ok(read)
+                            self.byte_start = 0;
-                    } else {
+                            self.byte_end = read;
-                        self.inner.read(buf).map(|read| {
+                            self.count_reader_bytes(read);
-                            self.data_read += read;
+                            &self.byte_buf[..read]
-                            self.count_reader_bytes(read)
                         })
-                    };
+                    } else {
+                        Ok(&self.byte_buf[self.byte_start..self.byte_end])
+                    }
+                }
+                Some(true) => {
+                    break if self.found_end {
+                        Ok(&[])
+                    } else if self.byte_start >= self.byte_end {
+                        if self.read_next_armor_line()? {
+                            Ok(&[])
+                        } else {
+                            Ok(&self.byte_buf[self.byte_start..self.byte_end])
+                        }
+                    } else {
+                        Ok(&self.byte_buf[self.byte_start..self.byte_end])
+                    }
                 }
-                Some(true) => break,
             }
         }
-        if self.found_end {
+    }
-            return Ok(0);
-        }
 
+    fn consume(&mut self, amt: usize) {
+        self.byte_start += amt;
+        self.data_read += amt;
+        assert!(self.byte_start <= self.byte_end);
+    }
+}
+
+impl<R: BufRead> ArmoredReader<R> {
+    /// Fills `self.byte_buf` with the next line of armored data.
+    ///
+    /// Returns `true` if this was the last line.
+    fn read_next_armor_line(&mut self) -> io::Result<bool> {
+        assert_eq!(self.is_armored, Some(true));
+
+        // Read the next line
+        self.inner
+            .read_line(&mut self.line_buf)
+            .map(|read| self.count_reader_bytes(read))?;
+
+        // Parse the line into bytes
+        if self.parse_armor_line()? {
+            // This was the last line! Check for trailing garbage.
+            loop {
+                let amt = match self.inner.fill_buf()? {
+                    &[] => break,
+                    buf => {
+                        if buf.iter().any(|b| !b.is_ascii_whitespace()) {
+                            return Err(io::Error::new(
+                                io::ErrorKind::InvalidData,
+                                ArmoredReadError::TrailingGarbage,
+                            ));
+                        }
+                        buf.len()
+                    }
+                };
+                self.inner.consume(amt);
+            }
+            Ok(true)
+        } else {
+            Ok(false)
+        }
+    }
+}
+
+impl<R: BufRead> Read for ArmoredReader<R> {
+    fn read(&mut self, mut buf: &mut [u8]) -> io::Result<usize> {
         let buf_len = buf.len();
 
-        // Output any remaining bytes from the previous line
-        if let Some(read) = self.read_cached_data(buf) {
-            buf = &mut buf[read..];
-        }
-
         while !buf.is_empty() {
-            // Read the next line
+            match self.fill_buf()? {
-            self.inner
+                [] => break,
-                .read_line(&mut self.line_buf)
+                next => {
-                .map(|read| self.count_reader_bytes(read))?;
+                    let read = cmp::min(next.len(), buf.len());
 
-            // Parse the line into bytes
+                    if next.len() < buf.len() {
-            if self.parse_armor_line()? {
+                        buf[..read].copy_from_slice(next);
-                // This was the last line! Check for trailing garbage.
+                    } else {
-                loop {
+                        buf.copy_from_slice(&next[..read]);
-                    let amt = match self.inner.fill_buf()? {
+                    }
-                        &[] => break,
+
-                        buf => {
+                    self.consume(read);
-                            if buf.iter().any(|b| !b.is_ascii_whitespace()) {
+                    buf = &mut buf[read..];
-                                return Err(io::Error::new(
-                                    io::ErrorKind::InvalidData,
-                                    ArmoredReadError::TrailingGarbage,
-                                ));
-                            }
-                            buf.len()
-                        }
-                    };
-                    self.inner.consume(amt);
                 }
-                break;
-            }
-
-            // Output as much as we can of this line
-            if let Some(read) = self.read_cached_data(buf) {
-                buf = &mut buf[read..];
             }
         }
 
@@ -1061,49 +1099,9 @@ impl<R: BufRead + Seek> Seek for ArmoredReader<R> {
                     self.inner.read_exact(&mut self.byte_buf[..MIN_ARMOR_LEN])?;
                     self.detect_armor()?
                 }
-                Some(false) => {
+                Some(armored) => {
-                    break if self.byte_start >= self.byte_end {
-                        // Map the data read onto the underlying stream.
-                        let start = self.start()?;
-                        let pos = match pos {
-                            SeekFrom::Start(offset) => SeekFrom::Start(start + offset),
-                            // Current and End positions don't need to be shifted.
-                            x => x,
-                        };
-                        self.inner.seek(pos)
-                    } else {
-                        // We are still inside the first line.
-                        match pos {
-                            SeekFrom::Start(offset) => self.byte_start = offset as usize,
-                            SeekFrom::Current(offset) => {
-                                let res = (self.byte_start as i64) + offset;
-                                if res >= 0 {
-                                    self.byte_start = res as usize;
-                                } else {
-                                    return Err(io::Error::new(
-                                        io::ErrorKind::InvalidData,
-                                        "cannot seek before the start",
-                                    ));
-                                }
-                            }
-                            SeekFrom::End(offset) => {
-                                let res = (self.line_buf.len() as i64) + offset;
-                                if res >= 0 {
-                                    self.byte_start = res as usize;
-                                } else {
-                                    return Err(io::Error::new(
-                                        io::ErrorKind::InvalidData,
-                                        "cannot seek before the start",
-                                    ));
-                                }
-                            }
-                        }
-                        Ok(self.byte_start as u64)
-                    };
-                }
-                Some(true) => {
                     // Convert the offset into the target position within the data inside
-                    // the armor.
+                    // the (maybe) armor.
                     let start = self.start()?;
                     let target_pos = match pos {
                         SeekFrom::Start(offset) => offset,
@@ -1146,6 +1144,15 @@ impl<R: BufRead + Seek> Seek for ArmoredReader<R> {
                         }
                     };
 
+                    if !armored {
+                        // We can seek directly on the inner reader.
+                        self.inner.seek(SeekFrom::Start(start + target_pos))?;
+                        self.byte_start = 0;
+                        self.byte_end = 0;
+                        self.data_read = target_pos as usize;
+                        break Ok(self.data_read as u64);
+                    }
+
                     // Jump back to the start of the armor data, and then read and drop
                     // until we reach the target position. This is very inefficient, but
                     // as armored files can have arbitrary line endings within the file,