diff --git a/src/pty.rs b/src/pty.rs index 1e5bba9..71b137b 100644 --- a/src/pty.rs +++ b/src/pty.rs @@ -19,26 +19,26 @@ use std::{env, fs}; type ExtraEnv = HashMap; -pub trait Recorder { - fn start(&mut self, size: TtySize) -> io::Result<()>; - fn output(&mut self, data: &[u8]); +pub trait Handler { + fn start(&mut self, size: TtySize); + fn output(&mut self, data: &[u8]) -> bool; fn input(&mut self, data: &[u8]) -> bool; - fn resize(&mut self, size: TtySize); + fn resize(&mut self, size: TtySize) -> bool; } -pub fn exec, T: Tty + ?Sized, R: Recorder>( +pub fn exec, T: Tty + ?Sized, H: Handler>( command: &[S], extra_env: &ExtraEnv, tty: &mut T, - recorder: &mut R, + handler: &mut H, ) -> Result { let winsize = tty.get_size(); - recorder.start(winsize.into())?; + handler.start(winsize.into()); let result = unsafe { pty::forkpty(Some(&winsize), None) }?; match result.fork_result { ForkResult::Parent { child } => { - handle_parent(result.master.as_raw_fd(), child, tty, recorder) + handle_parent(result.master.as_raw_fd(), child, tty, handler) } ForkResult::Child => { @@ -48,13 +48,13 @@ pub fn exec, T: Tty + ?Sized, R: Recorder>( } } -fn handle_parent( +fn handle_parent( master_fd: RawFd, child: unistd::Pid, tty: &mut T, - recorder: &mut R, + handler: &mut H, ) -> Result { - let wait_result = match copy(master_fd, child, tty, recorder) { + let wait_result = match copy(master_fd, child, tty, handler) { Ok(Some(status)) => Ok(status), Ok(None) => wait::waitpid(child, None), @@ -74,11 +74,11 @@ fn handle_parent( const BUF_SIZE: usize = 128 * 1024; -fn copy( +fn copy( master_raw_fd: RawFd, child: unistd::Pid, tty: &mut T, - recorder: &mut R, + handler: &mut H, ) -> Result> { let mut master = unsafe { fs::File::from_raw_fd(master_raw_fd) }; let mut buf = [0u8; BUF_SIZE]; @@ -146,8 +146,9 @@ fn copy( if master_read { while let Some(n) = read_non_blocking(&mut master, &mut buf)? { if n > 0 { - recorder.output(&buf[0..n]); - output.extend_from_slice(&buf[0..n]); + if handler.output(&buf[0..n]) { + output.extend_from_slice(&buf[0..n]); + } } else if output.is_empty() { return Ok(None); } else { @@ -204,7 +205,7 @@ fn copy( if tty_read { while let Some(n) = read_non_blocking(tty, &mut buf)? { if n > 0 { - if recorder.input(&buf[0..n]) { + if handler.input(&buf[0..n]) { input.extend_from_slice(&buf[0..n]); } } else { @@ -216,8 +217,10 @@ fn copy( if sigwinch_read { sigwinch_fd.flush(); let winsize = tty.get_size(); - set_pty_size(master_raw_fd, &winsize); - recorder.resize(winsize.into()); + + if handler.resize(winsize.into()) { + set_pty_size(master_raw_fd, &winsize); + } } let mut kill_the_child = false; @@ -366,34 +369,37 @@ impl Drop for SignalFd { #[cfg(test)] mod tests { - use super::Recorder; + use super::Handler; use crate::pty::ExtraEnv; use crate::tty::{FixedSizeTty, NullTty, TtySize}; #[derive(Default)] - struct TestRecorder { + struct TestHandler { tty_size: Option, output: Vec>, } - impl Recorder for TestRecorder { - fn start(&mut self, tty_size: TtySize) -> std::io::Result<()> { + impl Handler for TestHandler { + fn start(&mut self, tty_size: TtySize) { self.tty_size = Some(tty_size); - Ok(()) } - fn output(&mut self, data: &[u8]) { + fn output(&mut self, data: &[u8]) -> bool { self.output.push(data.into()); + + true } fn input(&mut self, _data: &[u8]) -> bool { true } - fn resize(&mut self, _size: TtySize) {} + fn resize(&mut self, _size: TtySize) -> bool { + true + } } - impl TestRecorder { + impl TestHandler { fn output(&self) -> Vec { self.output .iter() @@ -404,7 +410,7 @@ mod tests { #[test] fn exec_basic() { - let mut recorder = TestRecorder::default(); + let mut handler = TestHandler::default(); let code = r#" import sys; @@ -419,47 +425,47 @@ sys.stdout.write('bar'); &["python3", "-c", code], &ExtraEnv::new(), &mut NullTty::open().unwrap(), - &mut recorder, + &mut handler, ) .unwrap(); - assert_eq!(recorder.output(), vec!["foo", "bar"]); - assert_eq!(recorder.tty_size, Some(TtySize(80, 24))); + assert_eq!(handler.output(), vec!["foo", "bar"]); + assert_eq!(handler.tty_size, Some(TtySize(80, 24))); } #[test] fn exec_no_output() { - let mut recorder = TestRecorder::default(); + let mut handler = TestHandler::default(); super::exec( &["true"], &ExtraEnv::new(), &mut NullTty::open().unwrap(), - &mut recorder, + &mut handler, ) .unwrap(); - assert!(recorder.output().is_empty()); + assert!(handler.output().is_empty()); } #[test] fn exec_quick() { - let mut recorder = TestRecorder::default(); + let mut handler = TestHandler::default(); super::exec( &["printf", "hello world\n"], &ExtraEnv::new(), &mut NullTty::open().unwrap(), - &mut recorder, + &mut handler, ) .unwrap(); - assert!(!recorder.output().is_empty()); + assert!(!handler.output().is_empty()); } #[test] fn exec_extra_env() { - let mut recorder = TestRecorder::default(); + let mut handler = TestHandler::default(); let mut env = ExtraEnv::new(); env.insert("ASCIINEMA_TEST_FOO".to_owned(), "bar".to_owned()); @@ -468,25 +474,25 @@ sys.stdout.write('bar'); &["sh", "-c", "echo -n $ASCIINEMA_TEST_FOO"], &env, &mut NullTty::open().unwrap(), - &mut recorder, + &mut handler, ) .unwrap(); - assert_eq!(recorder.output(), vec!["bar"]); + assert_eq!(handler.output(), vec!["bar"]); } #[test] fn exec_winsize_override() { - let mut recorder = TestRecorder::default(); + let mut handler = TestHandler::default(); super::exec( &["true"], &ExtraEnv::new(), &mut FixedSizeTty::new(NullTty::open().unwrap(), Some(100), Some(50)), - &mut recorder, + &mut handler, ) .unwrap(); - assert_eq!(recorder.tty_size, Some(TtySize(100, 50))); + assert_eq!(handler.tty_size, Some(TtySize(100, 50))); } } diff --git a/src/recorder.rs b/src/recorder.rs index f14ab2b..7b8c059 100644 --- a/src/recorder.rs +++ b/src/recorder.rs @@ -81,10 +81,10 @@ impl Recorder { } } -impl pty::Recorder for Recorder { - fn start(&mut self, tty_size: tty::TtySize) -> io::Result<()> { +impl pty::Handler for Recorder { + fn start(&mut self, tty_size: tty::TtySize) { let mut output = self.output.take().unwrap(); - output.start(&tty_size)?; + let _ = output.start(&tty_size); let receiver = self.receiver.take().unwrap(); let mut notifier = self.notifier.take().unwrap(); @@ -134,17 +134,15 @@ impl pty::Recorder for Recorder { self.handle = Some(util::JoinHandle::new(handle)); self.start_time = Instant::now(); - - Ok(()) } - fn output(&mut self, data: &[u8]) { - if self.pause_time.is_some() { - return; + fn output(&mut self, data: &[u8]) -> bool { + if self.pause_time.is_none() { + let msg = Message::Output(self.elapsed_time(), data.into()); + self.sender.send(msg).expect("output send should succeed"); } - let msg = Message::Output(self.elapsed_time(), data.into()); - self.sender.send(msg).expect("output send should succeed"); + true } fn input(&mut self, data: &[u8]) -> bool { @@ -187,9 +185,11 @@ impl pty::Recorder for Recorder { true } - fn resize(&mut self, size: tty::TtySize) { + fn resize(&mut self, size: tty::TtySize) -> bool { let msg = Message::Resize(self.elapsed_time(), size); self.sender.send(msg).expect("resize send should succeed"); + + true } } diff --git a/src/streamer/mod.rs b/src/streamer/mod.rs index ff3b861..640e22d 100644 --- a/src/streamer/mod.rs +++ b/src/streamer/mod.rs @@ -7,7 +7,6 @@ use crate::notifier::Notifier; use crate::pty; use crate::tty; use crate::util; -use std::io; use std::net; use std::thread; use std::time::Duration; @@ -85,8 +84,8 @@ impl Streamer { } } -impl pty::Recorder for Streamer { - fn start(&mut self, tty_size: tty::TtySize) -> io::Result<()> { +impl pty::Handler for Streamer { + fn start(&mut self, tty_size: tty::TtySize) { let pty_rx = self.pty_rx.take().unwrap(); let (clients_tx, mut clients_rx) = mpsc::channel(1); let shutdown_token = tokio_util::sync::CancellationToken::new(); @@ -139,17 +138,15 @@ impl pty::Recorder for Streamer { })); self.start_time = Instant::now(); - - Ok(()) } - fn output(&mut self, raw: &[u8]) { - if self.paused { - return; + fn output(&mut self, raw: &[u8]) -> bool { + if !self.paused { + let event = Event::Output(self.elapsed_time(), raw.into()); + let _ = self.pty_tx.send(event); } - let event = Event::Output(self.elapsed_time(), raw.into()); - let _ = self.pty_tx.send(event); + true } fn input(&mut self, raw: &[u8]) -> bool { @@ -185,9 +182,11 @@ impl pty::Recorder for Streamer { true } - fn resize(&mut self, size: crate::tty::TtySize) { + fn resize(&mut self, size: crate::tty::TtySize) -> bool { let event = Event::Resize(self.elapsed_time(), size); let _ = self.pty_tx.send(event); + + true } }