shutdown function

This commit is contained in:
ssrlive 2023-04-14 17:27:37 +08:00 committed by B. Blechschmidt
parent fb3ad33b53
commit 62a04229db
2 changed files with 24 additions and 1 deletions

View file

@ -149,3 +149,7 @@ pub fn main_entry(
} }
ttp.run() ttp.run()
} }
pub fn shutdown() -> Result<(), Error> {
TunToProxy::shutdown()
}

View file

@ -244,6 +244,9 @@ pub(crate) trait ConnectionManager {
const TCP_TOKEN: Token = Token(0); const TCP_TOKEN: Token = Token(0);
const UDP_TOKEN: Token = Token(1); const UDP_TOKEN: Token = Token(1);
const EXIT_TOKEN: Token = Token(34255);
const EXIT_LISTENER: &str = "127.0.0.1:34255";
pub(crate) struct TunToProxy<'a> { pub(crate) struct TunToProxy<'a> {
tun: TunTapInterface, tun: TunTapInterface,
@ -257,6 +260,7 @@ pub(crate) struct TunToProxy<'a> {
device: VirtualTunDevice, device: VirtualTunDevice,
options: Options, options: Options,
write_sockets: HashSet<Token>, write_sockets: HashSet<Token>,
_exit_listener: mio::net::TcpListener,
} }
impl<'a> TunToProxy<'a> { impl<'a> TunToProxy<'a> {
@ -274,6 +278,10 @@ impl<'a> TunToProxy<'a> {
Interest::READABLE, Interest::READABLE,
)?; )?;
let mut _exit_listener = mio::net::TcpListener::bind(EXIT_LISTENER.parse()?)?;
poll.registry()
.register(&mut _exit_listener, EXIT_TOKEN, Interest::READABLE)?;
let config = match tun.capabilities().medium { let config = match tun.capabilities().medium {
Medium::Ethernet => Config::new( Medium::Ethernet => Config::new(
smoltcp::wire::EthernetAddress([0x02, 0x00, 0x00, 0x00, 0x00, 0x01]).into(), smoltcp::wire::EthernetAddress([0x02, 0x00, 0x00, 0x00, 0x00, 0x01]).into(),
@ -305,6 +313,7 @@ impl<'a> TunToProxy<'a> {
device: virt, device: virt,
options, options,
write_sockets: HashSet::default(), write_sockets: HashSet::default(),
_exit_listener,
}; };
Ok(tun) Ok(tun)
} }
@ -762,6 +771,10 @@ impl<'a> TunToProxy<'a> {
Ok(()) => { Ok(()) => {
for event in events.iter() { for event in events.iter() {
match event.token() { match event.token() {
EXIT_TOKEN => {
log::info!("exiting...");
return Ok(());
}
TCP_TOKEN => self.tun_event(event)?, TCP_TOKEN => self.tun_event(event)?,
UDP_TOKEN => self.udp_event(event), UDP_TOKEN => self.udp_event(event),
_ => self.mio_socket_event(event)?, _ => self.mio_socket_event(event)?,
@ -773,10 +786,16 @@ impl<'a> TunToProxy<'a> {
if e.kind() != std::io::ErrorKind::Interrupted { if e.kind() != std::io::ErrorKind::Interrupted {
return Err(e.into()); return Err(e.into());
} else { } else {
log::warn!("Poll interrupted") log::warn!("Poll interrupted: {e}")
} }
} }
} }
} }
} }
pub(crate) fn shutdown() -> Result<(), Error> {
let addr: SocketAddr = EXIT_LISTENER.parse()?;
let _ = std::net::TcpStream::connect(addr)?;
Ok(())
}
} }