luckyimnotanigger/tun2proxy
1
0
Fork 0
mirror of https://github.com/tun2proxy/tun2proxy.git synced 2025-04-23 15:29:10 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 3

exiting_event_handler

Browse source
This commit is contained in:
ssrlive 2023-09-25 14:11:18 +08:00
parent f67950e868
commit 7c507874e8
1 changed files with 59 additions and 51 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

68
src/tun2proxy.rs
View file

@ -208,8 +208,8 @@ pub(crate) trait ConnectionManager {
const TUN_TOKEN: Token = Token(0); const TUN_TOKEN: Token = Token(0);
const PIPE_TOKEN: Token = Token(1); const PIPE_TOKEN: Token = Token(1);
const EXIT_TOKEN_SENDER: Token = Token(2); const EXIT_TRIGGER_TOKEN: Token = Token(2);
const EXIT_TOKEN: Token = Token(99); const EXIT_TOKEN: Token = Token(10);
pub struct TunToProxy<'a> { pub struct TunToProxy<'a> {
#[cfg(any(target_os = "linux", target_os = "android"))] #[cfg(any(target_os = "linux", target_os = "android"))]
@ -230,11 +230,11 @@ pub struct TunToProxy<'a> {
#[cfg(target_family = "unix")] #[cfg(target_family = "unix")]
exit_receiver: mio::unix::pipe::Receiver, exit_receiver: mio::unix::pipe::Receiver,
#[cfg(target_family = "unix")] #[cfg(target_family = "unix")]
exit_sender: Option<mio::unix::pipe::Sender>, exit_trigger: Option<mio::unix::pipe::Sender>,
#[cfg(target_os = "windows")] #[cfg(target_os = "windows")]
exit_receiver: mio::windows::NamedPipe, exit_receiver: mio::windows::NamedPipe,
#[cfg(target_os = "windows")] #[cfg(target_os = "windows")]
exit_sender: Option<mio::windows::NamedPipe>, exit_trigger: Option<mio::windows::NamedPipe>,
} }
impl<'a> TunToProxy<'a> { impl<'a> TunToProxy<'a> {
@ -270,12 +270,12 @@ impl<'a> TunToProxy<'a> {
} }
#[cfg(target_family = "unix")] #[cfg(target_family = "unix")]
let (mut exit_sender, mut exit_receiver) = mio::unix::pipe::new()?; let (mut exit_trigger, mut exit_receiver) = mio::unix::pipe::new()?;
#[cfg(target_family = "windows")] #[cfg(target_family = "windows")]
let (mut exit_sender, mut exit_receiver) = wintuninterface::pipe()?; let (mut exit_trigger, mut exit_receiver) = wintuninterface::pipe()?;
poll.registry() poll.registry()
.register(&mut exit_sender, EXIT_TOKEN_SENDER, Interest::WRITABLE)?; .register(&mut exit_trigger, EXIT_TRIGGER_TOKEN, Interest::WRITABLE)?;
poll.registry() poll.registry()
.register(&mut exit_receiver, EXIT_TOKEN, Interest::READABLE)?; .register(&mut exit_receiver, EXIT_TOKEN, Interest::READABLE)?;
@ -310,7 +310,7 @@ impl<'a> TunToProxy<'a> {
options, options,
write_sockets: HashSet::default(), write_sockets: HashSet::default(),
exit_receiver, exit_receiver,
exit_sender: Some(exit_sender), exit_trigger: Some(exit_trigger),
}; };
Ok(tun) Ok(tun)
} }
@ -1129,16 +1129,15 @@ impl<'a> TunToProxy<'a> {
Ok(()) Ok(())
} }
pub fn run(&mut self) -> Result<(), Error> { #[cfg(any(target_os = "windows", target_os = "linux", target_os = "macos"))]
#[cfg(target_os = "windows")] fn prepare_exiting_signal_trigger(&mut self) -> Result<()> {
{ let mut exit_trigger = self.exit_trigger.take().ok_or("Already running")?;
let mut exit_sender = self.exit_sender.take().ok_or("Already running")?;
ctrlc::set_handler(move || { ctrlc::set_handler(move || {
let mut count = 0; let mut count = 0;
loop { loop {
match exit_sender.write(b"EX") { match exit_trigger.write(b"EXIT") {
Ok(_) => { Ok(_) => {
log::trace!("Send exit signal successfully"); log::trace!("Exit signal triggered successfully");
break; break;
} }
Err(err) if err.kind() == std::io::ErrorKind::WouldBlock => { Err(err) if err.kind() == std::io::ErrorKind::WouldBlock => {
@ -1157,8 +1156,13 @@ impl<'a> TunToProxy<'a> {
} }
} }
})?; })?;
Ok(())
} }
pub fn run(&mut self) -> Result<(), Error> {
#[cfg(any(target_os = "windows", target_os = "linux", target_os = "macos"))]
self.prepare_exiting_signal_trigger()?;
let mut events = Events::with_capacity(1024); let mut events = Events::with_capacity(1024);
loop { loop {
if let Err(err) = self.poll.poll(&mut events, None) { if let Err(err) = self.poll.poll(&mut events, None) {
@ -1171,24 +1175,12 @@ impl<'a> TunToProxy<'a> {
for event in events.iter() { for event in events.iter() {
match event.token() { match event.token() {
EXIT_TOKEN => { EXIT_TOKEN => {
let mut buffer = vec![0; 100]; if self.exiting_event_handler()? {
match self.exit_receiver.read(&mut buffer) {
Ok(size) => {
log::trace!("Received exit signal: {:?}", &buffer[..size]);
log::info!("Exiting tun2proxy...");
return Ok(()); return Ok(());
} }
Err(err) if err.kind() == std::io::ErrorKind::WouldBlock => {
log::trace!("Exiting reciever is ready");
} }
Err(err) => { EXIT_TRIGGER_TOKEN => {
log::error!("Exiting tun2proxy... {}", err); log::trace!("Exiting trigger is ready, {:?}", self.exit_trigger);
return Err(err.into());
}
}
}
EXIT_TOKEN_SENDER => {
log::trace!("Exiting sender is ready, {:?}", self.exit_sender);
} }
TUN_TOKEN => self.tun_event(event)?, TUN_TOKEN => self.tun_event(event)?,
PIPE_TOKEN => self.pipe_event(event)?, PIPE_TOKEN => self.pipe_event(event)?,
@ -1201,9 +1193,25 @@ impl<'a> TunToProxy<'a> {
} }
} }
fn exiting_event_handler(&mut self) -> Result<bool> {
let mut buffer = vec![0; 100];
match self.exit_receiver.read(&mut buffer) {
Ok(size) => {
log::trace!("Received exit signal: {:?}", &buffer[..size]);
log::info!("Exiting tun2proxy...");
Ok(true)
}
Err(err) if err.kind() == std::io::ErrorKind::WouldBlock => {
log::trace!("Exiting reciever is ready");
Ok(false)
}
Err(err) => Err(err.into()),
}
}
pub fn shutdown(&mut self) -> Result<(), Error> { pub fn shutdown(&mut self) -> Result<(), Error> {
log::debug!("Shutdown tun2proxy..."); log::debug!("Shutdown tun2proxy...");
self.exit_sender.as_mut().ok_or("Already shutdown")?.write_all(b"EX")?; _ = self.exit_trigger.as_mut().ok_or("Already triggered")?.write(b"EXIT")?;
Ok(()) Ok(())
} }
} }