From 4a48b297f31662955fac34c0e5de146283485c1f Mon Sep 17 00:00:00 2001
From: "B. Blechschmidt" <git@blechschmidt.io>
Date: Sat, 25 Mar 2023 21:41:40 +0100
Subject: [PATCH] Restore network configuration on Drop

---
 src/lib.rs   |  2 +-
 src/main.rs  |  9 +++++----
 src/setup.rs | 21 +++++++++++++++------
 3 files changed, 21 insertions(+), 11 deletions(-)

diff --git a/src/lib.rs b/src/lib.rs
index 5c4fe14..2d985ac 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -3,7 +3,7 @@ use crate::socks5::SocksVersion;
 use crate::{http::HttpManager, socks5::SocksManager, tun2proxy::TunToProxy};
 use std::net::{SocketAddr, ToSocketAddrs};
 
-mod error;
+pub mod error;
 mod http;
 pub mod setup;
 mod socks5;
diff --git a/src/main.rs b/src/main.rs
index 18550d9..38cbe22 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,6 +1,6 @@
 use clap::Parser;
 use env_logger::Env;
-use std::process::exit;
+use std::process::ExitCode;
 
 use tun2proxy::setup::{get_default_cidrs, Setup};
 use tun2proxy::Options;
@@ -44,7 +44,7 @@ enum ArgSetup {
     Auto,
 }
 
-fn main() {
+fn main() -> ExitCode {
     dotenvy::dotenv().ok();
     env_logger::Builder::from_env(Env::default().default_filter_or("info")).init();
     let args = Args::parse();
@@ -63,12 +63,13 @@ fn main() {
         setup = Setup::new(&args.tun, &args.proxy.addr.ip(), get_default_cidrs());
         if let Err(e) = setup.setup() {
             log::error!("{e}");
-            exit(1);
+            return ExitCode::FAILURE;
         }
     }
 
     if let Err(e) = main_entry(&args.tun, args.proxy, options) {
         log::error!("{e}");
-        exit(1);
+        return ExitCode::FAILURE;
     }
+    ExitCode::SUCCESS
 }
diff --git a/src/setup.rs b/src/setup.rs
index 513e7ed..805cae5 100644
--- a/src/setup.rs
+++ b/src/setup.rs
@@ -95,7 +95,7 @@ impl Setup {
             let mut f = std::fs::File::from_raw_fd(fd);
             f.write_all("nameserver 198.18.0.1\n".as_bytes())?;
             mem::forget(f);
-            if libc::fchmod(fd, 0o644) == -1 {
+            if libc::fchmod(fd, 0o444) == -1 {
                 return Err("Failed to change ownership of /etc/resolv.conf".into());
             }
             let fd_path = format!("/proc/self/fd/{}", fd);
@@ -133,10 +133,13 @@ impl Setup {
         Ok(())
     }
 
-    fn shutdown(tun_name: String) {
-        let _ = Command::new("ip")
-            .args(["link", "del", tun_name.as_str()])
-            .output();
+    fn shutdown(&self) {
+        Self::shutdown_with_args(&self.tun);
+    }
+
+    fn shutdown_with_args(tun_name: &str) {
+        log::info!("Restoring network configuration");
+        let _ = Command::new("ip").args(["link", "del", tun_name]).output();
         unsafe {
             let umount_path = CString::new("/etc/resolv.conf").unwrap();
             libc::umount(umount_path.as_ptr());
@@ -155,7 +158,7 @@ impl Setup {
         let tun_name = self.tun.clone();
         // TODO: This is not optimal.
         ctrlc::set_handler(move || {
-            Self::shutdown(tun_name.clone());
+            Self::shutdown_with_args(&tun_name);
             std::process::exit(0);
         })?;
 
@@ -178,3 +181,9 @@ impl Setup {
         Ok(())
     }
 }
+
+impl Drop for Setup {
+    fn drop(&mut self) {
+        self.shutdown();
+    }
+}