CLI

src/main.rs

@@ -1,19 +1,16 @@
 mod challenge;
+mod cli;
+mod config;
 mod http;
 mod policy;
 
 use http::HeaderLineIterator;
-use policy::{CompiledPolicies, Policy};
+use policy::CompiledPolicies;
 
 use rand::Rng;
 use realm_syscall::socket2::TcpKeepalive;
 use regex::bytes::Regex;
-use std::{
-	io::{Read, Write},
-	net::SocketAddr,
-	str::FromStr,
-	time::Duration,
-};
+use std::{io::Write, net::SocketAddr, time::Duration};
 use tokio::{
 	io::{AsyncWriteExt, ReadBuf},
 	net::{TcpSocket, TcpStream},
@@ -23,7 +20,6 @@ use tokio::{
 const SALT_LEN: usize = 16;
 const SECRET_LEN: usize = 32;
 const MAC_LEN: usize = 32;
-const CHALLENGE_TIMEOUT: u64 = 3600;
 const TARGET_ZEROS: u32 = 15;
 
 static CHALLENGE_BODY: &str = include_str!("challenge.html");
@@ -39,35 +35,9 @@ macro_rules! mk_static {
 
 #[tokio::main]
 async fn main() {
-	let mut config_file =
-		std::fs::File::open("example-config.yaml").expect("Cannot open config file");
-	let mut config_str = String::new();
-	config_file
-		.read_to_string(&mut config_str)
-		.expect("Cannot read config file");
-	let config_yaml = saphyr::Yaml::load_from_str(&config_str).expect("Error parsing config");
-	let config_doc = &config_yaml[0];
-	let listen_addr: SocketAddr = config_doc["listen"]
-		.as_str()
-		.expect("Missing listen address in config")
-		.parse()
-		.expect("Invalid listen address");
-	let pass_addr: SocketAddr = config_doc["pass"]
-		.as_str()
-		.expect("Missing pass address in config")
-		.parse()
-		.expect("Invalid pass address");
-	let default_action = policy::Action::from_str(
-		config_doc["default-action"]
-			.as_str()
-			.expect("Missing default action in config"),
-	)
-	.expect("Invalid default action");
-	let policy_groups: Vec<Vec<Policy>> = config_doc["policy-groups"].as_vec().expect("Missing policies in config").into_iter().map(|policy_group| policy_group.as_vec().expect("Missing policies in config").into_iter().map(|policy| Policy {
-		name: policy["name"].as_str().expect("Expected policy name string").to_string(),
-		first_line_regex: policy["first-line"].as_str().expect("Expected policy first line regex string").to_string(),
-		action: policy::Action::from_str(policy["action"].as_str().expect("Expected policy action string")).expect("Invalid policy action"),
-	}).collect()).collect();
+	let cli: cli::Cli = argp::parse_args_or_exit(argp::DEFAULT);
+
+	let config = config::Config::from_file(&cli.config);
 
 	let mut rng = rand::thread_rng();
 
@@ -75,17 +45,18 @@ async fn main() {
 
 	let policy_groups = &*mk_static!(
 		Vec<CompiledPolicies>,
-		policy_groups
+		config
+			.policy_groups
 			.into_iter()
 			.map(CompiledPolicies::new)
 			.collect()
 	);
 
-	let socket = realm_syscall::new_tcp_socket(&listen_addr).unwrap();
+	let socket = realm_syscall::new_tcp_socket(&config.listen_addr).unwrap();
 
 	socket.set_reuse_address(true).ok();
 
-	socket.bind(&listen_addr.into()).unwrap();
+	socket.bind(&config.listen_addr.into()).unwrap();
 	socket.listen(1024).unwrap();
 
 	let listener = tokio::net::TcpListener::from_std(socket.into()).unwrap();
@@ -144,7 +115,7 @@ async fn main() {
 				return;
 			};
 
-			let mut action = default_action;
+			let mut action = config.default_action;
 			for policy_group in policy_groups.iter() {
 				if let Some(policy) = policy_group.evaluate(first_line) {
 					println!("Applying policy {}", policy.name);
@@ -156,7 +127,7 @@ async fn main() {
 			match action {
 				policy::Action::Drop => {}
 				policy::Action::Allow => {
-					do_proxy(pass_addr, client_stream).await;
+					do_proxy(config.pass_addr, client_stream).await;
 				}
 				policy::Action::Challenge => {
 					let mut req_challenge = None;
@@ -187,6 +158,7 @@ async fn main() {
 							&secret,
 							req_user_agent,
 							req_ip,
+							config.challenge_timeout,
 						);
 						allow = dbg!(valid_challenge)
 							&& dbg!(challenge::check_challenge(
@@ -197,7 +169,7 @@ async fn main() {
 					}
 
 					if allow {
-						do_proxy(pass_addr, client_stream).await;
+						do_proxy(config.pass_addr, client_stream).await;
 					} else {
 						let salt: [u8; SALT_LEN] = rand::thread_rng().r#gen();