tricot/src/proxy_config.rs

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

use anyhow::Result;
use opentelemetry::{metrics, KeyValue};

use tokio::{select, sync::watch};
use tracing::*;

use crate::consul;

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

#[derive(Debug, Clone, Eq, PartialEq, Ord, PartialOrd, Hash)]
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,
		}
	}
}

impl std::fmt::Display for HostDescription {
	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
		match self {
			HostDescription::Hostname(h) => write!(f, "{}", h),
			HostDescription::Pattern(p) => write!(f, "[{}]", p.as_str()),
		}
	}
}

#[derive(Debug, Clone)]
pub struct UrlPrefix {
	/// Publicly exposed TLS hostnames for matching this rule
	pub host: HostDescription,

	/// Path prefix for matching this rule
	pub path_prefix: Option<String>,
}

impl PartialEq for UrlPrefix {
	fn eq(&self, other: &Self) -> bool {
		self.host == other.host && self.path_prefix == other.path_prefix
	}
}
impl Eq for UrlPrefix {}

impl UrlPrefix {
	fn new(raw_prefix: &str) -> Option<Self> {
		let (raw_host, path_prefix) = match raw_prefix.find('/') {
			Some(i) => {
				let (host, pp) = raw_prefix.split_at(i);
				(host, Some(pp.to_string()))
			}
			None => (raw_prefix, None),
		};

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

		Some(Self { host, path_prefix })
	}
}

#[derive(Debug)]
pub struct ProxyEntry {
	/// An Url prefix is made of a host and maybe a path prefix
	pub url_prefix: UrlPrefix,
	/// Priority with which this rule is considered (highest first)
	pub priority: u32,

	/// Consul service name
	pub service_name: String,
	/// 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,

	/// Flags for target selection
	pub flags: ProxyEntryFlags,

	/// Add the following headers to all responses returned
	/// when matching this rule
	pub add_headers: Vec<(String, String)>,

	/// Try to match all these redirection before forwarding to the backend
	/// when matching this rule
	pub redirects: Vec<(UrlPrefix, UrlPrefix, u16)>,

	/// Wether or not the domain must be validated before asking a certificate
	/// to let's encrypt (only for Glob patterns)
	pub on_demand_tls_ask: Option<String>,

	/// Number of calls in progress, used to deprioritize slow back-ends
	pub calls_in_progress: atomic::AtomicI64,
	/// Time of last call, used for round-robin selection
	pub last_call: atomic::AtomicI64,
}

impl PartialEq for ProxyEntry {
	fn eq(&self, other: &Self) -> bool {
		self.url_prefix == other.url_prefix
			&& self.priority == other.priority
			&& self.service_name == other.service_name
			&& self.target_addr == other.target_addr
			&& self.https_target == other.https_target
			&& self.flags == other.flags
			&& self.add_headers == other.add_headers
	}
}
impl Eq for ProxyEntry {}

impl ProxyEntry {
	fn new(
		service_name: String,
		frontend: MatchTag,
		target_addr: SocketAddr,
		middleware: &[ConfigTag],
		flags: ProxyEntryFlags,
	) -> Self {
		let (url_prefix, priority, https_target) = match frontend {
			MatchTag::Http(u, p) => (u, p, false),
			MatchTag::HttpWithTls(u, p) => (u, p, true),
		};

		let mut add_headers = vec![];
		let mut redirects = vec![];
		let mut on_demand_tls_ask: Option<String> = None;
		for mid in middleware.into_iter() {
			// LocalLb and GlobalLb are handled in the parent function
			match mid {
				ConfigTag::AddHeader(k, v) => add_headers.push((k.to_string(), v.clone())),
				ConfigTag::AddRedirect(m, r, c) => redirects.push(((*m).clone(), (*r).clone(), *c)),
				ConfigTag::OnDemandTlsAsk(url) => on_demand_tls_ask = Some(url.to_string()),
				ConfigTag::LocalLb | ConfigTag::GlobalLb => (),
			};
		}

		ProxyEntry {
			// id
			service_name,
			// frontend
			url_prefix,
			priority,
			// backend
			target_addr,
			https_target,
			// middleware
			flags,
			add_headers,
			redirects,
			on_demand_tls_ask,
			// internal
			last_call: atomic::AtomicI64::from(0),
			calls_in_progress: atomic::AtomicI64::from(0),
		}
	}
}

#[derive(Debug, Clone, Copy, Eq, PartialEq)]
pub struct ProxyEntryFlags {
	/// Is the target healthy?
	pub healthy: 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,

	/// Is site-wide load balancing enabled for this service?
	pub site_lb: bool,
	/// Is global load balancing enabled for this service?
	pub global_lb: bool,
}

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)?;
		write!(
			f,
			"{}{} {}",
			self.url_prefix.host,
			self.url_prefix.path_prefix.as_deref().unwrap_or_default(),
			self.priority
		)?;
		if !self.flags.healthy {
			write!(f, " UNHEALTHY")?;
		}
		if self.flags.same_node {
			write!(f, " OURSELF")?;
		} else if self.flags.same_site {
			write!(f, " SAME_SITE")?;
		}
		if self.flags.global_lb {
			write!(f, " GLOBAL-LB")?;
		} else if self.flags.site_lb {
			write!(f, " SITE-LB")?;
		}
		if !self.add_headers.is_empty() {
			write!(f, " +Headers: {:?}", self.add_headers)?;
		}
		Ok(())
	}
}

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

#[derive(Debug)]
enum ParsedTag<'a> {
	Frontend(MatchTag),
	Middleware(ConfigTag<'a>),
}

#[derive(Debug)]
enum MatchTag {
	/// HTTP backend (plain text)
	Http(UrlPrefix, u32),
	/// HTTPS backend (TLS encrypted)
	HttpWithTls(UrlPrefix, u32),
}

#[derive(Debug)]
enum ConfigTag<'a> {
	AddHeader(&'a str, String),
	AddRedirect(UrlPrefix, UrlPrefix, u16),
	OnDemandTlsAsk(&'a str),
	GlobalLb,
	LocalLb,
}

fn parse_tricot_tags(tag: &str) -> Option<ParsedTag> {
	let splits = tag.splitn(4, ' ').collect::<Vec<_>>();
	let parsed_tag = match splits.as_slice() {
		["tricot", raw_prefix, maybe_priority @ ..] => {
			// priority is set to 100 when value is invalid or missing
			let priority: u32 = maybe_priority
				.iter()
				.next()
				.map_or(Ok(100), |x| x.parse::<u32>())
				.unwrap_or(100);
			UrlPrefix::new(raw_prefix)
				.map(|prefix| ParsedTag::Frontend(MatchTag::Http(prefix, priority)))
		}
		["tricot-https", raw_prefix, maybe_priority @ ..] => {
			// priority is set to 100 when value is invalid or missing
			let priority: u32 = maybe_priority
				.iter()
				.next()
				.map_or(Ok(100), |x| x.parse::<u32>())
				.unwrap_or(100);
			UrlPrefix::new(raw_prefix)
				.map(|prefix| ParsedTag::Frontend(MatchTag::HttpWithTls(prefix, priority)))
		}
		["tricot-add-header", header_key, header_values @ ..] => Some(ParsedTag::Middleware(
			ConfigTag::AddHeader(header_key, header_values.join(" ")),
		)),
		["tricot-add-redirect", raw_match, raw_replace, maybe_raw_code @ ..] => {
			let (p_match, p_replace) =
				match (UrlPrefix::new(raw_match), UrlPrefix::new(raw_replace)) {
					(Some(m), Some(r)) => (m, r),
					_ => {
						debug!(
							"tag {} is ignored, one of the url prefix can't be parsed",
							tag
						);
						return None;
					}
				};

			if matches!(p_replace.host, HostDescription::Pattern(_)) {
				debug!(
					"tag {} ignored as redirect to a glob pattern is not supported",
					tag
				);
				return None;
			}

			let maybe_parsed_code = maybe_raw_code
				.iter()
				.next()
				.map(|c| c.parse::<u16>().ok())
				.flatten();
			let http_code = match maybe_parsed_code {
				Some(301) => 301,
				Some(302) => 302,
				Some(303) => 303,
				Some(307) => 307,
				_ => {
					debug!(
						"tag {} has a missing or invalid http code, setting it to 302",
						tag
					);
					302
				}
			};

			Some(ParsedTag::Middleware(ConfigTag::AddRedirect(
				p_match, p_replace, http_code,
			)))
		}
		["tricot-on-demand-tls-ask", url, ..] => {
			Some(ParsedTag::Middleware(ConfigTag::OnDemandTlsAsk(url)))
		}
		["tricot-global-lb", ..] => Some(ParsedTag::Middleware(ConfigTag::GlobalLb)),
		["tricot-local-lb", ..] => Some(ParsedTag::Middleware(ConfigTag::LocalLb)),
		_ => None,
	};

	trace!("tag {} parsed as {:?}", tag, parsed_tag);
	parsed_tag
}

fn parse_consul_service(
	s: &consul::catalog::HealthServiceNode,
	mut flags: ProxyEntryFlags,
) -> Vec<ProxyEntry> {
	trace!("Parsing service: {:#?}", s);

	let ip_addr = match s.service.address.parse() {
		Ok(ip) => ip,
		_ => match s.node.address.parse() {
			Ok(ip) => ip,
			_ => {
				warn!(
					"Could not get address for service {} at node {}",
					s.service.service, s.node.node
				);
				return vec![];
			}
		},
	};
	let addr = SocketAddr::new(ip_addr, s.service.port);

	// tag parsing
	let mut collected_middleware = vec![];
	let mut collected_frontends = vec![];
	for tag in s.service.tags.iter() {
		match parse_tricot_tags(tag) {
			Some(ParsedTag::Frontend(x)) => collected_frontends.push(x),
			Some(ParsedTag::Middleware(y)) => collected_middleware.push(y),
			_ => trace!(
				"service {}: tag '{}' could not be parsed",
				s.service.service,
				tag
			),
		}
	}

	// some legacy processing that would need a refactor later
	for mid in collected_middleware.iter() {
		match mid {
			ConfigTag::GlobalLb => flags.global_lb = true,
			ConfigTag::LocalLb => flags.site_lb = true,
			_ => (),
		};
	}

	// build proxy entries
	let entries = collected_frontends
		.into_iter()
		.map(|frt| {
			ProxyEntry::new(
				s.service.service.clone(),
				frt,
				addr,
				collected_middleware.as_ref(),
				flags,
			)
		})
		.collect::<Vec<_>>();

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

	entries
}

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

	let metrics = ProxyConfigMetrics::new(rx.clone());
	let consul = Arc::new(consul);

	tokio::spawn(async move {
		let mut catalog_rx = consul.watch_all_service_health(Duration::from_secs(300));
		let mut local_node_site = None;

		while !*must_exit.borrow() {
			select! {
				_ = catalog_rx.changed() => (),
				_ = must_exit.changed() => continue,
			};

			let services = catalog_rx.borrow_and_update().clone();
			if local_node_site.is_none() {
				for (_, svcnodes) in services.iter() {
					for svcnode in svcnodes.iter() {
						if svcnode.node.node == local_node {
							if let Some(site) = svcnode.node.meta.get("site") {
								local_node_site = Some(site.to_string());
							}
						}
					}
				}
			}

			let mut entries = vec![];

			for (_service, svcnodes) in services.iter() {
				for svcnode in svcnodes.iter() {
					let healthy = !svcnode.checks.iter().any(|x| x.status == "critical");

					let same_node = svcnode.node.node == local_node;
					let same_site = match (svcnode.node.meta.get("site"), local_node_site.as_ref())
					{
						(Some(s1), Some(s2)) => s1 == s2,
						_ => false,
					};

					let flags = ProxyEntryFlags {
						healthy,
						same_node,
						same_site,
						site_lb: false,
						global_lb: false,
					};

					entries.extend(parse_consul_service(&svcnode, flags));
				}
			}

			entries.sort_by_cached_key(|ent| ent.to_string());

			let config = ProxyConfig { entries };

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

		drop(metrics); // ensure Metrics lives up to here
	});

	rx
}

// ----

struct ProxyConfigMetrics {
	_proxy_config_entries: metrics::ObservableGauge<u64>,
}

impl ProxyConfigMetrics {
	fn new(rx: watch::Receiver<Arc<ProxyConfig>>) -> Self {
		let meter = opentelemetry::global::meter("tricot");
		Self {
			_proxy_config_entries: meter
				.u64_observable_gauge("proxy_config_entries")
				.with_callback(move |observer| {
					let mut patterns = HashMap::new();
					for ent in rx.borrow().entries.iter() {
						let attrs = (
							ent.url_prefix.host.to_string(),
							ent.url_prefix.path_prefix.clone().unwrap_or_default(),
							ent.service_name.clone(),
						);
						*patterns.entry(attrs).or_default() += 1;
					}
					for ((host, prefix, svc), num) in patterns {
						observer.observe(
							num,
							&[
								KeyValue::new("host", host),
								KeyValue::new("path_prefix", prefix),
								KeyValue::new("service", svc),
							],
						);
					}
				})
				.with_description("Number of proxy entries (back-ends) configured in Tricot")
				.init(),
		}
	}
}

// ----

#[cfg(test)]
mod tests {
	use super::*;

	#[test]
	fn test_parse_tricot_add_header_tag() {
		match parse_tricot_tags("tricot-add-header Content-Security-Policy default-src 'none'; img-src 'self'; script-src 'self'; style-src 'self'") {
          Some(ParsedTag::Middleware(ConfigTag::AddHeader(name, value))) => {
            assert_eq!(name, "Content-Security-Policy");
            assert_eq!(value, "default-src 'none'; img-src 'self'; script-src 'self'; style-src 'self'");
          }
          _ => panic!("Passed a valid tag but the function says it is not valid")
        }
	}
}