From f5db8d64828db756b80b6022322265a2b4f1c11b Mon Sep 17 00:00:00 2001 From: Rory Dudley Date: Sat, 6 Apr 2024 23:32:30 -0600 Subject: 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. --- src/poem/elements/rune.rs | 84 ++++++++++++++++++++++++++++------------------- 1 file changed, 51 insertions(+), 33 deletions(-) (limited to 'src/poem/elements') 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 { + pub fn incant_none(verse: &Verse, out: &mut Vec) -> Result { 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 { + pub fn incant_couplet(verse: &Verse, out: &mut Vec) -> Result { 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>>` - 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, pids: &mut Arc>>, ) -> Result { 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 { + pub fn incant_and(verse: &Verse, out: &mut Vec) -> Result { Rune::incant_none(verse, out) } /// Alias to [Rune::incant_none] - pub fn incant_continue(verse: &Verse, out: &mut String) -> Result { + pub fn incant_continue(verse: &Verse, out: &mut Vec) -> Result { 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, pids: &mut Arc>>, ) -> Result { // 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_ @@ -260,27 +256,27 @@ impl Rune { /// file paths in `next` pub fn incant_write( verse: &mut Verse, - out: &mut String, + out: &mut Vec, pids: &mut Arc>>, ) -> Result { // Split the verse from the paths - let paths = verse.split("<"); + let mut paths = Arc::new(Mutex::new(verse.split("<"))); // Alias incant_ // 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, pids: &mut Arc>>, ) -> Result { // Split the verse from the paths - let paths = verse.split("<"); + let mut paths = Arc::new(Mutex::new(verse.split("<"))); // Alias incant_ // 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, pids: &mut Arc>>, + paths: &mut Arc>>, ) -> Result { - 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 = 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(); } } -- cgit v1.2.3