Added inject response header functionality

This commit is contained in:
Ara Sadoyan
2025-03-24 18:54:38 +01:00
parent 2c87a330ef
commit f3bca5a001
7 changed files with 106 additions and 15 deletions

View File

@@ -21,29 +21,30 @@ pub struct APIUpstreamProvider {
#[async_trait]
pub trait Discovery {
async fn start(&self, tx: Sender<UpstreamsDashMap>);
async fn start(&self, tx: Sender<(UpstreamsDashMap, Headers)>);
}
#[async_trait]
impl Discovery for APIUpstreamProvider {
async fn start(&self, toreturn: Sender<UpstreamsDashMap>) {
async fn start(&self, toreturn: Sender<(UpstreamsDashMap, Headers)>) {
webserver::run_server(self.address.clone(), toreturn).await;
}
}
#[async_trait]
impl Discovery for FromFileProvider {
async fn start(&self, tx: Sender<UpstreamsDashMap>) {
async fn start(&self, tx: Sender<(UpstreamsDashMap, Headers)>) {
tokio::spawn(watch_file(self.path.clone(), tx.clone()));
}
}
pub async fn watch_file(fp: String, mut toreturn: Sender<UpstreamsDashMap>) {
pub async fn watch_file(fp: String, mut toreturn: Sender<(UpstreamsDashMap, Headers)>) {
sleep(Duration::from_millis(50)).await; // For having nice logs :-)
let file_path = fp.as_str();
let parent_dir = Path::new(file_path).parent().unwrap();
let (local_tx, mut local_rx) = tokio::sync::mpsc::channel::<notify::Result<Event>>(1);
info!("Watching for changes in {:?}", parent_dir);
let snd = load_yaml_to_dashmap(file_path, "filepath");
match snd {
Some(snd) => {
toreturn.send(snd).await.unwrap();

View File

@@ -21,10 +21,12 @@ struct HostConfig {
struct PathConfig {
ssl: bool,
servers: Vec<String>,
headers: Vec<String>,
}
pub fn load_yaml_to_dashmap(d: &str, kind: &str) -> Option<UpstreamsDashMap> {
pub fn load_yaml_to_dashmap(d: &str, kind: &str) -> Option<(UpstreamsDashMap, Headers)> {
let dashmap = UpstreamsDashMap::new();
let headers = DashMap::new();
let mut yaml_data = d.to_string();
match kind {
"filepath" => {
@@ -35,7 +37,7 @@ pub fn load_yaml_to_dashmap(d: &str, kind: &str) -> Option<UpstreamsDashMap> {
}
Err(e) => {
error!("Reading: {}: {:?}", d, e.to_string());
warn!("Running with empty upstreams list, chane it via API");
warn!("Running with empty upstreams list, update it via API");
return None;
}
};
@@ -51,8 +53,17 @@ pub fn load_yaml_to_dashmap(d: &str, kind: &str) -> Option<UpstreamsDashMap> {
Ok(parsed) => {
for (hostname, host_config) in parsed.upstreams {
let path_map = DashMap::new();
let header_list = DashMap::new();
for (path, path_config) in host_config.paths {
let mut server_list = Vec::new();
let mut hl = Vec::new();
for header in path_config.headers.iter().by_ref() {
if let Some((key, val)) = header.split_once(':') {
hl.push((key.to_string(), val.to_string()));
}
}
header_list.insert(path.clone(), hl);
// println!(" {:?} == {:?}", hostname, header_list);
for server in path_config.servers {
if let Some((ip, port_str)) = server.split_once(':') {
if let Ok(port) = port_str.parse::<u16>() {
@@ -62,9 +73,10 @@ pub fn load_yaml_to_dashmap(d: &str, kind: &str) -> Option<UpstreamsDashMap> {
}
path_map.insert(path, (server_list, AtomicUsize::new(0)));
}
headers.insert(hostname.clone(), header_list);
dashmap.insert(hostname, path_map);
}
Some(dashmap)
Some((dashmap, headers))
}
Err(e) => {
error!("Failed to parse upstreams file: {}", e);

View File

@@ -21,6 +21,8 @@ pub fn print_upstreams(upstreams: &UpstreamsDashMap) {
}
pub type UpstreamsDashMap = DashMap<String, DashMap<String, (Vec<(String, u16, bool)>, AtomicUsize)>>;
// pub type HeadersList = DashMap<String, Vec<(String, String)>>;
pub type Headers = DashMap<String, DashMap<String, Vec<(String, String)>>>;
// pub type UpstreamMap = DashMap<String, (Vec<(String, u16)>, AtomicUsize)>;
#[allow(dead_code)]