Fix recording in headless mode - when TTY is not available

This commit is contained in:
Marcin Kulik
2024-01-05 14:53:01 +01:00
parent e26c5f67dc
commit 477bc8c0f4
4 changed files with 150 additions and 39 deletions

View File

@@ -3,6 +3,7 @@ use crate::format::{asciicast, raw};
use crate::locale;
use crate::pty;
use crate::recorder;
use crate::tty;
use anyhow::Result;
use clap::Args;
use std::collections::{HashMap, HashSet};
@@ -116,9 +117,17 @@ impl Cli {
println!("asciinema: recording asciicast to {}", self.filename);
println!("asciinema: press <ctrl+d> or type \"exit\" when you're done");
let tty: Box<dyn tty::Tty> = if let Ok(dev_tty) = tty::DevTty::open() {
Box::new(dev_tty)
} else {
println!("asciinema: TTY not available, recording in headless mode");
Box::new(tty::DevNull::open()?)
};
pty::exec(
&exec_args,
&exec_extra_env,
tty,
(self.cols, self.rows),
&mut recorder,
)?;

View File

@@ -6,6 +6,7 @@ mod locale;
mod player;
mod pty;
mod recorder;
mod tty;
use crate::config::Config;
use anyhow::Result;
use clap::{Parser, Subcommand};

View File

@@ -1,4 +1,5 @@
use crate::io::set_non_blocking;
use crate::tty::Tty;
use anyhow::bail;
use mio::unix::SourceFd;
use nix::{libc, pty, sys::signal, sys::wait, unistd, unistd::ForkResult};
@@ -7,11 +8,9 @@ use signal_hook_mio::v0_8::Signals;
use std::collections::HashMap;
use std::ffi::{CString, NulError};
use std::io::{self, Read, Write};
use std::ops::Deref;
use std::os::fd::RawFd;
use std::os::unix::io::{AsRawFd, FromRawFd};
use std::{env, fs};
use termion::raw::IntoRawMode;
type ExtraEnv = HashMap<String, String>;
@@ -25,19 +24,19 @@ pub trait Recorder {
pub fn exec<S: AsRef<str>, R: Recorder>(
args: &[S],
extra_env: &ExtraEnv,
tty: Box<dyn Tty>,
winsize_override: (Option<u16>, Option<u16>),
recorder: &mut R,
) -> anyhow::Result<i32> {
let tty = open_tty()?;
let winsize = get_tty_size(tty.as_raw_fd(), winsize_override);
let winsize = get_tty_size(&*tty, winsize_override);
recorder.start((winsize.ws_col, winsize.ws_row))?;
let result = unsafe { pty::forkpty(Some(&winsize), None) }?;
match result.fork_result {
ForkResult::Parent { child } => handle_parent(
result.master.as_raw_fd(),
tty,
child,
tty,
winsize_override,
recorder,
),
@@ -51,12 +50,12 @@ pub fn exec<S: AsRef<str>, R: Recorder>(
fn handle_parent<R: Recorder>(
master_fd: RawFd,
tty: fs::File,
child: unistd::Pid,
tty: Box<dyn Tty>,
winsize_override: (Option<u16>, Option<u16>),
recorder: &mut R,
) -> anyhow::Result<i32> {
let copy_result = copy(master_fd, tty, child, winsize_override, recorder);
let copy_result = copy(master_fd, child, tty, winsize_override, recorder);
let wait_result = wait::waitpid(child, None);
copy_result?;
@@ -75,8 +74,8 @@ const BUF_SIZE: usize = 128 * 1024;
fn copy<R: Recorder>(
master_fd: RawFd,
tty: fs::File,
child: unistd::Pid,
mut tty: Box<dyn Tty>,
winsize_override: (Option<u16>, Option<u16>),
recorder: &mut R,
) -> anyhow::Result<()> {
@@ -84,7 +83,6 @@ fn copy<R: Recorder>(
let mut poll = mio::Poll::new()?;
let mut events = mio::Events::with_capacity(128);
let mut master_source = SourceFd(&master_fd);
let mut tty = tty.into_raw_mode()?;
let tty_fd = tty.as_raw_fd();
let mut tty_source = SourceFd(&tty_fd);
let mut signals = Signals::new([SIGWINCH, SIGINT, SIGTERM, SIGQUIT, SIGHUP])?;
@@ -94,7 +92,6 @@ fn copy<R: Recorder>(
let mut flush = false;
set_non_blocking(&master_fd)?;
set_non_blocking(&tty_fd)?;
poll.registry()
.register(&mut master_source, MASTER, mio::Interest::READABLE)?;
@@ -147,11 +144,11 @@ fn copy<R: Recorder>(
if event.is_read_closed() {
poll.registry().deregister(&mut master_source)?;
if !output.is_empty() {
flush = true;
} else {
if output.is_empty() {
return Ok(());
}
flush = true;
}
}
@@ -162,7 +159,8 @@ fn copy<R: Recorder>(
if left == 0 {
if flush {
return Ok(());
} else {
}
poll.registry().reregister(
&mut tty_source,
TTY,
@@ -170,11 +168,10 @@ fn copy<R: Recorder>(
)?;
}
}
}
if event.is_readable() {
let offset = input.len();
let read = read_all(&mut tty.deref(), &mut buf, &mut input)?;
let read = read_all(&mut tty, &mut buf, &mut input)?;
if read > 0 {
recorder.input(&input[offset..]);
@@ -188,7 +185,7 @@ fn copy<R: Recorder>(
}
if event.is_read_closed() {
poll.registry().deregister(&mut tty_source).unwrap();
poll.registry().deregister(&mut tty_source)?;
return Ok(());
}
}
@@ -197,7 +194,7 @@ fn copy<R: Recorder>(
for signal in signals.pending() {
match signal {
SIGWINCH => {
let winsize = get_tty_size(tty_fd, winsize_override);
let winsize = get_tty_size(&*tty, winsize_override);
set_pty_size(master_fd, &winsize);
recorder.resize((winsize.ws_col, winsize.ws_row));
}
@@ -237,22 +234,11 @@ fn handle_child<S: AsRef<str>>(args: &[S], extra_env: &ExtraEnv) -> anyhow::Resu
unsafe { libc::_exit(1) }
}
fn open_tty() -> io::Result<fs::File> {
fs::OpenOptions::new()
.read(true)
.write(true)
.open("/dev/tty")
}
fn get_tty_size(tty_fd: i32, winsize_override: (Option<u16>, Option<u16>)) -> pty::Winsize {
let mut winsize = pty::Winsize {
ws_row: 24,
ws_col: 80,
ws_xpixel: 0,
ws_ypixel: 0,
};
unsafe { libc::ioctl(tty_fd, libc::TIOCGWINSZ, &mut winsize) };
fn get_tty_size<T: Tty + ?Sized>(
tty: &T,
winsize_override: (Option<u16>, Option<u16>),
) -> pty::Winsize {
let mut winsize = tty.get_size();
if let Some(cols) = winsize_override.0 {
winsize.ws_col = cols;
@@ -327,6 +313,7 @@ fn write_all<W: Write>(sink: &mut W, data: &mut Vec<u8>) -> io::Result<usize> {
#[cfg(test)]
mod tests {
use crate::pty::ExtraEnv;
use crate::tty::DevNull;
#[derive(Default)]
struct TestRecorder {
@@ -373,12 +360,13 @@ sys.stdout.write('bar');
let result = super::exec(
&["python3", "-c", code],
&ExtraEnv::new(),
Box::new(DevNull::open().unwrap()),
(None, None),
&mut recorder,
);
assert_eq!(recorder.output(), vec!["foo", "bar"]);
assert!(result.is_ok());
assert_eq!(recorder.output(), vec!["foo", "bar"]);
assert!(recorder.size.is_some());
}
}

113
src/tty.rs Normal file
View File

@@ -0,0 +1,113 @@
use anyhow::Result;
use mio::unix::pipe;
use nix::{libc, pty};
use std::{
fs, io,
os::fd::{AsRawFd, RawFd},
};
use termion::raw::{IntoRawMode, RawTerminal};
pub trait Tty: io::Write + io::Read + AsRawFd {
fn get_size(&self) -> pty::Winsize;
}
pub struct DevTty {
file: RawTerminal<fs::File>,
}
impl DevTty {
pub fn open() -> Result<Self> {
let file = fs::OpenOptions::new()
.read(true)
.write(true)
.open("/dev/tty")?
.into_raw_mode()?;
crate::io::set_non_blocking(&file.as_raw_fd())?;
Ok(Self { file })
}
}
impl Tty for DevTty {
fn get_size(&self) -> pty::Winsize {
let mut winsize = pty::Winsize {
ws_row: 24,
ws_col: 80,
ws_xpixel: 0,
ws_ypixel: 0,
};
unsafe { libc::ioctl(self.file.as_raw_fd(), libc::TIOCGWINSZ, &mut winsize) };
winsize
}
}
impl io::Read for DevTty {
fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
self.file.read(buf)
}
}
impl io::Write for DevTty {
fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
self.file.write(buf)
}
fn flush(&mut self) -> io::Result<()> {
self.file.flush()
}
}
impl AsRawFd for DevTty {
fn as_raw_fd(&self) -> std::os::unix::prelude::RawFd {
self.file.as_raw_fd()
}
}
pub struct DevNull {
tx: pipe::Sender,
_rx: pipe::Receiver,
}
impl DevNull {
pub fn open() -> Result<Self> {
let (tx, rx) = pipe::new()?;
Ok(Self { tx, _rx: rx })
}
}
impl Tty for DevNull {
fn get_size(&self) -> pty::Winsize {
pty::Winsize {
ws_row: 24,
ws_col: 80,
ws_xpixel: 0,
ws_ypixel: 0,
}
}
}
impl io::Read for DevNull {
fn read(&mut self, _buf: &mut [u8]) -> io::Result<usize> {
panic!("read attempt from DevNull impl of Tty");
}
}
impl io::Write for DevNull {
fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
Ok(buf.len())
}
fn flush(&mut self) -> io::Result<()> {
Ok(())
}
}
impl AsRawFd for DevNull {
fn as_raw_fd(&self) -> RawFd {
self.tx.as_raw_fd()
}
}