|
| 1 | +use 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. |
| 5 | +pub fn stdin_stdout_to_console<F, T>(f: F) -> Result<T, Error> |
| 6 | +where |
| 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 | + |
| 19 | +enum Stdio { |
| 20 | + Stdin, |
| 21 | + Stdout, |
| 22 | +} |
| 23 | + |
| 24 | +#[cfg(windows)] |
| 25 | +mod 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)] |
| 83 | +mod 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)] |
| 137 | +mod 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::path::PathFixture::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 | +} |
0 commit comments