From 3223ca4e22003fd87e90f89ce352d84d38227f2e Mon Sep 17 00:00:00 2001
From: "B. Blechschmidt" <github@blechschmidt.pw>
Date: Wed, 22 Mar 2023 11:17:28 +0100
Subject: [PATCH] Pass proxy via URL

---
 Cargo.toml  |  1 +
 src/lib.rs  |  2 +-
 src/main.rs | 97 ++++++++++++++++++++++++++++++++---------------------
 3 files changed, 61 insertions(+), 39 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index 66c04ea..f499840 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -11,6 +11,7 @@ env_logger = "0.10"
 log = "0.4"
 mio = { version = "0.8", features = ["os-poll", "net", "os-ext"] }
 smoltcp = { version = "0.9", features = ["std"] }
+url = "2.3"
 
 [dev-dependencies]
 ctor = "0.1"
diff --git a/src/lib.rs b/src/lib.rs
index cf8113b..c68e142 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -7,7 +7,7 @@ pub mod socks5;
 pub mod tun2proxy;
 pub mod virtdevice;
 
-#[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord)]
+#[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Debug)]
 pub enum ProxyType {
     Socks5,
     Http,
diff --git a/src/main.rs b/src/main.rs
index b1fa359..e1bca13 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,6 +1,6 @@
-use std::net::SocketAddr;
+use std::net::{SocketAddr, ToSocketAddrs};
 
-use clap::{Parser, ValueEnum};
+use clap::Parser;
 use env_logger::Env;
 
 use tun2proxy::tun2proxy::Credentials;
@@ -11,55 +11,76 @@ use tun2proxy::{main_entry, ProxyType};
 #[command(author, version, about = "Tunnel interface to proxy.", long_about = None)]
 struct Args {
     /// Name of the tun interface
-    #[arg(short, long, value_name = "name")]
+    #[arg(short, long, value_name = "name", default_value = "tun0")]
     tun: String,
 
     /// What proxy type to run
-    #[arg(short, long = "proxy", value_name = "type", value_enum)]
-    proxy_type: ArgProxyType,
+    #[arg(short, long = "proxy", value_parser = proxy_url_parser)]
+    proxy: ArgProxy,
+}
 
-    /// Server address with format ip:port
-    #[clap(short, long, value_name = "ip:port")]
+#[derive(Clone)]
+struct ArgProxy {
+    proxy_type: ProxyType,
     addr: SocketAddr,
+    credentials: Credentials,
+}
 
-    /// Username for authentication
-    #[clap(long, value_name = "username")]
-    username: Option<String>,
+fn proxy_url_parser(s: &str) -> Result<ArgProxy, String> {
+    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"))?;
 
-    /// Password for authentication
-    #[clap(long, value_name = "password")]
-    password: Option<String>,
-}
+    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());
 
-#[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord, ValueEnum)]
-enum ArgProxyType {
-    /// SOCKS5 server to use
-    Socks5,
-    /// HTTP server to use
-    Http,
-}
+    let mut addr_iter = url_host
+        .to_socket_addrs()
+        .map_err(|_| format!("`{host}` could not be resolved"))?;
 
-fn main() {
-    env_logger::Builder::from_env(Env::default().default_filter_or("info")).init();
-    let args = Args::parse();
+    let addr = addr_iter
+        .next()
+        .ok_or(format!("`{host}` does not resolve to a usable IP address"))?;
 
-    let credentials = if args.username.is_some() || args.password.is_some() {
+    let credentials = if url.username() == "" && url.password().is_none() {
+        Credentials::none()
+    } else {
         Credentials::new(
-            args.username.unwrap_or(String::from("")),
-            args.password.unwrap_or(String::from("")),
+            String::from(url.username()),
+            String::from(url.password().unwrap_or("")),
         )
-    } else {
-        Credentials::none()
     };
 
-    match args.proxy_type {
-        ArgProxyType::Socks5 => {
-            log::info!("SOCKS5 server: {}", args.addr);
-            main_entry(&args.tun, args.addr, ProxyType::Socks5, credentials);
-        }
-        ArgProxyType::Http => {
-            log::info!("HTTP server: {}", args.addr);
-            main_entry(&args.tun, args.addr, ProxyType::Http, credentials);
-        }
+    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,
+    })
+}
+
+fn main() {
+    env_logger::Builder::from_env(Env::default().default_filter_or("info")).init();
+    let args = Args::parse();
+
+    let addr = args.proxy.addr;
+    log::info!("Proxy server: {addr}");
+
+    main_entry(
+        &args.tun,
+        args.proxy.addr,
+        args.proxy.proxy_type,
+        args.proxy.credentials,
+    );
 }