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