From 62a04229db7c9570d650f708961b7a8ae56d0545 Mon Sep 17 00:00:00 2001
From: ssrlive <30760636+ssrlive@users.noreply.github.com>
Date: Fri, 14 Apr 2023 17:27:37 +0800
Subject: [PATCH] shutdown function

---
 src/lib.rs       |  4 ++++
 src/tun2proxy.rs | 21 ++++++++++++++++++++-
 2 files changed, 24 insertions(+), 1 deletion(-)

diff --git a/src/lib.rs b/src/lib.rs
index 1f167a2..c0a5c30 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -149,3 +149,7 @@ pub fn main_entry(
     }
     ttp.run()
 }
+
+pub fn shutdown() -> Result<(), Error> {
+    TunToProxy::shutdown()
+}
diff --git a/src/tun2proxy.rs b/src/tun2proxy.rs
index 3940e3a..4e17711 100644
--- a/src/tun2proxy.rs
+++ b/src/tun2proxy.rs
@@ -244,6 +244,9 @@ pub(crate) trait ConnectionManager {
 
 const TCP_TOKEN: Token = Token(0);
 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> {
     tun: TunTapInterface,
@@ -257,6 +260,7 @@ pub(crate) struct TunToProxy<'a> {
     device: VirtualTunDevice,
     options: Options,
     write_sockets: HashSet<Token>,
+    _exit_listener: mio::net::TcpListener,
 }
 
 impl<'a> TunToProxy<'a> {
@@ -274,6 +278,10 @@ impl<'a> TunToProxy<'a> {
             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 {
             Medium::Ethernet => Config::new(
                 smoltcp::wire::EthernetAddress([0x02, 0x00, 0x00, 0x00, 0x00, 0x01]).into(),
@@ -305,6 +313,7 @@ impl<'a> TunToProxy<'a> {
             device: virt,
             options,
             write_sockets: HashSet::default(),
+            _exit_listener,
         };
         Ok(tun)
     }
@@ -762,6 +771,10 @@ impl<'a> TunToProxy<'a> {
                 Ok(()) => {
                     for event in events.iter() {
                         match event.token() {
+                            EXIT_TOKEN => {
+                                log::info!("exiting...");
+                                return Ok(());
+                            }
                             TCP_TOKEN => self.tun_event(event)?,
                             UDP_TOKEN => self.udp_event(event),
                             _ => self.mio_socket_event(event)?,
@@ -773,10 +786,16 @@ impl<'a> TunToProxy<'a> {
                     if e.kind() != std::io::ErrorKind::Interrupted {
                         return Err(e.into());
                     } 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(())
+    }
 }