Added config reloading

This commit is contained in:
Filip Znachor 2023-06-23 01:32:17 +02:00
parent b694c11d78
commit 161a006ebe
6 changed files with 49 additions and 22 deletions

11
Cargo.lock generated
View file

@ -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"

View file

@ -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"

View file

@ -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<Mutex<RootConf>> = 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 {

16
data.rs
View file

@ -8,7 +8,7 @@ use hyper::http::HeaderValue;
use crate::conf::{ProxyConf, self};
lazy_static! {
pub static ref HOST_MAP: HashMap<String, String> = generate_host_map();
pub static ref HOST_MAP: Arc<Mutex<HashMap<String, String>>> = Arc::new(Mutex::new(generate_host_map()));
pub static ref SERVICES: Arc<Mutex<HashMap<String, ServiceData>>> = Arc::new(Mutex::new(HashMap::from_iter(vec![])));
}
@ -28,23 +28,21 @@ impl ServiceData {
}
}
pub fn get_proxy(name: Option<&String>) -> Option<ProxyConf> {
pub fn get_proxy(name: Option<String>) -> Option<ProxyConf> {
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<String> {
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<String, String> {
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)
}

30
main.rs
View file

@ -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<Body>) -> Result<Response<Body>, 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<Body>) -> Result<Response<Body>, 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);
}
}
}
fn register_signals() -> Result<(), Box<dyn Error>> {
let mut signals = Signals::new(&[SIGHUP])?;
thread::spawn(move || {
signals.forever().for_each(|_| {
conf::reload()
});
});
Ok(())
}

View file

@ -16,7 +16,7 @@ fn target_to_address(target: &str) -> Option<SocketAddr> {
})
}
fn modify_service_data<F>(name: &String, modify_fn: F)
fn modify_service_data<F>(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| {