Skip to main content

compio_fs\stdio/
windows.rs

1use std::{
2    io::{self, IsTerminal, Read, Write},
3    os::windows::io::{AsRawHandle, BorrowedHandle, RawHandle},
4    sync::OnceLock,
5    task::Poll,
6};
7
8use compio_buf::{BufResult, IntoInner, IoBuf, IoBufMut};
9use compio_driver::{
10    AsFd, AsRawFd, BorrowedFd, BufferRef, OpCode, OpType, RawFd, ResultTakeBuffer, SharedFd,
11    op::{BufResultExt, Read as OpRead, ReadManaged, Write as OpWrite},
12};
13use compio_io::{AsyncRead, AsyncReadManaged, AsyncReadMulti, AsyncWrite};
14use compio_runtime::Runtime;
15use futures_util::{Stream, StreamExt};
16use windows_sys::Win32::System::IO::OVERLAPPED;
17
18#[cfg(doc)]
19use super::{stderr, stdin, stdout};
20
21struct StdRead<R: Read, B: IoBufMut> {
22    reader: R,
23    buffer: B,
24}
25
26impl<R: Read, B: IoBufMut> StdRead<R, B> {
27    pub fn new(reader: R, buffer: B) -> Self {
28        Self { reader, buffer }
29    }
30}
31
32unsafe impl<R: Read, B: IoBufMut> OpCode for StdRead<R, B> {
33    type Control = ();
34
35    unsafe fn init(&mut self, _: &mut Self::Control) {}
36
37    fn op_type(&self, _: &Self::Control) -> OpType {
38        OpType::Blocking
39    }
40
41    unsafe fn operate(
42        &mut self,
43        _: &mut Self::Control,
44        _optr: *mut OVERLAPPED,
45    ) -> Poll<io::Result<usize>> {
46        #[cfg(feature = "read_buf")]
47        {
48            // SAFETY: `read_buf` only writes initialized bytes.
49            let slice = unsafe { self.buffer.as_uninit() };
50            let mut buf = io::BorrowedBuf::from(slice);
51            let mut cursor = buf.unfilled();
52            self.reader.read_buf(cursor.reborrow())?;
53            Poll::Ready(Ok(cursor.written()))
54        }
55        #[cfg(not(feature = "read_buf"))]
56        {
57            use compio_buf::IoBufMutExt;
58
59            let slice = self.buffer.ensure_init();
60            self.reader.read(slice).into()
61        }
62    }
63}
64
65impl<R: Read, B: IoBufMut> IntoInner for StdRead<R, B> {
66    type Inner = B;
67
68    fn into_inner(self) -> Self::Inner {
69        self.buffer
70    }
71}
72
73struct StdWrite<W: Write, B: IoBuf> {
74    writer: W,
75    buffer: B,
76}
77
78impl<W: Write, B: IoBuf> StdWrite<W, B> {
79    pub fn new(writer: W, buffer: B) -> Self {
80        Self { writer, buffer }
81    }
82}
83
84unsafe impl<W: Write, B: IoBuf> OpCode for StdWrite<W, B> {
85    type Control = ();
86
87    unsafe fn init(&mut self, _: &mut Self::Control) {}
88
89    fn op_type(&self, _: &Self::Control) -> OpType {
90        OpType::Blocking
91    }
92
93    unsafe fn operate(
94        &mut self,
95        _: &mut Self::Control,
96        _optr: *mut OVERLAPPED,
97    ) -> Poll<io::Result<usize>> {
98        let slice = self.buffer.as_init();
99        self.writer.write(slice).into()
100    }
101}
102
103impl<W: Write, B: IoBuf> IntoInner for StdWrite<W, B> {
104    type Inner = B;
105
106    fn into_inner(self) -> Self::Inner {
107        self.buffer
108    }
109}
110
111#[derive(Debug)]
112struct StaticFd(RawHandle);
113
114impl AsFd for StaticFd {
115    fn as_fd(&self) -> BorrowedFd<'_> {
116        // SAFETY: we only use it for console handles.
117        BorrowedFd::File(unsafe { BorrowedHandle::borrow_raw(self.0) })
118    }
119}
120
121impl AsRawFd for StaticFd {
122    fn as_raw_fd(&self) -> RawFd {
123        self.0 as _
124    }
125}
126
127static STDIN_ISATTY: OnceLock<bool> = OnceLock::new();
128
129/// A handle to the standard input stream of a process.
130///
131/// See [`stdin`].
132#[derive(Debug, Clone)]
133pub struct Stdin {
134    fd: SharedFd<StaticFd>,
135    isatty: bool,
136}
137
138impl Stdin {
139    pub(crate) fn new() -> Self {
140        let stdin = io::stdin();
141        let isatty = *STDIN_ISATTY.get_or_init(|| {
142            stdin.is_terminal()
143                || Runtime::with_current(|r| r.attach(stdin.as_raw_handle() as _)).is_err()
144        });
145        Self {
146            fd: SharedFd::new(StaticFd(stdin.as_raw_handle())),
147            isatty,
148        }
149    }
150}
151
152impl AsyncRead for Stdin {
153    async fn read<B: IoBufMut>(&mut self, buf: B) -> BufResult<usize, B> {
154        let res = if self.isatty {
155            let op = StdRead::new(io::stdin(), buf);
156            compio_runtime::submit(op).await.into_inner()
157        } else {
158            let op = OpRead::new(self.fd.clone(), buf);
159            compio_runtime::submit(op).await.into_inner()
160        };
161        unsafe { res.map_advanced() }
162    }
163}
164
165impl AsyncReadManaged for Stdin {
166    type Buffer = BufferRef;
167
168    async fn read_managed(&mut self, len: usize) -> io::Result<Option<Self::Buffer>> {
169        (&*self).read_managed(len).await
170    }
171}
172
173impl AsyncReadManaged for &Stdin {
174    type Buffer = BufferRef;
175
176    async fn read_managed(&mut self, len: usize) -> io::Result<Option<Self::Buffer>> {
177        let runtime = Runtime::current();
178        let buffer_pool = runtime.buffer_pool()?;
179        if self.isatty {
180            let buf = buffer_pool.pop()?;
181            let op = StdRead::new(io::stdin(), buf);
182            unsafe { compio_runtime::submit(op).await.take_buffer() }
183        } else {
184            let op = ReadManaged::new(self.fd.clone(), &buffer_pool, len)?;
185            unsafe { compio_runtime::submit(op).await.take_buffer() }
186        }
187    }
188}
189
190impl AsyncReadMulti for Stdin {
191    fn read_multi(&mut self, len: usize) -> impl Stream<Item = io::Result<Self::Buffer>> {
192        futures_util::stream::once(self.read_managed(len))
193            .filter_map(|res| std::future::ready(res.transpose()))
194    }
195}
196
197impl AsyncReadMulti for &Stdin {
198    fn read_multi(&mut self, len: usize) -> impl Stream<Item = io::Result<Self::Buffer>> {
199        futures_util::stream::once(self.read_managed(len))
200            .filter_map(|res| std::future::ready(res.transpose()))
201    }
202}
203
204impl AsRawFd for Stdin {
205    fn as_raw_fd(&self) -> RawFd {
206        self.fd.as_raw_fd()
207    }
208}
209
210static STDOUT_ISATTY: OnceLock<bool> = OnceLock::new();
211
212/// A handle to the standard output stream of a process.
213///
214/// See [`stdout`].
215#[derive(Debug, Clone)]
216pub struct Stdout {
217    fd: SharedFd<StaticFd>,
218    isatty: bool,
219}
220
221impl Stdout {
222    pub(crate) fn new() -> Self {
223        let stdout = io::stdout();
224        let isatty = *STDOUT_ISATTY.get_or_init(|| {
225            stdout.is_terminal()
226                || Runtime::with_current(|r| r.attach(stdout.as_raw_handle() as _)).is_err()
227        });
228        Self {
229            fd: SharedFd::new(StaticFd(stdout.as_raw_handle())),
230            isatty,
231        }
232    }
233}
234
235impl AsyncWrite for Stdout {
236    async fn write<T: IoBuf>(&mut self, buf: T) -> BufResult<usize, T> {
237        if self.isatty {
238            let op = StdWrite::new(io::stdout(), buf);
239            compio_runtime::submit(op).await.into_inner()
240        } else {
241            let op = OpWrite::new(self.fd.clone(), buf);
242            compio_runtime::submit(op).await.into_inner()
243        }
244    }
245
246    async fn flush(&mut self) -> io::Result<()> {
247        Ok(())
248    }
249
250    async fn shutdown(&mut self) -> io::Result<()> {
251        self.flush().await
252    }
253}
254
255impl AsRawFd for Stdout {
256    fn as_raw_fd(&self) -> RawFd {
257        self.fd.as_raw_fd()
258    }
259}
260
261static STDERR_ISATTY: OnceLock<bool> = OnceLock::new();
262
263/// A handle to the standard output stream of a process.
264///
265/// See [`stderr`].
266#[derive(Debug, Clone)]
267pub struct Stderr {
268    fd: SharedFd<StaticFd>,
269    isatty: bool,
270}
271
272impl Stderr {
273    pub(crate) fn new() -> Self {
274        let stderr = io::stderr();
275        let isatty = *STDERR_ISATTY.get_or_init(|| {
276            stderr.is_terminal()
277                || Runtime::with_current(|r| r.attach(stderr.as_raw_handle() as _)).is_err()
278        });
279        Self {
280            fd: SharedFd::new(StaticFd(stderr.as_raw_handle())),
281            isatty,
282        }
283    }
284}
285
286impl AsyncWrite for Stderr {
287    async fn write<T: IoBuf>(&mut self, buf: T) -> BufResult<usize, T> {
288        if self.isatty {
289            let op = StdWrite::new(io::stderr(), buf);
290            compio_runtime::submit(op).await.into_inner()
291        } else {
292            let op = OpWrite::new(self.fd.clone(), buf);
293            compio_runtime::submit(op).await.into_inner()
294        }
295    }
296
297    async fn flush(&mut self) -> io::Result<()> {
298        Ok(())
299    }
300
301    async fn shutdown(&mut self) -> io::Result<()> {
302        self.flush().await
303    }
304}
305
306impl AsRawFd for Stderr {
307    fn as_raw_fd(&self) -> RawFd {
308        self.fd.as_raw_fd()
309    }
310}