From 161a006ebe680c5fde0ac28b22a464086789875d Mon Sep 17 00:00:00 2001 From: Filip Znachor Date: Fri, 23 Jun 2023 01:32:17 +0200 Subject: [PATCH] Added config reloading --- Cargo.lock | 11 +++++++++++ Cargo.toml | 1 + conf.rs | 3 +++ data.rs | 16 +++++++--------- main.rs | 30 ++++++++++++++++++++++-------- services.rs | 10 +++++----- 6 files changed, 49 insertions(+), 22 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index dc5fa53..ab10bdc 100755 --- a/Cargo.lock +++ b/Cargo.lock @@ -328,6 +328,7 @@ dependencies = [ "lazy_static", "serde", "serde_yaml", + "signal-hook", "tokio", "tower", "url", @@ -508,6 +509,16 @@ dependencies = [ "unsafe-libyaml", ] +[[package]] +name = "signal-hook" +version = "0.3.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "732768f1176d21d09e076c23a93123d40bba92d50c4058da34d45c8de8e682b9" +dependencies = [ + "libc", + "signal-hook-registry", +] + [[package]] name = "signal-hook-registry" version = "1.4.1" diff --git a/Cargo.toml b/Cargo.toml index eb94070..c371aac 100755 --- a/Cargo.toml +++ b/Cargo.toml @@ -14,6 +14,7 @@ serde = { version = "1.0", features = ["derive"] } serde_yaml = "0.9" lazy_static = "1.4.0" url = "2.3.0" +signal-hook = "0.3.15" [[bin]] name = "odproxy" diff --git a/conf.rs b/conf.rs index b264612..2861983 100755 --- a/conf.rs +++ b/conf.rs @@ -7,6 +7,8 @@ use serde::Deserialize; use lazy_static::lazy_static; use serde_yaml::from_str; +use crate::data::{HOST_MAP, generate_host_map}; + lazy_static! { pub static ref CONFIG: Arc> = Arc::new(Mutex::new(load())); } @@ -54,6 +56,7 @@ fn load() -> RootConf { pub fn reload() { let conf: RootConf = load(); *CONFIG.lock().unwrap() = conf; + *HOST_MAP.lock().unwrap() = generate_host_map(); } pub fn get() -> RootConf { diff --git a/data.rs b/data.rs index a51e914..714bfa0 100755 --- a/data.rs +++ b/data.rs @@ -8,7 +8,7 @@ use hyper::http::HeaderValue; use crate::conf::{ProxyConf, self}; lazy_static! { - pub static ref HOST_MAP: HashMap = generate_host_map(); + pub static ref HOST_MAP: Arc>> = Arc::new(Mutex::new(generate_host_map())); pub static ref SERVICES: Arc>> = Arc::new(Mutex::new(HashMap::from_iter(vec![]))); } @@ -28,23 +28,21 @@ impl ServiceData { } } -pub fn get_proxy(name: Option<&String>) -> Option { +pub fn get_proxy(name: Option) -> Option { let c = conf::get(); match name { - Some(name) => c.proxy.get(name).cloned(), + Some(name) => c.proxy.get(&name).cloned(), None => None } } -pub fn get_proxy_name(host: Option<&HeaderValue>) -> Option<&String> { +pub fn get_proxy_name(host: Option<&HeaderValue>) -> Option { match host { Some(host) => { let host_parts: Vec<&str> = host.to_str().unwrap().split(":").collect(); let domain = host_parts.get(0); - match domain { - Some(domain) => HOST_MAP.get(&domain.to_string()), - None => None - } + let host_map = HOST_MAP.lock().ok()?; + host_map.get(&domain?.to_string()).cloned() }, None => None } @@ -55,7 +53,7 @@ pub fn generate_host_map() -> HashMap { for (name, proxy) in conf::get().proxy.iter() { for host in proxy.hosts.iter() { hosts.push((host.to_string(), name.to_string())); - }; + } } HashMap::from_iter(hosts) } \ No newline at end of file diff --git a/main.rs b/main.rs index 6171eb1..9a7f5e1 100755 --- a/main.rs +++ b/main.rs @@ -2,11 +2,12 @@ mod conf; mod data; mod services; -use std::str::FromStr; +use std::{error::Error, thread, str::FromStr}; use data::HOST_MAP; use hyperlocal::UnixClientExt; use services::check_service; use tower::make::Shared; +use signal_hook::{iterator::Signals, consts::SIGHUP}; use hyper::{service::service_fn, Body, Client, Request, Response, Server}; @@ -16,11 +17,12 @@ async fn run(req: Request) -> Result, hyper::Error> { let host = req.headers().get("host"); let name = data::get_proxy_name(host); - let proxy = data::get_proxy(name); - match proxy { - Some(p) => { + let proxy = data::get_proxy(name.clone()); - check_service(name.unwrap(), &p).await; + match (name, proxy) { + (Some(name), Some(p)) => { + + check_service(&name, &p).await; // Create new Request let mut request_builder = Request::builder().method(req.method()); @@ -49,7 +51,7 @@ async fn run(req: Request) -> Result, hyper::Error> { } }, - None => { + _ => { println!("Unknown host accessed: {:?}", host.unwrap()); return Ok(Response::new(Body::empty())); } @@ -62,15 +64,27 @@ async fn main() { prepare_services().await; + _ = register_signals(); + let make_service = Shared::new(service_fn(run)); let server = Server::bind(&conf::get().listen).serve(make_service); - let host_count = HOST_MAP.len(); + let host_count = HOST_MAP.lock().ok().map(|m| m.len()).unwrap_or(0); let service_count = conf::get().proxy.len(); println!("odproxy is listening on {} with {} hosts and {} services", conf::get().listen, host_count, service_count); if let Err(e) = server.await { println!("error: {}", e); } -} \ No newline at end of file +} + +fn register_signals() -> Result<(), Box> { + let mut signals = Signals::new(&[SIGHUP])?; + thread::spawn(move || { + signals.forever().for_each(|_| { + conf::reload() + }); + }); + Ok(()) +} diff --git a/services.rs b/services.rs index 6fae356..e00e0ca 100755 --- a/services.rs +++ b/services.rs @@ -16,7 +16,7 @@ fn target_to_address(target: &str) -> Option { }) } -fn modify_service_data(name: &String, modify_fn: F) +fn modify_service_data(name: &str, modify_fn: F) where F: FnOnce(&mut ServiceData) { let mut hashmap = SERVICES.lock().unwrap(); @@ -31,7 +31,7 @@ pub async fn check_service(name: &String, proxy: &ProxyConf) { let mut ready = false; let mut running = false; - modify_service_data(name, |s| { + modify_service_data(&name, |s| { ready = s.child.is_some(); running = s.running; s.last_active = SystemTime::now().duration_since(UNIX_EPOCH).unwrap().as_secs(); @@ -45,16 +45,16 @@ pub async fn check_service(name: &String, proxy: &ProxyConf) { } if !running { - start_service(name, proxy); + start_service(&name, proxy); wait_for_service(proxy).await; - modify_service_data(name, |s| s.running = true); + modify_service_data(&name, |s| s.running = true); } } } -fn start_service(name: &String, proxy: &ProxyConf) -> bool { +fn start_service(name: &str, proxy: &ProxyConf) -> bool { let mut status = false; let spawn = proxy.spawn.as_ref().unwrap(); modify_service_data(name, |s| {