diff --git a/README.md b/README.md index c9ef51b..404c95f 100644 --- a/README.md +++ b/README.md @@ -49,6 +49,18 @@ usage: wireguard-proxy [options...] -h, --help print this usage text -st, --socket-timeout Socket timeout (time to wait for data) before terminating, default: 0 + + Environment variable support: + For every command line option, short and long, if you replace all + leading - with WGP_, and replace all remaining - with _, and uppercase + the whole thing, if you don't specify that command line option we will + read that environment variable for the argument. boolean arguments are + true if anything but unset, empty, 0, or false. + Examples: + --tcp-target ARG is WGP_TCP_TARGET=ARG + --socket-timeout 5 is WGP_SOCKET_TIMEOUT=5 + --tls is WGP_TLS=1 or WGP_TLS=true + WGP_TLS=0 or WGP_TLS=false would be like not sending --tls ``` Binaries: diff --git a/src/bin/udp-test.rs b/src/bin/udp-test.rs index b5be39e..db8d840 100644 --- a/src/bin/udp-test.rs +++ b/src/bin/udp-test.rs @@ -124,6 +124,18 @@ fn main() { sha256 hashes preceded by "sha256//" and separated by ";". Identical to curl's --pinnedpubkey and CURLOPT_PINNEDPUBLICKEY + + Environment variable support: + For every command line option, short and long, if you replace all + leading - with WGP_, and replace all remaining - with _, and uppercase + the whole thing, if you don't specify that command line option we will + read that environment variable for the argument. boolean arguments are + true if anything but unset, empty, 0, or false. + Examples: + --tcp-target ARG is WGP_TCP_TARGET=ARG + --socket-timeout 5 is WGP_SOCKET_TIMEOUT=5 + --tls is WGP_TLS=1 or WGP_TLS=true + WGP_TLS=0 or WGP_TLS=false would be like not sending --tls "#, default_udp_host_target, default_udp_host_target, default_socket_timeout); return; } else if args.flag("-s") || args.flag("--self-test") { @@ -137,10 +149,12 @@ fn main() { let mut proxyd_args = vec!["-th", tcp_host, "-ut", host]; + let tls_key = tls_key.as_ref().map(String::as_str); + let tls_cert = tls_cert.as_ref().map(String::as_str); if tls { let tls_key = tls_key.unwrap(); let tls_cert = tls_cert.unwrap(); - proxyd_args.extend(["-tk", tls_key, "-tc", tls_cert].iter().cloned()); + proxyd_args.extend(["-tk", &tls_key, "-tc", &tls_cert].iter().cloned()); } println!("executing: {} {}", proxy, proxyd_args.join(" ")); @@ -153,11 +167,12 @@ fn main() { let mut proxy_args = vec!["-tt", tcp_host]; + let pinnedpubkey = pinnedpubkey.as_ref().map(String::as_str); if tls { proxy_args.push("--tls"); if pinnedpubkey.is_some() { proxy_args.push("--pinnedpubkey"); - proxy_args.push(pinnedpubkey.unwrap()); + proxy_args.push(&pinnedpubkey.unwrap()); } } @@ -236,6 +251,7 @@ fn main() { if tls { let hostname = tcp_host.split(":").next(); + let pinnedpubkey = pinnedpubkey.as_ref().map(String::as_str); match pinnedpubkey { Some(pinnedpubkey) => println!("executing: wireguard-proxy -tt {} --tls --pinnedpubkey {}", tcp_host, pinnedpubkey), @@ -253,7 +269,7 @@ fn main() { println!("waiting: {:?} for wireguard-proxy client to come up.....", sleep); thread::sleep(sleep); - first_arg = host; + first_arg = host.to_owned(); } let server = Server::new( diff --git a/src/bin/wireguard-proxy.rs b/src/bin/wireguard-proxy.rs index 08e8e0d..5e4349c 100644 --- a/src/bin/wireguard-proxy.rs +++ b/src/bin/wireguard-proxy.rs @@ -54,6 +54,18 @@ fn main() { -h, --help print this usage text -st, --socket-timeout Socket timeout (time to wait for data) before terminating, default: {} + + Environment variable support: + For every command line option, short and long, if you replace all + leading - with WGP_, and replace all remaining - with _, and uppercase + the whole thing, if you don't specify that command line option we will + read that environment variable for the argument. boolean arguments are + true if anything but unset, empty, 0, or false. + Examples: + --tcp-target ARG is WGP_TCP_TARGET=ARG + --socket-timeout 5 is WGP_SOCKET_TIMEOUT=5 + --tls is WGP_TLS=1 or WGP_TLS=true + WGP_TLS=0 or WGP_TLS=false would be like not sending --tls "#, default_udp_host_target, default_udp_host_target, default_socket_timeout); return; } @@ -61,9 +73,9 @@ fn main() { let socket_timeout = args.get(&["-st", "--socket-timeout"], default_socket_timeout); if tcp_target.is_some() { - client(tcp_target.unwrap(), socket_timeout, args); + client(&tcp_target.unwrap(), socket_timeout, args); } else { - server(tcp_host.unwrap(), socket_timeout, args); + server(&tcp_host.unwrap(), socket_timeout, args); } } @@ -85,9 +97,9 @@ fn client(tcp_target: &str, socket_timeout: u64, args: Args) { ); if tls { - let hostname = args.get_option(&["--tls-hostname"]).or_else(|| tcp_target.split(":").next()); + let hostname = args.get_option(&["--tls-hostname"]).or_else(|| tcp_target.split(":").next().map(&str::to_owned)); let pinnedpubkey = args.get_option(&["--pinnedpubkey"]); - proxy_client.start_tls(hostname, pinnedpubkey).expect("error running tls proxy_client"); + proxy_client.start_tls(hostname.as_ref().map(String::as_str), pinnedpubkey.as_ref().map(String::as_str)).expect("error running tls proxy_client"); } else { proxy_client.start().expect("error running proxy_client"); } @@ -138,7 +150,7 @@ fn server(tcp_host: &str, socket_timeout: u64, args: Args) { ); if tls_key.is_some() && tls_cert.is_some() { - proxy_server.start_tls(tls_key.unwrap(), tls_cert.unwrap()).expect("error running TLS proxy_server"); + proxy_server.start_tls(&tls_key.unwrap(), &tls_cert.unwrap()).expect("error running TLS proxy_server"); } else if tls_key.is_none() && tls_cert.is_none() { proxy_server.start().expect("error running proxy_server"); } else { diff --git a/src/lib.rs b/src/lib.rs index ce39edb..6e309ba 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -24,6 +24,25 @@ mod tls { use tls::{TlsStream, TlsListener}; +fn arg_to_env<'a>(arg: &'a str) -> String { + let env = "WGP_".to_owned(); + let mut env = env + &arg.trim_matches('-').replace("-", "_"); + env.make_ascii_uppercase(); + env +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_arg_to_env() { + assert_eq!(arg_to_env("--tcp-host"), "WGP_TCP_HOST"); + assert_eq!(arg_to_env("--tls"), "WGP_TLS"); + assert_eq!(arg_to_env("-h"), "WGP_H"); + } +} + pub struct Args<'a> { args: &'a Vec, } @@ -33,26 +52,40 @@ impl<'a> Args<'a> { Args { args } } pub fn flag(&self, flag: &'a str) -> bool { - self.args.contains(&flag.to_owned()) + if self.args.contains(&flag.to_owned()) { + return true; + } + // because env we want slightly special handling of empty/0/false + match std::env::var(arg_to_env(flag)) { + Ok(env) => &env != "" && &env != "0" && &env != "false", + Err(_) => false, + } } - pub fn get_option(&self, flags: &[&'a str]) -> Option<&'a str> { + pub fn get_option(&self, flags: &[&'a str]) -> Option { for flag in flags.iter() { let mut found = false; for arg in self.args.iter() { if found { - return Some(arg); + return Some(arg.to_owned()); } if arg == flag { found = true; } } } + // no matching arguments are found, so check env variables as a fallback + for flag in flags.iter() { + let env = std::env::var(arg_to_env(flag)).ok(); + if env.is_some() { + return env; + } + } return None; } - pub fn get_str(&self, flags: &[&'a str], def: &'a str) -> &'a str { + pub fn get_str(&self, flags: &[&'a str], def: &'a str) -> String { match self.get_option(flags) { Some(ret) => ret, - None => def, + None => def.to_owned(), } } pub fn get(&self, flags: &[&'a str], def: T) -> T {