1 use std::convert::TryFrom;
2 use std::io;
3 use std::sync::Once;
4 
5 use crate::signal::registry::{globals, EventId, EventInfo, Init, Storage};
6 use crate::signal::RxFuture;
7 
8 use winapi::shared::minwindef::{BOOL, DWORD, FALSE, TRUE};
9 use winapi::um::consoleapi::SetConsoleCtrlHandler;
10 use winapi::um::wincon::{CTRL_BREAK_EVENT, CTRL_C_EVENT};
11 
ctrl_c() -> io::Result<RxFuture>12 pub(super) fn ctrl_c() -> io::Result<RxFuture> {
13     new(CTRL_C_EVENT)
14 }
15 
ctrl_break() -> io::Result<RxFuture>16 pub(super) fn ctrl_break() -> io::Result<RxFuture> {
17     new(CTRL_BREAK_EVENT)
18 }
19 
new(signum: DWORD) -> io::Result<RxFuture>20 fn new(signum: DWORD) -> io::Result<RxFuture> {
21     global_init()?;
22     let rx = globals().register_listener(signum as EventId);
23     Ok(RxFuture::new(rx))
24 }
25 
26 #[derive(Debug)]
27 pub(crate) struct OsStorage {
28     ctrl_c: EventInfo,
29     ctrl_break: EventInfo,
30 }
31 
32 impl Init for OsStorage {
init() -> Self33     fn init() -> Self {
34         Self {
35             ctrl_c: EventInfo::default(),
36             ctrl_break: EventInfo::default(),
37         }
38     }
39 }
40 
41 impl Storage for OsStorage {
event_info(&self, id: EventId) -> Option<&EventInfo>42     fn event_info(&self, id: EventId) -> Option<&EventInfo> {
43         match DWORD::try_from(id) {
44             Ok(CTRL_C_EVENT) => Some(&self.ctrl_c),
45             Ok(CTRL_BREAK_EVENT) => Some(&self.ctrl_break),
46             _ => None,
47         }
48     }
49 
for_each<'a, F>(&'a self, mut f: F) where F: FnMut(&'a EventInfo),50     fn for_each<'a, F>(&'a self, mut f: F)
51     where
52         F: FnMut(&'a EventInfo),
53     {
54         f(&self.ctrl_c);
55         f(&self.ctrl_break);
56     }
57 }
58 
59 #[derive(Debug)]
60 pub(crate) struct OsExtraData {}
61 
62 impl Init for OsExtraData {
init() -> Self63     fn init() -> Self {
64         Self {}
65     }
66 }
67 
global_init() -> io::Result<()>68 fn global_init() -> io::Result<()> {
69     static INIT: Once = Once::new();
70 
71     let mut init = None;
72 
73     INIT.call_once(|| unsafe {
74         let rc = SetConsoleCtrlHandler(Some(handler), TRUE);
75         let ret = if rc == 0 {
76             Err(io::Error::last_os_error())
77         } else {
78             Ok(())
79         };
80 
81         init = Some(ret);
82     });
83 
84     init.unwrap_or_else(|| Ok(()))
85 }
86 
handler(ty: DWORD) -> BOOL87 unsafe extern "system" fn handler(ty: DWORD) -> BOOL {
88     let globals = globals();
89     globals.record_event(ty as EventId);
90 
91     // According to https://docs.microsoft.com/en-us/windows/console/handlerroutine
92     // the handler routine is always invoked in a new thread, thus we don't
93     // have the same restrictions as in Unix signal handlers, meaning we can
94     // go ahead and perform the broadcast here.
95     if globals.broadcast() {
96         TRUE
97     } else {
98         // No one is listening for this notification any more
99         // let the OS fire the next (possibly the default) handler.
100         FALSE
101     }
102 }
103 
104 #[cfg(all(test, not(loom)))]
105 mod tests {
106     use super::*;
107     use crate::runtime::Runtime;
108 
109     use tokio_test::{assert_ok, assert_pending, assert_ready_ok, task};
110 
111     #[test]
ctrl_c()112     fn ctrl_c() {
113         let rt = rt();
114         let _enter = rt.enter();
115 
116         let mut ctrl_c = task::spawn(crate::signal::ctrl_c());
117 
118         assert_pending!(ctrl_c.poll());
119 
120         // Windows doesn't have a good programmatic way of sending events
121         // like sending signals on Unix, so we'll stub out the actual OS
122         // integration and test that our handling works.
123         unsafe {
124             super::handler(CTRL_C_EVENT);
125         }
126 
127         assert_ready_ok!(ctrl_c.poll());
128     }
129 
130     #[test]
ctrl_break()131     fn ctrl_break() {
132         let rt = rt();
133 
134         rt.block_on(async {
135             let mut ctrl_break = assert_ok!(crate::signal::windows::ctrl_break());
136 
137             // Windows doesn't have a good programmatic way of sending events
138             // like sending signals on Unix, so we'll stub out the actual OS
139             // integration and test that our handling works.
140             unsafe {
141                 super::handler(CTRL_BREAK_EVENT);
142             }
143 
144             ctrl_break.recv().await.unwrap();
145         });
146     }
147 
rt() -> Runtime148     fn rt() -> Runtime {
149         crate::runtime::Builder::new_current_thread()
150             .build()
151             .unwrap()
152     }
153 }
154