1 // Copyright (c) 2017 CtrlC developers
2 // Licensed under the Apache License, Version 2.0
3 // <LICENSE-APACHE or
4 // http://www.apache.org/licenses/LICENSE-2.0> or the MIT
5 // license <LICENSE-MIT or http://opensource.org/licenses/MIT>,
6 // at your option. All files in the project carrying such
7 // notice may not be copied, modified, or distributed except
8 // according to those terms.
9
10 #[cfg(unix)]
11 mod platform {
12 use std::io;
13
setup() -> io::Result<()>14 pub unsafe fn setup() -> io::Result<()> {
15 Ok(())
16 }
17
cleanup() -> io::Result<()>18 pub unsafe fn cleanup() -> io::Result<()> {
19 Ok(())
20 }
21
raise_ctrl_c()22 pub unsafe fn raise_ctrl_c() {
23 nix::sys::signal::raise(nix::sys::signal::SIGINT).unwrap();
24 }
25
print(fmt: ::std::fmt::Arguments)26 pub unsafe fn print(fmt: ::std::fmt::Arguments) {
27 use self::io::Write;
28 let stdout = ::std::io::stdout();
29 stdout.lock().write_fmt(fmt).unwrap();
30 }
31 }
32
33 #[cfg(windows)]
34 mod platform {
35 use std::io;
36 use std::ptr;
37
38 use winapi::shared::minwindef::DWORD;
39 use winapi::shared::ntdef::{CHAR, HANDLE};
40 use winapi::um::consoleapi::{AllocConsole, GetConsoleMode};
41 use winapi::um::fileapi::WriteFile;
42 use winapi::um::handleapi::INVALID_HANDLE_VALUE;
43 use winapi::um::processenv::{GetStdHandle, SetStdHandle};
44 use winapi::um::winbase::{STD_ERROR_HANDLE, STD_OUTPUT_HANDLE};
45 use winapi::um::wincon::{AttachConsole, FreeConsole, GenerateConsoleCtrlEvent};
46
47 /// Stores a piped stdout handle or a cache that gets
48 /// flushed when we reattached to the old console.
49 enum Output {
50 Pipe(HANDLE),
51 Cached(Vec<u8>),
52 }
53
54 static mut OLD_OUT: *mut Output = 0 as *mut Output;
55
56 impl io::Write for Output {
write(&mut self, buf: &[u8]) -> io::Result<usize>57 fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
58 match *self {
59 Output::Pipe(handle) => unsafe {
60 use winapi::shared::ntdef::VOID;
61
62 let mut n = 0u32;
63 if WriteFile(
64 handle,
65 buf.as_ptr() as *const VOID,
66 buf.len() as DWORD,
67 &mut n as *mut DWORD,
68 ptr::null_mut(),
69 ) == 0
70 {
71 Err(io::Error::last_os_error())
72 } else {
73 Ok(n as usize)
74 }
75 },
76 Output::Cached(ref mut s) => s.write(buf),
77 }
78 }
79
flush(&mut self) -> io::Result<()>80 fn flush(&mut self) -> io::Result<()> {
81 Ok(())
82 }
83 }
84
85 impl Output {
86 /// Stores current piped stdout or creates a new output cache that will
87 /// be written to stdout at a later time.
new() -> io::Result<Output>88 fn new() -> io::Result<Output> {
89 unsafe {
90 let stdout = GetStdHandle(STD_OUTPUT_HANDLE);
91 if stdout.is_null() || stdout == INVALID_HANDLE_VALUE {
92 return Err(io::Error::last_os_error());
93 }
94
95 let mut out = 0u32;
96 match GetConsoleMode(stdout, &mut out as *mut DWORD) {
97 0 => Ok(Output::Pipe(stdout)),
98 _ => Ok(Output::Cached(Vec::new())),
99 }
100 }
101 }
102
103 /// Set stdout/stderr and flush cache.
set_as_std(self) -> io::Result<()>104 unsafe fn set_as_std(self) -> io::Result<()> {
105 let stdout = match self {
106 Output::Pipe(h) => h,
107 Output::Cached(_) => get_stdout()?,
108 };
109
110 if SetStdHandle(STD_OUTPUT_HANDLE, stdout) == 0 {
111 return Err(io::Error::last_os_error());
112 }
113
114 if SetStdHandle(STD_ERROR_HANDLE, stdout) == 0 {
115 return Err(io::Error::last_os_error());
116 }
117
118 match self {
119 Output::Pipe(_) => Ok(()),
120 Output::Cached(ref s) => {
121 // Write cached output
122 use self::io::Write;
123 let out = io::stdout();
124 out.lock().write_all(&s[..])?;
125 Ok(())
126 }
127 }
128 }
129 }
130
get_stdout() -> io::Result<HANDLE>131 unsafe fn get_stdout() -> io::Result<HANDLE> {
132 use winapi::um::fileapi::{CreateFileA, OPEN_EXISTING};
133 use winapi::um::handleapi::INVALID_HANDLE_VALUE;
134 use winapi::um::winnt::{FILE_SHARE_WRITE, GENERIC_READ, GENERIC_WRITE};
135
136 let stdout = CreateFileA(
137 "CONOUT$\0".as_ptr() as *const CHAR,
138 GENERIC_READ | GENERIC_WRITE,
139 FILE_SHARE_WRITE,
140 ptr::null_mut(),
141 OPEN_EXISTING,
142 0,
143 ptr::null_mut(),
144 );
145
146 if stdout.is_null() || stdout == INVALID_HANDLE_VALUE {
147 Err(io::Error::last_os_error())
148 } else {
149 Ok(stdout)
150 }
151 }
152
153 /// Detach from the current console and create a new one,
154 /// We do this because GenerateConsoleCtrlEvent() sends ctrl-c events
155 /// to all processes on the same console. We want events to be received
156 /// only by our process.
157 ///
158 /// This breaks rust's stdout pre 1.18.0. Rust used to
159 /// [cache the std handles](https://github.com/rust-lang/rust/pull/40516)
160 ///
setup() -> io::Result<()>161 pub unsafe fn setup() -> io::Result<()> {
162 let old_out = Output::new()?;
163
164 if FreeConsole() == 0 {
165 return Err(io::Error::last_os_error());
166 }
167
168 if AllocConsole() == 0 {
169 return Err(io::Error::last_os_error());
170 }
171
172 // AllocConsole will not always set stdout/stderr to the to the console buffer
173 // of the new terminal.
174
175 let stdout = get_stdout()?;
176 if SetStdHandle(STD_OUTPUT_HANDLE, stdout) == 0 {
177 return Err(io::Error::last_os_error());
178 }
179
180 if SetStdHandle(STD_ERROR_HANDLE, stdout) == 0 {
181 return Err(io::Error::last_os_error());
182 }
183
184 OLD_OUT = Box::into_raw(Box::new(old_out));
185
186 Ok(())
187 }
188
189 /// Reattach to the old console.
cleanup() -> io::Result<()>190 pub unsafe fn cleanup() -> io::Result<()> {
191 if FreeConsole() == 0 {
192 return Err(io::Error::last_os_error());
193 }
194
195 if AttachConsole(winapi::um::wincon::ATTACH_PARENT_PROCESS) == 0 {
196 return Err(io::Error::last_os_error());
197 }
198
199 Box::from_raw(OLD_OUT).set_as_std()?;
200
201 Ok(())
202 }
203
204 /// This will signal the whole process group.
raise_ctrl_c()205 pub unsafe fn raise_ctrl_c() {
206 assert!(GenerateConsoleCtrlEvent(winapi::um::wincon::CTRL_C_EVENT, 0) != 0);
207 }
208
209 /// Print to both consoles, this is not thread safe.
print(fmt: ::std::fmt::Arguments)210 pub unsafe fn print(fmt: ::std::fmt::Arguments) {
211 use self::io::Write;
212 {
213 let stdout = io::stdout();
214 stdout.lock().write_fmt(fmt).unwrap();
215 }
216 {
217 assert!(!OLD_OUT.is_null());
218 (*OLD_OUT).write_fmt(fmt).unwrap();
219 }
220 }
221 }
222
test_set_handler()223 fn test_set_handler() {
224 let (tx, rx) = ::std::sync::mpsc::channel();
225 ctrlc::set_handler(move || {
226 tx.send(true).unwrap();
227 })
228 .unwrap();
229
230 unsafe {
231 platform::raise_ctrl_c();
232 }
233
234 rx.recv_timeout(::std::time::Duration::from_secs(10))
235 .unwrap();
236
237 match ctrlc::set_handler(|| {}) {
238 Err(ctrlc::Error::MultipleHandlers) => {}
239 ret => panic!("{:?}", ret),
240 }
241 }
242
243 macro_rules! run_tests {
244 ( $($test_fn:ident),* ) => {
245 unsafe {
246 platform::print(format_args!("\n"));
247 $(
248 platform::print(format_args!("test tests::{} ... ", stringify!($test_fn)));
249 $test_fn();
250 platform::print(format_args!("ok\n"));
251 )*
252 platform::print(format_args!("\n"));
253 }
254 }
255 }
256
main()257 fn main() {
258 unsafe {
259 platform::setup().unwrap();
260 }
261
262 let default = std::panic::take_hook();
263 std::panic::set_hook(Box::new(move |info| {
264 unsafe {
265 platform::cleanup().unwrap();
266 }
267 (default)(info);
268 }));
269
270 run_tests!(test_set_handler);
271
272 unsafe {
273 platform::cleanup().unwrap();
274 }
275 }
276