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, SetStdHandle, STD_HANDLE, STD_INPUT_HANDLE, STD_OUTPUT_HANDLE,
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::{close, dup, dup2, STDIN_FILENO, STDOUT_FILENO};
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::imp::ReplacementGuard;
142    use super::Stdio;
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}