openconnect_core/
command.rs

1use crate::VpnClient;
2use lazy_static::lazy_static;
3use openconnect_sys::{OC_CMD_CANCEL, OC_CMD_DETACH, OC_CMD_PAUSE, OC_CMD_STATS};
4use std::sync::{atomic::Ordering, Mutex, Weak};
5
6#[derive(Debug, Clone, Copy)]
7pub enum Command {
8    Cancel,
9    Detach,
10    Pause,
11    Stats,
12}
13
14impl From<Command> for u8 {
15    fn from(cmd: Command) -> u8 {
16        match cmd {
17            Command::Cancel => OC_CMD_CANCEL,
18            Command::Detach => OC_CMD_DETACH,
19            Command::Pause => OC_CMD_PAUSE,
20            Command::Stats => OC_CMD_STATS,
21        }
22    }
23}
24
25pub trait CmdPipe {
26    fn set_sock_block(&self, cmd_fd: i32);
27    fn send_command(&self, cmd: Command);
28}
29
30impl CmdPipe for VpnClient {
31    fn set_sock_block(&self, cmd_fd: i32) {
32        #[cfg(not(target_os = "windows"))]
33        {
34            unsafe {
35                libc::fcntl(
36                    cmd_fd,
37                    libc::F_SETFL,
38                    libc::fcntl(cmd_fd, libc::F_GETFL) & !libc::O_NONBLOCK,
39                );
40            }
41        }
42
43        #[cfg(target_os = "windows")]
44        {
45            let mut mode: u32 = 0;
46            unsafe {
47                windows_sys::Win32::Networking::WinSock::ioctlsocket(
48                    cmd_fd as usize,
49                    windows_sys::Win32::Networking::WinSock::FIONBIO,
50                    &mut mode,
51                );
52            }
53        }
54    }
55    fn send_command(&self, cmd: Command) {
56        let cmd: u8 = cmd.into();
57
58        #[cfg(not(target_os = "windows"))]
59        {
60            let cmd_fd = self.cmd_fd.load(Ordering::SeqCst);
61            if cmd != 0 && cmd_fd >= 0 {
62                let ret = unsafe { libc::write(cmd_fd, std::ptr::from_ref(&cmd) as *const _, 1) };
63
64                if ret < 0 {
65                    // TODO: log error
66                }
67            }
68        }
69
70        #[cfg(target_os = "windows")]
71        {
72            let cmd_fd = self.cmd_fd.load(Ordering::SeqCst);
73            if cmd_fd >= 0 {
74                let ret = {
75                    unsafe {
76                        windows_sys::Win32::Networking::WinSock::send(
77                            cmd_fd as usize,
78                            std::ptr::from_ref(&cmd) as *const _,
79                            1,
80                            0,
81                        )
82                    }
83                };
84
85                if ret < 0 {
86                    // TODO: log error
87                }
88            }
89        }
90    }
91}
92
93pub struct SignalHandle {
94    client: Mutex<Weak<VpnClient>>,
95}
96
97lazy_static! {
98    pub static ref SIGNAL_HANDLE: SignalHandle = {
99        let sig_handle = SignalHandle {
100            client: Mutex::new(Weak::new()),
101        };
102        sig_handle.set_sig_handler();
103        sig_handle
104    };
105}
106
107impl SignalHandle {
108    /// Set the current client singleton to the given client.
109    /// This is used when signal handler is called to send command to the client.
110    pub fn update_client_singleton(&self, client: Weak<VpnClient>) {
111        let saved_client = self.client.lock();
112        if let Ok(mut saved_client) = saved_client {
113            *saved_client = client;
114        }
115    }
116
117    /// Set the signal handler for the current process.
118    fn set_sig_handler(&self) {
119        #[cfg(not(target_os = "windows"))]
120        {
121            use signal_hook::{
122                consts::{SIGHUP, SIGINT, SIGTERM, SIGUSR1, SIGUSR2},
123                iterator::Signals,
124            };
125
126            let mut signals = Signals::new([SIGINT, SIGTERM, SIGHUP, SIGUSR1, SIGUSR2])
127                .expect("Failed to register signal handler");
128
129            std::thread::spawn(move || {
130                for sig in signals.forever() {
131                    let cmd = match sig {
132                        SIGINT | SIGTERM => {
133                            println!("Received SIGINT or SIGTERM");
134                            Command::Cancel
135                        }
136                        SIGHUP => {
137                            println!("Received SIGHUP");
138                            Command::Detach
139                        }
140                        SIGUSR2 => {
141                            println!("Received SIGUSR2");
142                            Command::Pause
143                        }
144                        SIGUSR1 => {
145                            println!("Received SIGUSR1");
146                            Command::Stats
147                        }
148                        _ => {
149                            println!("Received unknown signal");
150                            unreachable!()
151                        }
152                    };
153
154                    {
155                        let this = SIGNAL_HANDLE
156                            .client
157                            .lock()
158                            .ok()
159                            .and_then(|this| this.upgrade());
160
161                        if let Some(this) = this {
162                            this.send_command(cmd);
163                        }
164                    }
165
166                    if sig == SIGINT || sig == SIGTERM {
167                        // Exit the signal handler thread since the process is going to exit
168                        break;
169                    }
170                }
171            });
172        }
173
174        #[cfg(target_os = "windows")]
175        {
176            use windows_sys::Win32::System::Console::{
177                SetConsoleCtrlHandler, CTRL_BREAK_EVENT, CTRL_CLOSE_EVENT, CTRL_C_EVENT,
178                CTRL_LOGOFF_EVENT, CTRL_SHUTDOWN_EVENT,
179            };
180
181            unsafe extern "system" fn console_control_handle(dw_ctrl_type: u32) -> i32 {
182                let cmd = match dw_ctrl_type {
183                    CTRL_C_EVENT | CTRL_CLOSE_EVENT | CTRL_LOGOFF_EVENT | CTRL_SHUTDOWN_EVENT => {
184                        Command::Cancel
185                    }
186                    CTRL_BREAK_EVENT => Command::Detach,
187                    _ => unreachable!(),
188                };
189
190                {
191                    let this = SIGNAL_HANDLE
192                        .client
193                        .lock()
194                        .ok()
195                        .and_then(|this| this.upgrade());
196
197                    if let Some(this) = this {
198                        this.send_command(cmd);
199                    }
200                }
201
202                1
203            }
204
205            unsafe {
206                SetConsoleCtrlHandler(Some(console_control_handle), 1);
207            }
208        }
209    }
210}