Add set_nonblocking method on Stdin

Author: GuillaumeGomez

src/libstd/io/stdio.rs

@@ -97,6 +97,13 @@ impl Read for StdinRaw {
         Initializer::nop()
     }
 }
+impl StdinRaw {
+    #[cfg(any(unix, windows))]
+    #[stable(since = "1.46.0", feature = "stdin_nonblocking")]
+    pub unsafe fn set_nonblocking(&self, nonblocking: bool) -> io::Result<()> {
+        self.0.set_nonblocking(nonblocking)
+    }
+}
 impl Write for StdoutRaw {
     fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
         self.0.write(buf)
@@ -365,6 +372,33 @@ impl Stdin {
     pub fn read_line(&self, buf: &mut String) -> io::Result<usize> {
         self.lock().read_line(buf)
     }
+
+    /// Set the `stdin` in non-blocking mode.
+    ///
+    /// It is useful in case you're playing with termcaps. However, please note that it is
+    /// unsafe because it'll modify the behaviour of all other stdin.
+    ///
+    /// # Example
+    ///
+    /// ```no_run
+    /// use std::io::{self, Read};
+    ///
+    /// fn main() -> io::Result<()> {
+    ///     let mut buffer = String::new();
+    ///     let stdin = io::stdin();
+    ///     let mut handle = stdin.lock();
+    ///
+    ///     handle.set_nonblocking(true)?;
+    ///     handle.read_to_string(&mut buffer)?;
+    ///
+    ///     Ok(())
+    /// }
+    /// ```
+    #[cfg(any(unix, windows))]
+    #[stable(since = "1.46.0", feature = "stdin_nonblocking")]
+    pub unsafe fn set_nonblocking(&self, nonblocking: bool) -> io::Result<()> {
+        self.lock().set_nonblocking(nonblocking)
+    }
 }
 
 #[stable(feature = "std_debug", since = "1.16.0")]
@@ -439,6 +473,38 @@ impl fmt::Debug for StdinLock<'_> {
     }
 }
 
+impl StdinLock<'_> {
+    /// Set the `stdin` in non-blocking mode.
+    ///
+    /// It is useful in case you're playing with termcaps. However, please note that it is
+    /// unsafe because it'll modify the behaviour of all other stdin.
+    ///
+    /// # Example
+    ///
+    /// ```no_run
+    /// use std::io::{self, Read};
+    ///
+    /// fn main() -> io::Result<()> {
+    ///     let mut buffer = String::new();
+    ///     let stdin = io::stdin();
+    ///     let mut handle = stdin.lock();
+    ///
+    ///     handle.set_nonblocking(true)?;
+    ///     handle.read_to_string(&mut buffer)?;
+    ///
+    ///     Ok(())
+    /// }
+    /// ```
+    #[cfg(any(unix, windows))]
+    #[stable(since = "1.46.0", feature = "stdin_nonblocking")]
+    pub unsafe fn set_nonblocking(&self, nonblocking: bool) -> io::Result<()> {
+        match self.inner.get_ref() {
+            Maybe::Real(ref stdin) => stdin.set_nonblocking(nonblocking),
+            Maybe::Fake => Ok(()),
+        }
+    }
+}
+
 /// A handle to the global standard output stream of the current process.
 ///
 /// Each handle shares a global buffer of data to be written to the standard

src/libstd/sys/unix/stdio.rs

@@ -1,5 +1,6 @@
 use crate::io::{self, IoSlice, IoSliceMut};
 use crate::mem::ManuallyDrop;
+use crate::sys::cvt;
 use crate::sys::fd::FileDesc;
 
 pub struct Stdin(());
@@ -10,6 +11,16 @@ impl Stdin {
     pub fn new() -> io::Result<Stdin> {
         Ok(Stdin(()))
     }
+    pub unsafe fn set_nonblocking(&self, nonblocking: bool) -> io::Result<()> {
+        let mut flags = libc::fcntl(libc::STDIN_FILENO, libc::F_GETFL);
+        if nonblocking {
+            flags |= libc::O_NONBLOCK;
+        } else {
+            flags &= !libc::O_NONBLOCK;
+        }
+        cvt(libc::fcntl(libc::STDIN_FILENO, libc::F_SETFL, flags))?;
+        Ok(())
+    }
 }
 
 impl io::Read for Stdin {

src/libstd/sys/windows/c.rs

@@ -669,6 +669,7 @@ if #[cfg(not(target_vendor = "uwp"))] {
     pub const HANDLE_FLAG_INHERIT: DWORD = 0x00000001;
 
     pub const TOKEN_READ: DWORD = 0x20008;
+    pub const ENABLE_LINE_INPUT: DWORD = 0x0002;
 
     extern "system" {
         #[link_name = "SystemFunction036"]
@@ -688,6 +689,8 @@ if #[cfg(not(target_vendor = "uwp"))] {
 
         pub fn GetConsoleMode(hConsoleHandle: HANDLE,
                               lpMode: LPDWORD) -> BOOL;
+        pub fn SetConsoleMode(hConsoleHandle: HANDLE,
+                              lpMode: DWORD) -> BOOL;
         // Allowed but unused by UWP
         pub fn OpenProcessToken(ProcessHandle: HANDLE,
                                 DesiredAccess: DWORD,

src/libstd/sys/windows/stdio.rs

@@ -134,6 +134,28 @@ impl Stdin {
     pub fn new() -> io::Result<Stdin> {
         Ok(Stdin { surrogate: 0 })
     }
+    pub unsafe fn set_nonblocking(&self, nonblocking: bool) -> io::Result<()> {
+        let handle = get_handle(c::STD_INPUT_HANDLE)?;
+        let mut mode = 0;
+        if c::GetConsoleMode(handle, &mut mode) == 0 {
+            return Err(io::Error::new(
+                io::ErrorKind::InvalidInput,
+                "Cannot get console mode for the given handle",
+            ));
+        }
+        if nonblocking {
+            mode &= !c::ENABLE_LINE_INPUT;
+        } else {
+            mode |= c::ENABLE_LINE_INPUT;
+        }
+        if c::SetConsoleMode(handle, mode) == 0 {
+            return Err(io::Error::new(
+                io::ErrorKind::InvalidInput,
+                "Cannot set console mode for the given handle",
+            ));
+        }
+        Ok(())
+    }
 }
 
 impl io::Read for Stdin {