tricot/src/proxy_config.rs

use std::collections::HashMap;
use std::net::SocketAddr;
use std::sync::{atomic, Arc};
use std::{cmp, time::Duration};

use anyhow::Result;

use futures::future::BoxFuture;
use futures::stream::{FuturesUnordered, StreamExt};

use log::*;
use tokio::{select, sync::watch, time::sleep};

use crate::consul::*;

// ---- Extract proxy config from Consul catalog ----

#[derive(Debug)]
pub enum HostDescription {
	Hostname(String),
	Pattern(glob::Pattern),
}

impl HostDescription {
	fn new(desc: &str) -> Result<Self> {
		if desc.chars().any(|x| matches!(x, '*' | '?' | '[' | ']')) {
			Ok(Self::Pattern(glob::Pattern::new(desc)?))
		} else {
			Ok(Self::Hostname(desc.to_string()))
		}
	}

	pub fn matches(&self, v: &str) -> bool {
		match self {
			Self::Pattern(p) => p.matches(v),
			Self::Hostname(s) => s == v,
		}
	}
}

#[derive(Debug)]
pub struct ProxyEntry {
	/// Publicly exposed TLS hostnames for matching this rule
	pub host: HostDescription,
	/// Path prefix for matching this rule
	pub path_prefix: Option<String>,
	/// Priority with which this rule is considered (highest first)
	pub priority: u32,

	/// Node address (ip+port) to handle requests that match this entry
	pub target_addr: SocketAddr,
	/// Is the target serving HTTPS instead of HTTP?
	pub https_target: bool,

	/// Is the target the same node as we are running on?
	/// (if yes priorize it over other matching targets)
	pub same_node: bool,
	/// Is the target the same site as this node?
	/// (if yes priorize it over other matching targets)
	pub same_site: bool,

	/// Add the following headers to all responses returned
	/// when matching this rule
	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
	// entries for the same host and same path prefix.
	pub calls: atomic::AtomicI64,
}

impl std::fmt::Display for ProxyEntry {
	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
		if self.https_target {
			write!(f, "https://")?;
		}
		write!(f, "{} ", self.target_addr)?;
		match &self.host {
			HostDescription::Hostname(h) => write!(f, "{}", h)?,
			HostDescription::Pattern(p) => write!(f, "Pattern('{}')", p.as_str())?,
		}
		write!(
			f,
			"{} {}",
			self.path_prefix.as_ref().unwrap_or(&String::new()),
			self.priority
		)?;
		if self.same_node {
			write!(f, " OURSELF")?;
		} else if self.same_site {
			write!(f, " SAME_SITE")?;
		}
		if !self.add_headers.is_empty() {
			write!(f, " +Headers: {:?}", self.add_headers)?;
		}
		write!(f, " ({})", self.calls.load(atomic::Ordering::Relaxed))
	}
}

#[derive(Debug)]
pub struct ProxyConfig {
	pub entries: Vec<ProxyEntry>,
}

fn retry_to_time(retries: u32, max_time: Duration) -> Duration {
	// 1.2^x seems to be a good value to exponentially increase time at a good pace
	// eg. 1.2^32 = 341 seconds ~= 5 minutes - ie. after 32 retries we wait 5
	// minutes
	Duration::from_secs(cmp::min(
		max_time.as_secs(),
		1.2f64.powf(retries as f64) as u64,
	))
}

fn parse_tricot_tag(
	tag: &str,
	target_addr: SocketAddr,
	add_headers: &[(String, String)],
	same_node: bool,
	same_site: bool,
) -> Option<ProxyEntry> {
	let splits = tag.split(' ').collect::<Vec<_>>();
	if (splits.len() != 2 && splits.len() != 3)
		|| (splits[0] != "tricot" && splits[0] != "tricot-https")
	{
		return None;
	}

	let (host, path_prefix) = match splits[1].find('/') {
		Some(i) => {
			let (host, pp) = splits[1].split_at(i);
			(host, Some(pp.to_string()))
		}
		None => (splits[1], None),
	};

	let priority = match splits.len() {
		3 => splits[2].parse().ok()?,
		_ => 100,
	};

	let host = match HostDescription::new(host) {
		Ok(h) => h,
		Err(e) => {
			warn!("Invalid hostname pattern {}: {}", host, e);
			return None;
		}
	};

	Some(ProxyEntry {
		target_addr,
		https_target: (splits[0] == "tricot-https"),
		host,
		same_node,
		same_site,
		path_prefix,
		priority,
		add_headers: add_headers.to_vec(),
		calls: atomic::AtomicI64::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,
	same_node: bool,
	same_site: bool,
) -> Vec<ProxyEntry> {
	trace!("Parsing node catalog: {:#?}", catalog);

	let mut entries = vec![];

	for (_, svc) in catalog.services.iter() {
		let ip_addr = match svc.address.parse() {
			Ok(ip) => ip,
			_ => match catalog.node.address.parse() {
				Ok(ip) => ip,
				_ => {
					warn!(
						"Could not get address for service {} at node {}",
						svc.service, catalog.node.node
					);
					continue;
				}
			},
		};
		let addr = SocketAddr::new(ip_addr, svc.port);

		let (same_node, same_site) = if svc.tags.contains(&"tricot-global-lb".into()) {
			(false, false)
		} else if svc.tags.contains(&"tricot-site-lb".into()) {
			(false, same_site)
		} else {
			(same_node, same_site)
		};

		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(tag, addr, &add_headers[..], same_node, same_site) {
				entries.push(ent);
			}
		}
	}

	trace!("Result of parsing catalog:");
	for ent in entries.iter() {
		trace!("    {}", ent);
	}

	entries
}

#[derive(Default)]
struct NodeWatchState {
	last_idx: Option<usize>,
	last_catalog: Option<ConsulNodeCatalog>,
	retries: u32,
}

pub fn spawn_proxy_config_task(
	consul: Consul,
	mut must_exit: watch::Receiver<bool>,
) -> watch::Receiver<Arc<ProxyConfig>> {
	let (tx, rx) = watch::channel(Arc::new(ProxyConfig {
		entries: Vec::new(),
	}));

	let consul = Arc::new(consul);

	tokio::spawn(async move {
		let mut nodes = HashMap::new();
		let mut watches = FuturesUnordered::<BoxFuture<'static, (String, Result<_>)>>::new();

		let mut node_site = HashMap::new();

		while !*must_exit.borrow() {
			let list_nodes = select! {
				ln = consul.list_nodes() => ln,
				_ = must_exit.changed() => continue,
			};

			match list_nodes {
				Ok(consul_nodes) => {
					info!("Watched consul nodes: {:?}", consul_nodes);
					for consul_node in consul_nodes {
						let node = &consul_node.node;
						if !nodes.contains_key(node) {
							nodes.insert(node.clone(), NodeWatchState::default());

							let node = node.to_string();
							let consul = consul.clone();

							watches.push(Box::pin(async move {
								let res = consul.watch_node(&node, None).await;
								(node, res)
							}));
						}
						if let Some(site) = consul_node.meta.get("site") {
							node_site.insert(node.clone(), site.clone());
						}
					}
				}
				Err(e) => {
					warn!("Could not get Consul node list: {}", e);
				}
			}

			let next_watch = select! {
				nw = watches.next() => nw,
				_ = must_exit.changed() => continue,
			};

			let (node, res): (String, Result<_>) = match next_watch {
				Some(v) => v,
				None => {
					warn!("No nodes currently watched in proxy_config.rs");
					sleep(Duration::from_secs(10)).await;
					continue;
				}
			};

			match res {
				Ok((catalog, new_idx)) => {
					let mut watch_state = nodes.get_mut(&node).unwrap();
					watch_state.last_idx = Some(new_idx);
					watch_state.last_catalog = Some(catalog);
					watch_state.retries = 0;

					let idx = watch_state.last_idx;
					let consul = consul.clone();
					watches.push(Box::pin(async move {
						let res = consul.watch_node(&node, idx).await;
						(node, res)
					}));
				}
				Err(e) => {
					let mut watch_state = nodes.get_mut(&node).unwrap();
					watch_state.retries += 1;
					watch_state.last_idx = None;

					let will_retry_in =
						retry_to_time(watch_state.retries, Duration::from_secs(600));
					error!(
						"Failed to query consul for node {}. Will retry in {}s. {}",
						node,
						will_retry_in.as_secs(),
						e
					);

					let consul = consul.clone();
					watches.push(Box::pin(async move {
						sleep(will_retry_in).await;
						let res = consul.watch_node(&node, None).await;
						(node, res)
					}));
					continue;
				}
			}

			let mut entries = vec![];
			for (node_name, watch_state) in nodes.iter() {
				if let Some(catalog) = &watch_state.last_catalog {
					let same_node = *node_name == consul.local_node;
					let same_site =
						match (node_site.get(node_name), node_site.get(&consul.local_node)) {
							(Some(s1), Some(s2)) => s1 == s2,
							_ => false,
						};

					entries.extend(parse_consul_catalog(catalog, same_node, same_site));
				}
			}
			let config = ProxyConfig { entries };

			tx.send(Arc::new(config)).expect("Internal error");
		}
	});

	rx
}