90 lines
2.2 KiB
Rust
90 lines
2.2 KiB
Rust
use std::{fs::File, path::Path};
|
|
|
|
use serde::Deserialize;
|
|
use simple_error::bail;
|
|
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>,
|
|
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).expect("Failed to parse!"); //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).expect("Failed to parse!"); //TODO: Print path
|
|
|
|
return Ok(config);
|
|
}
|
|
|
|
pub fn load_config(&self) -> Result<Config> {
|
|
for path in [
|
|
"config.yaml",
|
|
"config.json",
|
|
"/etc/rustocat.yaml",
|
|
"/etc/rustocat.json",
|
|
]
|
|
.iter()
|
|
{
|
|
// if(p)
|
|
let config = if path.ends_with(".yaml") {
|
|
self.load_yaml(Path::new(path))
|
|
} else {
|
|
self.load_json(Path::new(path))
|
|
};
|
|
if config.is_ok() {
|
|
return config;
|
|
}
|
|
}
|
|
bail!("No config file found");
|
|
}
|
|
}
|
|
|
|
#[async_trait::async_trait]
|
|
impl ConfigProvider for FileConfigProvider {
|
|
async fn get_targets(&self) -> Result<Vec<Target>> {
|
|
let config = self.load_config()?;
|
|
return Ok(config.mappings);
|
|
}
|
|
|
|
async fn wait_for_change(&mut self) -> Result<()> {
|
|
self.sighup_stream.recv().await;
|
|
|
|
return Ok(());
|
|
}
|
|
}
|