Add support for custom headers

1 files changed, 21 insertions, 3 deletions

diff --git a/src/proxy_config.rs b/src/proxy_config.rs
index 31a2659..9d07604 100644
--- a/src/proxy_config.rs
+++ b/src/proxy_config.rs
@@ -22,6 +22,7 @@ pub struct ProxyEntry {
 	pub host: String,
 	pub path_prefix: Option<String>,
 	pub priority: u32,
+	pub add_headers: Vec<(String, String)>,
 
 	// Counts the number of times this proxy server has been called to
 	// This implements a round-robin load balancer if there are multiple
@@ -44,7 +45,7 @@ fn retry_to_time(retries: u32, max_time: Duration) -> Duration {
 	));
 }
 
-fn parse_tricot_tag(target_addr: SocketAddr, tag: &str) -> Option<ProxyEntry> {
+fn parse_tricot_tag(tag: &str, target_addr: SocketAddr, add_headers: &[(String, String)]) -> Option<ProxyEntry> {
 	let splits = tag.split(' ').collect::<Vec<_>>();
 	if (splits.len() != 2 && splits.len() != 3) || splits[0] != "tricot" {
 		return None;
@@ -65,10 +66,20 @@ fn parse_tricot_tag(target_addr: SocketAddr, tag: &str) -> Option<ProxyEntry> {
 		host: host.to_string(),
 		path_prefix,
 		priority,
+		add_headers: add_headers.to_vec(),
 		calls: atomic::AtomicU64::from(0),
 	})
 }
 
+fn parse_tricot_add_header_tag(tag: &str) -> Option<(String, String)> {
+	let splits = tag.split(' ').collect::<Vec<_>>();
+	if splits.len() == 3 && splits[0] == "tricot-add-header" {
+		Some((splits[1].to_string(), splits[2].to_string()))
+	} else {
+		None
+	}
+}
+
 fn parse_consul_catalog(catalog: &ConsulNodeCatalog) -> Vec<ProxyEntry> {
 	let mut entries = vec![];
 
@@ -78,8 +89,16 @@ fn parse_consul_catalog(catalog: &ConsulNodeCatalog) -> Vec<ProxyEntry> {
 			_ => continue,
 		};
 		let addr = SocketAddr::new(ip_addr, svc.port);
+
+		let mut add_headers = vec![];
+		for tag in svc.tags.iter() {
+			if let Some(pair) = parse_tricot_add_header_tag(tag) {
+				add_headers.push(pair);
+			}
+		}
+
 		for tag in svc.tags.iter() {
-			if let Some(ent) = parse_tricot_tag(addr, tag) {
+			if let Some(ent) = parse_tricot_tag(tag, addr, &add_headers[..]) {
 				entries.push(ent);
 			}
 		}
@@ -181,7 +200,6 @@ pub fn spawn_proxy_config_task(consul: Consul) -> watch::Receiver<Arc<ProxyConfi
 				}
 			}
 			let config = ProxyConfig { entries };
-			debug!("Extracted configuration: {:#?}", config);
 
 			tx.send(Arc::new(config)).expect("Internal error");
 		}