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, 18 insertions, 8 deletions

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)
     }
 }