use crate::utils::structs::InnerMap; use crate::web::proxyhttp::LB; use async_trait::async_trait; use std::sync::atomic::Ordering; #[derive(Debug, Clone)] pub struct GetHostsReturHeaders { pub client_headers: Option>, pub server_headers: Option>, } #[async_trait] pub trait GetHost { fn get_host(&self, peer: &str, path: &str, backend_id: Option<&str>) -> Option; fn get_header(&self, peer: &str, path: &str) -> Option; } #[async_trait] impl GetHost for LB { fn get_host(&self, peer: &str, path: &str, backend_id: Option<&str>) -> Option { if let Some(b) = backend_id { if let Some(bb) = self.ump_byid.get(b) { return Some(bb.value().clone()); } } let host_entry = self.ump_upst.get(peer)?; let mut current_path = path.to_string(); let mut best_match: Option = None; loop { if let Some(entry) = host_entry.get(¤t_path) { let (servers, index) = entry.value(); if !servers.is_empty() { let idx = index.fetch_add(1, Ordering::Relaxed) % servers.len(); best_match = Some(servers[idx].clone()); break; } } if let Some(pos) = current_path.rfind('/') { current_path.truncate(pos); } else { break; } } if best_match.is_none() { if let Some(entry) = host_entry.get("/") { let (servers, index) = entry.value(); if !servers.is_empty() { let idx = index.fetch_add(1, Ordering::Relaxed) % servers.len(); best_match = Some(servers[idx].clone()); } } } best_match } fn get_header(&self, peer: &str, path: &str) -> Option { let client_entry = self.client_headers.get(peer)?; let server_entry = self.server_headers.get(peer)?; let mut current_path = path; let mut best_match = None; loop { if let Some(entry) = client_entry.get(current_path) { if !entry.value().is_empty() { best_match = Some(entry.value().clone()); break; } } if let Some(pos) = current_path.rfind('/') { current_path = if pos == 0 { "/" } else { ¤t_path[..pos] }; } else { break; } } current_path = path; let mut serv_match = None; loop { if let Some(entry) = server_entry.get(current_path) { if !entry.value().is_empty() { serv_match = Some(entry.value().clone()); break; } } if let Some(pos) = current_path.rfind('/') { current_path = if pos == 0 { "/" } else { ¤t_path[..pos] }; } else { break; } if best_match.is_none() { if let Some(entry) = server_entry.get("/") { if !entry.value().is_empty() { best_match = Some(entry.value().clone()); break; } } } } let result = GetHostsReturHeaders { client_headers: best_match, server_headers: serv_match, }; Some(result) } }