110 lines
2.7 KiB
Rust
110 lines
2.7 KiB
Rust
use std::{fs::File, path::Path};
|
|
|
|
use log::info;
|
|
use serde::Deserialize;
|
|
use tokio::signal::unix::Signal;
|
|
|
|
use crate::Result;
|
|
|
|
#[derive(Debug, Deserialize)]
|
|
pub struct Target {
|
|
pub udp: Option<bool>,
|
|
pub source: String,
|
|
pub targets: Vec<String>,
|
|
}
|
|
|
|
#[derive(Debug, Deserialize)]
|
|
pub struct Config {
|
|
pub consul: Option<bool>,
|
|
pub consul_http_addr: Option<String>,
|
|
pub consul_http_token: Option<String>,
|
|
mappings: Vec<Target>,
|
|
}
|
|
|
|
#[async_trait::async_trait]
|
|
pub trait ConfigProvider {
|
|
async fn get_targets(&self) -> Result<Vec<Target>>;
|
|
async fn wait_for_change(&mut self) -> Result<()>;
|
|
}
|
|
|
|
pub struct FileConfigProvider {
|
|
sighup_stream: Signal,
|
|
}
|
|
|
|
impl FileConfigProvider {
|
|
pub fn new() -> Self {
|
|
Self {
|
|
sighup_stream: tokio::signal::unix::signal(tokio::signal::unix::SignalKind::hangup())
|
|
.expect("Failed to create sighup stream"),
|
|
}
|
|
}
|
|
|
|
fn load_yaml(&self, path: &Path) -> Result<Config> {
|
|
let file = File::open(path)?;
|
|
let config: Config = serde_yaml::from_reader(file)?; //TODO: Print path
|
|
|
|
return Ok(config);
|
|
}
|
|
|
|
fn load_json(&self, path: &Path) -> Result<Config> {
|
|
let file = File::open(path)?;
|
|
let config: Config = serde_json::from_reader(file)?; //TODO: Print path
|
|
|
|
return Ok(config);
|
|
}
|
|
|
|
pub fn load_config(&self) -> Result<Config> {
|
|
let configs = [
|
|
"config.yaml",
|
|
"config.yml",
|
|
"config.json",
|
|
"/etc/rustocat.yaml",
|
|
"/etc/rustocat.json",
|
|
]
|
|
.iter()
|
|
.map(|path| {
|
|
if std::path::Path::new(path).exists() {
|
|
return Some(path);
|
|
} else {
|
|
return None;
|
|
}
|
|
})
|
|
.filter(|p| p.is_some())
|
|
.map(|p| p.unwrap())
|
|
.map(|path| {
|
|
if path.ends_with(".yaml") {
|
|
self.load_yaml(Path::new(path))
|
|
} else {
|
|
self.load_json(Path::new(path))
|
|
}
|
|
});
|
|
|
|
for config in configs {
|
|
match config {
|
|
Ok(c) => return Ok(c),
|
|
Err(e) => {
|
|
info!("Error loading config: {}", e);
|
|
}
|
|
}
|
|
}
|
|
|
|
Err("No (valid) config file found".into())
|
|
}
|
|
}
|
|
|
|
#[async_trait::async_trait]
|
|
impl ConfigProvider for FileConfigProvider {
|
|
async fn get_targets(&self) -> Result<Vec<Target>> {
|
|
info!("Getting targets from file");
|
|
let config = self.load_config()?;
|
|
return Ok(config.mappings);
|
|
}
|
|
|
|
async fn wait_for_change(&mut self) -> Result<()> {
|
|
info!("Waiting for file config change (SIGHUP)");
|
|
self.sighup_stream.recv().await;
|
|
|
|
return Ok(());
|
|
}
|
|
}
|