#
cargo_util/
read2.rs

    1pub use self::imp::read2;
    2
    3#[cfg(unix)]
    4mod imp {
    5    use libc::{F_GETFL, F_SETFL, O_NONBLOCK, c_int, fcntl};
    6    use std::io;
    7    use std::io::prelude::*;
    8    use std::mem;
    9    use std::os::unix::prelude::*;
    10    use std::process::{ChildStderr, ChildStdout};
    11
    12    fn set_nonblock(fd: c_int) -> io::Result<()> {
    13        let flags = unsafe { fcntl(fd, F_GETFL) };
    14        if flags == -1 || unsafe { fcntl(fd, F_SETFL, flags | O_NONBLOCK) } == -1 {
    15            return Err(io::Error::last_os_error());
    16        }
    17        Ok(())
    18    }
    19
    20    pub fn read2(
    21        mut out_pipe: ChildStdout,
    22        mut err_pipe: ChildStderr,
    23        data: &mut dyn FnMut(bool, &mut Vec<u8>, bool),
    24    ) -> io::Result<()> {
    25        set_nonblock(out_pipe.as_raw_fd())?;
    26        set_nonblock(err_pipe.as_raw_fd())?;
    27
    28        let mut out_done = false;
    29        let mut err_done = false;
    30        let mut out = Vec::new();
    31        let mut err = Vec::new();
    32
    33        let mut fds: [libc::pollfd; 2] = unsafe { mem::zeroed() };
    34        fds[0].fd = out_pipe.as_raw_fd();
    35        fds[0].events = libc::POLLIN;
    36        fds[1].fd = err_pipe.as_raw_fd();
    37        fds[1].events = libc::POLLIN;
    38        let mut nfds = 2;
    39        let mut errfd = 1;
    40
    41        while nfds > 0 {
    42            // wait for either pipe to become readable using `poll`
    43            let r = unsafe { libc::poll(fds.as_mut_ptr(), nfds, -1) };
    44            if r == -1 {
    45                let err = io::Error::last_os_error();
    46                if err.kind() == io::ErrorKind::Interrupted {
    47                    continue;
    48                }
    49                return Err(err);
    50            }
    51
    52            // Read as much as we can from each pipe, ignoring EWOULDBLOCK or
    53            // EAGAIN. If we hit EOF, then this will happen because the underlying
    54            // reader will return Ok(0), in which case we'll see `Ok` ourselves. In
    55            // this case we flip the other fd back into blocking mode and read
    56            // whatever's leftover on that file descriptor.
    57            let handle = |res: io::Result<_>| match res {
    58                Ok(_) => Ok(true),
    59                Err(e) => {
    60                    if e.kind() == io::ErrorKind::WouldBlock {
    61                        Ok(false)
    62                    } else {
    63                        Err(e)
    64                    }
    65                }
    66            };
    67            if !err_done && fds[errfd].revents != 0 && handle(err_pipe.read_to_end(&mut err))? {
    68                err_done = true;
    69                nfds -= 1;
    70            }
    71            data(false, &mut err, err_done);
    72            if !out_done && fds[0].revents != 0 && handle(out_pipe.read_to_end(&mut out))? {
    73                out_done = true;
    74                fds[0].fd = err_pipe.as_raw_fd();
    75                errfd = 0;
    76                nfds -= 1;
    77            }
    78            data(true, &mut out, out_done);
    79        }
    80        Ok(())
    81    }
    82}
    83
    84#[cfg(windows)]
    85mod imp {
    86    use std::io;
    87    use std::os::windows::prelude::*;
    88    use std::process::{ChildStderr, ChildStdout};
    89    use std::slice;
    90
    91    use miow::Overlapped;
    92    use miow::iocp::{CompletionPort, CompletionStatus};
    93    use miow::pipe::NamedPipe;
    94    use windows_sys::Win32::Foundation::ERROR_BROKEN_PIPE;
    95
    96    struct Pipe<'a> {
    97        dst: &'a mut Vec<u8>,
    98        overlapped: Overlapped,
    99        pipe: NamedPipe,
    100        done: bool,
    101    }
    102
    103    pub fn read2(
    104        out_pipe: ChildStdout,
    105        err_pipe: ChildStderr,
    106        data: &mut dyn FnMut(bool, &mut Vec<u8>, bool),
    107    ) -> io::Result<()> {
    108        let mut out = Vec::new();
    109        let mut err = Vec::new();
    110
    111        let port = CompletionPort::new(1)?;
    112        port.add_handle(0, &out_pipe)?;
    113        port.add_handle(1, &err_pipe)?;
    114
    115        unsafe {
    116            let mut out_pipe = Pipe::new(out_pipe, &mut out);
    117            let mut err_pipe = Pipe::new(err_pipe, &mut err);
    118
    119            out_pipe.read()?;
    120            err_pipe.read()?;
    121
    122            let mut status = [CompletionStatus::zero(), CompletionStatus::zero()];
    123
    124            while !out_pipe.done || !err_pipe.done {
    125                for status in port.get_many(&mut status, None)? {
    126                    if status.token() == 0 {
    127                        out_pipe.complete(status);
    128                        data(true, out_pipe.dst, out_pipe.done);
    129                        out_pipe.read()?;
    130                    } else {
    131                        err_pipe.complete(status);
    132                        data(false, err_pipe.dst, err_pipe.done);
    133                        err_pipe.read()?;
    134                    }
    135                }
    136            }
    137
    138            Ok(())
    139        }
    140    }
    141
    142    impl<'a> Pipe<'a> {
    143        unsafe fn new<P: IntoRawHandle>(p: P, dst: &'a mut Vec<u8>) -> Pipe<'a> {
    144            // SAFETY: Handle must be owned, open, and closeable with CloseHandle.
    145            let pipe = unsafe { NamedPipe::from_raw_handle(p.into_raw_handle()) };
    146            Pipe {
    147                dst,
    148                pipe,
    149                overlapped: Overlapped::zero(),
    150                done: false,
    151            }
    152        }
    153
    154        unsafe fn read(&mut self) -> io::Result<()> {
    155            let dst = unsafe { slice_to_end(self.dst) };
    156            // SAFETY: The buffer must be valid until the end of the I/O,
    157            // which is handled in `read2`.
    158            match unsafe { self.pipe.read_overlapped(dst, self.overlapped.raw()) } {
    159                Ok(_) => Ok(()),
    160                Err(e) => {
    161                    if e.raw_os_error() == Some(ERROR_BROKEN_PIPE as i32) {
    162                        self.done = true;
    163                        Ok(())
    164                    } else {
    165                        Err(e)
    166                    }
    167                }
    168            }
    169        }
    170
    171        unsafe fn complete(&mut self, status: &CompletionStatus) {
    172            let prev = self.dst.len();
    173            unsafe { self.dst.set_len(prev + status.bytes_transferred() as usize) };
    174            if status.bytes_transferred() == 0 {
    175                self.done = true;
    176            }
    177        }
    178    }
    179
    180    unsafe fn slice_to_end(v: &mut Vec<u8>) -> &mut [u8] {
    181        if v.capacity() == 0 {
    182            v.reserve(16);
    183        }
    184        if v.capacity() == v.len() {
    185            v.reserve(1);
    186        }
    187        unsafe { slice::from_raw_parts_mut(v.as_mut_ptr().add(v.len()), v.capacity() - v.len()) }
    188    }
    189}