pub trait Locate: Send + Sync {
fn base_for<'a>(
&'a self,
service: &'a str,
) -> std::pin::Pin<Box<dyn Future<Output = Option<String>> + Send + 'a>>;
}
#[derive(Clone)]
pub struct Upstream {
base: String,
strip: Option<String>,
host: Option<String>,
locate: Option<(String, std::sync::Arc<dyn Locate>)>,
open: bool,
}
impl std::fmt::Debug for Upstream {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("Upstream")
.field("base", &self.base)
.field("strip", &self.strip)
.field("host", &self.host)
.field("discovered", &self.locate.is_some())
.field("open", &self.open)
.finish()
}
}
impl Upstream {
pub fn at(base: impl Into<String>) -> Upstream {
let base = base.into();
Upstream {
base: base.trim_end_matches('/').to_string(),
strip: None,
host: None,
locate: None,
open: false,
}
}
pub fn service(name: impl Into<String>, locate: std::sync::Arc<dyn Locate>) -> Upstream {
let name = name.into().to_ascii_uppercase();
Upstream {
base: format!("service://{name}"),
strip: None,
host: None,
locate: Some((name, locate)),
open: false,
}
}
pub fn strip(mut self, prefix: impl Into<String>) -> Upstream {
self.strip = Some(prefix.into());
self
}
pub fn host(mut self, host: impl Into<String>) -> Upstream {
self.host = Some(host.into());
self
}
pub fn open(mut self) -> Upstream {
self.open = true;
self
}
pub fn is_open(&self) -> bool {
self.open
}
pub fn base(&self) -> &str {
&self.base
}
pub fn forwarded_host(&self) -> Option<&str> {
self.host.as_deref()
}
pub fn service_name(&self) -> Option<&str> {
self.locate.as_ref().map(|(name, _)| name.as_str())
}
pub async fn resolved(&self) -> Option<std::borrow::Cow<'_, Upstream>> {
let Some((name, locate)) = &self.locate else {
return Some(std::borrow::Cow::Borrowed(self));
};
let base = locate.base_for(name).await?;
let mut resolved = self.clone();
resolved.base = base.trim_end_matches('/').to_string();
resolved.locate = None;
Some(std::borrow::Cow::Owned(resolved))
}
pub fn url_for(&self, target: &str) -> String {
let path = match &self.strip {
Some(prefix) => target.strip_prefix(prefix.as_str()).unwrap_or(target),
None => target,
};
let path = if path.is_empty() { "/" } else { path };
let separator = if path.starts_with('/') { "" } else { "/" };
format!("{}{separator}{path}", self.base)
}
}
#[derive(Debug, Clone, Default)]
pub struct Routes {
entries: Vec<(String, Upstream)>,
}
impl Routes {
pub fn new() -> Routes {
Routes::default()
}
pub fn add(&mut self, pattern: impl Into<String>, upstream: Upstream) {
self.entries.push((pattern.into(), upstream));
self.entries.sort_by_key(|(pattern, _)| std::cmp::Reverse(pattern.len()));
}
pub fn find(&self, path: &str) -> Option<&Upstream> {
self.entries.iter().find_map(|(pattern, upstream)| {
match pattern.strip_suffix('*') {
Some(prefix) => path.starts_with(prefix).then_some(upstream),
None => (path == pattern).then_some(upstream),
}
})
}
pub fn is_empty(&self) -> bool {
self.entries.is_empty()
}
pub fn upstreams(&self) -> impl Iterator<Item = (&str, &Upstream)> {
self.entries.iter().map(|(pattern, upstream)| (pattern.as_str(), upstream))
}
}
pub const HOP_BY_HOP: &[&str] = &[
"connection",
"keep-alive",
"proxy-authenticate",
"proxy-authorization",
"te",
"trailer",
"transfer-encoding",
"upgrade",
"content-length",
"host",
];
pub fn is_forwardable(name: &str) -> bool {
let name = name.to_ascii_lowercase();
!HOP_BY_HOP.contains(&name.as_str())
}
#[cfg(test)]
mod tests {
use super::*;
struct Fixed(Option<String>);
impl Locate for Fixed {
fn base_for<'a>(
&'a self,
_service: &'a str,
) -> std::pin::Pin<Box<dyn Future<Output = Option<String>> + Send + 'a>> {
let answer = self.0.clone();
Box::pin(async move { answer })
}
}
#[tokio::test]
async fn a_fixed_upstream_resolves_to_itself_without_copying() {
let upstream = Upstream::at("http://api:9000");
let resolved = upstream.resolved().await.expect("a fixed upstream always resolves");
assert!(matches!(resolved, std::borrow::Cow::Borrowed(_)), "a fixed route was cloned");
assert_eq!(resolved.base(), "http://api:9000");
}
#[tokio::test]
async fn a_named_service_takes_its_address_from_the_resolver() {
let upstream = Upstream::service(
"orders",
std::sync::Arc::new(Fixed(Some("http://10.0.0.7:8080/".into()))),
)
.strip("/api");
assert_eq!(upstream.service_name(), Some("ORDERS"));
assert_eq!(upstream.base(), "service://ORDERS");
let resolved = upstream.resolved().await.expect("an instance was available");
assert_eq!(resolved.base(), "http://10.0.0.7:8080", "the trailing slash survived");
assert_eq!(resolved.url_for("/api/orders/7"), "http://10.0.0.7:8080/orders/7");
}
#[tokio::test]
async fn a_service_with_no_instances_resolves_to_nothing() {
let upstream = Upstream::service("orders", std::sync::Arc::new(Fixed(None)));
assert!(upstream.resolved().await.is_none());
}
#[tokio::test]
async fn resolving_is_not_repeatable() {
let upstream = Upstream::service(
"orders",
std::sync::Arc::new(Fixed(Some("http://10.0.0.7:8080".into()))),
);
let resolved = upstream.resolved().await.unwrap().into_owned();
assert_eq!(resolved.service_name(), None);
assert_eq!(resolved.resolved().await.unwrap().base(), "http://10.0.0.7:8080");
}
#[test]
fn a_wildcard_matches_everything_under_it_and_an_exact_route_only_itself() {
let mut routes = Routes::new();
routes.add("/api/*", Upstream::at("http://api:9002"));
routes.add("/health", Upstream::at("http://health:9003"));
assert_eq!(routes.find("/api/orders").map(Upstream::base), Some("http://api:9002"));
assert_eq!(routes.find("/health").map(Upstream::base), Some("http://health:9003"));
assert!(routes.find("/health/deep").is_none(), "an exact route matched a longer path");
assert!(routes.find("/nothing").is_none());
}
#[test]
fn the_more_specific_route_wins_regardless_of_order() {
let mut general = Routes::new();
general.add("/api/*", Upstream::at("http://general"));
general.add("/api/billing/*", Upstream::at("http://billing"));
let mut reversed = Routes::new();
reversed.add("/api/billing/*", Upstream::at("http://billing"));
reversed.add("/api/*", Upstream::at("http://general"));
for routes in [general, reversed] {
assert_eq!(routes.find("/api/billing/invoices").map(Upstream::base), Some("http://billing"));
assert_eq!(routes.find("/api/orders").map(Upstream::base), Some("http://general"));
}
}
#[test]
fn a_stripped_prefix_does_not_reach_the_upstream() {
let upstream = Upstream::at("https://api.example.com").strip("/maps");
assert_eq!(upstream.url_for("/maps/roads?zoom=3"), "https://api.example.com/roads?zoom=3");
}
#[test]
fn a_path_that_strips_to_nothing_becomes_a_slash() {
let upstream = Upstream::at("http://service/").strip("/maps");
assert_eq!(upstream.url_for("/maps"), "http://service/");
assert_eq!(upstream.url_for("/maps/"), "http://service/");
}
#[test]
fn connection_headers_do_not_cross_the_gateway() {
for refused in ["Connection", "transfer-encoding", "Keep-Alive", "content-length", "Host"] {
assert!(!is_forwardable(refused), "{refused} would have been passed on");
}
for allowed in ["authorization", "content-type", "accept", "x-request-id"] {
assert!(is_forwardable(allowed), "{allowed} was dropped");
}
}
}