From 3e7341a98ca95a6a64e5b696034219c2844d940a Mon Sep 17 00:00:00 2001
From: Zhaofeng Li <hello@zhaofeng.li>
Date: Mon, 25 Apr 2022 00:03:57 -0700
Subject: [PATCH] utils.rs: Tolerate invalid UTF-8 in streams

This will substitute any invalid UTF-8 sequence with `\u{fffd}`. Such
substitutions are okay for human-readable log outputs.

Correctness: A side-effect of this is that it also affects the code
that captures filesystem paths from stdout, in which case any substitution
is unacceptable. Currently we only capture store paths from Nix, and Nix
restricts characters allowed in derivation names. Nevertheless we will
refactor this for correctness.

Fixes #75.
---
 src/util.rs | 37 +++++++++++++++++++++++++++++++++++--
 1 file changed, 35 insertions(+), 2 deletions(-)

diff --git a/src/util.rs b/src/util.rs
index 4b75f7d..04a3f1c 100644
--- a/src/util.rs
+++ b/src/util.rs
@@ -14,6 +14,8 @@ use super::nix::{Flake, Hive, HivePath, StorePath};
 use super::nix::deployment::TargetNodeMap;
 use super::job::JobHandle;
 
+const NEWLINE: u8 = 0xa;
+
 /// Non-interactive execution of an arbitrary command.
 pub struct CommandExecution {
     command: Command,
@@ -288,8 +290,9 @@ pub async fn capture_stream<R>(mut stream: BufReader<R>, job: Option<JobHandle>,
     let mut log = String::new();
 
     loop {
-        let mut line = String::new();
-        let len = stream.read_line(&mut line).await?;
+        let mut line = Vec::new();
+        let len = stream.read_until(NEWLINE, &mut line).await?;
+        let line = String::from_utf8_lossy(&line);
 
         if len == 0 {
             break;
@@ -315,3 +318,33 @@ pub async fn capture_stream<R>(mut stream: BufReader<R>, job: Option<JobHandle>,
 pub fn get_label_width(targets: &TargetNodeMap) -> Option<usize> {
     targets.keys().map(|n| n.len()).max()
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    use tokio::io::BufReader;
+    use tokio_test::block_on;
+
+    #[test]
+    fn test_capture_stream() {
+        let expected = "Hello\nWorld\n";
+
+        let stream = BufReader::new(expected.as_bytes());
+        let captured = block_on(async {
+            capture_stream(stream, None, false).await.unwrap()
+        });
+
+        assert_eq!(expected, captured);
+    }
+
+    #[test]
+    fn test_capture_stream_with_invalid_utf8() {
+        let stream = BufReader::new([0x80, 0xa].as_slice());
+        let captured = block_on(async {
+            capture_stream(stream, None, false).await.unwrap()
+        });
+
+        assert_eq!("\u{fffd}\n", captured);
+    }
+}