proc_reader/
lib.rs

1//! A std::io::Read implementation for stdout/stderr of other process.
2//!
3//! # Examples
4//!
5//! ```
6//! # use proc_reader::ProcReader;
7//! # use std::process::Command;
8//! # use std::io::Read;
9//! # use std::time::Duration;
10//! # use std::thread;
11//! # fn main() {
12//! // Create a process for reading stdout
13//! let mut child = Command::new("sh").arg("-c").arg("sleep 1; echo aaa").spawn().unwrap();
14//!
15//! // Create ProcReader from pid
16//! let mut reader = ProcReader::from_stdout(child.id());
17//!
18//! // Wait the end of process
19//! thread::sleep(Duration::from_secs(2));
20//!
21//! // Read from ProcReader
22//! let mut line = String::new();
23//! let _ = reader.read_to_string(&mut line);
24//! assert_eq!( "aaa\n", line);
25//! # }
26//! ```
27
28use libc::user_regs_struct;
29use nix::sys::ptrace::{self, AddressType, Options};
30use nix::sys::signal::Signal;
31use nix::sys::wait::{self, WaitStatus};
32use nix::unistd::Pid;
33use std::io::Read;
34use std::mem;
35use std::sync::mpsc::{channel, Receiver, Sender};
36use std::thread::{self, JoinHandle};
37use thiserror::Error;
38
39// -------------------------------------------------------------------------------------------------
40// Error
41// -------------------------------------------------------------------------------------------------
42
43#[derive(Debug, Error)]
44pub enum Error {
45    #[error("failed to access process ({pid})")]
46    ProcAccessFailed { pid: Pid },
47    #[error(transparent)]
48    Nix(#[from] nix::Error),
49    #[error(transparent)]
50    Recv(#[from] std::sync::mpsc::RecvError),
51    #[error(transparent)]
52    Send(#[from] std::sync::mpsc::SendError<Vec<u8>>),
53}
54
55// -------------------------------------------------------------------------------------------------
56// Type per arch
57// -------------------------------------------------------------------------------------------------
58
59#[cfg(target_arch = "x86_64")]
60type Word = u64;
61
62#[cfg(target_arch = "x86")]
63type Word = u32;
64
65#[cfg(target_arch = "x86_64")]
66const WORD_BYTES: u64 = 8;
67
68#[cfg(target_arch = "x86")]
69const WORD_BYTES: u32 = 4;
70
71/// Represents all possible ptrace-accessible registers on x86_64
72#[cfg(target_arch = "x86_64")]
73#[derive(Clone, Copy, Debug)]
74struct UserRegs(user_regs_struct);
75
76/// Represents all possible ptrace-accessible registers on x86
77#[cfg(target_arch = "x86")]
78#[derive(Clone, Copy, Debug)]
79struct UserRegs(user_regs_struct);
80
81#[cfg(target_arch = "x86_64")]
82#[allow(dead_code)]
83impl UserRegs {
84    fn syscall(self) -> Word {
85        self.0.orig_rax
86    }
87
88    fn ret(self) -> Word {
89        self.0.rax
90    }
91
92    fn arg1(self) -> Word {
93        self.0.rdi
94    }
95
96    fn arg2(self) -> Word {
97        self.0.rsi
98    }
99
100    fn arg3(self) -> Word {
101        self.0.rdx
102    }
103
104    fn arg4(self) -> Word {
105        self.0.r10
106    }
107
108    fn arg5(self) -> Word {
109        self.0.r8
110    }
111
112    fn arg6(self) -> Word {
113        self.0.r9
114    }
115}
116
117#[cfg(target_arch = "x86")]
118#[allow(dead_code)]
119impl UserRegs {
120    fn syscall(self) -> Word {
121        self.0.orig_eax
122    }
123
124    fn ret(self) -> Word {
125        self.0.eax
126    }
127
128    fn arg1(self) -> Word {
129        self.0.ebx
130    }
131
132    fn arg2(self) -> Word {
133        self.0.ecx
134    }
135
136    fn arg3(self) -> Word {
137        self.0.edx
138    }
139
140    fn arg4(self) -> Word {
141        self.0.esi
142    }
143
144    fn arg5(self) -> Word {
145        self.0.edi
146    }
147
148    fn arg6(self) -> Word {
149        self.0.ebp
150    }
151}
152
153// -------------------------------------------------------------------------------------------------
154// ProcReader
155// -------------------------------------------------------------------------------------------------
156
157enum ProcReaderMsg {
158    Stop,
159    Redirect,
160}
161
162/// The struct `ProcReader` provide reader from stdout/stderr of other process.
163pub struct ProcReader {
164    ctl: Sender<ProcReaderMsg>,
165    buf: Receiver<Vec<u8>>,
166    err: Receiver<Error>,
167    child: Option<JoinHandle<()>>,
168    rest: Vec<u8>,
169}
170
171#[derive(PartialEq)]
172enum StdType {
173    Any,
174    Out,
175    Err,
176}
177
178impl ProcReader {
179    /// Create a new `ProcReader` for stdout/stderr of the process specified by `pid`
180    pub fn from_stdany(pid: u32) -> Self {
181        let pid = Pid::from_raw(pid as i32);
182        ProcReader::new(pid, StdType::Any)
183    }
184
185    /// Create a new `ProcReader` for stdout of the process specified by `pid`
186    pub fn from_stdout(pid: u32) -> Self {
187        let pid = Pid::from_raw(pid as i32);
188        ProcReader::new(pid, StdType::Out)
189    }
190
191    /// Create a new `ProcReader` for stderr of the process specified by `pid`
192    pub fn from_stderr(pid: u32) -> Self {
193        let pid = Pid::from_raw(pid as i32);
194        ProcReader::new(pid, StdType::Err)
195    }
196
197    /// Enable redirect trace
198    pub fn with_redirect(self) -> Self {
199        let _ = self.ctl.send(ProcReaderMsg::Redirect);
200        self
201    }
202
203    fn new(pid: Pid, typ: StdType) -> Self {
204        let (ctl_tx, ctl_rx) = channel();
205        let (buf_tx, buf_rx) = channel();
206        let (err_tx, err_rx) = channel();
207
208        let child = thread::spawn(
209            move || match ProcReader::collect(pid, typ, ctl_rx, buf_tx) {
210                Err(x) => {
211                    let _ = err_tx.send(x);
212                }
213                _ => (),
214            },
215        );
216
217        ProcReader {
218            ctl: ctl_tx,
219            buf: buf_rx,
220            err: err_rx,
221            child: Some(child),
222            rest: Vec::new(),
223        }
224    }
225
226    fn collect(
227        pid: Pid,
228        typ: StdType,
229        ctl_rx: Receiver<ProcReaderMsg>,
230        buf_tx: Sender<Vec<u8>>,
231    ) -> Result<(), Error> {
232        ptrace::attach(pid).map_err(|_| Error::ProcAccessFailed { pid })?;
233        ProcReader::set_tracesysgood(pid).map_err(|_| Error::ProcAccessFailed { pid })?;
234
235        // pid stack
236        let mut pids = Vec::new();
237        pids.push(pid);
238
239        // fd stack
240        let mut fd = [0; 1024];
241        fd[1] = 1;
242        fd[2] = 2;
243        let mut fds = Vec::new();
244        fds.push(fd.clone());
245
246        let mut enable_redirect = false;
247        let mut is_syscall_before = false;
248        let mut prev_orig_rax = 0;
249
250        loop {
251            let mut pid = *pids.last().unwrap();
252            match wait::waitpid(pid, None).map_err(|_| Error::ProcAccessFailed { pid })? {
253                WaitStatus::PtraceSyscall(_) => {
254                    let regs =
255                        ProcReader::get_regs(pid).map_err(|_| Error::ProcAccessFailed { pid })?;
256
257                    is_syscall_before = if prev_orig_rax == regs.syscall() {
258                        !is_syscall_before
259                    } else {
260                        true
261                    };
262                    prev_orig_rax = regs.syscall();
263
264                    if !is_syscall_before && enable_redirect {
265                        fd = ProcReader::update_fd(fd, regs);
266                    }
267
268                    let sys_clone = regs.syscall() == libc::SYS_clone as Word;
269                    let sys_fork = regs.syscall() == libc::SYS_fork as Word;
270                    let sys_vfork = regs.syscall() == libc::SYS_vfork as Word;
271
272                    if (sys_clone || sys_fork || sys_vfork) && !is_syscall_before {
273                        pid = Pid::from_raw(regs.ret() as i32);
274                        pids.push(pid);
275                        fds.push(fd.clone());
276                        continue;
277                    }
278
279                    if regs.syscall() == libc::SYS_write as Word && is_syscall_before {
280                        let out = ProcReader::peek_bytes(pid, regs.arg2(), regs.arg3());
281                        let out_typ = regs.arg1();
282
283                        let send_stdout = fd[out_typ as usize] == 1
284                            && (typ == StdType::Any || typ == StdType::Out);
285                        let send_stderr = fd[out_typ as usize] == 2
286                            && (typ == StdType::Any || typ == StdType::Err);
287
288                        if send_stdout || send_stderr {
289                            buf_tx.send(out)?;
290                        }
291                    }
292                }
293                WaitStatus::Exited(_, _) => {
294                    pids.pop();
295                    if pids.is_empty() {
296                        break;
297                    } else {
298                        pid = *pids.last().unwrap();
299                        fd = fds.pop().unwrap();
300                    }
301                }
302                _ => (),
303            }
304
305            match ctl_rx.try_recv() {
306                Ok(ProcReaderMsg::Stop) => {
307                    ptrace::detach(pid, None).map_err(|_| Error::ProcAccessFailed { pid })?;
308                    break;
309                }
310                Ok(ProcReaderMsg::Redirect) => {
311                    enable_redirect = true;
312                    ptrace::syscall(pid, None).map_err(|_| Error::ProcAccessFailed { pid })?;
313                }
314                _ => {
315                    ptrace::syscall(pid, None).map_err(|_| Error::ProcAccessFailed { pid })?;
316                }
317            }
318        }
319        Ok(())
320    }
321
322    fn set_tracesysgood(pid: Pid) -> Result<(), Error> {
323        loop {
324            match wait::waitpid(pid, None).map_err(|_| Error::ProcAccessFailed { pid })? {
325                // setoptions must be called at stopped condition
326                WaitStatus::Stopped(_, Signal::SIGSTOP) => {
327                    // set TRACESYSGOOD to enable PtraceSyscall
328                    // set TRACECLONE/FORK/VFORK to trace chile process
329                    let opt = Options::PTRACE_O_TRACESYSGOOD
330                        | Options::PTRACE_O_TRACECLONE
331                        | Options::PTRACE_O_TRACEFORK
332                        | Options::PTRACE_O_TRACEVFORK;
333                    ptrace::setoptions(pid, opt).map_err(|_| Error::ProcAccessFailed { pid })?;
334                    ptrace::syscall(pid, None).map_err(|_| Error::ProcAccessFailed { pid })?;
335                    break;
336                }
337                _ => {
338                    ptrace::syscall(pid, None).map_err(|_| Error::ProcAccessFailed { pid })?;
339                }
340            }
341        }
342
343        Ok(())
344    }
345
346    fn get_regs(pid: Pid) -> Result<UserRegs, Error> {
347        let regs = ptrace::getregs(pid).map(|x| UserRegs(x))?;
348        Ok(regs)
349    }
350
351    fn peek_bytes(pid: Pid, addr: Word, size: Word) -> Vec<u8> {
352        let mut vec = (0..(size + WORD_BYTES - 1) / WORD_BYTES)
353            .filter_map(|i| {
354                ptrace::read(pid, (addr + WORD_BYTES * i) as AddressType)
355                    .map(|l| unsafe { mem::transmute(l) })
356                    .ok()
357            })
358            .collect::<Vec<[u8; WORD_BYTES as usize]>>()
359            .concat();
360        vec.truncate(size as usize);
361        vec
362    }
363
364    fn update_fd(mut fd: [Word; 1024], regs: UserRegs) -> [Word; 1024] {
365        // detect dup2 for redirect
366        if regs.syscall() == libc::SYS_dup2 as Word {
367            let src = regs.arg1();
368            let dst = regs.arg2();
369            fd[dst as usize] = fd[src as usize];
370        }
371
372        // detect fcntl for fd backup
373        if regs.syscall() == libc::SYS_fcntl as Word {
374            if regs.arg2() == libc::F_DUPFD as Word {
375                let src = regs.arg1();
376                let dst = regs.ret();
377                fd[dst as usize] = fd[src as usize];
378            }
379        }
380        fd
381    }
382}
383
384impl Read for ProcReader {
385    fn read(&mut self, buf: &mut [u8]) -> std::result::Result<usize, std::io::Error> {
386        let err = match self.err.try_recv() {
387            Ok(x) => Some(x),
388            _ => None,
389        };
390
391        loop {
392            match self.buf.try_recv() {
393                Ok(mut x) => self.rest.append(&mut x),
394                Err(_) => break,
395            }
396        }
397
398        let len = if buf.len() >= self.rest.len() {
399            let len = self.rest.len();
400            self.rest.resize(buf.len(), 0);
401            buf.copy_from_slice(&self.rest);
402            self.rest.clear();
403            len
404        } else {
405            let len = buf.len();
406            let rest = self.rest.split_off(len);
407            buf.copy_from_slice(&self.rest);
408            self.rest = rest;
409            len
410        };
411
412        if len != 0 {
413            Ok(len)
414        } else if let Some(err) = err {
415            Err(std::io::Error::new(
416                std::io::ErrorKind::Other,
417                format!("{}", err),
418            ))
419        } else {
420            Ok(len)
421        }
422    }
423}
424
425impl Drop for ProcReader {
426    fn drop(&mut self) {
427        if self.child.is_some() {
428            let _ = self.ctl.send(ProcReaderMsg::Stop);
429            let _ = self.child.take().unwrap().join();
430        }
431    }
432}
433
434// -------------------------------------------------------------------------------------------------
435// Test
436// -------------------------------------------------------------------------------------------------
437
438#[cfg(test)]
439mod tests {
440    use super::*;
441    use std::io::BufReader;
442    use std::process::Command;
443    use std::thread;
444    use std::time::Duration;
445
446    static SCRIPT: &'static str = r#"
447        sleep 1;
448        print "aaa\n";
449        sleep 1;
450        print "bbb\n";
451        sleep 1;
452        print "ccc\n";
453    "#;
454
455    static SCRIPT_WITH_ERR: &'static str = r#"
456        sleep 1;
457        print "aaa\n";
458        warn "eee\n";
459    "#;
460
461    static SCRIPT_REDIRECT: &'static str = r#"
462        sleep 1;
463        echo 'aaa';
464        echo 'bbb' > /dev/null 1>&2;
465        perl -e 'warn "ccc\n"' 2>&1;
466        perl -e 'warn "ddd\n"';
467    "#;
468
469    #[test]
470    fn test_bufreader() {
471        let child = Command::new("perl").arg("-e").arg(SCRIPT).spawn().unwrap();
472        let mut reader = BufReader::new(ProcReader::from_stdout(child.id()));
473
474        thread::sleep(Duration::from_secs(4));
475
476        let mut line = String::new();
477        let _ = reader.read_to_string(&mut line);
478        assert_eq!("aaa\nbbb\nccc\n", line);
479    }
480
481    #[test]
482    fn test_short_array() {
483        let child = Command::new("perl").arg("-e").arg(SCRIPT).spawn().unwrap();
484        let mut reader = ProcReader::from_stdout(child.id());
485
486        thread::sleep(Duration::from_secs(4));
487
488        let mut buf = [0; 10];
489        let _ = reader.read_exact(&mut buf);
490        assert_eq!("aaa\nbbb\ncc", String::from_utf8(buf.to_vec()).unwrap());
491    }
492
493    #[test]
494    fn test_kill() {
495        let mut child = Command::new("perl").arg("-e").arg(SCRIPT).spawn().unwrap();
496        let mut reader = ProcReader::from_stdout(child.id());
497        let _ = child.kill();
498
499        thread::sleep(Duration::from_secs(4));
500
501        let mut buf = [0; 10];
502        let ret = reader.read_exact(&mut buf);
503        assert_eq!(
504            &format!("{:?}", ret)[0..60],
505            "Err(Custom { kind: Other, error: \"failed to access process ("
506        );
507    }
508
509    #[test]
510    fn test_kill2() {
511        let mut child = Command::new("perl").arg("-e").arg(SCRIPT).spawn().unwrap();
512        let mut reader = ProcReader::from_stdout(child.id());
513
514        thread::sleep(Duration::from_secs(2));
515        let _ = child.kill();
516        thread::sleep(Duration::from_secs(2));
517
518        let mut buf = [0; 10];
519        let ret = reader.read_exact(&mut buf);
520        assert_eq!(
521            &format!("{:?}", ret)[0..60],
522            "Err(Error { kind: UnexpectedEof, message: \"failed to fill wh"
523        );
524    }
525
526    #[test]
527    fn test_stderr() {
528        let child = Command::new("perl")
529            .arg("-e")
530            .arg(SCRIPT_WITH_ERR)
531            .spawn()
532            .unwrap();
533        let mut reader = BufReader::new(ProcReader::from_stderr(child.id()));
534
535        thread::sleep(Duration::from_secs(2));
536
537        let mut line = String::new();
538        let _ = reader.read_to_string(&mut line);
539        assert_eq!("eee\n", line);
540    }
541
542    #[test]
543    fn test_both() {
544        let child = Command::new("perl")
545            .arg("-e")
546            .arg(SCRIPT_WITH_ERR)
547            .spawn()
548            .unwrap();
549        let mut reader = BufReader::new(ProcReader::from_stdany(child.id()));
550
551        thread::sleep(Duration::from_secs(2));
552
553        let mut line = String::new();
554        let _ = reader.read_to_string(&mut line);
555        assert_eq!("aaa\neee\n", line);
556    }
557
558    #[test]
559    fn test_stdout_without_redirect() {
560        let child = Command::new("sh")
561            .arg("-c")
562            .arg(SCRIPT_REDIRECT)
563            .spawn()
564            .unwrap();
565        let mut reader = BufReader::new(ProcReader::from_stdout(child.id()));
566
567        thread::sleep(Duration::from_secs(4));
568
569        let mut line = String::new();
570        let _ = reader.read_to_string(&mut line);
571        assert_eq!("aaa\nbbb\n", line);
572    }
573
574    #[test]
575    #[cfg(target_arch = "x86_64")]
576    fn test_stdout_with_redirect() {
577        let child = Command::new("sh")
578            .arg("-c")
579            .arg(SCRIPT_REDIRECT)
580            .spawn()
581            .unwrap();
582        let mut reader = BufReader::new(ProcReader::from_stdout(child.id()).with_redirect());
583
584        thread::sleep(Duration::from_secs(4));
585
586        let mut line = String::new();
587        let _ = reader.read_to_string(&mut line);
588        assert_eq!("aaa\nccc\n", line);
589    }
590
591    #[test]
592    fn test_stderr_without_redirect() {
593        let child = Command::new("sh")
594            .arg("-c")
595            .arg(SCRIPT_REDIRECT)
596            .spawn()
597            .unwrap();
598        let mut reader = BufReader::new(ProcReader::from_stderr(child.id()));
599
600        thread::sleep(Duration::from_secs(4));
601
602        let mut line = String::new();
603        let _ = reader.read_to_string(&mut line);
604        assert_eq!("ccc\nddd\n", line);
605    }
606
607    #[test]
608    #[cfg(target_arch = "x86_64")]
609    fn test_stderr_with_redirect() {
610        let child = Command::new("sh")
611            .arg("-c")
612            .arg(SCRIPT_REDIRECT)
613            .spawn()
614            .unwrap();
615        let mut reader = BufReader::new(ProcReader::from_stderr(child.id()).with_redirect());
616
617        thread::sleep(Duration::from_secs(4));
618
619        let mut line = String::new();
620        let _ = reader.read_to_string(&mut line);
621        assert_eq!("bbb\nddd\n", line);
622    }
623}
proc_reader/lib.rs

proc_reader/
lib.rs
