cargo/util/
read2.rs

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