Fix append_paths when input does not end in NUL

1 files changed, 60 insertions, 5 deletions

diff --git a/src/path_list_editor.rs b/src/path_list_editor.rs
index e25b0db..9addb7f 100644
--- a/src/path_list_editor.rs
+++ b/src/path_list_editor.rs
@@ -5,19 +5,44 @@ use std::{
     os::unix::ffi::OsStrExt,
 };
 
+const NUL: &[u8] = b"\0";
+
 pub fn append_paths<R: Read, W: Write>(
-    mut input: R,
+    input: R,
     mut output: W,
     new_paths: &[OsString],
 ) -> std::io::Result<()> {
-    std::io::copy(&mut input, &mut output)?;
+    let last_byte_in_input = copy_while_tracking_last_byte(input, &mut output)?;
+    // Add NUL if there is input and it does not end in NUL.
+    if matches!(last_byte_in_input, Some(byte) if byte != 0) {
+        output.write_all(NUL)?;
+    }
     for path in new_paths {
         output.write_all(path.as_bytes())?;
-        output.write_all(b"\0")?;
+        output.write_all(NUL)?;
     }
     Ok(())
 }
 
+fn copy_while_tracking_last_byte<R: Read, W: Write>(
+    mut input: R,
+    mut output: W,
+) -> std::io::Result<Option<u8>> {
+    let mut buf = [0; 8192];
+    let mut last = None;
+
+    loop {
+        let read_bytes = input.read(&mut buf)?;
+        if read_bytes == 0 {
+            break;
+        }
+
+        last = Some(buf[read_bytes - 1]);
+        output.write_all(&buf[..read_bytes])?;
+    }
+    Ok(last)
+}
+
 pub fn prepend_paths<R: Read, W: Write>(
     mut input: R,
     mut output: W,
@@ -25,7 +50,7 @@ pub fn prepend_paths<R: Read, W: Write>(
 ) -> std::io::Result<()> {
     for path in new_paths {
         output.write_all(path.as_bytes())?;
-        output.write_all(b"\0")?;
+        output.write_all(NUL)?;
     }
     std::io::copy(&mut input, &mut output)?;
     Ok(())
@@ -46,7 +71,7 @@ pub fn remove_paths<R: Read, W: Write>(
         let os_string_segment = OsStr::from_bytes(segment).to_os_string();
         if !unwanted_paths.contains(&os_string_segment) {
             output.write_all(segment)?;
-            output.write_all(b"\0")?;
+            output.write_all(NUL)?;
         }
     }
     Ok(())
@@ -109,6 +134,36 @@ mod tests {
     }
 
     #[test]
+    fn append_paths_inserts_separator_when_stdin_is_not_null_terminated() {
+        let input = b"a.rs";
+        let mut output = Vec::new();
+
+        append_paths(&input[..], &mut output, &[OsString::from("b.rs")]).unwrap();
+
+        assert_eq!(output, b"a.rs\0b.rs\0");
+    }
+
+    #[test]
+    fn append_paths_does_not_insert_extra_separator_when_stdin_is_null_terminated() {
+        let input = b"a.rs\0";
+        let mut output = Vec::new();
+
+        append_paths(&input[..], &mut output, &[OsString::from("b.rs")]).unwrap();
+
+        assert_eq!(output, b"a.rs\0b.rs\0");
+    }
+
+    #[test]
+    fn append_paths_handles_partially_terminated_input_sequence() {
+        let input = b"a.rs\0b.rs";
+        let mut output = Vec::new();
+
+        append_paths(&input[..], &mut output, &[OsString::from("c.rs")]).unwrap();
+
+        assert_eq!(output, b"a.rs\0b.rs\0c.rs\0");
+    }
+
+    #[test]
     fn prepend_paths_works_with_empty_stdin_and_args() {
         let input = b"";
         let mut output = Vec::new();