Support falling back to environment variables if command line arguments don't exist

README.md

@@ -49,6 +49,18 @@ usage: wireguard-proxy [options...]
  -h, --help                      print this usage text
  -st, --socket-timeout <seconds> 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:

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(

src/bin/wireguard-proxy.rs

@@ -54,6 +54,18 @@ fn main() {
  -h, --help                      print this usage text
  -st, --socket-timeout <seconds> 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 {

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<String>,
 }
@@ -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<String> {
         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<T: FromStr>(&self, flags: &[&'a str], def: T) -> T {