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.

6 files changed, 77 insertions, 49 deletions

diff --git a/src/poem/anthology.rs b/src/poem/anthology.rs
index 79f48f2..7766f62 100644
--- a/src/poem/anthology.rs
+++ b/src/poem/anthology.rs
@@ -43,7 +43,7 @@ pub fn lookup(verb: &str) -> Option<usize> {
 ///     ...
 /// }
 /// ```
-pub fn incant(verse: &Verse, out: &mut String, index: usize, env: &mut Environment) -> i32 {
+pub fn incant(verse: &Verse, out: &mut Vec<u8>, index: usize, env: &mut Environment) -> i32 {
     let verb = INDEX[index];
     match verb {
         "alias" => alias::incant(verse, out, &mut env.aliases),
diff --git a/src/poem/anthology/alias.rs b/src/poem/anthology/alias.rs
index 4c9b7b6..248342e 100644
--- a/src/poem/anthology/alias.rs
+++ b/src/poem/anthology/alias.rs
@@ -10,7 +10,7 @@ use std::collections::HashMap;
 /// ```sh
 /// alias vim=nvim
 /// ```
-pub fn incant(verse: &Verse, out: &mut String, aliases: &mut HashMap<String, String>) -> i32 {
+pub fn incant(verse: &Verse, out: &mut Vec<u8>, aliases: &mut HashMap<String, String>) -> i32 {
     match verse.clause() {
         Some(clause) => {
             for stanza in clause {
@@ -39,7 +39,7 @@ pub fn incant(verse: &Verse, out: &mut String, aliases: &mut HashMap<String, Str
             }
 
             if verse.couplet {
-                *out = format!("{}\n", lines.join("\n"));
+                *out = format!("{}\n", lines.join("\n")).as_bytes().to_vec();
             } else {
                 println!("{}", lines.join("\n"));
             }
diff --git a/src/poem/anthology/source.rs b/src/poem/anthology/source.rs
index 182fef4..43d6204 100644
--- a/src/poem/anthology/source.rs
+++ b/src/poem/anthology/source.rs
@@ -12,7 +12,7 @@ use std::fs;
 /// ```sh
 /// source ~/.dwvshrc
 /// ```
-pub fn incant(verse: &Verse, out: &mut String, env: &mut Environment) -> i32 {
+pub fn incant(verse: &Verse, out: &mut Vec<u8>, env: &mut Environment) -> i32 {
     let files = match verse.clause() {
         Some(clause) => clause,
         None => {
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();
                     }
                 }
diff --git a/src/poem/recite.rs b/src/poem/recite.rs
index 4038e37..60b7857 100644
--- a/src/poem/recite.rs
+++ b/src/poem/recite.rs
@@ -4,6 +4,7 @@ use crate::compose::Environment;
 use crate::path;
 use crate::poem::anthology;
 use crate::poem::elements::rune::Rune;
+use crate::poem::elements::stanza::Stanza;
 use std::env;
 use std::{
     io,
@@ -11,16 +12,16 @@ use std::{
 };
 
 pub trait Reciteable {
-    fn recite(&self, env: &mut Environment, stdout: Option<bool>) -> Result<String, io::Error>;
+    fn recite(&self, env: &mut Environment, stdout: Option<bool>) -> Result<Vec<u8>, io::Error>;
 }
 
 impl Reciteable for Poem {
-    fn recite(&self, env: &mut Environment, stdout: Option<bool>) -> Result<String, io::Error> {
+    fn recite(&self, env: &mut Environment, stdout: Option<bool>) -> Result<Vec<u8>, io::Error> {
         // Should we print to stdout or always capture it
         let stdout = stdout.unwrap_or(true);
 
         // Variable for storing the output of a piped verse
-        let mut out: String = String::new();
+        let mut out: Vec<u8> = Vec::new();
 
         // Keep track of pids for background processes
         let mut pids: Arc<Mutex<Vec<i32>>> = Arc::new(Mutex::new(Vec::new()));
@@ -64,7 +65,7 @@ impl Reciteable for Poem {
 
             // Run interal poems
             let v = verse.clone();
-            let mut new_stanza = None;
+            let mut new_stanza: Option<Stanza> = None;
             if verse.poems() {
                 // Collect all the words that have vertical tabs
                 let mut wordp_indicies = vec![];
@@ -93,7 +94,17 @@ impl Reciteable for Poem {
                             Some(poem) => poem,
                             None => break, // TODO: Return an error
                         };
-                        let out = poem.recite(env, Some(false))?;
+                        let mut out = poem.recite(env, Some(false))?;
+                        match out.last() {
+                            Some(last) => {
+                                if *last == b'\n' {
+                                    out.remove(out.len() - 1);
+                                }
+                            }
+                            None => {}
+                        }
+
+                        let out = String::from_utf8_lossy(&out);
                         if out.contains("\n") && index.is_none() {
                             let mut out = out.split("\n");
                             let next = out.next().unwrap_or("").trim();
@@ -112,9 +123,8 @@ impl Reciteable for Poem {
                             left.append(&mut right);
                             new_stanza = Some(left.clone());
                         } else {
-                            *wordp = wordp.replacen("\x0b", out.as_str(), 1).to_string();
+                            *wordp = wordp.replacen("\x0b", &out, 1).to_string();
                         }
-                        *wordp = wordp.replacen("\x0b", out.as_str(), 1).to_string();
                     }
                     j += 1;
                 }
@@ -190,6 +200,6 @@ impl Reciteable for Poem {
 
         // If we've successfully exited the loop, then all verses were properly
         // recited
-        Ok(out.trim().to_string())
+        Ok(out)
     }
 }
diff --git a/src/poem/recite/ps.rs b/src/poem/recite/ps.rs
index 61ed66d..5700ae8 100644
--- a/src/poem/recite/ps.rs
+++ b/src/poem/recite/ps.rs
@@ -17,7 +17,7 @@ macro_rules! task {
                 .spawn()?;
 
             let stdin = child.stdin.as_mut().ok_or(io::ErrorKind::BrokenPipe)?;
-            stdin.write_all(&$out.as_bytes())?;
+            stdin.write_all(&$out)?;
             $out.clear();
 
             child
@@ -50,7 +50,7 @@ macro_rules! ctask {
                 .spawn()?;
 
             let stdin = child.stdin.as_mut().ok_or(io::ErrorKind::BrokenPipe)?;
-            stdin.write_all(&$out.as_bytes())?;
+            stdin.write_all(&$out)?;
             $out.clear();
 
             child
@@ -83,7 +83,7 @@ macro_rules! btask {
                 .spawn()?;
 
             let stdin = child.stdin.as_mut().ok_or(io::ErrorKind::BrokenPipe)?;
-            stdin.write_all(&$out.as_bytes())?;
+            stdin.write_all(&$out)?;
             $out.clear();
 
             child
@@ -119,7 +119,7 @@ macro_rules! iobtask {
                 .spawn()?;
 
             let stdin = child.stdin.as_mut().ok_or(io::ErrorKind::BrokenPipe)?;
-            stdin.write_all(&$out.as_bytes())?;
+            stdin.write_all(&$out)?;
             $out.clear();
 
             child