#
cargo_credential/
stdio.rs

    1use std::{fs::File, io::Error};
    2
    3/// Reset stdin and stdout to the attached console / tty for the duration of the closure.
    4/// If no console is available, stdin and stdout will be redirected to null.
    5pub fn stdin_stdout_to_console<F, T>(f: F) -> Result<T, Error>
    6where
    7    F: FnOnce() -> T,
    8{
    9    let open_write = |f| std::fs::OpenOptions::new().write(true).open(f);
    10
    11    let mut stdin = File::open(imp::IN_DEVICE).or_else(|_| File::open(imp::NULL_DEVICE))?;
    12    let mut stdout = open_write(imp::OUT_DEVICE).or_else(|_| open_write(imp::NULL_DEVICE))?;
    13
    14    let _stdin_guard = imp::ReplacementGuard::new(Stdio::Stdin, &mut stdin)?;
    15    let _stdout_guard = imp::ReplacementGuard::new(Stdio::Stdout, &mut stdout)?;
    16    Ok(f())
    17}
    18
    19enum Stdio {
    20    Stdin,
    21    Stdout,
    22}
    23
    24#[cfg(windows)]
    25mod imp {
    26    use super::Stdio;
    27    use std::{fs::File, io::Error, os::windows::prelude::AsRawHandle};
    28    use windows_sys::Win32::{
    29        Foundation::{HANDLE, INVALID_HANDLE_VALUE},
    30        System::Console::{
    31            GetStdHandle, STD_HANDLE, STD_INPUT_HANDLE, STD_OUTPUT_HANDLE, SetStdHandle,
    32        },
    33    };
    34    pub const OUT_DEVICE: &str = "CONOUT$";
    35    pub const IN_DEVICE: &str = "CONIN$";
    36    pub const NULL_DEVICE: &str = "NUL";
    37
    38    /// Restores previous stdio when dropped.
    39    pub struct ReplacementGuard {
    40        std_handle: STD_HANDLE,
    41        previous: HANDLE,
    42    }
    43
    44    impl ReplacementGuard {
    45        pub(super) fn new(stdio: Stdio, replacement: &mut File) -> Result<ReplacementGuard, Error> {
    46            let std_handle = match stdio {
    47                Stdio::Stdin => STD_INPUT_HANDLE,
    48                Stdio::Stdout => STD_OUTPUT_HANDLE,
    49            };
    50
    51            let previous;
    52            unsafe {
    53                // Make a copy of the current handle
    54                previous = GetStdHandle(std_handle);
    55                if previous == INVALID_HANDLE_VALUE {
    56                    return Err(std::io::Error::last_os_error());
    57                }
    58
    59                // Replace stdin with the replacement handle
    60                if SetStdHandle(std_handle, replacement.as_raw_handle() as HANDLE) == 0 {
    61                    return Err(std::io::Error::last_os_error());
    62                }
    63            }
    64
    65            Ok(ReplacementGuard {
    66                previous,
    67                std_handle,
    68            })
    69        }
    70    }
    71
    72    impl Drop for ReplacementGuard {
    73        fn drop(&mut self) {
    74            unsafe {
    75                // Put previous handle back in to stdin
    76                SetStdHandle(self.std_handle, self.previous);
    77            }
    78        }
    79    }
    80}
    81
    82#[cfg(unix)]
    83mod imp {
    84    use super::Stdio;
    85    use libc::{STDIN_FILENO, STDOUT_FILENO, close, dup, dup2};
    86    use std::{fs::File, io::Error, os::fd::AsRawFd};
    87    pub const IN_DEVICE: &str = "/dev/tty";
    88    pub const OUT_DEVICE: &str = "/dev/tty";
    89    pub const NULL_DEVICE: &str = "/dev/null";
    90
    91    /// Restores previous stdio when dropped.
    92    pub struct ReplacementGuard {
    93        std_fileno: i32,
    94        previous: i32,
    95    }
    96
    97    impl ReplacementGuard {
    98        pub(super) fn new(stdio: Stdio, replacement: &mut File) -> Result<ReplacementGuard, Error> {
    99            let std_fileno = match stdio {
    100                Stdio::Stdin => STDIN_FILENO,
    101                Stdio::Stdout => STDOUT_FILENO,
    102            };
    103
    104            let previous;
    105            unsafe {
    106                // Duplicate the existing stdin file to a new descriptor
    107                previous = dup(std_fileno);
    108                if previous == -1 {
    109                    return Err(std::io::Error::last_os_error());
    110                }
    111                // Replace stdin with the replacement file
    112                if dup2(replacement.as_raw_fd(), std_fileno) == -1 {
    113                    return Err(std::io::Error::last_os_error());
    114                }
    115            }
    116
    117            Ok(ReplacementGuard {
    118                previous,
    119                std_fileno,
    120            })
    121        }
    122    }
    123
    124    impl Drop for ReplacementGuard {
    125        fn drop(&mut self) {
    126            unsafe {
    127                // Put previous file back in to stdin
    128                dup2(self.previous, self.std_fileno);
    129                // Close the file descriptor we used as a backup
    130                close(self.previous);
    131            }
    132        }
    133    }
    134}
    135
    136#[cfg(test)]
    137mod test {
    138    use std::fs::OpenOptions;
    139    use std::io::{Seek, Write};
    140
    141    use super::Stdio;
    142    use super::imp::ReplacementGuard;
    143
    144    #[test]
    145    fn stdin() {
    146        let tempdir = snapbox::dir::DirRoot::mutable_temp().unwrap();
    147        let file = tempdir.path().unwrap().join("stdin");
    148        let mut file = OpenOptions::new()
    149            .read(true)
    150            .write(true)
    151            .create(true)
    152            .open(file)
    153            .unwrap();
    154
    155        writeln!(&mut file, "hello").unwrap();
    156        file.seek(std::io::SeekFrom::Start(0)).unwrap();
    157        {
    158            let _guard = ReplacementGuard::new(Stdio::Stdin, &mut file).unwrap();
    159            let line = std::io::stdin().lines().next().unwrap().unwrap();
    160            assert_eq!(line, "hello");
    161        }
    162    }
    163}