From 780a0076255afb192ffd55d2c3ec46e4b93b0700 Mon Sep 17 00:00:00 2001 From: Hristo Venev Date: Fri, 17 May 2019 19:29:04 +0300 Subject: Implement command-line config parsing. --- src/main.rs | 179 ++++++++++++++++++++++++++++++++++++++++++++++++++---------- 1 file changed, 150 insertions(+), 29 deletions(-) (limited to 'src/main.rs') diff --git a/src/main.rs b/src/main.rs index d56ba39..3d74674 100644 --- a/src/main.rs +++ b/src/main.rs @@ -9,7 +9,11 @@ extern crate arrayref; use std::ffi::{OsStr, OsString}; use std::time::Instant; -use std::{env, fs, io, process, thread}; +use std::{env, process, thread}; + +#[cfg(feature = "toml")] +use std::{fs, io}; +#[cfg(feature = "toml")] use toml; mod builder; @@ -19,7 +23,8 @@ mod model; mod proto; mod wg; -fn load_config(path: &OsStr) -> io::Result { +#[cfg(feature = "toml")] +fn file_config(path: OsString) -> io::Result { let mut data = String::new(); { use io::Read; @@ -31,6 +36,82 @@ fn load_config(path: &OsStr) -> io::Result { .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e)) } +fn cli_config(args: &mut impl Iterator) -> Option { + use std::str::FromStr; + + let mut cfg = config::Config::default(); + + let mut cur_src: Option<&mut config::Source> = None; + while let Some(key) = args.next() { + let arg; + + if let Some(ref mut s) = cur_src { + if key == "psk" { + arg = args.next()?; + let arg = arg.to_str()?; + s.psk = Some(model::Key::from_str(arg).ok()?); + continue; + } + if key == "ipv4" { + arg = args.next()?; + let arg = arg.to_str()?; + for arg in arg.split(',') { + s.ipv4.insert(model::Ipv4Net::from_str(arg).ok()?); + } + continue; + } + if key == "ipv6" { + arg = args.next()?; + let arg = arg.to_str()?; + for arg in arg.split(',') { + s.ipv6.insert(model::Ipv6Net::from_str(arg).ok()?); + } + continue; + } + if key == "required" { + s.required = true; + continue; + } + } + cur_src = None; + + if key == "min_keepalive" { + arg = args.next()?; + let arg = arg.to_str()?; + cfg.peer_config.min_keepalive = u32::from_str(arg).ok()?; + continue; + } + if key == "max_keepalive" { + arg = args.next()?; + let arg = arg.to_str()?; + cfg.peer_config.max_keepalive = u32::from_str(arg).ok()?; + continue; + } + if key == "refresh_sec" { + arg = args.next()?; + let arg = arg.to_str()?; + cfg.update_config.refresh_sec = u32::from_str(arg).ok()?; + continue; + } + if key == "source" { + let name = args.next()?.into_string().ok()?; + let url = args.next()?.into_string().ok()?; + cur_src = Some(cfg.sources.entry(name).or_insert(config::Source { + url, + psk: None, + ipv4: model::Ipv4Set::new(), + ipv6: model::Ipv6Set::new(), + required: false, + })); + continue; + } + + return None; + } + + Some(cfg) +} + fn usage(argv0: &str) -> i32 { eprintln!( "<1>Invalid arguments. See `{} --help` for more information", @@ -39,13 +120,16 @@ fn usage(argv0: &str) -> i32 { 1 } -fn help(argv0: &str) -> i32 { - println!("Usage:"); - println!( - " {} IFNAME CONFIG - run daemon on iterface", - argv0 +fn help(argv0: &str, _args: &mut impl Iterator) -> i32 { + print!( + "\ +Usage: + {} IFNAME CONFIG - run daemon on iterface + {} --check-source PATH - validate source JSON + {} --cmdline IFNAME ... - run daemon using config passed as arguments +", + argv0, argv0, argv0 ); - println!(" {} --check-source PATH - validate source JSON", argv0); 1 } @@ -57,23 +141,53 @@ fn maybe_get_var(out: &mut Option>, var: impl AsRef) } } -fn run_daemon(argv0: String, args: Vec) -> i32 { - if args.len() != 2 { - return usage(&argv0); +#[cfg(feature = "toml")] +fn run_with_file(argv0: &str, args: &mut impl Iterator) -> i32 { + let ifname = match args.next() { + Some(v) => v, + None => return usage(argv0), + }; + let path = match args.next() { + Some(v) => v, + None => return usage(argv0), + }; + if args.next().is_some() { + return usage(argv0); } - let mut args = args.into_iter(); - let ifname = args.next().unwrap(); - let config_path = args.next().unwrap(); - assert!(args.next().is_none()); - let mut config = match load_config(&config_path) { + let config = match file_config(path) { Ok(c) => c, Err(e) => { eprintln!("<1>Failed to load config: {}", e); - process::exit(1); + return 1; + } + }; + run_daemon(ifname, config) +} + +#[cfg(not(feature = "toml"))] +fn run_with_file(_: &str, _: &mut impl Iterator) -> i32 { + eprintln!("<1>Config loading not supported"); + 1 +} + +fn run_with_cmdline(argv0: &str, args: &mut impl Iterator) -> i32 { + let ifname = match args.next() { + Some(v) => v, + None => return usage(argv0), + }; + + let config = match cli_config(args) { + Some(c) => c, + None => { + eprintln!("<1>Invalid config"); + return 1; } }; + run_daemon(ifname, config) +} +fn run_daemon(ifname: OsString, mut config: config::Config) -> i32 { maybe_get_var(&mut config.cache_directory, "CACHE_DIRECTORY"); maybe_get_var(&mut config.runtime_directory, "RUNTIME_DIRECTORY"); @@ -81,7 +195,7 @@ fn run_daemon(argv0: String, args: Vec) -> i32 { Ok(m) => m, Err(e) => { eprintln!("<1>Failed to open device: {}", e); - process::exit(1); + return 1; } }; @@ -90,7 +204,7 @@ fn run_daemon(argv0: String, args: Vec) -> i32 { Ok(t) => t, Err(e) => { eprintln!("<1>{}", e); - process::exit(1); + return 1; } }; let now = Instant::now(); @@ -101,13 +215,14 @@ fn run_daemon(argv0: String, args: Vec) -> i32 { } } -fn run_check_source(argv0: String, args: Vec) -> i32 { - if args.len() != 1 { - usage(&argv0); +fn run_check_source(argv0: &str, args: &mut impl Iterator) -> i32 { + let path = match args.next() { + Some(v) => v, + None => return usage(argv0), + }; + if args.next().is_some() { + return usage(argv0); } - let mut args = args.into_iter(); - let path = args.next().unwrap(); - assert!(args.next().is_none()); match manager::load_source(&path) { Ok(_) => { @@ -123,10 +238,11 @@ fn run_check_source(argv0: String, args: Vec) -> i32 { fn main() { let mut iter_args = env::args_os(); - let argv0 = iter_args.next().unwrap().to_string_lossy().into_owned(); + let argv0 = iter_args.next().unwrap(); + let argv0 = argv0.to_string_lossy(); let mut args = Vec::new(); - let mut run: for<'a> fn(String, Vec) -> i32 = run_daemon; + let mut run: for<'a> fn(&'a str, &'a mut std::vec::IntoIter) -> i32 = run_with_file; let mut parse_args = true; for arg in iter_args { if !parse_args || !arg.to_string_lossy().starts_with('-') { @@ -134,14 +250,19 @@ fn main() { } else if arg == "--" { parse_args = false; } else if arg == "-h" || arg == "--help" { - process::exit(help(&argv0)); + run = help; + break; } else if arg == "--check-source" { run = run_check_source; parse_args = false; + } else if arg == "--cmdline" { + run = run_with_cmdline; + parse_args = false; } else { usage(&argv0); } } - process::exit(run(argv0, args)); + let mut args = args.into_iter(); + process::exit(run(&argv0, &mut args)); } -- cgit