don't panic on ip resolution

2 weeks ago · 82724ad9b1
parent 9d84d3defe
commit 82724ad9b1
5 changed files with 307 additions and 263 deletions
--- a/src/config.rs
+++ b/src/config.rs
@ -0,0 +1,2 @@
 pub const HOST: &str = "ortem-tun";
 pub const BASE_CONFIG: &str = concat!(env!("HOME"), "/.ssh/config");
--- a/src/main.rs
+++ b/src/main.rs
@ -1,266 +1,14 @@
-use nix::{
+mod config;
-    sys::signal::{SigHandler, Signal, signal},
+mod monitoring;
    unistd::{ForkResult, close as fdclose, fork, getppid, setsid},
 };
 use reqwest::{Proxy, blocking::Client, header};
 use std::{
    env,
    io::Result as IOResult,
    net::ToSocketAddrs,
    panic,
    process::{Child, Command, Stdio, exit},
    sync::{Mutex, MutexGuard, OnceLock},
    thread, time,
 };
 mod ssh_config;
-use ssh_config::TinySSHConfig;
+mod trigger;
-
+mod utils;
 const HOST: &str = "ortem-tun";
 const BASE_CONFIG: &str = concat!(env!("HOME"), "/.ssh/config");
 static mut DAEMONIZED: bool = false;
 struct ThresholdCaller {
    // call Fn once after threshold is exceeded
    cmd: Box<dyn Fn(&str) + Send + 'static>,
    arg: Option<String>,
    triggered: bool,
    retries: usize,
    max_retries: usize,
 }
 impl ThresholdCaller {
    fn new(cmd: impl Fn(&str) + Send + 'static, max_retries: usize) -> Self {
        ThresholdCaller {
            cmd: Box::new(cmd),
            arg: None,
            triggered: false,
            retries: 0,
            max_retries,
        }
    }
    pub fn get<'m>() -> MutexGuard<'m, Self> {
        static FAIL_NOTIFIER: OnceLock<Mutex<ThresholdCaller>> = OnceLock::new();
        FAIL_NOTIFIER
            .get_or_init(|| {
                Mutex::new(ThresholdCaller::new(
                    |info| notify(format!("Can't setup socks proxy: {}", info), false),
                    5,
                ))
            })
            .lock()
            .unwrap()
    }
    pub fn call_once(&mut self, info: Option<&String>) {
        //dbg!(self.retries, self.triggered);
        if !self.triggered && self.retries >= self.max_retries {
            let arg = info.or(self.arg.as_ref());
            (self.cmd)(arg.unwrap_or(&String::new()));
            self.triggered = true;
        }
    }
    pub fn set_arg(&mut self, arg: Option<String>) {
        self.arg = arg;
    }
    pub fn append_fail<S: Into<String>>(&mut self, msg: Option<S>) {
        self.inc();
        if let Some(m) = msg {
            let msg = m.into();
            unsafe {
                if !DAEMONIZED {
                    eprintln!("{}", &msg);
                }
            }
            self.set_arg(Some(msg));
        }
    }
    pub fn was_triggered(&self) -> bool {
        self.triggered
    }
    pub fn reset(&mut self) {
        self.retries = 0;
        self.triggered = false;
    }
    pub fn inc(&mut self) {
        self.retries += 1;
    }
 }
 pub fn daemonize() -> Result<i32, String> {
    if getppid().as_raw() != 1 {
        setsig(Signal::SIGTTOU, SigHandler::SigIgn);
        setsig(Signal::SIGTTIN, SigHandler::SigIgn);
        setsig(Signal::SIGTSTP, SigHandler::SigIgn);
    }
    for fd in 0..=2 {
        fdclose(fd).ok();
    }
    unsafe {
        match fork() {
            Ok(ForkResult::Parent { .. }) => {
                exit(0);
            }
            Ok(ForkResult::Child) => match setsid() {
                Ok(pid) => Ok(pid.as_raw()),
                Err(e) => Err(e.to_string()),
            },
            Err(_) => exit(255),
        }
    }
 }
 pub fn setsig(sig: Signal, hnd: SigHandler) {
    unsafe {
        signal(sig, hnd).unwrap();
    }
 }
 fn spawn(cmd: &str, args: &[&str]) -> IOResult<Child> {
    Command::new(cmd).args(args).spawn()
 }
 fn _dbg() {
    spawn("touch", &["/tmp/asdqwe"]).ok();
 }
 fn notify<S: Into<String>>(msg: S, critical: bool) {
    let (urg, prefix) = if critical {
        ("critical", "Отвалилась жопа: ")
    } else {
        ("normal", "")
    };
    spawn(
        "notify-send",
        &[
            "-a",
            env!("CARGO_PKG_NAME"),
            "-t",
            "5000",
            "-u",
            urg,
            &format!("{prefix}{}", msg.into()),
        ],
    )
    .unwrap()
    .wait()
    .unwrap();
 }
 fn run_proxy(host: &str) -> IOResult<Child> {
    Command::new("/usr/bin/ssh")
        .args(&["-f", "-N", host, "-o", "ConnectTimeout=3"])
        .stderr(Stdio::piped())
        .spawn()
 }
 fn pgrep(cmd: &str) -> Option<i32> {
    let grep = Command::new("pgrep").args(&["-f", cmd]).output().unwrap();
    if grep.status.success() {
        let output = String::from_utf8(grep.stdout).unwrap();
        output.lines().next().unwrap().parse().ok()
    } else {
        None
    }
 }
 fn process_present(cmd: &str) -> bool {
    pgrep(cmd).is_some()
 }
-fn kill_by_cmd(cmd: &str) {
+use std::{env, panic, process::exit, thread};
    if let Some(pid) = pgrep(cmd) {
        Command::new("kill")
            .args(&["-9", &pid.to_string()])
            .status()
            .ok();
    }
 }
-fn monitor_connection(ssh_host: &str) -> ! {
+use crate::{config::HOST, utils::notify};
    let config = TinySSHConfig::parse(BASE_CONFIG).expect("ssh config file has wrong format");
    let host = config.get_host(ssh_host).expect("no host found in config");
    assert!(host.can_be_proxy(), "expected DynamicForward entry");
    let proxy = Proxy::all(&format!(
        "socks5://127.0.0.1:{}",
        host.get("DynamicForward").unwrap()
    ))
    .unwrap();
    let mut headers = header::HeaderMap::new();
    headers.insert(
        header::USER_AGENT,
        header::HeaderValue::from_static("curl/7.74.0"),
    );
    let client = Client::builder()
        .proxy(proxy)
        .timeout(time::Duration::from_secs(10))
        .default_headers(headers)
        .build()
        .unwrap();
    let proxy_ip = {
        let host = host.get("HostName").unwrap();
        let host = format!("{host}:80");
        host.to_socket_addrs().unwrap().next().unwrap().ip()
    };
    let killer = |msg| {
        ThresholdCaller::get().append_fail(Some(msg));
        kill_by_cmd(ssh_host);
    };
    loop {
        if process_present(HOST) {
            match client.get("https://2ip.ru/").send() {
                Ok(resp) => match resp.text() {
                    Ok(received_ip) => {
                        if received_ip.trim() != proxy_ip.to_string() {
                            killer("wrong ip");
                        } else {
                            if ThresholdCaller::get().was_triggered() {
                                notify("Соединение восстановлено", false);
                            }
                            ThresholdCaller::get().reset();
                        }
                    }
                    Err(_) => {
                        killer("can't receive response");
                    }
                },
                Err(_) => {
                    killer("can't send request");
                }
            }
        }
        thread::sleep(time::Duration::from_secs(10));
    }
 }
-fn monitor_process() -> ! {
+static mut DAEMONIZED: bool = false;
    loop {
        if !process_present(HOST) {
            match run_proxy(HOST) {
                Ok(result) => {
                    let exited = result.wait_with_output().unwrap();
                    if !exited.status.success() {
                        ThresholdCaller::get()
                            .append_fail(Some(String::from_utf8(exited.stderr).unwrap()));
                    }
                }
                Err(e) => {
                    ThresholdCaller::get().append_fail(Some(e.to_string()));
                }
            }
            ThresholdCaller::get().call_once(None);
        }
        thread::sleep(time::Duration::from_secs(1));
    }
 }
 fn main() {
    if !cfg!(debug_assertions) {
@ -272,19 +20,20 @@ fn main() {
    let mut args = env::args();
    if let Some(arg) = args.nth(1) {
        if arg == "-d" {
-            daemonize().unwrap();
+            utils::daemonize().unwrap();
            unsafe {
                DAEMONIZED = true;
            }
        }
    }
-    thread::spawn(|| monitor_connection(HOST));
+    thread::spawn(|| monitoring::monitor_connection(HOST));
-    monitor_process();
+    monitoring::monitor_process();
 }
 #[cfg(test)]
 mod tests {
-    use super::*;
+    use crate::{config::BASE_CONFIG, ssh_config::TinySSHConfig};
    use std::error::Error;
    type TestResult = Result<(), Box<dyn Error>>;
--- a/src/monitoring.rs
+++ b/src/monitoring.rs
@ -0,0 +1,117 @@
 use std::{net::ToSocketAddrs as _, thread, time::Duration};
 use crate::{
    config::{BASE_CONFIG, HOST},
    ssh_config::TinySSHConfig,
    trigger::Trigger,
    utils::*,
 };
 use reqwest::{Proxy, blocking::Client, header};
 pub fn monitor_connection(ssh_host: &str) -> ! {
    let config = TinySSHConfig::parse(BASE_CONFIG).expect("ssh config file has wrong format");
    let host = config
        .get_host(ssh_host)
        .expect("no host found in ssh config");
    assert!(host.can_be_proxy(), "expected DynamicForward entry");
    let proxy = Proxy::all(&format!(
        "socks5://127.0.0.1:{}",
        host.get("DynamicForward").unwrap()
    ))
    .unwrap();
    let mut headers = header::HeaderMap::new();
    headers.insert(
        header::USER_AGENT,
        header::HeaderValue::from_static("curl/7.74.0"),
    );
    let client = Client::builder()
        .proxy(proxy)
        .timeout(Duration::from_secs(10))
        .default_headers(headers)
        .build()
        .unwrap();
    let kill_with = |msg| {
        Trigger::get().append_fail(Some(msg));
        kill_by_cmd(ssh_host);
    };
    let proxy_host = {
        let host = host.get("HostName").unwrap();
        format!("{host}:80")
    };
    let mut proxy_ip = None;
    'mon: loop {
        let Some(proxy_ip) = &proxy_ip else {
            let ip = match resolve_host(&proxy_host) {
                Ok(ip) => ip,
                Err(err) => {
                    kill_with(err);
                    continue 'mon;
                }
            };
            proxy_ip = Some(ip);
            continue 'mon;
        };
        if process_present(HOST) {
            if let Err(err) = verify_proxy_ip(&client, &proxy_ip) {
                kill_with(err)
            } else {
                if Trigger::get().was_triggered() {
                    notify("Соединение восстановлено", false);
                }
                Trigger::get().reset();
            }
        }
        thread::sleep(Duration::from_secs(5));
    }
 }
 fn verify_proxy_ip(client: &Client, expected_proxy_ip: &str) -> Result<(), String> {
    let received_ip = client
        .get("https://2ip.ru/")
        .send()
        .map_err(|e| format!("can't send request while checking ip: {}", e.to_string()))?
        .text()
        .map_err(|e| {
            format!(
                "can't receive response while checking ip: {}",
                e.to_string()
            )
        })?;
    let received_ip = received_ip.trim();
    if received_ip != expected_proxy_ip {
        return Err(format!(
            "wrong proxy ip: {received_ip}, expected: {expected_proxy_ip}"
        ));
    }
    Ok(())
 }
 fn resolve_host(host: &str) -> Result<String, String> {
    host.to_socket_addrs()
        .map_err(|e| e.to_string())?
        .next()
        .map(|ip| ip.ip().to_string())
        .ok_or_else(|| "No IPs returned".to_string())
 }
 pub fn monitor_process() -> ! {
    loop {
        if !process_present(HOST) {
            match run_proxy(HOST) {
                Ok(result) => {
                    let exited = result.wait_with_output().unwrap();
                    if !exited.status.success() {
                        Trigger::get().append_fail(Some(String::from_utf8(exited.stderr).unwrap()));
                    }
                }
                Err(e) => {
                    Trigger::get().append_fail(Some(e.to_string()));
                }
            }
            Trigger::get().call_once(None);
        }
        thread::sleep(Duration::from_secs(1));
    }
 }
--- a/src/trigger.rs
+++ b/src/trigger.rs
@ -0,0 +1,76 @@
 use std::sync::{Mutex, MutexGuard, OnceLock};
 use crate::{DAEMONIZED, utils::notify};
 pub struct Trigger {
    // call Fn once after threshold is exceeded
    cmd: Box<dyn Fn(&str) + Send + 'static>,
    arg: Option<String>,
    triggered: bool,
    retries: usize,
    max_retries: usize,
 }
 impl Trigger {
    fn new(cmd: impl Fn(&str) + Send + 'static, max_retries: usize) -> Self {
        Trigger {
            cmd: Box::new(cmd),
            arg: None,
            triggered: false,
            retries: 0,
            max_retries,
        }
    }
    pub fn get() -> MutexGuard<'static, Self> {
        static FAIL_NOTIFIER: OnceLock<Mutex<Trigger>> = OnceLock::new();
        FAIL_NOTIFIER
            .get_or_init(|| {
                Mutex::new(Trigger::new(
                    |info| notify(format!("Can't setup socks proxy: {}", info), false),
                    5,
                ))
            })
            .lock()
            .unwrap()
    }
    pub fn call_once(&mut self, info: Option<&str>) {
        //dbg!(self.retries, self.triggered);
        if !self.triggered && self.retries >= self.max_retries {
            let arg = info.or(self.arg.as_deref());
            (self.cmd)(arg.unwrap_or(""));
            self.triggered = true;
        }
    }
    pub fn set_arg(&mut self, arg: Option<String>) {
        self.arg = arg;
    }
    pub fn append_fail<S: Into<String>>(&mut self, msg: Option<S>) {
        self.inc();
        if let Some(m) = msg {
            let msg = m.into();
            unsafe {
                if !DAEMONIZED {
                    eprintln!("{}", &msg);
                }
            }
            self.set_arg(Some(msg));
        }
    }
    pub fn was_triggered(&self) -> bool {
        self.triggered
    }
    pub fn reset(&mut self) {
        self.retries = 0;
        self.triggered = false;
    }
    pub fn inc(&mut self) {
        self.retries += 1;
    }
 }
--- a/src/utils.rs
+++ b/src/utils.rs
@ -0,0 +1,100 @@
 use std::{
    io,
    process::{Child, Command, Stdio, exit},
 };
 use nix::{
    sys::signal::{SigHandler, Signal, signal},
    unistd::{ForkResult, close as fdclose, fork, getppid, setsid},
 };
 pub fn daemonize() -> Result<i32, String> {
    if getppid().as_raw() != 1 {
        setsig(Signal::SIGTTOU, SigHandler::SigIgn);
        setsig(Signal::SIGTTIN, SigHandler::SigIgn);
        setsig(Signal::SIGTSTP, SigHandler::SigIgn);
    }
    for fd in 0..=2 {
        fdclose(fd).ok();
    }
    unsafe {
        match fork() {
            Ok(ForkResult::Parent { .. }) => {
                exit(0);
            }
            Ok(ForkResult::Child) => match setsid() {
                Ok(pid) => Ok(pid.as_raw()),
                Err(e) => Err(e.to_string()),
            },
            Err(_) => exit(255),
        }
    }
 }
 pub fn setsig(sig: Signal, hnd: SigHandler) {
    unsafe {
        signal(sig, hnd).unwrap();
    }
 }
 pub fn spawn(cmd: &str, args: &[&str]) -> io::Result<Child> {
    Command::new(cmd).args(args).spawn()
 }
 pub fn _dbg() {
    spawn("touch", &["/tmp/asdqwe"]).ok();
 }
 pub fn notify<S: Into<String>>(msg: S, critical: bool) {
    let (urg, prefix) = if critical {
        ("critical", "Отвалилась жопа: ")
    } else {
        ("normal", "")
    };
    spawn(
        "notify-send",
        &[
            "-a",
            env!("CARGO_PKG_NAME"),
            "-t",
            "5000",
            "-u",
            urg,
            &format!("{prefix}{}", msg.into()),
        ],
    )
    .unwrap()
    .wait()
    .unwrap();
 }
 pub fn run_proxy(host: &str) -> io::Result<Child> {
    Command::new("/usr/bin/ssh")
        .args(&["-f", "-N", host, "-o", "ConnectTimeout=3"])
        .stderr(Stdio::piped())
        .spawn()
 }
 pub fn pgrep(cmd: &str) -> Option<i32> {
    let grep = Command::new("pgrep").args(&["-f", cmd]).output().unwrap();
    if grep.status.success() {
        let output = String::from_utf8(grep.stdout).unwrap();
        output.lines().next().unwrap().parse().ok()
    } else {
        None
    }
 }
 pub fn process_present(cmd: &str) -> bool {
    pgrep(cmd).is_some()
 }
 pub fn kill_by_cmd(cmd: &str) {
    if let Some(pid) = pgrep(cmd) {
        Command::new("kill")
            .args(&["-9", &pid.to_string()])
            .status()
            .ok();
    }
 }
		`@ -0,0 +1,2 @@`
							`pub const HOST: &str = "ortem-tun";`
							`pub const BASE_CONFIG: &str = concat!(env!("HOME"), "/.ssh/config");`