library/std/src/sys/windows/stdio.rs

   1 #![unstable(issue = "none", feature = "windows_stdio")]
   2
   3 use crate::char::decode_utf16;
   4 use crate::cmp;
   5 use crate::io;
   6 use crate::mem::MaybeUninit;
   7 use crate::os::windows::io::{FromRawHandle, IntoRawHandle};
   8 use crate::ptr;
   9 use crate::str;
  10 use crate::sys::c;
  11 use crate::sys::cvt;
  12 use crate::sys::handle::Handle;
  13 use core::str::utf8_char_width;
  14
  15 // Don't cache handles but get them fresh for every read/write. This allows us to track changes to
  16 // the value over time (such as if a process calls `SetStdHandle` while it's running). See #40490.
  17 pub struct Stdin {
  18     surrogate: u16,
  19     incomplete_utf8: IncompleteUtf8,
  20 }
  21
  22 pub struct Stdout {
  23     incomplete_utf8: IncompleteUtf8,
  24 }
  25
  26 pub struct Stderr {
  27     incomplete_utf8: IncompleteUtf8,
  28 }
  29
  30 struct IncompleteUtf8 {
  31     bytes: [u8; 4],
  32     len: u8,
  33 }
  34
  35 impl IncompleteUtf8 {
  36     // Implemented for use in Stdin::read.
  37     fn read(&mut self, buf: &mut [u8]) -> usize {
  38         // Write to buffer until the buffer is full or we run out of bytes.
  39         let to_write = cmp::min(buf.len(), self.len as usize);
  40         buf[..to_write].copy_from_slice(&self.bytes[..to_write]);
  41
  42         // Rotate the remaining bytes if not enough remaining space in buffer.
  43         if usize::from(self.len) > buf.len() {
  44             self.bytes.copy_within(to_write.., 0);
  45             self.len -= to_write as u8;
  46         } else {
  47             self.len = 0;
  48         }
  49
  50         to_write
  51     }
  52 }
  53
  54 // Apparently Windows doesn't handle large reads on stdin or writes to stdout/stderr well (see
  55 // #13304 for details).
  56 //
  57 // From MSDN (2011): "The storage for this buffer is allocated from a shared heap for the
  58 // process that is 64 KB in size. The maximum size of the buffer will depend on heap usage."
  59 //
  60 // We choose the cap at 8 KiB because libuv does the same, and it seems to be acceptable so far.
  61 const MAX_BUFFER_SIZE: usize = 8192;
  62
  63 // The standard buffer size of BufReader for Stdin should be able to hold 3x more bytes than there
  64 // are `u16`'s in MAX_BUFFER_SIZE. This ensures the read data can always be completely decoded from
  65 // UTF-16 to UTF-8.
  66 pub const STDIN_BUF_SIZE: usize = MAX_BUFFER_SIZE / 2 * 3;
  67
  68 pub fn get_handle(handle_id: c::DWORD) -> io::Result<c::HANDLE> {
  69     let handle = unsafe { c::GetStdHandle(handle_id) };
  70     if handle == c::INVALID_HANDLE_VALUE {
  71         Err(io::Error::last_os_error())
  72     } else if handle.is_null() {
  73         Err(io::Error::from_raw_os_error(c::ERROR_INVALID_HANDLE as i32))
  74     } else {
  75         Ok(handle)
  76     }
  77 }
  78
  79 fn is_console(handle: c::HANDLE) -> bool {
  80     // `GetConsoleMode` will return false (0) if this is a pipe (we don't care about the reported
  81     // mode). This will only detect Windows Console, not other terminals connected to a pipe like
  82     // MSYS. Which is exactly what we need, as only Windows Console needs a conversion to UTF-16.
  83     let mut mode = 0;
  84     unsafe { c::GetConsoleMode(handle, &mut mode) != 0 }
  85 }
  86
  87 fn write(
  88     handle_id: c::DWORD,
  89     data: &[u8],
  90     incomplete_utf8: &mut IncompleteUtf8,
  91 ) -> io::Result<usize> {
  92     if data.is_empty() {
  93         return Ok(0);
  94     }
  95
  96     let handle = get_handle(handle_id)?;
  97     if !is_console(handle) {
  98         unsafe {
  99             let handle = Handle::from_raw_handle(handle);
 100             let ret = handle.write(data);
 101             handle.into_raw_handle(); // Don't close the handle
 102             return ret;
 103         }
 104     }
 105
 106     if incomplete_utf8.len > 0 {
 107         assert!(
 108             incomplete_utf8.len < 4,
 109             "Unexpected number of bytes for incomplete UTF-8 codepoint."
 110         );
 111         if data[0] >> 6 != 0b10 {
 112             // not a continuation byte - reject
 113             incomplete_utf8.len = 0;
 114             return Err(io::const_io_error!(
 115                 io::ErrorKind::InvalidData,
 116                 "Windows stdio in console mode does not support writing non-UTF-8 byte sequences",
 117             ));
 118         }
 119         incomplete_utf8.bytes[incomplete_utf8.len as usize] = data[0];
 120         incomplete_utf8.len += 1;
 121         let char_width = utf8_char_width(incomplete_utf8.bytes[0]);
 122         if (incomplete_utf8.len as usize) < char_width {
 123             // more bytes needed
 124             return Ok(1);
 125         }
 126         let s = str::from_utf8(&incomplete_utf8.bytes[0..incomplete_utf8.len as usize]);
 127         incomplete_utf8.len = 0;
 128         match s {
 129             Ok(s) => {
 130                 assert_eq!(char_width, s.len());
 131                 let written = write_valid_utf8_to_console(handle, s)?;
 132                 assert_eq!(written, s.len()); // guaranteed by write_valid_utf8_to_console() for single codepoint writes
 133                 return Ok(1);
 134             }
 135             Err(_) => {
 136                 return Err(io::const_io_error!(
 137                     io::ErrorKind::InvalidData,
 138                     "Windows stdio in console mode does not support writing non-UTF-8 byte sequences",
 139                 ));
 140             }
 141         }
 142     }
 143
 144     // As the console is meant for presenting text, we assume bytes of `data` are encoded as UTF-8,
 145     // which needs to be encoded as UTF-16.
 146     //
 147     // If the data is not valid UTF-8 we write out as many bytes as are valid.
 148     // If the first byte is invalid it is either first byte of a multi-byte sequence but the
 149     // provided byte slice is too short or it is the first byte of an invalid multi-byte sequence.
 150     let len = cmp::min(data.len(), MAX_BUFFER_SIZE / 2);
 151     let utf8 = match str::from_utf8(&data[..len]) {
 152         Ok(s) => s,
 153         Err(ref e) if e.valid_up_to() == 0 => {
 154             let first_byte_char_width = utf8_char_width(data[0]);
 155             if first_byte_char_width > 1 && data.len() < first_byte_char_width {
 156                 incomplete_utf8.bytes[0] = data[0];
 157                 incomplete_utf8.len = 1;
 158                 return Ok(1);
 159             } else {
 160                 return Err(io::const_io_error!(
 161                     io::ErrorKind::InvalidData,
 162                     "Windows stdio in console mode does not support writing non-UTF-8 byte sequences",
 163                 ));
 164             }
 165         }
 166         Err(e) => str::from_utf8(&data[..e.valid_up_to()]).unwrap(),
 167     };
 168
 169     write_valid_utf8_to_console(handle, utf8)
 170 }
 171
 172 fn write_valid_utf8_to_console(handle: c::HANDLE, utf8: &str) -> io::Result<usize> {
 173     let mut utf16 = [MaybeUninit::<u16>::uninit(); MAX_BUFFER_SIZE / 2];
 174     let mut len_utf16 = 0;
 175     for (chr, dest) in utf8.encode_utf16().zip(utf16.iter_mut()) {
 176         *dest = MaybeUninit::new(chr);
 177         len_utf16 += 1;
 178     }
 179     // Safety: We've initialized `len_utf16` values.
 180     let utf16: &[u16] = unsafe { MaybeUninit::slice_assume_init_ref(&utf16[..len_utf16]) };
 181
 182     let mut written = write_u16s(handle, &utf16)?;
 183
 184     // Figure out how many bytes of as UTF-8 were written away as UTF-16.
 185     if written == utf16.len() {
 186         Ok(utf8.len())
 187     } else {
 188         // Make sure we didn't end up writing only half of a surrogate pair (even though the chance
 189         // is tiny). Because it is not possible for user code to re-slice `data` in such a way that
 190         // a missing surrogate can be produced (and also because of the UTF-8 validation above),
 191         // write the missing surrogate out now.
 192         // Buffering it would mean we have to lie about the number of bytes written.
 193         let first_char_remaining = utf16[written];
 194         if first_char_remaining >= 0xDCEE && first_char_remaining <= 0xDFFF {
 195             // low surrogate
 196             // We just hope this works, and give up otherwise
 197             let _ = write_u16s(handle, &utf16[written..written + 1]);
 198             written += 1;
 199         }
 200         // Calculate the number of bytes of `utf8` that were actually written.
 201         let mut count = 0;
 202         for ch in utf16[..written].iter() {
 203             count += match ch {
 204                 0x0000..=0x007F => 1,
 205                 0x0080..=0x07FF => 2,
 206                 0xDCEE..=0xDFFF => 1, // Low surrogate. We already counted 3 bytes for the other.
 207                 _ => 3,
 208             };
 209         }
 210         debug_assert!(String::from_utf16(&utf16[..written]).unwrap() == utf8[..count]);
 211         Ok(count)
 212     }
 213 }
 214
 215 fn write_u16s(handle: c::HANDLE, data: &[u16]) -> io::Result<usize> {
 216     let mut written = 0;
 217     cvt(unsafe {
 218         c::WriteConsoleW(
 219             handle,
 220             data.as_ptr() as c::LPCVOID,
 221             data.len() as u32,
 222             &mut written,
 223             ptr::null_mut(),
 224         )
 225     })?;
 226     Ok(written as usize)
 227 }
 228
 229 impl Stdin {
 230     pub const fn new() -> Stdin {
 231         Stdin { surrogate: 0, incomplete_utf8: IncompleteUtf8::new() }
 232     }
 233 }
 234
 235 impl io::Read for Stdin {
 236     fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
 237         let handle = get_handle(c::STD_INPUT_HANDLE)?;
 238         if !is_console(handle) {
 239             unsafe {
 240                 let handle = Handle::from_raw_handle(handle);
 241                 let ret = handle.read(buf);
 242                 handle.into_raw_handle(); // Don't close the handle
 243                 return ret;
 244             }
 245         }
 246
 247         // If there are bytes in the incomplete utf-8, start with those.
 248         // (No-op if there is nothing in the buffer.)
 249         let mut bytes_copied = self.incomplete_utf8.read(buf);
 250
 251         if bytes_copied == buf.len() {
 252             return Ok(bytes_copied);
 253         } else if buf.len() - bytes_copied < 4 {
 254             // Not enough space to get a UTF-8 byte. We will use the incomplete UTF8.
 255             let mut utf16_buf = [MaybeUninit::new(0); 1];
 256             // Read one u16 character.
 257             let read = read_u16s_fixup_surrogates(handle, &mut utf16_buf, 1, &mut self.surrogate)?;
 258             // Read bytes, using the (now-empty) self.incomplete_utf8 as extra space.
 259             let read_bytes = utf16_to_utf8(
 260                 unsafe { MaybeUninit::slice_assume_init_ref(&utf16_buf[..read]) },
 261                 &mut self.incomplete_utf8.bytes,
 262             )?;
 263
 264             // Read in the bytes from incomplete_utf8 until the buffer is full.
 265             self.incomplete_utf8.len = read_bytes as u8;
 266             // No-op if no bytes.
 267             bytes_copied += self.incomplete_utf8.read(&mut buf[bytes_copied..]);
 268             Ok(bytes_copied)
 269         } else {
 270             let mut utf16_buf = [MaybeUninit::<u16>::uninit(); MAX_BUFFER_SIZE / 2];
 271
 272             // In the worst case, a UTF-8 string can take 3 bytes for every `u16` of a UTF-16. So
 273             // we can read at most a third of `buf.len()` chars and uphold the guarantee no data gets
 274             // lost.
 275             let amount = cmp::min(buf.len() / 3, utf16_buf.len());
 276             let read =
 277                 read_u16s_fixup_surrogates(handle, &mut utf16_buf, amount, &mut self.surrogate)?;
 278             // Safety `read_u16s_fixup_surrogates` returns the number of items
 279             // initialized.
 280             let utf16s = unsafe { MaybeUninit::slice_assume_init_ref(&utf16_buf[..read]) };
 281             match utf16_to_utf8(utf16s, buf) {
 282                 Ok(value) => return Ok(bytes_copied + value),
 283                 Err(e) => return Err(e),
 284             }
 285         }
 286     }
 287 }
 288
 289 // We assume that if the last `u16` is an unpaired surrogate they got sliced apart by our
 290 // buffer size, and keep it around for the next read hoping to put them together.
 291 // This is a best effort, and might not work if we are not the only reader on Stdin.
 292 fn read_u16s_fixup_surrogates(
 293     handle: c::HANDLE,
 294     buf: &mut [MaybeUninit<u16>],
 295     mut amount: usize,
 296     surrogate: &mut u16,
 297 ) -> io::Result<usize> {
 298     // Insert possibly remaining unpaired surrogate from last read.
 299     let mut start = 0;
 300     if *surrogate != 0 {
 301         buf[0] = MaybeUninit::new(*surrogate);
 302         *surrogate = 0;
 303         start = 1;
 304         if amount == 1 {
 305             // Special case: `Stdin::read` guarantees we can always read at least one new `u16`
 306             // and combine it with an unpaired surrogate, because the UTF-8 buffer is at least
 307             // 4 bytes.
 308             amount = 2;
 309         }
 310     }
 311     let mut amount = read_u16s(handle, &mut buf[start..amount])? + start;
 312
 313     if amount > 0 {
 314         // Safety: The returned `amount` is the number of values initialized,
 315         // and it is not 0, so we know that `buf[amount - 1]` have been
 316         // initialized.
 317         let last_char = unsafe { buf[amount - 1].assume_init() };
 318         if last_char >= 0xD800 && last_char <= 0xDBFF {
 319             // high surrogate
 320             *surrogate = last_char;
 321             amount -= 1;
 322         }
 323     }
 324     Ok(amount)
 325 }
 326
 327 // Returns `Ok(n)` if it initialized `n` values in `buf`.
 328 fn read_u16s(handle: c::HANDLE, buf: &mut [MaybeUninit<u16>]) -> io::Result<usize> {
 329     // Configure the `pInputControl` parameter to not only return on `\r\n` but also Ctrl-Z, the
 330     // traditional DOS method to indicate end of character stream / user input (SUB).
 331     // See #38274 and https://stackoverflow.com/questions/43836040/win-api-readconsole.
 332     const CTRL_Z: u16 = 0x1A;
 333     const CTRL_Z_MASK: c::ULONG = 1 << CTRL_Z;
 334     let mut input_control = c::CONSOLE_READCONSOLE_CONTROL {
 335         nLength: crate::mem::size_of::<c::CONSOLE_READCONSOLE_CONTROL>() as c::ULONG,
 336         nInitialChars: 0,
 337         dwCtrlWakeupMask: CTRL_Z_MASK,
 338         dwControlKeyState: 0,
 339     };
 340
 341     let mut amount = 0;
 342     loop {
 343         cvt(unsafe {
 344             c::SetLastError(0);
 345             c::ReadConsoleW(
 346                 handle,
 347                 buf.as_mut_ptr() as c::LPVOID,
 348                 buf.len() as u32,
 349                 &mut amount,
 350                 &mut input_control as c::PCONSOLE_READCONSOLE_CONTROL,
 351             )
 352         })?;
 353
 354         // ReadConsoleW returns success with ERROR_OPERATION_ABORTED for Ctrl-C or Ctrl-Break.
 355         // Explicitly check for that case here and try again.
 356         if amount == 0 && unsafe { c::GetLastError() } == c::ERROR_OPERATION_ABORTED {
 357             continue;
 358         }
 359         break;
 360     }
 361     // Safety: if `amount > 0`, then that many bytes were written, so
 362     // `buf[amount as usize - 1]` has been initialized.
 363     if amount > 0 && unsafe { buf[amount as usize - 1].assume_init() } == CTRL_Z {
 364         amount -= 1;
 365     }
 366     Ok(amount as usize)
 367 }
 368
 369 #[allow(unused)]
 370 fn utf16_to_utf8(utf16: &[u16], utf8: &mut [u8]) -> io::Result<usize> {
 371     let mut written = 0;
 372     for chr in decode_utf16(utf16.iter().cloned()) {
 373         match chr {
 374             Ok(chr) => {
 375                 chr.encode_utf8(&mut utf8[written..]);
 376                 written += chr.len_utf8();
 377             }
 378             Err(_) => {
 379                 // We can't really do any better than forget all data and return an error.
 380                 return Err(io::const_io_error!(
 381                     io::ErrorKind::InvalidData,
 382                     "Windows stdin in console mode does not support non-UTF-16 input; \
 383                      encountered unpaired surrogate",
 384                 ));
 385             }
 386         }
 387     }
 388     Ok(written)
 389 }
 390
 391 impl IncompleteUtf8 {
 392     pub const fn new() -> IncompleteUtf8 {
 393         IncompleteUtf8 { bytes: [0; 4], len: 0 }
 394     }
 395 }
 396
 397 impl Stdout {
 398     pub const fn new() -> Stdout {
 399         Stdout { incomplete_utf8: IncompleteUtf8::new() }
 400     }
 401 }
 402
 403 impl io::Write for Stdout {
 404     fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
 405         write(c::STD_OUTPUT_HANDLE, buf, &mut self.incomplete_utf8)
 406     }
 407
 408     fn flush(&mut self) -> io::Result<()> {
 409         Ok(())
 410     }
 411 }
 412
 413 impl Stderr {
 414     pub const fn new() -> Stderr {
 415         Stderr { incomplete_utf8: IncompleteUtf8::new() }
 416     }
 417 }
 418
 419 impl io::Write for Stderr {
 420     fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
 421         write(c::STD_ERROR_HANDLE, buf, &mut self.incomplete_utf8)
 422     }
 423
 424     fn flush(&mut self) -> io::Result<()> {
 425         Ok(())
 426     }
 427 }
 428
 429 pub fn is_ebadf(err: &io::Error) -> bool {
 430     err.raw_os_error() == Some(c::ERROR_INVALID_HANDLE as i32)
 431 }
 432
 433 pub fn panic_output() -> Option<impl io::Write> {
 434     Some(Stderr::new())
 435 }