diff --git a/src/input.rs b/src/input.rs index facd64bf..d79c8198 100644 --- a/src/input.rs +++ b/src/input.rs @@ -9,6 +9,8 @@ use content_inspector::{self, ContentType}; use crate::error::*; +const CONTENT_INSPECTION_LIMIT: usize = 1024; + /// A description of an Input source. /// This tells bat how to refer to the input. #[derive(Clone)] @@ -263,10 +265,29 @@ impl<'a> InputReader<'a> { } pub(crate) fn try_new(mut reader: R) -> io::Result> { - let mut first_line = vec![]; - reader.read_until(b'\n', &mut first_line)?; + // content_inspector scans at most 1024 bytes. Capture the already-buffered + // prefix before splitting out the first line so an early newline in binary + // data does not shorten the inspected content. This does not consume input + // or perform an additional read beyond the one read_until needs anyway. + let mut inspection_prefix = { + let buffered = reader.fill_buf()?; + buffered[..buffered.len().min(CONTENT_INSPECTION_LIMIT)].to_vec() + }; - let content_type = inspect_content_type(&first_line); + let mut first_line = vec![]; + if !inspection_prefix.is_empty() { + reader.read_until(b'\n', &mut first_line)?; + } + + // A custom BufRead implementation may expose less than 1024 bytes at a + // time. Keep the old behavior for long first lines in that case. + let first_line_prefix_len = first_line.len().min(CONTENT_INSPECTION_LIMIT); + if first_line_prefix_len > inspection_prefix.len() { + inspection_prefix.clear(); + inspection_prefix.extend_from_slice(&first_line[..first_line_prefix_len]); + } + + let content_type = inspect_content_type(&inspection_prefix); if content_type == Some(ContentType::UTF_16LE) { read_utf16_line(&mut reader, &mut first_line, 0x00, 0x0A)?; @@ -410,6 +431,53 @@ fn non_zip_pk_prefix_is_not_treated_as_binary() { ); } +#[test] +fn binary_detection_scans_beyond_first_line_and_preserves_input() { + let mut content = vec![b'a'; CONTENT_INSPECTION_LIMIT + 1]; + content[1] = b'\n'; + content[CONTENT_INSPECTION_LIMIT - 1] = 0; + + let mut reader = InputReader::new(&content[..]); + assert_eq!(Some(ContentType::BINARY), reader.content_type); + + let mut replayed = Vec::new(); + while reader.read_line(&mut replayed).unwrap() {} + assert_eq!(content, replayed); + + drop(reader); + + content[CONTENT_INSPECTION_LIMIT - 1] = b'a'; + content[CONTENT_INSPECTION_LIMIT] = 0; + + let reader = InputReader::new(&content[..]); + assert_eq!(Some(ContentType::UTF_8), reader.content_type); +} + +#[test] +fn input_detection_does_not_read_twice() { + struct OneRead(Option<&'static [u8]>); + + impl Read for OneRead { + fn read(&mut self, buf: &mut [u8]) -> io::Result { + match self.0.take() { + None => Err(io::Error::new( + io::ErrorKind::WouldBlock, + "no more data is available yet", + )), + Some(content) => { + buf[..content.len()].copy_from_slice(content); + Ok(content.len()) + } + } + } + } + + for (content, expected) in [(&b"text\n"[..], Some(ContentType::UTF_8)), (&b""[..], None)] { + let input = InputReader::try_new(BufReader::new(OneRead(Some(content)))).unwrap(); + assert_eq!(expected, input.content_type); + } +} + #[test] fn input_open_returns_initial_read_errors() { struct FailingRead; diff --git a/tests/integration_tests.rs b/tests/integration_tests.rs index 124a9e5b..a600e5b3 100644 --- a/tests/integration_tests.rs +++ b/tests/integration_tests.rs @@ -2164,6 +2164,26 @@ fn header_binary() { .stderr(""); } +// Regression test for https://github.com/sharkdp/bat/issues/3554 +#[test] +fn header_binary_with_null_after_first_line() { + let tmp_dir = tempdir().expect("can create temporary directory"); + let tmp_path = tmp_dir.path().join("encrypted.gpg"); + std::fs::write(&tmp_path, b"packet-header\npayload\0bytes\n") + .expect("can write temporary file"); + + bat() + .arg(&tmp_path) + .arg("--decorations=always") + .arg("--style=header") + .arg("--line-range=0:0") + .arg("--file-name=encrypted.gpg") + .assert() + .success() + .stdout("File: encrypted.gpg \n") + .stderr(""); +} + #[test] fn header_zip_file_is_binary() { let tmp_dir = tempdir().expect("can create temporary directory");