Capture STDOUT as bytes, and convert to string when necessary

Previously, the recite() function created the 'out' variable, which was
a String, that got passed to the various incant functions, in order to
capture STDOUT in certain situations. In cases where STDOUT was
captured, it was first converted to a String, and then appended to the
'out' variable, by means of String::from_utf8_lossy(). This works for
basic text, however, does NOT work for binary data. This becomes
problematic, when for example, downling a tar file with curl/wget, that
is then piped ('|') to the tar program. Using from_utf8_lossy() in this
case can corrupt the tar file. This patch makes it so that out is stored
as bytes by default, and only converted to a String when necessary.

1 files changed, 51 insertions, 33 deletions

diff --git a/src/poem/elements/rune.rs b/src/poem/elements/rune.rs
index 1742778..6db69c9 100644
--- a/src/poem/elements/rune.rs
+++ b/src/poem/elements/rune.rs
@@ -106,7 +106,7 @@ impl Rune {
     /// # Arguments
     /// * `verse: &Verse` - The verse to recite
     /// * `out: &mut String` - A string that may have output from the last command
-    pub fn incant_none(verse: &Verse, out: &mut String) -> Result<i32, io::Error> {
+    pub fn incant_none(verse: &Verse, out: &mut Vec<u8>) -> Result<i32, io::Error> {
         let child = task!(verse, out);
 
         let output = child.wait_with_output()?;
@@ -132,20 +132,16 @@ impl Rune {
     /// # Arguments
     /// * `verse: &Verse` - The verse to recite
     /// * `out: &mut String` - A string that may have output from the last command
-    pub fn incant_couplet(verse: &Verse, out: &mut String) -> Result<i32, io::Error> {
+    pub fn incant_couplet(verse: &Verse, out: &mut Vec<u8>) -> Result<i32, io::Error> {
         let child = ctask!(verse, out);
 
-        let output = child.wait_with_output()?;
+        let mut output = child.wait_with_output()?;
 
         if !output.status.success() {
             return Ok(output.status.code().unwrap_or(-1));
         }
 
-        out.push_str(
-            String::from_utf8_lossy(&output.stdout)
-                .into_owned()
-                .as_str(),
-        );
+        out.append(&mut output.stdout);
 
         Ok(output.status.code().unwrap_or(0))
     }
@@ -166,7 +162,7 @@ impl Rune {
     /// * `pids: Arc<Mutex<Vec<i32>>>` - A vector that stores the PIDs of all background processes that belong to the shell
     pub fn incant_quiet(
         verse: &Verse,
-        out: &mut String,
+        out: &mut Vec<u8>,
         pids: &mut Arc<Mutex<Vec<i32>>>,
     ) -> Result<i32, io::Error> {
         let child = btask!(verse, out);
@@ -195,12 +191,12 @@ impl Rune {
     }
 
     /// Alias to [Rune::incant_none]
-    pub fn incant_and(verse: &Verse, out: &mut String) -> Result<i32, io::Error> {
+    pub fn incant_and(verse: &Verse, out: &mut Vec<u8>) -> Result<i32, io::Error> {
         Rune::incant_none(verse, out)
     }
 
     /// Alias to [Rune::incant_none]
-    pub fn incant_continue(verse: &Verse, out: &mut String) -> Result<i32, io::Error> {
+    pub fn incant_continue(verse: &Verse, out: &mut Vec<u8>) -> Result<i32, io::Error> {
         Rune::incant_none(verse, out)
     }
 
@@ -220,7 +216,7 @@ impl Rune {
     ///                        file paths in `next`
     pub fn incant_read(
         verse: &mut Verse,
-        out: &mut String,
+        out: &mut Vec<u8>,
         pids: &mut Arc<Mutex<Vec<i32>>>,
     ) -> Result<i32, io::Error> {
         // Split the verse from the paths
@@ -232,7 +228,7 @@ impl Rune {
             let mut file = OpenOptions::new().read(true).open(path)?;
             let mut contents = String::new();
             file.read_to_string(&mut contents)?;
-            out.push_str(contents.as_str());
+            out.append(&mut contents.as_bytes().to_vec());
         }
 
         // Alias incant_<meter>
@@ -260,27 +256,27 @@ impl Rune {
     ///                        file paths in `next`
     pub fn incant_write(
         verse: &mut Verse,
-        out: &mut String,
+        out: &mut Vec<u8>,
         pids: &mut Arc<Mutex<Vec<i32>>>,
     ) -> Result<i32, io::Error> {
         // Split the verse from the paths
-        let paths = verse.split("<");
+        let mut paths = Arc::new(Mutex::new(verse.split("<")));
 
         // Alias incant_<meter>
         // let status = Rune::incant_couplet(&verse, out)?;
         let status = match verse.meter {
-            Rune::None => Rune::incant_none(&verse, out)?,
+            Rune::None => Rune::incant_couplet(&verse, out)?,
             Rune::Couplet => Rune::incant_couplet(&verse, out)?,
-            Rune::Quiet => Rune::incant_quiet_io(&verse, out, pids)?,
-            Rune::And => Rune::incant_and(&verse, out)?,
-            Rune::Continue => Rune::incant_continue(&verse, out)?,
+            Rune::Quiet => Rune::incant_quiet_io(&verse, out, pids, &mut paths)?,
+            Rune::And => Rune::incant_couplet(&verse, out)?,
+            Rune::Continue => Rune::incant_couplet(&verse, out)?,
             _ => unreachable!(),
         };
 
         // Write output to each file specified in the next verse
-        for path in paths.iter() {
+        for path in paths.lock().unwrap().iter() {
             let mut file = OpenOptions::new().create(true).write(true).open(path)?;
-            file.write(out.as_bytes())?;
+            file.write(out)?;
         }
 
         // Clear out
@@ -303,27 +299,27 @@ impl Rune {
     ///                        file paths in `next`
     pub fn incant_addendum(
         verse: &mut Verse,
-        out: &mut String,
+        out: &mut Vec<u8>,
         pids: &mut Arc<Mutex<Vec<i32>>>,
     ) -> Result<i32, io::Error> {
         // Split the verse from the paths
-        let paths = verse.split("<");
+        let mut paths = Arc::new(Mutex::new(verse.split("<")));
 
         // Alias incant_<meter>
         // let status = Rune::incant_couplet(&verse, out)?;
         let status = match verse.meter {
-            Rune::None => Rune::incant_none(&verse, out)?,
+            Rune::None => Rune::incant_couplet(&verse, out)?,
             Rune::Couplet => Rune::incant_couplet(&verse, out)?,
-            Rune::Quiet => Rune::incant_quiet_io(&verse, out, pids)?,
-            Rune::And => Rune::incant_and(&verse, out)?,
-            Rune::Continue => Rune::incant_continue(&verse, out)?,
+            Rune::Quiet => Rune::incant_quiet_io(&verse, out, pids, &mut paths)?,
+            Rune::And => Rune::incant_couplet(&verse, out)?,
+            Rune::Continue => Rune::incant_couplet(&verse, out)?,
             _ => unreachable!(),
         };
 
         // Write output to each file specified in the next verse
-        for path in paths.iter() {
+        for path in paths.lock().unwrap().iter() {
             let mut file = OpenOptions::new().create(true).append(true).open(path)?;
-            file.write(out.as_bytes())?;
+            file.write(out)?;
         }
 
         // Clear out
@@ -336,15 +332,18 @@ impl Rune {
     /// Same as incant_quiet, except capture STDOUT into `out`
     pub fn incant_quiet_io(
         verse: &Verse,
-        out: &mut String,
+        out: &mut Vec<u8>,
         pids: &mut Arc<Mutex<Vec<i32>>>,
+        paths: &mut Arc<Mutex<Vec<String>>>,
     ) -> Result<i32, io::Error> {
-        let child = iobtask!(verse, out);
-        println!("[&]  {}", child.id());
+        let child = Arc::new(Mutex::new(iobtask!(verse, out)));
+        println!("[&]  {}", child.lock().unwrap().id());
 
-        pids.lock().unwrap().push(child.id() as i32);
+        pids.lock().unwrap().push(child.lock().unwrap().id() as i32);
         let stanza = verse.stanza.join(" ").to_string();
         let pids = Arc::clone(pids);
+        let paths = Arc::clone(paths);
+        let io = Arc::new(verse.io);
 
         unsafe {
             signal_hook::low_level::register(signal_hook::consts::SIGCHLD, move || {
@@ -354,6 +353,25 @@ impl Rune {
                     pid = waitpid(pid, &mut status, WNOHANG);
                     if pid > 0 {
                         print!("\n[&]  + done    {}", stanza);
+                        let mut bytes: Vec<u8> = Vec::new();
+                        let mut child = child.lock().unwrap();
+                        child
+                            .stdout
+                            .as_mut()
+                            .unwrap()
+                            .read_to_end(&mut bytes)
+                            .unwrap();
+                        for path in paths.lock().unwrap().iter() {
+                            let file = if io == Rune::Write.into() {
+                                OpenOptions::new().create(true).write(true).open(path)
+                            } else if io == Rune::Addendum.into() {
+                                OpenOptions::new().create(true).append(true).open(path)
+                            } else {
+                                unreachable!();
+                            };
+
+                            let _ = file.unwrap().write(&bytes);
+                        }
                         io::stdout().flush().unwrap();
                     }
                 }