tricot/src/proxy_config.rs

355 lines
8.6 KiB
Rust
Raw Normal View History

use std::collections::HashMap;
use std::net::SocketAddr;
2021-12-07 15:37:22 +00:00
use std::sync::{atomic, Arc};
use std::{cmp, time::Duration};
2021-12-07 16:56:15 +00:00
use anyhow::Result;
use futures::future::BoxFuture;
2021-12-07 16:56:15 +00:00
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
2021-12-07 17:19:51 +00:00
pub add_headers: Vec<(String, String)>,
2021-12-07 15:37:22 +00:00
// 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.
2021-12-08 21:29:08 +00:00
pub calls: atomic::AtomicI64,
}
2021-12-08 10:24:25 +00:00
impl std::fmt::Display for ProxyEntry {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
2021-12-08 21:58:19 +00:00
if self.https_target {
write!(f, "https://")?;
}
2021-12-08 10:24:25 +00:00
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")?;
}
2021-12-08 10:24:25 +00:00
if !self.add_headers.is_empty() {
2021-12-08 11:27:47 +00:00
write!(f, " +Headers: {:?}", self.add_headers)?;
2021-12-08 10:24:25 +00:00
}
2021-12-08 21:29:08 +00:00
write!(f, " ({})", self.calls.load(atomic::Ordering::Relaxed))
2021-12-08 10:24:25 +00:00
}
}
#[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
2021-12-09 11:20:37 +00:00
Duration::from_secs(cmp::min(
max_time.as_secs(),
1.2f64.powf(retries as f64) as u64,
2021-12-09 11:20:37 +00:00
))
}
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<_>>();
2021-12-08 21:58:19 +00:00
if (splits.len() != 2 && splits.len() != 3)
|| (splits[0] != "tricot" && splits[0] != "tricot-https")
{
return None;
}
2021-12-08 16:36:01 +00:00
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,
2021-12-08 21:58:19 +00:00
https_target: (splits[0] == "tricot-https"),
host,
same_node,
same_site,
path_prefix,
priority,
2021-12-07 17:19:51 +00:00
add_headers: add_headers.to_vec(),
2021-12-08 21:29:08 +00:00
calls: atomic::AtomicI64::from(0),
})
}
2021-12-07 17:19:51 +00:00
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> {
2021-12-08 12:28:07 +00:00
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);
2021-12-07 17:19:51 +00:00
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)
};
2021-12-07 17:19:51 +00:00
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);
}
}
}
2021-12-08 12:28:07 +00:00
trace!("Result of parsing catalog:");
for ent in entries.iter() {
trace!(" {}", ent);
}
2021-12-07 16:56:15 +00:00
entries
}
2021-12-07 16:56:15 +00:00
#[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>> {
2021-12-07 12:50:44 +00:00
let (tx, rx) = watch::channel(Arc::new(ProxyConfig {
entries: Vec::new(),
2021-12-07 12:50:44 +00:00
}));
2021-12-07 16:56:15 +00:00
let consul = Arc::new(consul);
tokio::spawn(async move {
2021-12-07 16:56:15 +00:00
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 {
2021-12-07 16:56:15 +00:00
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) {
2021-12-07 16:56:15 +00:00
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());
}
2021-12-07 16:56:15 +00:00
}
}
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 {
2021-12-07 16:56:15 +00:00
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) => {
2021-12-07 16:56:15 +00:00
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
);
2021-12-07 16:56:15 +00:00
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;
}
2021-12-07 16:56:15 +00:00
}
2021-12-07 16:56:15 +00:00
let mut entries = vec![];
for (node_name, watch_state) in nodes.iter() {
2021-12-07 16:56:15 +00:00
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));
2021-12-07 16:56:15 +00:00
}
}
let config = ProxyConfig { entries };
2021-12-07 12:50:44 +00:00
tx.send(Arc::new(config)).expect("Internal error");
}
});
rx
}