aboutsummaryrefslogtreecommitdiffstats
path: root/src/bin/ttyrec/main.rs
blob: 1db150456ee31bcac52ec1eceaca0766a2e2aa02 (plain) (blame)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
use async_std::io::{ReadExt as _, WriteExt as _};
use async_std::prelude::FutureExt as _;
use pty_process::Command as _;

#[derive(Debug, structopt::StructOpt)]
#[structopt(about = "ttyrec")]
struct Opt {
    #[structopt(short, long, default_value = "ttyrec")]
    file: std::ffi::OsString,
    #[structopt(short, long)]
    cmd: Option<std::ffi::OsString>,
}

fn get_cmd(
    cmd: Option<std::ffi::OsString>,
) -> (std::ffi::OsString, Vec<std::ffi::OsString>) {
    if let Some(cmd) = cmd {
        ("/bin/sh".into(), vec!["-c".into(), cmd])
    } else {
        let shell =
            std::env::var_os("SHELL").unwrap_or_else(|| "/bin/sh".into());
        (shell, vec![])
    }
}

enum Event {
    Key(textmode::Result<Option<textmode::Key>>),
    Stdout(std::io::Result<Vec<u8>>),
}

async fn async_main(opt: Opt) -> anyhow::Result<()> {
    let Opt { cmd, file } = opt;
    let (cmd, args) = get_cmd(cmd);

    let size = terminal_size::terminal_size().map_or(
        (24, 80),
        |(terminal_size::Width(w), terminal_size::Height(h))| (h, w),
    );
    let fh = async_std::fs::File::create(file).await?;
    let mut input = textmode::Input::new().await?;
    let _input_guard = input.take_raw_guard();
    let mut stdout = async_std::io::stdout();
    let child = async_std::process::Command::new(cmd)
        .args(args)
        .spawn_pty(Some(&pty_process::Size::new(size.0, size.1)))?;

    let (event_w, event_r) = async_std::channel::unbounded();
    let (input_w, input_r) = async_std::channel::unbounded();

    {
        let event_w = event_w.clone();
        async_std::task::spawn(async move {
            loop {
                event_w
                    .send(Event::Key(input.read_key().await))
                    .await
                    .unwrap();
            }
        });
    }

    {
        let event_w = event_w.clone();
        async_std::task::spawn(async move {
            loop {
                enum Res {
                    Read(Result<usize, std::io::Error>),
                    Write(Result<Vec<u8>, async_std::channel::RecvError>),
                }
                let mut buf = [0_u8; 4096];
                let mut pty = child.pty();
                let read = async { Res::Read(pty.read(&mut buf).await) };
                let write = async { Res::Write(input_r.recv().await) };
                match read.race(write).await {
                    Res::Read(res) => {
                        let res = res.map(|n| buf[..n].to_vec());
                        let err = res.is_err();
                        event_w.send(Event::Stdout(res)).await.unwrap();
                        if err {
                            break;
                        }
                    }
                    Res::Write(res) => {
                        let bytes = res.unwrap();
                        pty.write(&bytes).await.unwrap();
                    }
                }
            }
        });
    }

    let mut writer = ttyrec::Writer::new(fh);
    loop {
        match event_r.recv().await? {
            Event::Key(key) => {
                let key = key?;
                if let Some(key) = key {
                    input_w.send(key.into_bytes()).await?;
                } else {
                    break;
                }
            }
            Event::Stdout(bytes) => match bytes {
                Ok(bytes) => {
                    writer.frame(&bytes).await?;
                    stdout.write_all(&bytes).await?;
                    stdout.flush().await?;
                }
                Err(e) => {
                    if e.raw_os_error() == Some(libc::EIO) {
                        break;
                    } else {
                        anyhow::bail!(
                            "failed to read from child process: {}",
                            e
                        );
                    }
                }
            },
        }
    }

    Ok(())
}

#[paw::main]
fn main(opt: Opt) {
    match async_std::task::block_on(async_main(opt)) {
        Ok(_) => (),
        Err(e) => {
            eprintln!("ttyrec: {}", e);
            std::process::exit(1);
        }
    };
}