diff --git a/Cargo.toml b/Cargo.toml index f499840..2053e90 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -20,3 +20,4 @@ nix = { version = "0.26", features = ["process", "signal"] } prctl = "1.0" reqwest = { version = "0.11", features = ["blocking", "json"] } serial_test = "1.0" +test-with = "0.9" diff --git a/src/lib.rs b/src/lib.rs index 256586f..40ce19c 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,12 +1,64 @@ use crate::tun2proxy::Credentials; use crate::{http::HttpManager, socks5::Socks5Manager, tun2proxy::TunToProxy}; -use std::net::SocketAddr; +use std::net::{SocketAddr, ToSocketAddrs}; pub mod http; pub mod socks5; pub mod tun2proxy; pub mod virtdevice; +#[derive(Clone, Debug)] +pub struct Proxy { + pub proxy_type: ProxyType, + pub addr: SocketAddr, + pub credentials: Option, +} + +impl Proxy { + pub fn from_url(s: &str) -> Result { + let url = url::Url::parse(s).map_err(|_| format!("`{s}` is not a valid proxy URL"))?; + let host = url + .host_str() + .ok_or(format!("`{s}` does not contain a host"))?; + + let mut url_host = String::from(host); + let port = url.port().ok_or(format!("`{s}` does not contain a port"))?; + url_host.push(':'); + url_host.push_str(port.to_string().as_str()); + + let mut addr_iter = url_host + .to_socket_addrs() + .map_err(|_| format!("`{host}` could not be resolved"))?; + + let addr = addr_iter + .next() + .ok_or(format!("`{host}` does not resolve to a usable IP address"))?; + + let credentials = if url.username() == "" && url.password().is_none() { + None + } else { + let username = String::from(url.username()); + let password = String::from(url.password().unwrap_or("")); + Some(Credentials::new(&username, &password)) + }; + + let scheme = url.scheme(); + + let proxy_type = match url.scheme().to_ascii_lowercase().as_str() { + "socks5" => Some(ProxyType::Socks5), + "http" => Some(ProxyType::Http), + _ => None, + } + .ok_or(format!("`{scheme}` is an invalid proxy type"))?; + + Ok(Proxy { + proxy_type, + addr, + credentials, + }) + } +} + #[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Debug)] pub enum ProxyType { Socks5, @@ -22,19 +74,14 @@ impl std::fmt::Display for ProxyType { } } -pub fn main_entry( - tun: &str, - addr: SocketAddr, - proxy_type: ProxyType, - credentials: Option, -) { +pub fn main_entry(tun: &str, proxy: Proxy) { let mut ttp = TunToProxy::new(tun); - match proxy_type { + match proxy.proxy_type { ProxyType::Socks5 => { - ttp.add_connection_manager(Socks5Manager::new(addr, credentials)); + ttp.add_connection_manager(Socks5Manager::new(proxy.addr, proxy.credentials)); } ProxyType::Http => { - ttp.add_connection_manager(HttpManager::new(addr, credentials)); + ttp.add_connection_manager(HttpManager::new(proxy.addr, proxy.credentials)); } } ttp.run(); diff --git a/src/main.rs b/src/main.rs index 521c978..87dab34 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,10 +1,7 @@ -use std::net::{SocketAddr, ToSocketAddrs}; - use clap::Parser; use env_logger::Env; -use tun2proxy::tun2proxy::Credentials; -use tun2proxy::{main_entry, ProxyType}; +use tun2proxy::{main_entry, Proxy}; /// Tunnel interface to proxy #[derive(Parser)] @@ -15,58 +12,8 @@ struct Args { tun: String, /// The proxy URL in the form proto://[username[:password]@]host:port - #[arg(short, long = "proxy", value_parser = proxy_url_parser, value_name = "URL")] - proxy: ArgProxy, -} - -#[derive(Clone)] -struct ArgProxy { - proxy_type: ProxyType, - addr: SocketAddr, - credentials: Option, -} - -fn proxy_url_parser(s: &str) -> Result { - let url = url::Url::parse(s).map_err(|_| format!("`{s}` is not a valid proxy URL"))?; - let host = url - .host_str() - .ok_or(format!("`{s}` does not contain a host"))?; - - let mut url_host = String::from(host); - let port = url.port().ok_or(format!("`{s}` does not contain a port"))?; - url_host.push(':'); - url_host.push_str(port.to_string().as_str()); - - let mut addr_iter = url_host - .to_socket_addrs() - .map_err(|_| format!("`{host}` could not be resolved"))?; - - let addr = addr_iter - .next() - .ok_or(format!("`{host}` does not resolve to a usable IP address"))?; - - let credentials = if url.username() == "" && url.password().is_none() { - None - } else { - let username = String::from(url.username()); - let password = String::from(url.password().unwrap_or("")); - Some(Credentials::new(&username, &password)) - }; - - let scheme = url.scheme(); - - let proxy_type = match url.scheme().to_ascii_lowercase().as_str() { - "socks5" => Some(ProxyType::Socks5), - "http" => Some(ProxyType::Http), - _ => None, - } - .ok_or(format!("`{scheme}` is an invalid proxy type"))?; - - Ok(ArgProxy { - proxy_type, - addr, - credentials, - }) + #[arg(short, long = "proxy", value_parser = Proxy::from_url, value_name = "URL")] + proxy: Proxy, } fn main() { @@ -77,5 +24,5 @@ fn main() { let proxy_type = args.proxy.proxy_type; log::info!("Proxy {proxy_type} server: {addr}"); - main_entry(&args.tun, addr, proxy_type, args.proxy.credentials); + main_entry(&args.tun, args.proxy); } diff --git a/src/tun2proxy.rs b/src/tun2proxy.rs index a2940ec..a2f8da4 100644 --- a/src/tun2proxy.rs +++ b/src/tun2proxy.rs @@ -162,7 +162,7 @@ struct ConnectionState { handler: std::boxed::Box, } -#[derive(Default, Clone)] +#[derive(Default, Clone, Debug)] pub struct Credentials { pub(crate) username: Vec, pub(crate) password: Vec, diff --git a/tests/proxy.rs b/tests/proxy.rs index 3b7127b..ba2ab96 100644 --- a/tests/proxy.rs +++ b/tests/proxy.rs @@ -2,36 +2,41 @@ mod tests { extern crate reqwest; + use std::env; + use std::io::BufRead; + use std::net::SocketAddr; + use std::process::Command; + use std::string::ToString; + use fork::Fork; use nix::sys::signal; use nix::unistd::Pid; use serial_test::serial; - use std::env; - use std::io::BufRead; - use std::net::{SocketAddr, ToSocketAddrs}; - use std::process::Command; - use std::string::ToString; - use tun2proxy::{main_entry, ProxyType}; + + use tun2proxy::{main_entry, Proxy, ProxyType}; static TUN_TEST_DEVICE: &str = "tun0"; static ALL_ROUTES: [&str; 4] = ["0.0.0.0/1", "128.0.0.0/1", "::/1", "8000::/1"]; - #[derive(Clone, Copy)] + #[derive(Clone, Debug)] struct Test { - env: &'static str, - proxy_type: ProxyType, + proxy: Proxy, } - static TESTS: [Test; 2] = [ - Test { - env: "SOCKS5_SERVER", - proxy_type: ProxyType::Socks5, - }, - Test { - env: "HTTP_SERVER", - proxy_type: ProxyType::Http, - }, - ]; + fn proxy_from_env(env_var: &str) -> Result { + let url = + env::var(env_var).map_err(|_| format!("{env_var} environment variable not found"))?; + Proxy::from_url(url.as_str()).map_err(|_| format!("{env_var} URL cannot be parsed")) + } + + fn test_from_env(env_var: &str) -> Result { + let proxy = proxy_from_env(env_var)?; + Ok(Test { proxy }) + } + + fn tests() -> [Result; 2] { + [test_from_env("SOCKS5_SERVER"), test_from_env("HTTP_SERVER")] + } #[cfg(test)] #[ctor::ctor] @@ -48,17 +53,14 @@ mod tests { .expect("failed to delete tun device"); } - fn parse_server_addr(string: String) -> SocketAddr { - return string.to_socket_addrs().unwrap().next().unwrap(); - } - fn routes_setup() { let mut all_servers: Vec = Vec::new(); - for test in TESTS { - if let Ok(server) = env::var(test.env) { - all_servers.push(parse_server_addr(server)); + for test in tests() { + if test.is_err() { + continue; } + all_servers.push(test.unwrap().proxy.addr); } Command::new("ip") @@ -116,46 +118,45 @@ mod tests { where F: Fn(&Test) -> bool, { - for test in TESTS { - if !filter(&test) { - continue; - } - let env_var = env::var(test.env).expect( - format!( - "this test requires the {} environment variable to be set", - test.env - ) - .as_str(), - ); - let address = parse_server_addr(env_var); + for potential_test in tests() { + match potential_test { + Ok(test) => { + if filter(&test) { + continue; + } - match fork::fork() { - Ok(Fork::Parent(child)) => { - reqwest::blocking::get("https://1.1.1.1") - .expect("failed to issue HTTP request"); - signal::kill(Pid::from_raw(child), signal::SIGKILL) - .expect("failed to kill child"); - nix::sys::wait::waitpid(Pid::from_raw(child), None) - .expect("failed to wait for child"); + match fork::fork() { + Ok(Fork::Parent(child)) => { + reqwest::blocking::get("https://1.1.1.1") + .expect("failed to issue HTTP request"); + signal::kill(Pid::from_raw(child), signal::SIGKILL) + .expect("failed to kill child"); + nix::sys::wait::waitpid(Pid::from_raw(child), None) + .expect("failed to wait for child"); + } + Ok(Fork::Child) => { + prctl::set_death_signal(signal::SIGKILL as isize).unwrap(); // 9 == SIGKILL + main_entry(TUN_TEST_DEVICE, test.proxy); + } + Err(_) => assert!(false), + } } - Ok(Fork::Child) => { - prctl::set_death_signal(signal::SIGKILL as isize).unwrap(); // 9 == SIGKILL - main_entry(TUN_TEST_DEVICE, address, ProxyType::Socks5, None); + Err(_) => { + continue; } - Err(_) => assert!(false), } } } - #[test] #[serial] + #[test_with::env(SOCKS5_SERVER)] fn test_socks5() { - run_test(|test| test.proxy_type == ProxyType::Socks5) + run_test(|test| test.proxy.proxy_type == ProxyType::Socks5) } - #[test] #[serial] + #[test_with::env(HTTP_SERVER)] fn test_http() { - run_test(|test| test.proxy_type == ProxyType::Http) + run_test(|test| test.proxy.proxy_type == ProxyType::Http) } }