use std::{
cell::RefCell,
collections::{BTreeMap, HashMap},
net::{IpAddr, SocketAddr},
rc::Rc,
time::Duration,
};
use kawa::Block;
use mio::Token;
use sozu_command::{
logging::ansi_palette,
proto::command::{
Cluster, ListenerType, LoadBalancingAlgorithms, RedirectPolicy, RedirectScheme,
},
state::ClusterId,
};
#[cfg(debug_assertions)]
use super::DebugEvent;
use super::{
BackendId, BackendStatus, Connection, Context, GenericHttpStream, GlobalStreamId, Position,
StreamState,
};
use crate::{
BackendConnectionError, L7ListenerHandler, ListenerHandler, Readiness, RetrieveClusterError,
backends::BackendError,
load_balancing::{affinity_key_from_ip, affinity_key_from_value},
protocol::http::{
editor::{HeaderEditMode, HeaderEditSnapshot, HttpContext, emits_x_forwarded},
parser::compare_no_case,
},
router::{HeaderEdit, RouteResult},
socket::SessionTcpStream,
};
use crate::metrics::names;
macro_rules! log_module_context {
() => {{
let (open, reset, _, _, _) = ansi_palette();
format!("{open}MUX-ROUTER{reset}\t >>>", open = open, reset = reset)
}};
($http_context:expr) => {{
let (open, reset, grey, gray, white) = ansi_palette();
let http_ctx: &HttpContext = &$http_context;
let ctx = http_ctx.log_context();
format!(
"{gray}{ctx}{reset}\t{open}MUX-ROUTER{reset}\t{grey}Session{reset}({gray}frontend{reset}={white}{frontend:?}{reset}, {gray}method{reset}={white}{method:?}{reset}, {gray}authority_bytes{reset}={white}{authority_bytes:?}{reset})\t >>>",
open = open,
reset = reset,
grey = grey,
gray = gray,
white = white,
ctx = ctx,
frontend = http_ctx.session_address,
method = http_ctx.method,
authority_bytes = http_ctx.authority.as_ref().map(String::len),
)
}};
}
fn log_coalescing_accepted(context: &HttpContext, authority: &str, sni: &str, matched_name: &str) {
debug!(
"{} accepted coalesced authority (authority_bytes={}, sni_bytes={}, matched_san_kind={}, matched_san_bytes={})",
log_module_context!(context),
authority.len(),
sni.len(),
if matched_name.starts_with("*.") {
"wildcard"
} else {
"exact"
},
matched_name.len(),
);
}
fn log_sni_authority_mismatch(
context: &HttpContext,
authority: &str,
sni: &str,
certificate_names: Option<&[String]>,
) {
let certificate_sans_count = certificate_names.map_or(0, <[String]>::len);
let certificate_sans_bytes = certificate_names
.into_iter()
.flatten()
.map(String::len)
.fold(0usize, usize::saturating_add);
warn!(
"{} rejecting request: TLS cert SANs do not cover authority (authority_bytes={}, sni_bytes={}, certificate_sans_count={}, certificate_sans_bytes={})",
log_module_context!(context),
authority.len(),
sni.len(),
certificate_sans_count,
certificate_sans_bytes,
);
}
pub(super) struct RoutingView<'a> {
clusters: &'a HashMap<ClusterId, Cluster>,
listener_kind: ListenerType,
backend_retired: Option<&'a dyn Fn(&BackendId) -> bool>,
max_connection_attempts: u32,
}
impl<'a> RoutingView<'a> {
pub(super) fn new(
clusters: &'a HashMap<ClusterId, Cluster>,
listener_kind: ListenerType,
) -> Self {
Self {
clusters,
listener_kind,
backend_retired: None,
max_connection_attempts: sozu_command::config::DEFAULT_MAX_CONNECTION_ATTEMPTS,
}
}
pub(super) fn with_max_connection_attempts(mut self, attempts: u32) -> Self {
self.max_connection_attempts = attempts;
self
}
fn max_connection_attempts(&self, cluster_id: Option<&str>) -> u32 {
cluster_id
.and_then(|cluster_id| self.cluster(cluster_id))
.and_then(|cluster| cluster.max_connection_attempts)
.unwrap_or(self.max_connection_attempts)
}
pub(super) fn with_backend_retired(mut self, retired: &'a dyn Fn(&BackendId) -> bool) -> Self {
self.backend_retired = Some(retired);
self
}
fn is_retired(&self, backend: &BackendId) -> bool {
self.backend_retired.is_some_and(|retired| retired(backend))
}
fn cluster(&self, cluster_id: &str) -> Option<&Cluster> {
self.clusters.get(cluster_id)
}
fn is_http_listener(&self) -> bool {
matches!(self.listener_kind, ListenerType::Http)
}
}
#[derive(Debug)]
pub(super) enum ConnectStep {
Decided(ConnectPlan),
CheckIpLimit(ConnectResume),
}
#[derive(Debug)]
pub(super) struct ConnectResume {
h2: bool,
frontend_should_stick: bool,
ip: IpAddr,
max_connections_per_ip: Option<u64>,
cluster_retry_after: Option<u32>,
max_connections_per_subnet: Option<u64>,
}
impl ConnectResume {
pub(super) fn ip(&self) -> IpAddr {
self.ip
}
pub(super) fn max_connections_per_ip(&self) -> Option<u64> {
self.max_connections_per_ip
}
pub(super) fn cluster_retry_after(&self) -> Option<u32> {
self.cluster_retry_after
}
pub(super) fn max_connections_per_subnet(&self) -> Option<u64> {
self.max_connections_per_subnet
}
}
#[derive(Debug)]
pub(super) enum IpGateVerdict {
Admitted,
AtLimit { retry_after: u32 },
}
#[derive(Debug)]
pub(super) enum ConnectPlan {
Attached,
Dial {
cluster_id: ClusterId,
h2: bool,
frontend_should_stick: bool,
},
}
pub struct BackendConnections {
first: Option<(Token, Connection<SessionTcpStream>)>,
rest: BTreeMap<Token, Connection<SessionTcpStream>>,
}
impl BackendConnections {
pub const fn new() -> Self {
Self {
first: None,
rest: BTreeMap::new(),
}
}
pub fn len(&self) -> usize {
usize::from(self.first.is_some()) + self.rest.len()
}
pub fn is_empty(&self) -> bool {
self.first.is_none()
}
pub fn contains_key(&self, token: &Token) -> bool {
self.get(token).is_some()
}
pub fn get(&self, token: &Token) -> Option<&Connection<SessionTcpStream>> {
match &self.first {
Some((first, connection)) if first == token => Some(connection),
_ => self.rest.get(token),
}
}
pub fn get_mut(&mut self, token: &Token) -> Option<&mut Connection<SessionTcpStream>> {
match &mut self.first {
Some((first, connection)) if first == token => Some(connection),
_ => self.rest.get_mut(token),
}
}
pub fn insert(
&mut self,
token: Token,
connection: Connection<SessionTcpStream>,
) -> Option<Connection<SessionTcpStream>> {
match &mut self.first {
None => {
self.first = Some((token, connection));
None
}
Some((first, held)) if *first == token => Some(std::mem::replace(held, connection)),
Some((first, held)) if token < *first => {
let displaced_token = std::mem::replace(first, token);
let displaced = std::mem::replace(held, connection);
self.rest.insert(displaced_token, displaced)
}
Some(_) => self.rest.insert(token, connection),
}
}
pub fn remove(&mut self, token: &Token) -> Option<Connection<SessionTcpStream>> {
match &self.first {
Some((first, _)) if first == token => {
let removed = std::mem::replace(&mut self.first, self.rest.pop_first());
removed.map(|(_, connection)| connection)
}
_ => self.rest.remove(token),
}
}
pub fn iter(&self) -> BackendConnectionsIter<'_> {
let split: fn(&BackendEntry) -> (&Token, &Connection<SessionTcpStream>) =
|(token, connection)| (token, connection);
self.first.iter().map(split).chain(self.rest.iter())
}
pub fn iter_mut(&mut self) -> BackendConnectionsIterMut<'_> {
let split: fn(&mut BackendEntry) -> (&Token, &mut Connection<SessionTcpStream>) =
|(token, connection)| (&*token, connection);
self.first.iter_mut().map(split).chain(self.rest.iter_mut())
}
pub fn keys(&self) -> impl Iterator<Item = &Token> {
self.iter().map(|(token, _)| token)
}
pub fn values_mut(&mut self) -> impl Iterator<Item = &mut Connection<SessionTcpStream>> {
self.iter_mut().map(|(_, connection)| connection)
}
}
type BackendEntry = (Token, Connection<SessionTcpStream>);
pub type BackendConnectionsIter<'a> = std::iter::Chain<
std::iter::Map<
std::option::Iter<'a, BackendEntry>,
fn(&'a BackendEntry) -> (&'a Token, &'a Connection<SessionTcpStream>),
>,
std::collections::btree_map::Iter<'a, Token, Connection<SessionTcpStream>>,
>;
pub type BackendConnectionsIterMut<'a> = std::iter::Chain<
std::iter::Map<
std::option::IterMut<'a, BackendEntry>,
fn(&'a mut BackendEntry) -> (&'a Token, &'a mut Connection<SessionTcpStream>),
>,
std::collections::btree_map::IterMut<'a, Token, Connection<SessionTcpStream>>,
>;
impl<'a> IntoIterator for &'a BackendConnections {
type Item = (&'a Token, &'a Connection<SessionTcpStream>);
type IntoIter = BackendConnectionsIter<'a>;
fn into_iter(self) -> Self::IntoIter {
self.iter()
}
}
impl<'a> IntoIterator for &'a mut BackendConnections {
type Item = (&'a Token, &'a mut Connection<SessionTcpStream>);
type IntoIter = BackendConnectionsIterMut<'a>;
fn into_iter(self) -> Self::IntoIter {
self.iter_mut()
}
}
impl Default for BackendConnections {
fn default() -> Self {
Self::new()
}
}
impl std::fmt::Debug for BackendConnections {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_map().entries(self.iter()).finish()
}
}
#[derive(Debug)]
pub struct Router {
pub backends: BackendConnections,
pub configured_backend_timeout: Duration,
pub configured_connect_timeout: Duration,
pub(super) fallback_readiness: Readiness,
pub(super) dead_backends: Vec<Token>,
}
impl Router {
pub fn new(configured_backend_timeout: Duration, configured_connect_timeout: Duration) -> Self {
Self {
backends: BackendConnections::new(),
configured_backend_timeout,
configured_connect_timeout,
fallback_readiness: Readiness::new(),
dead_backends: Vec::new(),
}
}
pub(super) fn plan_connect<L: ListenerHandler + L7ListenerHandler>(
&mut self,
stream_id: GlobalStreamId,
context: &mut Context<L>,
view: &RoutingView<'_>,
) -> Result<ConnectStep, BackendConnectionError> {
let stream = &mut context.streams[stream_id];
if !matches!(stream.state, StreamState::Link) {
error!(
"{} stream {} expected to be in Link state, got {:?}",
log_module_context!(stream.context),
stream_id,
stream.state
);
return Err(BackendConnectionError::MaxSessionsMemory);
}
#[cfg(debug_assertions)]
context
.debug
.push(DebugEvent::Str(stream.context.get_route()));
let max_attempts = view.max_connection_attempts(stream.context.cluster_id.as_deref());
if u32::from(stream.attempts) >= max_attempts {
incr!(
names::backend::CONNECT_RETRIES_EXHAUSTED,
stream.context.cluster_id.as_deref(),
stream.context.backend_id.as_deref(),
stream.context.cluster_metrics_incarnation()
);
return Err(BackendConnectionError::MaxConnectionRetries(
stream.context.cluster_id.clone(),
));
}
stream.attempts += 1;
let replaying = stream.front.consumed;
if !replaying {
let (front_ref, stream_context_ref) = {
let stream_split = &mut *stream;
(&mut stream_split.front, &mut stream_split.context)
};
let routed = self
.route_from_request(stream_context_ref, front_ref, &context.listener, view)
.map_err(BackendConnectionError::RetrieveClusterError)?;
let context = &stream.context;
let key = view.cluster(&routed).and_then(|cluster| {
affinity_key(
cluster,
&stream.front,
context.session_address,
StickyCookie {
name: &context.sticky_name,
value: context.sticky_session_found.as_deref(),
},
)
});
stream.context.affinity_key = key;
stream.context.set_cluster_metrics_incarnation(
crate::metrics::capture_cluster_metrics(routed.as_ref()),
);
stream.context.cluster_id = Some(routed);
}
let stream_context = &stream.context;
let Some(cluster_id) = stream_context.cluster_id.as_deref() else {
return Err(BackendConnectionError::ReplayRefused(
"no cluster_id survived the first attempt",
));
};
let (
frontend_should_stick,
frontend_should_redirect_https,
h2,
cluster_max_connections_per_ip,
cluster_retry_after,
cluster_max_connections_per_subnet,
) = view
.cluster(cluster_id)
.map(|cluster| {
(
cluster.sticky_session,
cluster.https_redirect,
cluster.http2.unwrap_or(false),
cluster.max_connections_per_ip,
cluster.retry_after,
cluster.max_connections_per_subnet,
)
})
.unwrap_or((false, false, false, None, None, None));
if replaying && h2 {
return Err(BackendConnectionError::ReplayRefused(
"the cluster switched to HTTP/2 between attempts",
));
}
if frontend_should_redirect_https && view.is_http_listener() {
return Err(BackendConnectionError::RetrieveClusterError(
RetrieveClusterError::HttpsRedirect,
));
}
let Some(ip) = stream_context.session_address.map(|sa| sa.ip()) else {
return self
.decide_after_gate(stream_id, context, view, h2, frontend_should_stick)
.map(ConnectStep::Decided);
};
Ok(ConnectStep::CheckIpLimit(ConnectResume {
h2,
frontend_should_stick,
ip,
max_connections_per_ip: cluster_max_connections_per_ip,
cluster_retry_after,
max_connections_per_subnet: cluster_max_connections_per_subnet,
}))
}
pub(super) fn plan_connect_resume<L: ListenerHandler + L7ListenerHandler>(
&mut self,
stream_id: GlobalStreamId,
context: &mut Context<L>,
view: &RoutingView<'_>,
resume: ConnectResume,
verdict: IpGateVerdict,
) -> Result<ConnectPlan, BackendConnectionError> {
if let IpGateVerdict::AtLimit { retry_after } = verdict {
context.streams[stream_id].context.retry_after_seconds =
Some(retry_after).filter(|v| *v > 0);
return Err(BackendConnectionError::TooManyConnectionsPerIp {
cluster_id: context.streams[stream_id]
.context
.cluster_id
.clone()
.unwrap_or_default(),
});
}
self.decide_after_gate(
stream_id,
context,
view,
resume.h2,
resume.frontend_should_stick,
)
}
fn decide_after_gate<L: ListenerHandler + L7ListenerHandler>(
&mut self,
stream_id: GlobalStreamId,
context: &mut Context<L>,
view: &RoutingView<'_>,
h2: bool,
frontend_should_stick: bool,
) -> Result<ConnectPlan, BackendConnectionError> {
let stream_context = &context.streams[stream_id].context;
let Some(cluster_id) = stream_context.cluster_id.as_ref() else {
error!(
"{} stream {} reached the backend decision without a routed cluster",
log_module_context!(stream_context),
stream_id
);
return Err(BackendConnectionError::MaxSessionsMemory);
};
let mut reuse_token = None;
let mut best_h2_stream_count = usize::MAX;
for (token, backend) in &self.backends {
match (h2, backend.position()) {
(_, Position::Server) => {
error!(
"{} Backend connection unexpectedly behaves like a server",
log_module_context!(stream_context)
);
continue;
}
(_, Position::Client(_, _, BackendStatus::Disconnecting)) => {}
(_, Position::Client(other_cluster_id, backend, _))
if other_cluster_id == cluster_id && view.is_retired(backend) => {}
(true, Position::Client(other_cluster_id, _, BackendStatus::Connected)) => {
if other_cluster_id == cluster_id && !backend.is_draining() {
let Connection::H2(h2c) = backend else {
continue;
};
let stream_count = h2c.core.stream_count();
if stream_count
>= h2c.core.peer_settings.settings_max_concurrent_streams as usize
{
continue;
}
if stream_count < best_h2_stream_count {
best_h2_stream_count = stream_count;
reuse_token = Some(*token);
}
}
}
(true, Position::Client(other_cluster_id, _, BackendStatus::Connecting(_))) => {
if other_cluster_id == cluster_id
&& best_h2_stream_count == usize::MAX
&& matches!(backend, Connection::H2(_))
{
reuse_token = Some(*token)
}
}
(true, Position::Client(other_cluster_id, _, BackendStatus::KeepAlive)) => {
if other_cluster_id == cluster_id && matches!(backend, Connection::H2(_)) {
error!(
"{} ConnectionH2 unexpectedly behaves like H1 with KeepAlive",
log_module_context!(stream_context)
);
}
}
(false, Position::Client(old_cluster_id, _, BackendStatus::KeepAlive)) => {
if old_cluster_id == cluster_id {
reuse_token = Some(*token);
break;
}
}
(false, Position::Client(_, _, BackendStatus::Connected))
| (false, Position::Client(_, _, BackendStatus::Connecting(_))) => {}
}
}
trace!(
"{} connect: (stick={}, h2={}) -> (reuse={:?})",
log_module_context!(stream_context),
frontend_should_stick,
h2,
reuse_token
);
if let Some(token) = reuse_token {
incr!(names::backend::POOL_HIT);
trace!(
"{} reused backend: {:#?}",
log_module_context!(stream_context),
self.backends.get(&token)
);
let Some(backend_conn) = self.backends.get_mut(&token) else {
error!(
"{} reused backend token {:?} missing from backends map",
log_module_context!(stream_context),
token
);
return Err(BackendConnectionError::MaxSessionsMemory);
};
if !backend_conn.start_stream(stream_id, context) {
error!(
"{} Backend rejected stream start (max concurrent streams reached)",
log_module_context!(context.http_context(stream_id))
);
return Err(BackendConnectionError::MaxSessionsMemory);
}
if let Some(backend_conn) = self.backends.get(&token)
&& let Position::Client(_, backend, _) = backend_conn.position()
{
let stream = &mut context.streams[stream_id];
stream.context.backend_id = Some(Rc::clone(&backend.backend_id));
stream.context.backend_address = Some(backend.address);
stream.metrics.backend_id = Some(Rc::clone(&backend.backend_id));
stream.metrics.backend_start();
stream.metrics.backend_connected();
}
context.link_stream(stream_id, token);
return Ok(ConnectPlan::Attached);
}
incr!(names::backend::POOL_MISS);
Ok(ConnectPlan::Dial {
cluster_id: cluster_id.clone(),
h2,
frontend_should_stick,
})
}
pub(super) fn commit_dialed<L: ListenerHandler + L7ListenerHandler>(
&mut self,
stream_id: GlobalStreamId,
context: &mut Context<L>,
token: Token,
connection: Connection<SessionTcpStream>,
) {
self.backends.insert(token, connection);
context.link_stream(stream_id, token);
}
fn route_from_request<L: ListenerHandler + L7ListenerHandler>(
&mut self,
context: &mut HttpContext,
front: &mut super::GenericHttpStream,
listener: &Rc<RefCell<L>>,
view: &RoutingView<'_>,
) -> Result<ClusterId, RetrieveClusterError> {
let (host, uri, method) = match context.extract_route() {
Ok(tuple) => tuple,
Err(cluster_error) => {
error!(
"{} Malformed request in connect (should be caught at parsing) {:?}: {}",
log_module_context!(context),
context,
cluster_error
);
return Err(cluster_error);
}
};
if let Some(sni) = context
.tls_server_name
.as_deref()
.filter(|_| context.strict_sni_binding)
{
let matched: Option<&str> = match context.tls_cert_names.as_deref() {
Some(cert_names) => authority_matched_cert_name(host, cert_names),
None => {
if authority_matches_sni(host, sni) {
Some(sni)
} else {
None
}
}
};
match matched {
Some(matched_name) => {
if !authority_matches_sni(host, sni) && context.tls_alpn == Some("h2") {
incr!(names::h2::COALESCING_ACCEPTED);
log_coalescing_accepted(context, host, sni, matched_name);
}
}
None => {
incr!(names::http::SNI_AUTHORITY_MISMATCH);
log_sni_authority_mismatch(
context,
host,
sni,
context.tls_cert_names.as_deref().map(Vec::as_slice),
);
return Err(RetrieveClusterError::SniAuthorityMismatch {
sni: sni.to_owned(),
authority: host.to_owned(),
});
}
}
}
let route_result = listener.borrow().frontend_from_request(host, uri, method);
let route = match route_result {
Ok(route) => route,
Err(frontend_error) => {
trace!("{} {}", log_module_context!(context), frontend_error);
return Err(RetrieveClusterError::RetrieveFrontend(frontend_error));
}
};
let RouteResult {
cluster_id,
redirect,
redirect_scheme,
redirect_template,
rewritten_host,
rewritten_path,
rewritten_port,
headers_request,
headers_response,
required_auth: frontend_required_auth,
tags,
} = route;
context.tags = tags;
if matches!(context.protocol, crate::Protocol::HTTPS) {
snapshot_response_edits(&mut context.headers_response, &headers_response, |e| {
matches!(e.mode, HeaderEditMode::SetIfAbsent | HeaderEditMode::Set)
});
}
let (legacy_https_redirect, https_redirect_port, authorized_hashes, www_authenticate) =
match cluster_id.as_deref() {
Some(id) => view
.cluster(id)
.map(|c| {
(
c.https_redirect,
c.https_redirect_port,
c.authorized_hashes.clone(),
c.www_authenticate.clone(),
)
})
.unwrap_or((false, None, Vec::new(), None)),
None => (false, None, Vec::new(), None),
};
let redirect_status = match redirect {
RedirectPolicy::Permanent => Some(301u16),
RedirectPolicy::Found => Some(302u16),
RedirectPolicy::PermanentRedirect => Some(308u16),
RedirectPolicy::Forward | RedirectPolicy::Unauthorized => None,
};
if let Some(status_code) = redirect_status {
let scheme = resolve_redirect_scheme(redirect_scheme, context);
let port = rewritten_port.map(|p| p as u32).or(https_redirect_port);
context.redirect_location = Some(build_redirect_location(
scheme,
context,
port,
rewritten_host.as_deref(),
rewritten_path.as_deref(),
));
context.frontend_redirect_template = redirect_template;
context.redirect_status = Some(status_code);
return Err(RetrieveClusterError::HttpsRedirect);
}
if matches!(redirect, RedirectPolicy::Unauthorized) || cluster_id.is_none() {
context.www_authenticate = www_authenticate.clone();
trace!("{} RouteResult::deny", log_module_context!(context));
return Err(RetrieveClusterError::UnauthorizedRoute);
}
let Some(cluster_id) = cluster_id else {
unreachable!("cluster_id was checked Some above")
};
if legacy_https_redirect && view.is_http_listener() {
let port = https_redirect_port;
context.redirect_location =
Some(build_redirect_location("https", context, port, None, None));
}
if frontend_required_auth
&& !crate::protocol::mux::auth::check_basic(front, &authorized_hashes)
{
context.www_authenticate = www_authenticate.clone();
trace!(
"{} basic-auth check failed; emitting 401",
log_module_context!(context)
);
return Err(RetrieveClusterError::UnauthorizedRoute);
}
apply_request_rewrites_and_headers(
front,
context,
rewritten_host.as_deref(),
rewritten_path.as_deref(),
&headers_request,
);
snapshot_response_edits(&mut context.headers_response, &headers_response, |_| true);
Ok(cluster_id)
}
pub fn backend_from_request(
&mut self,
cluster_id: &str,
frontend_should_stick: bool,
context: &mut HttpContext,
selector: &mut dyn BackendSelector,
) -> Result<BackendId, BackendConnectionError> {
let affinity = if frontend_should_stick {
Affinity::Sticky(context.sticky_session_found.as_deref())
} else {
Affinity::Unpinned
};
let selected = selector
.select(cluster_id, affinity, context.affinity_key)
.map_err(|backend_error| {
trace!("{} {}", log_module_context!(context), backend_error);
BackendConnectionError::Backend(backend_error)
})?;
debug_assert_eq!(
selected.sticky_session.is_some(),
frontend_should_stick,
"a selector resolves a sticky cookie exactly when the frontend sticks"
);
if let Some(sticky_session) = selected.sticky_session {
context.sticky_session = Some(sticky_session);
}
context.backend_id = Some(Rc::clone(&selected.backend.backend_id));
context.backend_address = Some(selected.backend.address);
Ok(selected.backend)
}
}
pub trait BackendSelector {
fn select(
&mut self,
cluster_id: &str,
affinity: Affinity<'_>,
key: Option<u64>,
) -> Result<SelectedBackend, BackendError>;
}
pub(crate) fn cluster_reads_affinity_key(cluster: &Cluster) -> bool {
cluster.shard_percent.is_some()
|| matches!(
cluster.load_balancing(),
LoadBalancingAlgorithms::Hrw | LoadBalancingAlgorithms::Maglev
)
}
pub(super) fn affinity_key(
cluster: &Cluster,
front: &GenericHttpStream,
session_address: Option<SocketAddr>,
sticky: StickyCookie<'_>,
) -> Option<u64> {
if !cluster_reads_affinity_key(cluster) {
return None;
}
let buf = front.storage.buffer();
let value = if let Some(name) = cluster.affinity_header.as_deref() {
front.blocks.iter().find_map(|block| match block {
Block::Header(header)
if !header.is_elided()
&& compare_no_case(header.key.data(buf), name.as_bytes()) =>
{
Some(header.val.data(buf))
}
_ => None,
})
} else if let Some(name) = cluster.affinity_cookie.as_deref() {
if name == sticky.name {
sticky.value.map(str::as_bytes)
} else {
front.detached.jar.iter().find_map(|cookie| {
(!cookie.is_elided() && cookie.key.data(buf) == name.as_bytes())
.then(|| cookie.val.data(buf))
})
}
} else {
None
};
debug_assert!(
value.is_none() || cluster.affinity_header.is_some() || cluster.affinity_cookie.is_some(),
"a header or cookie value keys only a cluster that configured its name"
);
let key = match value.filter(|value| !value.is_empty()) {
Some(value) => Some(affinity_key_from_value(value)),
None => session_address.map(|address| affinity_key_from_ip(address.ip())),
};
debug_assert!(
key.is_some() || session_address.is_none(),
"a request with a source address must always be keyed"
);
key
}
#[derive(Clone, Copy, Debug)]
pub(super) struct StickyCookie<'a> {
pub(super) name: &'a str,
pub(super) value: Option<&'a str>,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Affinity<'a> {
Unpinned,
Sticky(Option<&'a str>),
}
#[derive(Debug)]
pub struct SelectedBackend {
pub backend: BackendId,
pub sticky_session: Option<String>,
}
fn apply_request_rewrites_and_headers(
kawa: &mut super::GenericHttpStream,
context: &mut HttpContext,
rewritten_host: Option<&str>,
rewritten_path: Option<&str>,
headers_request: &[HeaderEdit],
) {
use kawa::{Block, Pair, Store};
if rewritten_host.is_none() && rewritten_path.is_none() && headers_request.is_empty() {
return;
}
let original_authority: Option<String> = if rewritten_host.is_some() {
context.authority.clone()
} else {
None
};
if (rewritten_host.is_some() || rewritten_path.is_some())
&& let kawa::StatusLine::Request {
authority,
path,
uri,
..
} = &mut kawa.detached.status_line
{
if let Some(new_host) = rewritten_host {
*authority = Store::from_string(new_host.to_owned());
}
if let Some(new_path) = rewritten_path {
*path = Store::from_string(new_path.to_owned());
*uri = Store::from_string(new_path.to_owned());
}
}
let host_lower = b"host";
let xfh_lower = b"x-forwarded-host";
let rewriting_host = rewritten_host.is_some();
let synthesises_xfh = rewriting_host && emits_x_forwarded(context.forwarded_headers);
let mut keys_to_drop: Vec<Vec<u8>> = Vec::with_capacity(headers_request.len() + 2);
let mut to_insert: Vec<Block> = Vec::with_capacity(headers_request.len() + 2);
let mut operator_overrides_host = false;
let mut operator_overrides_xfh = false;
for edit in headers_request {
let key_is_host = edit.key.eq_ignore_ascii_case(host_lower);
let key_is_xfh = edit.key.eq_ignore_ascii_case(xfh_lower);
operator_overrides_host |= key_is_host;
operator_overrides_xfh |= key_is_xfh;
if edit.val.is_empty() {
keys_to_drop.push(edit.key.iter().map(u8::to_ascii_lowercase).collect());
} else {
to_insert.push(Block::Header(Pair {
key: Store::from_slice(&edit.key),
val: Store::from_slice(&edit.val),
}));
}
}
if rewriting_host || operator_overrides_host {
keys_to_drop.push(host_lower.to_vec());
}
if synthesises_xfh || operator_overrides_xfh {
keys_to_drop.push(xfh_lower.to_vec());
}
let buf_ptr = kawa.storage.buffer();
if !keys_to_drop.is_empty() {
let buf = buf_ptr;
kawa.blocks.retain(|block| {
if let Block::Header(Pair { key, val: _ }) = block {
if matches!(key, Store::Empty) {
return true;
}
let key_bytes = key.data(buf);
let key_lower: Vec<u8> = key_bytes.iter().map(u8::to_ascii_lowercase).collect();
!keys_to_drop
.iter()
.any(|k| k.as_slice() == key_lower.as_slice())
} else {
true
}
});
}
let end_header_idx = super::shared::end_of_headers_index(kawa);
if rewriting_host {
let mut synth: Vec<Block> = Vec::with_capacity(2);
if let Some(new_host) = rewritten_host {
synth.push(Block::Header(Pair {
key: Store::Static(b"Host"),
val: Store::from_string(new_host.to_owned()),
}));
}
if synthesises_xfh && let Some(orig) = original_authority.as_deref() {
synth.push(Block::Header(Pair {
key: Store::Static(b"X-Forwarded-Host"),
val: Store::from_string(orig.to_owned()),
}));
}
context.original_authority = original_authority;
synth.append(&mut to_insert);
to_insert = synth;
}
if !to_insert.is_empty() {
let insert_at = end_header_idx.unwrap_or(kawa.blocks.len());
for (offset, block) in to_insert.into_iter().enumerate() {
kawa.blocks.insert(insert_at + offset, block);
}
}
}
fn snapshot_response_edits<F>(target: &mut Vec<HeaderEditSnapshot>, src: &[HeaderEdit], filter: F)
where
F: Fn(&HeaderEdit) -> bool,
{
target.clear();
for edit in src.iter().filter(|e| filter(e)) {
target.push(HeaderEditSnapshot {
key: edit.key.to_vec(),
val: edit.val.to_vec(),
mode: edit.mode,
});
}
}
fn resolve_redirect_scheme(scheme: RedirectScheme, context: &HttpContext) -> &'static str {
match scheme {
RedirectScheme::UseHttps => "https",
RedirectScheme::UseHttp => "http",
RedirectScheme::UseSame => {
if context.tls_server_name.is_some() {
"https"
} else {
"http"
}
}
}
}
fn build_redirect_location(
scheme: &str,
context: &HttpContext,
port: Option<u32>,
host_override: Option<&str>,
path_override: Option<&str>,
) -> String {
let authority = host_override
.or(context.authority.as_deref())
.unwrap_or_default();
let path = path_override.or(context.path.as_deref()).unwrap_or("/");
let host_only = match authority.rsplit_once(':') {
Some((host, port_part))
if !port_part.is_empty() && port_part.bytes().all(|b| b.is_ascii_digit()) =>
{
host
}
_ => authority,
};
let port_suffix = match port {
Some(80) if scheme == "http" => String::new(),
Some(443) if scheme == "https" => String::new(),
Some(p) => format!(":{p}"),
None => String::new(),
};
format!("{scheme}://{host_only}{port_suffix}{path}")
}
pub(crate) fn authority_matches_sni(authority: &str, sni_lowercased: &str) -> bool {
let host = strip_authority_port(authority);
if host.len() != sni_lowercased.len() {
return false;
}
host.as_bytes()
.iter()
.zip(sni_lowercased.as_bytes())
.all(|(a, b)| a.to_ascii_lowercase() == *b)
}
fn strip_authority_port(authority: &str) -> &str {
match authority.rsplit_once(':') {
Some((h, port)) if !port.is_empty() && port.bytes().all(|b| b.is_ascii_digit()) => h,
_ => authority,
}
}
pub(crate) fn authority_matched_cert_name<'a>(
authority: &str,
names: &'a [String],
) -> Option<&'a str> {
let mut host = strip_authority_port(authority);
if let Some(trimmed) = host.strip_suffix('.') {
host = trimmed;
}
if host.is_empty() {
return None;
}
for entry in names {
if let Some(suffix) = entry.strip_prefix("*.") {
if suffix.contains('*') {
continue;
}
let Some((leftmost, rest)) = host.split_once('.') else {
continue;
};
if leftmost.is_empty() {
continue;
}
if rest.eq_ignore_ascii_case(suffix) {
return Some(entry);
}
continue;
}
if entry.contains('*') {
continue;
}
if host.eq_ignore_ascii_case(entry) {
return Some(entry);
}
}
None
}
#[cfg(test)]
mod tests {
use std::{cell::RefCell, collections::BTreeMap, rc::Rc, time::Duration};
use sozu_command::{
config::ListenerBuilder,
proto::command::{
ForwardedHeaders, PathRule, RequestHttpFrontend, RulePosition, SocketAddress,
},
};
use super::{Router, RoutingView, authority_matches_sni};
use crate::{
L7Proxy,
http::HttpProxy,
protocol::{
http::{editor::HttpContext, parser::Method},
mux::{buffer_source::PoolBufferSource, stream::Stream},
},
};
#[test]
fn a_wildcard_frontends_tags_reach_the_access_log_of_a_concrete_authority() {
const TAG: &str = "owner=team-wildcard-1379";
let output = crate::capture_test_logs(|| {
let port = crate::testing::provide_port();
let address = SocketAddress::new_v4(127, 0, 0, 1, port);
let config = ListenerBuilder::new_http(address)
.to_http(None)
.expect("test http listener config must build");
let parts = crate::testing::prebuild_server(10, 16_384, false)
.expect("test server parts must build");
let pool = parts.pool.clone();
let mut proxy =
HttpProxy::new(parts.registry, parts.sessions, parts.pool, parts.backends);
let token = mio::Token(0);
proxy
.add_listener(config, token)
.expect("test listener must register");
proxy
.add_http_frontend(RequestHttpFrontend {
cluster_id: Some("cluster-1379".to_owned()),
address,
hostname: "*.example.com".to_owned(),
path: PathRule::prefix("/".to_owned()),
position: RulePosition::Tree.into(),
tags: BTreeMap::from([("owner".to_owned(), "team-wildcard-1379".to_owned())]),
..Default::default()
})
.expect("the wildcard frontend must register");
let listener = proxy
.get_listener(&token)
.expect("the registered listener must be reachable");
let proxy: Rc<RefCell<dyn L7Proxy>> = Rc::new(RefCell::new(proxy));
let mut context = HttpContext::new(
rusty_ulid::Ulid::generate(),
rusty_ulid::Ulid::generate(),
crate::Protocol::HTTP,
"127.0.0.1:80"
.parse()
.expect("test public address must parse"),
Some(
"127.0.0.1:12345"
.parse()
.expect("test session address must parse"),
),
"SOZUBALANCEID".to_owned(),
"Sozu-Id".to_owned(),
false,
false,
);
context.authority = Some("foo.example.com".to_owned());
context.path = Some("/".to_owned());
context.method = Some(Method::Get);
let mut stream = Stream::new(
&mut PoolBufferSource::new(Rc::downgrade(&pool)),
context,
crate::protocol::mux::test_support::test_answers(),
65_535,
)
.expect("test stream must check out its buffers");
let mut router = Router::new(Duration::from_secs(10), Duration::from_secs(10));
let (front, stream_context) = {
let split = &mut stream;
(&mut split.front, &mut split.context)
};
let proxy_ref = proxy.borrow();
let view = RoutingView::new(proxy_ref.clusters(), proxy_ref.kind());
router
.route_from_request(stream_context, front, &listener, &view)
.expect("the wildcard frontend must route foo.example.com");
drop(proxy_ref);
let _ = stream.generate_access_log(false, None, listener, None, None);
});
assert!(
output.contains(TAG),
"the access log of a request routed by `*.example.com` must carry that frontend's tags, got: {output}"
);
}
#[test]
fn only_a_host_rewrite_records_the_original_authority() {
let port = crate::testing::provide_port();
let address = SocketAddress::new_v4(127, 0, 0, 1, port);
let config = ListenerBuilder::new_http(address)
.to_http(None)
.expect("test http listener config must build");
let parts = crate::testing::prebuild_server(10, 16_384, false)
.expect("test server parts must build");
let pool = parts.pool.clone();
let mut proxy = HttpProxy::new(parts.registry, parts.sessions, parts.pool, parts.backends);
let token = mio::Token(0);
proxy
.add_listener(config, token)
.expect("test listener must register");
for (hostname, rewrite_host) in [
("rewrite.example.com", Some("internal.example.com")),
("plain.example.com", None),
] {
proxy
.add_http_frontend(RequestHttpFrontend {
cluster_id: Some("cluster-1589".to_owned()),
address,
hostname: hostname.to_owned(),
path: PathRule::prefix("/".to_owned()),
position: RulePosition::Tree.into(),
rewrite_host: rewrite_host.map(str::to_owned),
..Default::default()
})
.unwrap_or_else(|error| panic!("{hostname} must register: {error}"));
}
let listener = proxy
.get_listener(&token)
.expect("the registered listener must be reachable");
let proxy: Rc<RefCell<dyn L7Proxy>> = Rc::new(RefCell::new(proxy));
let mut router = Router::new(Duration::from_secs(10), Duration::from_secs(10));
for (authority, expected) in [
("rewrite.example.com", Some("rewrite.example.com")),
("plain.example.com", None),
] {
let mut context = HttpContext::new(
rusty_ulid::Ulid::generate(),
rusty_ulid::Ulid::generate(),
crate::Protocol::HTTP,
"127.0.0.1:80"
.parse()
.expect("test public address must parse"),
None,
"SOZUBALANCEID".to_owned(),
"Sozu-Id".to_owned(),
false,
false,
);
context.authority = Some(authority.to_owned());
context.path = Some("/".to_owned());
context.method = Some(Method::Get);
let mut stream = Stream::new(
&mut PoolBufferSource::new(Rc::downgrade(&pool)),
context,
crate::protocol::mux::test_support::test_answers(),
65_535,
)
.expect("test stream must check out its buffers");
let split = &mut stream;
let proxy_ref = proxy.borrow();
let view = RoutingView::new(proxy_ref.clusters(), proxy_ref.kind());
router
.route_from_request(&mut split.context, &mut split.front, &listener, &view)
.unwrap_or_else(|error| panic!("{authority} must route: {error}"));
assert_eq!(
split.context.original_authority.as_deref(),
expected,
"{authority}: original_authority is recorded exactly when the host is rewritten",
);
let forwarded_host = split.front.blocks.iter().find_map(|block| match block {
kawa::Block::Header(pair)
if pair.key.data(split.front.storage.buffer()) == b"X-Forwarded-Host" =>
{
Some(pair.val.data(split.front.storage.buffer()).to_vec())
}
_ => None,
});
assert_eq!(
forwarded_host.as_deref(),
expected.map(str::as_bytes),
"{authority}: X-Forwarded-Host carries the authority the client sent, on a rewrite only",
);
}
}
#[test]
fn a_host_rewrite_discloses_x_forwarded_host_only_in_x_forwarded_modes() {
let port = crate::testing::provide_port();
let address = SocketAddress::new_v4(127, 0, 0, 1, port);
let config = ListenerBuilder::new_http(address)
.to_http(None)
.expect("test http listener config must build");
let parts = crate::testing::prebuild_server(10, 16_384, false)
.expect("test server parts must build");
let pool = parts.pool.clone();
let mut proxy = HttpProxy::new(parts.registry, parts.sessions, parts.pool, parts.backends);
let token = mio::Token(0);
proxy
.add_listener(config, token)
.expect("test listener must register");
proxy
.add_http_frontend(RequestHttpFrontend {
cluster_id: Some("cluster-322".to_owned()),
address,
hostname: "rewrite.example.com".to_owned(),
path: PathRule::prefix("/".to_owned()),
position: RulePosition::Tree.into(),
rewrite_host: Some("internal.example.com".to_owned()),
..Default::default()
})
.expect("the rewriting frontend must register");
let listener = proxy
.get_listener(&token)
.expect("the registered listener must be reachable");
let proxy: Rc<RefCell<dyn L7Proxy>> = Rc::new(RefCell::new(proxy));
let mut router = Router::new(Duration::from_secs(10), Duration::from_secs(10));
for (mode, client_xfh, expected) in [
(
ForwardedHeaders::Both,
true,
vec![b"rewrite.example.com".to_vec()],
),
(
ForwardedHeaders::XForwarded,
true,
vec![b"rewrite.example.com".to_vec()],
),
(ForwardedHeaders::Rfc7239, false, vec![]),
(
ForwardedHeaders::None,
true,
vec![b"client.example".to_vec()],
),
] {
let mut context = HttpContext::new(
rusty_ulid::Ulid::generate(),
rusty_ulid::Ulid::generate(),
crate::Protocol::HTTP,
"127.0.0.1:80"
.parse()
.expect("test public address must parse"),
None,
"SOZUBALANCEID".to_owned(),
"Sozu-Id".to_owned(),
false,
false,
);
context.forwarded_headers = mode;
context.authority = Some("rewrite.example.com".to_owned());
context.path = Some("/".to_owned());
context.method = Some(Method::Get);
let mut stream = Stream::new(
&mut PoolBufferSource::new(Rc::downgrade(&pool)),
context,
crate::protocol::mux::test_support::test_answers(),
65_535,
)
.expect("test stream must check out its buffers");
if client_xfh {
stream.front.push_block(kawa::Block::Header(kawa::Pair {
key: kawa::Store::Static(b"X-Forwarded-Host"),
val: kawa::Store::Static(b"client.example"),
}));
}
let split = &mut stream;
let proxy_ref = proxy.borrow();
let view = RoutingView::new(proxy_ref.clusters(), proxy_ref.kind());
router
.route_from_request(&mut split.context, &mut split.front, &listener, &view)
.unwrap_or_else(|error| panic!("{mode:?} must route: {error}"));
let buffer = split.front.storage.buffer();
let forwarded_hosts: Vec<Vec<u8>> = split
.front
.blocks
.iter()
.filter_map(|block| match block {
kawa::Block::Header(pair)
if !pair.is_elided()
&& pair
.key
.data(buffer)
.eq_ignore_ascii_case(b"X-Forwarded-Host") =>
{
Some(pair.val.data(buffer).to_vec())
}
_ => None,
})
.collect();
assert_eq!(
forwarded_hosts, expected,
"{mode:?}: X-Forwarded-Host after a host rewrite"
);
}
}
#[cfg(any(debug_assertions, feature = "logs-debug", feature = "logs-trace"))]
#[test]
fn routing_runtime_logs_bound_method_authority_sni_and_certificate_names() {
use super::{log_coalescing_accepted, log_sni_authority_mismatch};
const METHOD_SECRET: &str = "MUX_ROUTE_METHOD_SECRET_SENTINEL";
const AUTHORITY_SECRET: &str = "MUX_ROUTE_AUTHORITY_SECRET_SENTINEL";
const SNI_SECRET: &str = "MUX_ROUTE_SNI_SECRET_SENTINEL";
const SAN_SECRET: &str = "MUX_ROUTE_SAN_SECRET_SENTINEL";
let long_value = |marker: &str| format!("{marker}{}", "x".repeat(4096));
let method = long_value(METHOD_SECRET);
let authority = long_value(AUTHORITY_SECRET);
let sni = long_value(SNI_SECRET);
let san = long_value(SAN_SECRET);
let method_len = method.len();
let authority_len = authority.len();
let sni_len = sni.len();
let san_len = san.len();
let output = crate::capture_test_logs_at_level("debug", move || {
let mut context = HttpContext::new(
rusty_ulid::Ulid::generate(),
rusty_ulid::Ulid::generate(),
crate::Protocol::HTTPS,
"127.0.0.1:443"
.parse()
.expect("test public address must parse"),
Some(
"127.0.0.1:12345"
.parse()
.expect("test session address must parse"),
),
String::new(),
String::new(),
false,
false,
);
context.method = Some(crate::protocol::http::parser::Method::Custom(method));
context.authority = Some(authority.clone());
context.tls_cert_names = Some(std::sync::Arc::new(vec![san.clone()]));
log_coalescing_accepted(&context, &authority, &sni, &san);
log_sni_authority_mismatch(
&context,
&authority,
&sni,
context.tls_cert_names.as_deref().map(Vec::as_slice),
);
});
for secret in [METHOD_SECRET, AUTHORITY_SECRET, SNI_SECRET, SAN_SECRET] {
assert!(
!output.contains(secret),
"mux routing runtime log leaked {secret}: {output}"
);
}
for metadata in [
format!("bytes={method_len}"),
format!("authority_bytes=Some({authority_len})"),
format!("authority_bytes={authority_len}"),
format!("sni_bytes={sni_len}"),
format!("matched_san_bytes={san_len}"),
"certificate_sans_count=1".to_owned(),
format!("certificate_sans_bytes={san_len}"),
] {
assert!(
output.contains(&metadata),
"mux routing runtime log omitted {metadata}: {output}"
);
}
assert!(
output.len() <= 2048,
"mux routing runtime capture is not bounded: {} bytes",
output.len()
);
}
#[test]
fn match_exact() {
assert!(authority_matches_sni("example.com", "example.com"));
}
#[test]
fn match_different_case() {
assert!(authority_matches_sni("Example.COM", "example.com"));
}
#[test]
fn match_authority_with_port() {
assert!(authority_matches_sni("example.com:8443", "example.com"));
}
#[test]
fn reject_different_host() {
assert!(!authority_matches_sni(
"tenant-b.example.com",
"tenant-a.example.com"
));
}
#[test]
fn reject_substring_attack() {
assert!(!authority_matches_sni("example.co", "example.com"));
assert!(!authority_matches_sni("example.commons", "example.com"));
}
#[test]
fn reject_wildcard_not_expanded() {
assert!(!authority_matches_sni("foo.example.com", "*.example.com"));
}
#[test]
fn ipv6_bracketed_literal_with_port() {
assert!(authority_matches_sni("[::1]:8443", "[::1]"));
}
#[test]
fn ipv6_bracketed_without_port() {
assert!(authority_matches_sni("[::1]", "[::1]"));
}
}
#[cfg(test)]
mod authority_matched_cert_name_tests {
use super::authority_matched_cert_name;
#[test]
fn cert_name_match_exact_single_san() {
let names = vec!["example.com".to_owned()];
assert_eq!(
authority_matched_cert_name("example.com", &names),
Some("example.com"),
);
}
#[test]
fn cert_name_match_wildcard_left_most() {
let names = vec!["*.cleverapps.io".to_owned()];
assert_eq!(
authority_matched_cert_name("staging-3.cleverapps.io", &names),
Some("*.cleverapps.io"),
);
}
#[test]
fn cert_name_reject_wildcard_apex() {
let names = vec!["*.example.com".to_owned()];
assert_eq!(authority_matched_cert_name("example.com", &names), None);
}
#[test]
fn cert_name_reject_wildcard_two_labels() {
let names = vec!["*.example.com".to_owned()];
assert_eq!(authority_matched_cert_name("a.b.example.com", &names), None,);
}
#[test]
fn cert_name_reject_wildcard_not_left_most() {
let names = vec!["foo.*.example.com".to_owned()];
assert_eq!(
authority_matched_cert_name("foo.bar.example.com", &names),
None,
);
}
#[test]
fn cert_name_match_case_insensitive() {
let names = vec!["EXAMPLE.com".to_owned()];
assert!(authority_matched_cert_name("Example.COM", &names).is_some());
}
#[test]
fn cert_name_match_with_port() {
let names = vec!["example.com".to_owned()];
assert!(authority_matched_cert_name("example.com:8443", &names).is_some());
}
#[test]
fn cert_name_match_absolute_form_trailing_dot() {
let names = vec!["example.com".to_owned()];
assert!(authority_matched_cert_name("example.com.", &names).is_some());
assert!(authority_matched_cert_name("example.com.:8443", &names).is_some());
let wildcard = vec!["*.example.com".to_owned()];
assert!(authority_matched_cert_name("foo.example.com.", &wildcard).is_some());
}
#[test]
fn cert_name_match_idn_a_label() {
let names = vec!["xn--bcher-kva.example.com".to_owned()];
assert!(authority_matched_cert_name("xn--bcher-kva.example.com", &names).is_some());
}
#[test]
fn cert_name_reject_empty_names() {
assert_eq!(authority_matched_cert_name("example.com", &[]), None);
}
#[test]
fn cert_name_match_multi_san_one_hit() {
let names = vec!["foo.com".to_owned(), "*.example.org".to_owned()];
assert_eq!(
authority_matched_cert_name("bar.example.org", &names),
Some("*.example.org"),
);
}
#[test]
fn cert_name_reject_substring_attack() {
let names = vec!["*.example.com".to_owned()];
assert_eq!(authority_matched_cert_name("example.commons", &names), None,);
}
#[test]
fn cert_name_ipv6_bracketed_literal_with_port() {
let names = vec!["[::1]".to_owned()];
assert!(authority_matched_cert_name("[::1]:8443", &names).is_some());
}
}
#[cfg(test)]
mod backend_selection_order_tests {
use std::{cell::RefCell, collections::HashMap, net::SocketAddr, rc::Rc, time::Duration};
use mio::Token;
use rusty_ulid::Ulid;
use sozu_command::{
config::ListenerBuilder,
proto::command::{
Cluster, ListenerType, PathRule, RequestHttpFrontend, RulePosition, SocketAddress,
},
};
use super::Router;
use crate::{
BackendConnectionError, L7Proxy,
backends::Backend,
http::{HttpListener, HttpProxy},
pool::Pool,
protocol::{
http::parser::Method,
mux::{
BackendId, BackendRegistry, BackendStatus, Connection, Context, Position,
StreamState,
buffer_source::PoolBufferSource,
h2::H2ConnectionConfig,
h2_flood_detector::H2FloodConfig,
router::{ConnectPlan, ConnectStep, IpGateVerdict, RoutingView},
},
},
socket::SessionTcpStream,
};
fn decided(step: ConnectStep) -> ConnectPlan {
match step {
ConnectStep::Decided(plan) => plan,
ConnectStep::CheckIpLimit(_) => panic!(
"a stream with no session address must not reach the \
per-(cluster, source-IP) gate"
),
}
}
const STAGED_TOKENS: [usize; 4] = [23, 7, 41, 13];
const LOWEST_TOKEN: usize = 7;
const HIGHEST_TOKEN: usize = 41;
const ROUNDS: usize = 24;
const H2_CLUSTER: &str = "cluster-1338-h2";
const H1_CLUSTER: &str = "cluster-1338-h1";
const H2_AUTHORITY: &str = "h2.backend-order.example.com";
const H1_AUTHORITY: &str = "h1.backend-order.example.com";
enum Staged {
ConnectedH2,
ConnectingH2,
KeepAliveH1,
}
impl Staged {
fn cluster(&self) -> &'static str {
match self {
Staged::ConnectedH2 | Staged::ConnectingH2 => H2_CLUSTER,
Staged::KeepAliveH1 => H1_CLUSTER,
}
}
fn authority(&self) -> &'static str {
match self {
Staged::ConnectedH2 | Staged::ConnectingH2 => H2_AUTHORITY,
Staged::KeepAliveH1 => H1_AUTHORITY,
}
}
}
struct RoutingFixture {
proxy: Rc<RefCell<dyn L7Proxy>>,
backends: Rc<RefCell<crate::backends::BackendMap>>,
listener: Rc<RefCell<HttpListener>>,
pool: Rc<RefCell<Pool>>,
}
fn routing_fixture() -> RoutingFixture {
let port = crate::testing::provide_port();
let address = SocketAddress::new_v4(127, 0, 0, 1, port);
let config = ListenerBuilder::new_http(address)
.to_http(None)
.expect("test http listener config must build");
let parts =
crate::testing::prebuild_server(32, 65_536, false).expect("test server must build");
let pool = parts.pool.clone();
let backends = Rc::clone(&parts.backends);
let mut proxy = HttpProxy::new(parts.registry, parts.sessions, parts.pool, parts.backends);
let listener_token = Token(0);
proxy
.add_listener(config, listener_token)
.expect("test listener must register");
for (cluster_id, authority, http2) in [
(H2_CLUSTER, H2_AUTHORITY, Some(true)),
(H1_CLUSTER, H1_AUTHORITY, None),
] {
proxy
.add_cluster(Cluster {
cluster_id: cluster_id.to_owned(),
http2,
..Default::default()
})
.expect("the test cluster must register");
proxy
.add_http_frontend(RequestHttpFrontend {
cluster_id: Some(cluster_id.to_owned()),
address,
hostname: authority.to_owned(),
path: PathRule::prefix("/".to_owned()),
position: RulePosition::Tree.into(),
..Default::default()
})
.expect("the test frontend must register");
}
let listener = proxy
.get_listener(&listener_token)
.expect("the registered listener must be reachable");
RoutingFixture {
proxy: Rc::new(RefCell::new(proxy)),
backends,
listener,
pool,
}
}
fn staged_backend(
pool: &Rc<RefCell<Pool>>,
staged: &Staged,
backend_registry: &mut BackendRegistry,
) -> (Connection<SessionTcpStream>, std::net::TcpStream) {
let (socket, peer) = super::super::test_support::connected_socket();
let backend_address: SocketAddr =
"127.0.0.1:2".parse().expect("backend address must parse");
let backend = Rc::new(RefCell::new(Backend::new(
"test-backend",
backend_address,
None,
None,
None,
)));
let backend = backend_registry.id_for(&backend);
let session_ulid = Ulid::generate();
let socket = SessionTcpStream::new(socket, session_ulid, Some(backend_address));
let mut connection = match staged {
Staged::ConnectedH2 | Staged::ConnectingH2 => Connection::new_h2_client(
session_ulid,
socket,
staged.cluster().into(),
backend,
&mut PoolBufferSource::new(Rc::downgrade(pool)),
Duration::from_secs(30),
H2FloodConfig::default(),
H2ConnectionConfig::default(),
Duration::from_secs(30),
None,
)
.expect("the test pool must hand out a buffer for the H2 backend"),
Staged::KeepAliveH1 => Connection::new_h1_client(
session_ulid,
socket,
staged.cluster().into(),
backend,
Duration::from_secs(30),
),
};
if let Position::Client(_, _, status) = connection.position_mut() {
match staged {
Staged::ConnectedH2 => *status = BackendStatus::Connected,
Staged::ConnectingH2 => {}
Staged::KeepAliveH1 => *status = BackendStatus::KeepAlive,
}
}
(connection, peer)
}
fn selected_backend_token(fixture: &RoutingFixture, staged: &Staged) -> Token {
let pool = &fixture.pool;
let mut context = Context::new(
Ulid::generate(),
Rc::downgrade(pool),
fixture.listener.clone(),
None,
"127.0.0.1:80"
.parse()
.expect("test public address must parse"),
);
let stream_id = context
.create_stream(Ulid::generate(), 65_535)
.expect("the test pool must hand out a stream");
{
let stream = &mut context.streams[stream_id];
stream.state = StreamState::Link;
stream.context.authority = Some(staged.authority().to_owned());
stream.context.path = Some("/".to_owned());
stream.context.method = Some(Method::Get);
}
let mut router = Router::new(Duration::from_secs(10), Duration::from_secs(10));
let mut backend_registry = BackendRegistry::default();
let mut peers = Vec::with_capacity(STAGED_TOKENS.len());
for token in STAGED_TOKENS {
let (connection, peer) = staged_backend(pool, staged, &mut backend_registry);
peers.push(peer);
router.backends.insert(Token(token), connection);
}
let proxy_ref = fixture.proxy.borrow();
let view = RoutingView::new(proxy_ref.clusters(), proxy_ref.kind());
let plan = router
.plan_connect(stream_id, &mut context, &view)
.map(decided)
.expect("the router must reuse one of the staged backends");
assert!(
matches!(plan, ConnectPlan::Attached),
"the staged backends are all reusable, so no dial may be requested: got {plan:?}"
);
match context.streams[stream_id].state {
StreamState::Linked(token) => token,
other => panic!("connect must link the stream to a backend, got {other:?}"),
}
}
fn plan_against_one_staged(
fixture: &RoutingFixture,
staged: &Staged,
retired: bool,
) -> (ConnectPlan, StreamState) {
let mut context = Context::new(
Ulid::generate(),
Rc::downgrade(&fixture.pool),
fixture.listener.clone(),
None,
"127.0.0.1:80"
.parse()
.expect("test public address must parse"),
);
let stream_id = context
.create_stream(Ulid::generate(), 65_535)
.expect("the test pool must hand out a stream");
{
let stream = &mut context.streams[stream_id];
stream.state = StreamState::Link;
stream.context.authority = Some(staged.authority().to_owned());
stream.context.path = Some("/".to_owned());
stream.context.method = Some(Method::Get);
}
let mut router = Router::new(Duration::from_secs(10), Duration::from_secs(10));
let mut backend_registry = BackendRegistry::default();
let (connection, _peer) = staged_backend(&fixture.pool, staged, &mut backend_registry);
if retired {
let Position::Client(_, backend, _) = connection.position() else {
unreachable!("a staged backend is a client connection")
};
backend_registry
.handle(backend)
.expect("the staged backend was interned")
.borrow_mut()
.set_closing();
}
router.backends.insert(Token(LOWEST_TOKEN), connection);
let is_retired = |backend: &BackendId| backend_registry.is_retired(backend);
let proxy_ref = fixture.proxy.borrow();
let view = RoutingView::new(proxy_ref.clusters(), proxy_ref.kind())
.with_backend_retired(&is_retired);
let plan = router
.plan_connect(stream_id, &mut context, &view)
.map(decided)
.expect("routing must resolve the staged cluster");
(plan, context.streams[stream_id].state)
}
#[test]
fn a_retired_backend_connection_is_never_reused() {
let fixture = routing_fixture();
for (staged, arm) in [
(Staged::ConnectedH2, "H2 Connected"),
(Staged::ConnectingH2, "H2 Connecting"),
(Staged::KeepAliveH1, "H1 KeepAlive"),
] {
let (plan, state) = plan_against_one_staged(&fixture, &staged, false);
assert!(
matches!(plan, ConnectPlan::Attached),
"premise, {arm}: a connection to a live backend is reused, got {plan:?}"
);
assert_eq!(state, StreamState::Linked(Token(LOWEST_TOKEN)));
let (plan, state) = plan_against_one_staged(&fixture, &staged, true);
assert!(
matches!(&plan, ConnectPlan::Dial { cluster_id, .. } if &**cluster_id == staged.cluster()),
"{arm}: a connection to a retired backend must not be reused, got {plan:?}"
);
assert_eq!(
state,
StreamState::Link,
"{arm}: the stream must not be linked to the retired connection"
);
}
}
#[test]
fn idle_pooled_backend_names_only_streamless_pooled_connections() {
let fixture = routing_fixture();
for (staged, arm, idle) in [
(Staged::ConnectedH2, "H2 Connected", true),
(Staged::ConnectingH2, "H2 Connecting", false),
(Staged::KeepAliveH1, "H1 KeepAlive", true),
] {
let mut backend_registry = BackendRegistry::default();
let (connection, _peer) = staged_backend(&fixture.pool, &staged, &mut backend_registry);
assert_eq!(
connection.idle_pooled_backend().is_some(),
idle,
"{arm}: idle_pooled_backend must answer {idle}"
);
}
}
#[test]
fn an_admitted_stream_is_counted_before_the_decision_resumes() {
let fixture = routing_fixture();
let sessions = fixture.proxy.borrow().sessions();
let source: SocketAddr = "203.0.113.7:51000"
.parse()
.expect("test source address must parse");
let mut clusters = HashMap::new();
clusters.insert(
H2_CLUSTER.into(),
Cluster {
cluster_id: H2_CLUSTER.to_owned(),
max_connections_per_ip: Some(1),
..Default::default()
},
);
let view = RoutingView::new(&clusters, ListenerType::Http);
let run = |token: Token| {
let mut context = Context::new(
Ulid::generate(),
Rc::downgrade(&fixture.pool),
fixture.listener.clone(),
Some(source),
"127.0.0.1:80"
.parse()
.expect("test public address must parse"),
);
let stream_id = context
.create_stream(Ulid::generate(), 65_535)
.expect("the test pool must hand out a stream");
{
let stream = &mut context.streams[stream_id];
stream.state = StreamState::Link;
stream.context.authority = Some(H2_AUTHORITY.to_owned());
stream.context.path = Some("/".to_owned());
stream.context.method = Some(Method::Get);
}
let mut router = Router::new(Duration::from_secs(10), Duration::from_secs(10));
let step = router
.plan_connect(stream_id, &mut context, &view)
.expect("routing must resolve the staged cluster");
let ConnectStep::CheckIpLimit(resume) = step else {
panic!("a stream WITH a session address must reach the per-IP gate")
};
let verdict = super::super::consult_ip_gate(
&sessions,
token,
context
.http_context(stream_id)
.cluster_id
.as_ref()
.expect("plan_connect stamps the routed cluster"),
&resume,
);
let admitted = matches!(verdict, IpGateVerdict::Admitted);
let outcome =
router.plan_connect_resume(stream_id, &mut context, &view, resume, verdict);
let retry_after = context.streams[stream_id].context.retry_after_seconds;
(admitted, outcome, retry_after)
};
let (admitted, outcome, _) = run(Token(1));
assert!(
admitted,
"premise: the first token must be under the limit, or the second proves nothing"
);
assert!(
matches!(outcome, Ok(ConnectPlan::Dial { .. })),
"an admitted stream must go on to request a dial, got {outcome:?}"
);
let (admitted, outcome, retry_after) = run(Token(2));
assert!(
!admitted,
"a SECOND token from the same source IP must find the slot already counted — \
an Admitted verdict here means the first stream resumed without being tracked"
);
assert!(
matches!(
outcome,
Err(BackendConnectionError::TooManyConnectionsPerIp { .. })
),
"the refused stream must surface TooManyConnectionsPerIp, got {outcome:?}"
);
assert_eq!(
retry_after, None,
"with no cluster or global override the resolved Retry-After is 0, which is \
stashed as None so the 429 mapping elides the header rather than sending 0"
);
}
#[test]
fn a_subnet_cap_refuses_a_fresh_address_from_an_already_counted_subnet() {
let fixture = routing_fixture();
let sessions = fixture.proxy.borrow().sessions();
{
let mut manager = sessions.borrow_mut();
manager.subnet_ipv4_prefix = 24;
manager.max_connections_per_subnet = 3;
manager.max_connections_per_ip = 10;
}
let mut clusters = HashMap::new();
clusters.insert(
H2_CLUSTER.into(),
Cluster {
cluster_id: H2_CLUSTER.to_owned(),
..Default::default()
},
);
let view = RoutingView::new(&clusters, ListenerType::Http);
let run = |token: Token, source: SocketAddr| {
let mut context = Context::new(
Ulid::generate(),
Rc::downgrade(&fixture.pool),
fixture.listener.clone(),
Some(source),
"127.0.0.1:80"
.parse()
.expect("test public address must parse"),
);
let stream_id = context
.create_stream(Ulid::generate(), 65_535)
.expect("the test pool must hand out a stream");
{
let stream = &mut context.streams[stream_id];
stream.state = StreamState::Link;
stream.context.authority = Some(H2_AUTHORITY.to_owned());
stream.context.path = Some("/".to_owned());
stream.context.method = Some(Method::Get);
}
let mut router = Router::new(Duration::from_secs(10), Duration::from_secs(10));
let step = router
.plan_connect(stream_id, &mut context, &view)
.expect("routing must resolve the staged cluster");
let ConnectStep::CheckIpLimit(resume) = step else {
panic!("a stream WITH a session address must reach the connection gate")
};
let verdict = super::super::consult_ip_gate(
&sessions,
token,
context
.http_context(stream_id)
.cluster_id
.as_ref()
.expect("plan_connect stamps the routed cluster"),
&resume,
);
let admitted = matches!(verdict, IpGateVerdict::Admitted);
let outcome =
router.plan_connect_resume(stream_id, &mut context, &view, resume, verdict);
(admitted, outcome)
};
let address = |host: u8| -> SocketAddr {
format!("203.0.113.{host}:51000")
.parse()
.expect("test source address must parse")
};
for (index, host) in [1u8, 2, 3].into_iter().enumerate() {
let (admitted, outcome) = run(Token(index + 1), address(host));
assert!(
admitted,
"premise: 203.0.113.{host} is the {} of 3 allowed in the subnet and must be \
admitted, or the refusal below proves nothing",
index + 1
);
assert!(
matches!(outcome, Ok(ConnectPlan::Dial { .. })),
"an admitted stream must go on to request a dial, got {outcome:?}"
);
}
{
let manager = sessions.borrow();
for host in [1u8, 2, 3] {
let ip = address(host).ip();
assert!(
!manager.cluster_ip_at_limit(Token(99), H2_CLUSTER, &ip, None),
"premise: 203.0.113.{host} holds 1 of 10 per-IP slots and must be under \
its own cap — otherwise the refusal below is the per-IP cap, not the subnet"
);
}
}
let (admitted, outcome) = run(Token(4), address(4));
assert!(
!admitted,
"a fresh address from an already-counted /24 must be refused — an Admitted \
verdict here is the #1270 bypass: a new address per connection defeats the cap"
);
assert!(
matches!(
outcome,
Err(BackendConnectionError::TooManyConnectionsPerIp { .. })
),
"the refused stream must surface the too-many-connections error, got {outcome:?}"
);
let (admitted, _) = run(
Token(5),
"198.51.100.4:51000"
.parse()
.expect("test source address must parse"),
);
assert!(
admitted,
"an address in an UNCOUNTED subnet must still be admitted — otherwise the cap \
is refusing on something other than the subnet key"
);
}
#[test]
fn an_unset_subnet_cap_leaves_the_admission_path_untouched() {
let fixture = routing_fixture();
let sessions = fixture.proxy.borrow().sessions();
{
let mut manager = sessions.borrow_mut();
manager.subnet_ipv4_prefix = 24;
manager.max_connections_per_subnet = 0;
manager.max_connections_per_ip = 1;
}
let mut clusters = HashMap::new();
clusters.insert(
H2_CLUSTER.into(),
Cluster {
cluster_id: H2_CLUSTER.to_owned(),
..Default::default()
},
);
let view = RoutingView::new(&clusters, ListenerType::Http);
let run = |token: Token, source: SocketAddr| {
let mut context = Context::new(
Ulid::generate(),
Rc::downgrade(&fixture.pool),
fixture.listener.clone(),
Some(source),
"127.0.0.1:80"
.parse()
.expect("test public address must parse"),
);
let stream_id = context
.create_stream(Ulid::generate(), 65_535)
.expect("the test pool must hand out a stream");
{
let stream = &mut context.streams[stream_id];
stream.state = StreamState::Link;
stream.context.authority = Some(H2_AUTHORITY.to_owned());
stream.context.path = Some("/".to_owned());
stream.context.method = Some(Method::Get);
}
let mut router = Router::new(Duration::from_secs(10), Duration::from_secs(10));
let step = router
.plan_connect(stream_id, &mut context, &view)
.expect("routing must resolve the staged cluster");
let ConnectStep::CheckIpLimit(resume) = step else {
panic!("a stream WITH a session address must reach the connection gate")
};
let verdict = super::super::consult_ip_gate(
&sessions,
token,
context
.http_context(stream_id)
.cluster_id
.as_ref()
.expect("plan_connect stamps the routed cluster"),
&resume,
);
matches!(verdict, IpGateVerdict::Admitted)
};
for (index, host) in [1u8, 2, 3, 4].into_iter().enumerate() {
let source: SocketAddr = format!("203.0.113.{host}:51000")
.parse()
.expect("test source address must parse");
assert!(
run(Token(index + 1), source),
"with max_connections_per_subnet = 0, 203.0.113.{host} must be admitted — \
the default must not gate traffic the previous release allowed"
);
}
assert!(
sessions.borrow().subnet_tracking_is_empty(),
"a disabled subnet limiter must leave both subnet maps empty — a populated one \
means every deployment that never enables this feature pays for it anyway"
);
let repeat: SocketAddr = "203.0.113.1:51000"
.parse()
.expect("test source address must parse");
assert!(
!run(Token(90), repeat),
"the per-IP cap must still refuse a second connection from the same address — \
this feature must not have relaxed the limiter that already existed"
);
}
#[test]
fn the_routing_decision_reads_cluster_config_from_the_view() {
let fixture = routing_fixture();
let pool = &fixture.pool;
let mut context = Context::new(
Ulid::generate(),
Rc::downgrade(pool),
fixture.listener.clone(),
None,
"127.0.0.1:80"
.parse()
.expect("test public address must parse"),
);
let stream_id = context
.create_stream(Ulid::generate(), 65_535)
.expect("the test pool must hand out a stream");
{
let stream = &mut context.streams[stream_id];
stream.state = StreamState::Link;
stream.context.authority = Some(H2_AUTHORITY.to_owned());
stream.context.path = Some("/".to_owned());
stream.context.method = Some(Method::Get);
}
let mut clusters = HashMap::new();
clusters.insert(
H2_CLUSTER.into(),
Cluster {
cluster_id: H2_CLUSTER.to_owned(),
http2: Some(false),
..Default::default()
},
);
assert_eq!(
fixture
.proxy
.borrow()
.clusters()
.get(H2_CLUSTER)
.and_then(|c| c.http2),
Some(true),
"premise: the proxy and the view must disagree, or this proves nothing"
);
let view = RoutingView::new(&clusters, ListenerType::Http);
let mut router = Router::new(Duration::from_secs(10), Duration::from_secs(10));
let plan = router
.plan_connect(stream_id, &mut context, &view)
.map(decided)
.expect("routing must resolve the staged cluster");
match plan {
ConnectPlan::Dial {
ref cluster_id, h2, ..
} => {
assert_eq!(
&**cluster_id, H2_CLUSTER,
"premise: routing resolved a cluster"
);
assert!(
!h2,
"the plan must carry the view's `http2`, not the proxy's — \
`h2: true` means the decision read `L7Proxy::clusters` behind the view"
);
}
ConnectPlan::Attached => {
panic!("premise: an empty router has nothing to reuse, so it must ask for a dial")
}
}
}
#[test]
fn a_dial_plan_performs_no_effect_of_its_own() {
let fixture = routing_fixture();
let pool = &fixture.pool;
let mut context = Context::new(
Ulid::generate(),
Rc::downgrade(pool),
fixture.listener.clone(),
None,
"127.0.0.1:80"
.parse()
.expect("test public address must parse"),
);
let stream_id = context
.create_stream(Ulid::generate(), 65_535)
.expect("the test pool must hand out a stream");
{
let stream = &mut context.streams[stream_id];
stream.state = StreamState::Link;
stream.context.authority = Some(H2_AUTHORITY.to_owned());
stream.context.path = Some("/".to_owned());
stream.context.method = Some(Method::Get);
}
let mut router = Router::new(Duration::from_secs(10), Duration::from_secs(10));
let proxy_ref = fixture.proxy.borrow();
let view = RoutingView::new(proxy_ref.clusters(), proxy_ref.kind());
let plan = router
.plan_connect(stream_id, &mut context, &view)
.map(decided)
.expect("routing must resolve the staged cluster");
match &plan {
ConnectPlan::Dial { cluster_id, .. } => assert_eq!(
&**cluster_id, H2_CLUSTER,
"the dial request must name the cluster routing resolved"
),
ConnectPlan::Attached => {
panic!("premise: an empty router has nothing to reuse, so it must ask for a dial")
}
}
assert!(
router.backends.is_empty(),
"a plan must not have taken a backend connection into the router: {:?}",
router.backends.keys().collect::<Vec<_>>()
);
assert_eq!(
context.streams[stream_id].state,
StreamState::Link,
"a plan must not have linked the stream — the embedder has not dialled yet"
);
assert!(
context.backend_streams.is_empty(),
"a plan must not have touched the backend-stream reverse index"
);
assert!(
context.backend_deltas.is_empty(),
"a plan must not have charged any backend accounting: {:?}",
context.backend_deltas
);
}
#[test]
fn an_h2_stream_count_tie_always_picks_the_lowest_backend_token() {
let fixture = routing_fixture();
for round in 0..ROUNDS {
let chosen = selected_backend_token(&fixture, &Staged::ConnectedH2);
assert_eq!(
chosen,
Token(LOWEST_TOKEN),
"round {round}: a stream-count tie must resolve to the lowest token, \
not to whatever the map iterated first"
);
}
}
#[test]
fn an_all_connecting_h2_fallback_always_picks_the_highest_backend_token() {
let fixture = routing_fixture();
for round in 0..ROUNDS {
let chosen = selected_backend_token(&fixture, &Staged::ConnectingH2);
assert_eq!(
chosen,
Token(HIGHEST_TOKEN),
"round {round}: the connecting-backend fallback assigns last-wins, \
so the total order must land on the highest token"
);
}
}
#[test]
fn an_h1_keep_alive_reuse_always_picks_the_lowest_backend_token() {
let fixture = routing_fixture();
for round in 0..ROUNDS {
let chosen = selected_backend_token(&fixture, &Staged::KeepAliveH1);
assert_eq!(
chosen,
Token(LOWEST_TOKEN),
"round {round}: the first matching keep-alive socket wins, so the total \
order must land on the lowest token"
);
}
}
#[test]
fn a_request_on_a_reused_backend_connection_allocates_nothing() {
use std::hint::black_box;
use crate::test_allocations::allocations;
const REQUESTS: usize = 64;
let backend_token = Token(LOWEST_TOKEN);
let fixture = routing_fixture();
let mut context = Context::new(
Ulid::generate(),
Rc::downgrade(&fixture.pool),
fixture.listener.clone(),
None,
"127.0.0.1:80"
.parse()
.expect("test public address must parse"),
);
let stream_id = context
.create_stream(Ulid::generate(), 65_535)
.expect("the test pool must hand out a stream");
context.streams[stream_id].context.method = Some(Method::Get);
context.streams[stream_id].context.cluster_id = Some(H1_CLUSTER.into());
let mut router = Router::new(Duration::from_secs(10), Duration::from_secs(10));
let mut backend_registry = BackendRegistry::default();
let (connection, _peer) =
staged_backend(&fixture.pool, &Staged::KeepAliveH1, &mut backend_registry);
router.backends.insert(backend_token, connection);
let backend_registry = &backend_registry;
let release = |router: &mut Router, context: &mut Context<HttpListener>| {
assert_eq!(
context.streams[stream_id].state,
StreamState::Linked(backend_token)
);
context.unlink_stream(stream_id);
backend_registry.apply_all(&mut context.backend_deltas, &fixture.backends);
context.streams[stream_id].forget_upstream_replay();
context.streams[stream_id].state = StreamState::Link;
let Some(Connection::H1(h1)) = router.backends.get_mut(&backend_token) else {
unreachable!("the staged backend is an H1 connection")
};
h1.stream = None;
if let Position::Client(_, _, status) = &mut h1.position {
*status = BackendStatus::KeepAlive;
}
};
let clusters = HashMap::new();
let retired = |backend: &super::BackendId| backend_registry.is_retired(backend);
let view = RoutingView::new(&clusters, ListenerType::Http).with_backend_retired(&retired);
let decision = |router: &mut Router, context: &mut Context<HttpListener>| -> usize {
let before = allocations();
let plan =
router.decide_after_gate(stream_id, black_box(&mut *context), &view, false, false);
let allocated = allocations() - before;
assert!(
matches!(plan, Ok(ConnectPlan::Attached)),
"the keep-alive socket must be reused, got {plan:?}"
);
release(router, context);
allocated
};
let control = |router: &mut Router, context: &mut Context<HttpListener>| -> usize {
let before = allocations();
let started = router
.backends
.get_mut(&backend_token)
.expect("the staged backend is registered")
.start_stream(stream_id, black_box(&mut *context));
context.link_stream(stream_id, backend_token);
let allocated = allocations() - before;
assert!(started, "the keep-alive socket must accept the stream");
release(router, context);
allocated
};
decision(&mut router, &mut context);
control(&mut router, &mut context);
let (mut decided, mut controlled) = (0, 0);
for _ in 0..REQUESTS {
decided += decision(&mut router, &mut context);
controlled += control(&mut router, &mut context);
}
decision(&mut router, &mut context);
let stream = &context.streams[stream_id];
assert_eq!(stream.context.backend_id.as_deref(), Some("test-backend"));
assert_eq!(stream.metrics.backend_id.as_deref(), Some("test-backend"));
assert_eq!(
decided.saturating_sub(controlled),
0,
"{REQUESTS} requests on a reused keep-alive backend made {decided} \
heap allocations in the routing decision against {controlled} for \
the mux bookkeeping alone, expected no difference"
);
assert_eq!(
(decided, controlled),
(0, 0),
"{REQUESTS} requests on a reused keep-alive backend must allocate \
nothing in the decision nor in the mux bookkeeping it triggers"
);
}
#[test]
fn a_routed_request_on_a_reused_backend_connection_copies_no_cluster_id() {
use std::hint::black_box;
use crate::test_allocations::allocations;
const REQUESTS: usize = 64;
let backend_token = Token(LOWEST_TOKEN);
let fixture = routing_fixture();
let sessions = fixture.proxy.borrow().sessions();
for gated in [false, true] {
let mut context = Context::new(
Ulid::generate(),
Rc::downgrade(&fixture.pool),
fixture.listener.clone(),
gated.then(|| {
"127.0.0.1:4242"
.parse()
.expect("test session address must parse")
}),
"127.0.0.1:80"
.parse()
.expect("test public address must parse"),
);
let stream_id = context
.create_stream(Ulid::generate(), 65_535)
.expect("the test pool must hand out a stream");
{
let stream = &mut context.streams[stream_id];
stream.state = StreamState::Link;
stream.context.authority = Some(H1_AUTHORITY.to_owned());
stream.context.path = Some("/".to_owned());
stream.context.method = Some(Method::Get);
}
let mut router = Router::new(Duration::from_secs(10), Duration::from_secs(10));
let mut backend_registry = BackendRegistry::default();
let (connection, _peer) =
staged_backend(&fixture.pool, &Staged::KeepAliveH1, &mut backend_registry);
router.backends.insert(backend_token, connection);
let proxy_ref = fixture.proxy.borrow();
let view = RoutingView::new(proxy_ref.clusters(), proxy_ref.kind());
let branch = |router: &mut Router, context: &mut Context<HttpListener>| -> usize {
let before = allocations();
let step = router
.plan_connect(stream_id, black_box(&mut *context), &view)
.expect("the reused request must be planned");
let mut allocated = allocations() - before;
let plan = match step {
ConnectStep::Decided(plan) => {
assert!(!gated, "a stream with a source address must be gated");
plan
}
ConnectStep::CheckIpLimit(resume) => {
assert!(gated, "a stream without a source address must not be gated");
let before = allocations();
let verdict = super::super::consult_ip_gate(
&sessions,
Token(0),
context
.http_context(stream_id)
.cluster_id
.as_ref()
.expect("plan_connect stamps the routed cluster"),
&resume,
);
let plan = router.plan_connect_resume(
stream_id,
&mut *context,
&view,
resume,
verdict,
);
allocated += allocations() - before;
plan.expect("the gate is unlimited, so the stream is admitted")
}
};
assert!(
matches!(plan, ConnectPlan::Attached),
"the keep-alive socket must be reused, got {plan:?}"
);
let before = allocations();
context.unlink_stream(stream_id);
backend_registry.apply_all(&mut context.backend_deltas, &fixture.backends);
allocated += allocations() - before;
let stream = &mut context.streams[stream_id];
stream.forget_upstream_replay();
stream.state = StreamState::Link;
stream.attempts = 0;
let Some(Connection::H1(h1)) = router.backends.get_mut(&backend_token) else {
unreachable!("the staged backend is an H1 connection")
};
h1.stream = None;
if let Position::Client(_, _, status) = &mut h1.position {
*status = BackendStatus::KeepAlive;
}
allocated
};
let control = |context: &mut Context<HttpListener>| -> usize {
let before = allocations();
#[cfg(debug_assertions)]
{
use crate::protocol::mux::DebugEvent;
let route = context.streams[stream_id].context.get_route();
context.debug.push(DebugEvent::Str(route));
}
#[cfg(not(debug_assertions))]
let _ = &context;
allocations() - before
};
branch(&mut router, &mut context);
control(&mut context);
let (mut branched, mut controlled) = (0, 0);
for _ in 0..REQUESTS {
branched += branch(&mut router, &mut context);
controlled += control(&mut context);
}
assert_eq!(
context.http_context(stream_id).cluster_id.as_deref(),
Some(H1_CLUSTER)
);
assert_eq!(
branched, controlled,
"gated={gated}: {REQUESTS} requests on a reused keep-alive backend \
made {branched} heap allocations through the reuse branch against \
{controlled} for the debug history, expected no difference"
);
}
}
#[test]
fn a_dial_moves_the_planned_cluster_id_into_the_backend_connection() {
use crate::{Protocol, ProxySession, protocol::mux::Mux, server::ListenSession};
let fixture = routing_fixture();
let upstream =
std::net::TcpListener::bind("127.0.0.1:0").expect("a loopback listener must bind");
fixture.backends.borrow_mut().add_backend(
H1_CLUSTER,
Backend::new(
"dial-backend",
upstream
.local_addr()
.expect("a bound listener has an address"),
None,
None,
None,
),
);
let mut context = Context::new(
Ulid::generate(),
Rc::downgrade(&fixture.pool),
fixture.listener.clone(),
None,
"127.0.0.1:80"
.parse()
.expect("test public address must parse"),
);
let stream_id = context
.create_stream(Ulid::generate(), 65_535)
.expect("the test pool must hand out a stream");
{
let stream = &mut context.streams[stream_id];
stream.state = StreamState::Link;
stream.context.authority = Some(H1_AUTHORITY.to_owned());
stream.context.path = Some("/".to_owned());
stream.context.method = Some(Method::Get);
}
let mut router = Router::new(Duration::from_secs(10), Duration::from_secs(10));
let mut backend_registry = BackendRegistry::default();
let plan = {
let proxy_ref = fixture.proxy.borrow();
let view = RoutingView::new(proxy_ref.clusters(), proxy_ref.kind());
router
.plan_connect(stream_id, &mut context, &view)
.map(decided)
.expect("an empty router must ask for a dial")
};
let ConnectPlan::Dial {
cluster_id,
h2,
frontend_should_stick,
} = plan
else {
panic!("premise: an empty router has nothing to reuse, so it must ask for a dial")
};
let planned = cluster_id.as_ptr();
let session: Rc<RefCell<dyn ProxySession>> = Rc::new(RefCell::new(ListenSession {
protocol: Protocol::HTTPListen,
}));
Mux::<mio::net::TcpStream, HttpListener>::dial_backend(
&mut router,
&mut backend_registry,
&fixture.backends,
stream_id,
&mut context,
&session,
&fixture.proxy,
cluster_id,
h2,
frontend_should_stick,
)
.expect("the loopback backend must dial");
let token = match context.streams[stream_id].state {
StreamState::Linked(token) => token,
other => panic!("the dial must link the stream, got {other:?}"),
};
let Some(Position::Client(owned, _, _)) =
router.backends.get(&token).map(Connection::position)
else {
panic!("the dialled connection must be a client")
};
assert_eq!(&**owned, H1_CLUSTER);
assert_eq!(
context.http_context(stream_id).cluster_id.as_deref(),
Some(H1_CLUSTER)
);
assert!(
std::ptr::eq(owned.as_ptr(), planned),
"the connection must own the id the plan carried, not a copy of it"
);
let routed = context
.http_context(stream_id)
.cluster_id
.as_ref()
.expect("plan_connect stamps the routed cluster");
assert!(
std::ptr::eq(owned.as_ptr(), routed.as_ptr()),
"the connection must share the stream's routed id, not a copy of it"
);
}
fn dial_one(
address: std::net::SocketAddr,
) -> (Result<(), BackendConnectionError>, Rc<RefCell<Backend>>) {
dial_with(address, DialFault::None)
}
enum DialFault {
None,
PoolExhausted,
StartRefused,
RegisterRefused,
}
struct RegisterRefuses(Rc<RefCell<dyn L7Proxy>>);
impl L7Proxy for RegisterRefuses {
fn kind(&self) -> ListenerType {
self.0.borrow().kind()
}
fn register_socket(
&self,
_socket: &mut mio::net::TcpStream,
_token: Token,
_interest: mio::Interest,
) -> Result<(), std::io::Error> {
Err(std::io::Error::from_raw_os_error(24)) }
fn add_session(&self, session: Rc<RefCell<dyn crate::ProxySession>>) -> Token {
self.0.borrow().add_session(session)
}
fn remove_session(&self, token: Token) -> bool {
self.0.borrow().remove_session(token)
}
fn clusters(&self) -> &HashMap<sozu_command::state::ClusterId, Cluster> {
unreachable!("Mux::dial_backend never reads the cluster table")
}
fn sessions(&self) -> Rc<RefCell<crate::server::SessionManager>> {
self.0.borrow().sessions()
}
}
fn dial_with(
address: std::net::SocketAddr,
fault: DialFault,
) -> (Result<(), BackendConnectionError>, Rc<RefCell<Backend>>) {
use crate::{Protocol, ProxySession, protocol::mux::Mux, server::ListenSession};
let (cluster, authority) = match fault {
DialFault::PoolExhausted | DialFault::StartRefused => (H2_CLUSTER, H2_AUTHORITY),
DialFault::None | DialFault::RegisterRefused => (H1_CLUSTER, H1_AUTHORITY),
};
let fixture = routing_fixture();
fixture.backends.borrow_mut().add_backend(
cluster,
Backend::new("dial-once", address, None, None, None),
);
let backend = fixture
.backends
.borrow_mut()
.backends
.get_mut(cluster)
.and_then(|list| list.find_backend(&address).cloned())
.expect("the backend was just added");
let mut context = Context::new(
Ulid::generate(),
Rc::downgrade(&fixture.pool),
fixture.listener.clone(),
None,
"127.0.0.1:80"
.parse()
.expect("test public address must parse"),
);
let stream_id = context
.create_stream(Ulid::generate(), 65_535)
.expect("the test pool must hand out a stream");
{
let stream = &mut context.streams[stream_id];
stream.state = StreamState::Link;
stream.context.authority = Some(authority.to_owned());
stream.context.path = Some("/".to_owned());
stream.context.method = Some(Method::Get);
}
let mut router = Router::new(Duration::from_secs(10), Duration::from_secs(10));
let mut backend_registry = BackendRegistry::default();
let plan = {
let proxy_ref = fixture.proxy.borrow();
let view = RoutingView::new(proxy_ref.clusters(), proxy_ref.kind());
router
.plan_connect(stream_id, &mut context, &view)
.map(decided)
.expect("an empty router must ask for a dial")
};
let ConnectPlan::Dial {
cluster_id,
h2,
frontend_should_stick,
} = plan
else {
panic!("premise: an empty router has nothing to reuse, so it must ask for a dial")
};
assert_eq!(
h2,
cluster == H2_CLUSTER,
"premise: the plan must dial the protocol the cluster declares"
);
let mut exhausted = Vec::new();
let mut proxy = fixture.proxy.clone();
match fault {
DialFault::None => {}
DialFault::PoolExhausted => {
while let Some(buffer) = fixture.pool.borrow_mut().checkout() {
exhausted.push(buffer);
}
}
DialFault::StartRefused => {
context.streams[stream_id]
.front
.out
.push_back(kawa::OutBlock::Store(kawa::Store::Static(b"encoded")));
}
DialFault::RegisterRefused => {
proxy = Rc::new(RefCell::new(RegisterRefuses(fixture.proxy.clone())));
}
}
let session: Rc<RefCell<dyn ProxySession>> = Rc::new(RefCell::new(ListenSession {
protocol: Protocol::HTTPListen,
}));
let dialed = Mux::<mio::net::TcpStream, HttpListener>::dial_backend(
&mut router,
&mut backend_registry,
&fixture.backends,
stream_id,
&mut context,
&session,
&proxy,
cluster_id,
h2,
frontend_should_stick,
);
drop(exhausted);
assert!(
dialed.is_ok() || router.backends.is_empty(),
"a dial that failed must keep no backend connection, got {dialed:?}"
);
backend_registry.apply_all(&mut context.backend_deltas, &fixture.backends);
(dialed, backend)
}
#[test]
fn a_successful_dial_counts_its_connection_once() {
let upstream =
std::net::TcpListener::bind("127.0.0.1:0").expect("a loopback listener must bind");
let (dialed, backend) = dial_one(
upstream
.local_addr()
.expect("a bound listener has an address"),
);
dialed.expect("the loopback backend must dial");
let backend = backend.borrow();
assert_eq!(
backend.active_connections, 1,
"one dial must hold exactly one connection on its backend"
);
assert_eq!(backend.failures, 0, "a successful dial records no failure");
}
#[test]
fn a_synchronously_failed_dial_releases_its_reservation() {
let (dialed, backend) = dial_one("255.255.255.255:80".parse().expect("a literal address"));
assert!(
matches!(
dialed,
Err(BackendConnectionError::Backend(
crate::backends::BackendError::ConnectionFailures { failures: 1, .. }
))
),
"a connect to the broadcast address must fail synchronously, got {dialed:?}"
);
let backend = backend.borrow();
assert_eq!(
backend.active_connections, 0,
"the failed dial must release the reservation selection took"
);
assert_eq!(
backend.failures, 1,
"the failed dial must count one failure"
);
}
fn assert_a_local_failure_releases_the_reservation(
fault: DialFault,
expected: fn(&BackendConnectionError) -> bool,
) {
use crate::retry::{RetryAction, RetryPolicy};
let upstream =
std::net::TcpListener::bind("127.0.0.1:0").expect("a loopback listener must bind");
let (dialed, backend) = dial_with(
upstream
.local_addr()
.expect("a bound listener has an address"),
fault,
);
assert!(
dialed.as_ref().is_err_and(expected),
"premise: the staged fault must fail the dial at its exit, got {dialed:?}"
);
let backend = backend.borrow();
assert_eq!(
backend.active_connections, 0,
"a dial abandoned after its connect must release the reservation \
selection took"
);
assert_eq!(
backend.failures, 0,
"a local failure must not count against the backend"
);
assert_eq!(
backend.retry_policy.can_try(std::time::Instant::now()),
Some(RetryAction::OKAY),
"a local failure must not arm the backend's backoff"
);
}
#[test]
fn an_h2_client_the_pool_cannot_build_releases_its_reservation() {
assert_a_local_failure_releases_the_reservation(DialFault::PoolExhausted, |error| {
matches!(error, BackendConnectionError::MaxBuffers)
});
}
#[test]
fn a_refused_stream_start_releases_its_reservation() {
assert_a_local_failure_releases_the_reservation(DialFault::StartRefused, |error| {
matches!(error, BackendConnectionError::MaxSessionsMemory)
});
}
#[test]
fn a_socket_the_poller_refuses_releases_its_reservation() {
assert_a_local_failure_releases_the_reservation(DialFault::RegisterRefused, |error| {
matches!(error, BackendConnectionError::MaxSessionsMemory)
});
}
#[test]
fn the_first_backend_connection_of_a_session_allocates_nothing() {
use crate::test_allocations::allocations;
let fixture = routing_fixture();
let mut backend_registry = BackendRegistry::default();
let mut staged = |router: &mut Router, token: usize| {
let (connection, peer) =
staged_backend(&fixture.pool, &Staged::KeepAliveH1, &mut backend_registry);
let before = allocations();
let replaced = router.backends.insert(Token(token), connection);
let allocated = allocations() - before;
assert!(replaced.is_none(), "Token({token}) was not held yet");
(allocated, peer)
};
let mut router = Router::new(Duration::from_secs(10), Duration::from_secs(10));
let (first, _first_peer) = staged(&mut router, LOWEST_TOKEN + 2);
assert_eq!(
first, 0,
"the first backend connection of a session must not allocate"
);
let (_, _lower_peer) = staged(&mut router, LOWEST_TOKEN);
let (_, _higher_peer) = staged(&mut router, LOWEST_TOKEN + 5);
let order = |router: &Router| router.backends.keys().map(|t| t.0).collect::<Vec<_>>();
assert_eq!(
order(&router),
[LOWEST_TOKEN, LOWEST_TOKEN + 2, LOWEST_TOKEN + 5]
);
assert_eq!(router.backends.len(), 3);
assert!(router.backends.remove(&Token(LOWEST_TOKEN)).is_some());
assert_eq!(order(&router), [LOWEST_TOKEN + 2, LOWEST_TOKEN + 5]);
assert!(router.backends.contains_key(&Token(LOWEST_TOKEN + 2)));
assert!(!router.backends.contains_key(&Token(LOWEST_TOKEN)));
assert!(router.backends.remove(&Token(LOWEST_TOKEN + 5)).is_some());
assert!(router.backends.remove(&Token(LOWEST_TOKEN + 2)).is_some());
assert!(router.backends.is_empty());
}
}
#[cfg(test)]
mod backend_selector_tests {
use std::{
cell::RefCell,
net::SocketAddr,
time::{Duration, Instant},
};
use sozu_command::proto::command::{LoadBalancingAlgorithms, LoadMetric};
use super::Router;
use crate::{
backends::{Backend, BackendMap},
protocol::{
http::editor::HttpContext,
mux::{BackendChange, BackendDelta, BackendId, BackendRegistry, RegistrySelector},
},
retry::{RetryAction, RetryPolicy},
};
const CLUSTER: &str = "cluster-1340-q12";
const STICKY_B: &str = "sticky-b";
struct Fixture {
backends: RefCell<BackendMap>,
registry: BackendRegistry,
router: Router,
now: Instant,
}
fn fixture_with(policy: LoadBalancingAlgorithms) -> Fixture {
let now = Instant::now();
let address = |port: u16| SocketAddr::from(([127, 0, 0, 1], port));
let mut backends = BackendMap::with_seed(0x1684);
backends.set_load_balancing_policy_for_cluster(
CLUSTER,
policy,
Some(LoadMetric::Connections),
);
backends.add_backend(
CLUSTER,
Backend::new_at("backend-a", address(9501), None, None, None, now),
);
backends.add_backend(
CLUSTER,
Backend::new_at(
"backend-b",
address(9502),
Some(STICKY_B.to_owned()),
None,
None,
now,
),
);
Fixture {
backends: RefCell::new(backends),
registry: BackendRegistry::default(),
router: Router::new(Duration::from_secs(10), Duration::from_secs(10)),
now,
}
}
fn fixture() -> Fixture {
fixture_with(LoadBalancingAlgorithms::LeastLoaded)
}
fn request(cookie: Option<&str>) -> HttpContext {
let mut context = HttpContext::new(
rusty_ulid::Ulid::generate(),
rusty_ulid::Ulid::generate(),
crate::Protocol::HTTP,
"127.0.0.1:80"
.parse()
.expect("test public address must parse"),
None,
"SOZUBALANCEID".to_owned(),
"Sozu-Id".to_owned(),
false,
false,
);
context.sticky_session_found = cookie.map(ToOwned::to_owned);
context
}
impl Fixture {
fn select(&mut self, should_stick: bool, context: &mut HttpContext) -> BackendId {
let mut backends = self.backends.borrow_mut();
let mut selector = RegistrySelector {
backends: &mut backends,
registry: &mut self.registry,
now: self.now,
exclude: &[],
};
let backend = self
.router
.backend_from_request(CLUSTER, should_stick, context, &mut selector)
.expect("a cluster with two available backends must select");
assert_eq!(
context.backend_id.as_deref(),
Some(&*backend.backend_id),
"the request must be stamped with the backend it was reserved on"
);
backend
}
fn backend(&self, id: &BackendId) -> std::cell::Ref<'_, Backend> {
self.registry
.handle(id)
.expect("an id this fixture minted must resolve")
.borrow()
}
fn active_connections(&self, backend: &BackendId) -> usize {
self.backend(backend).active_connections
}
}
#[test]
fn a_second_selection_observes_the_first_selections_reservation() {
let mut fixture = fixture();
let first = fixture.select(false, &mut request(None));
let second = fixture.select(false, &mut request(None));
assert_eq!(
&*first.backend_id, "backend-a",
"with both backends idle the first minimum wins"
);
assert_ne!(
first.backend_id,
second.backend_id,
"the second selection landed on the first one's backend: the first \
selection did not reserve its connection \
(first={} active={}, second={} active={})",
first.backend_id,
fixture.active_connections(&first),
second.backend_id,
fixture.active_connections(&second),
);
assert_eq!(
(
fixture.active_connections(&first),
fixture.active_connections(&second)
),
(1, 1),
"each selection must reserve exactly one connection on the backend it chose"
);
}
#[test]
fn power_of_two_weighs_every_reservation() {
const SELECTIONS: usize = 64;
let mut fixture = fixture_with(LoadBalancingAlgorithms::PowerOfTwo);
let first = fixture.select(false, &mut request(None));
let mut ids = vec![first];
for selections in 2..=SELECTIONS {
let before: Vec<usize> = ids
.iter()
.map(|id| fixture.active_connections(id))
.collect();
let picked = fixture.select(false, &mut request(None));
if !ids.iter().any(|id| id.backend_id == picked.backend_id) {
ids.push(picked.clone());
}
let loads: Vec<usize> = ids
.iter()
.map(|id| fixture.active_connections(id))
.collect();
assert_eq!(
loads.iter().sum::<usize>(),
selections,
"{selections} selections must hold {selections} reservations, got {loads:?}"
);
if ids.len() == 2 {
let (low, high) = (loads[0].min(loads[1]), loads[0].max(loads[1]));
assert!(
high - low <= 1,
"power-of-two let the reservations drift apart: {loads:?} (before {before:?})"
);
}
}
assert_eq!(ids.len(), 2, "both backends must be used");
}
#[test]
fn a_failed_dial_releases_its_reservation_and_records_the_failure() {
let mut fixture = fixture();
let backend = fixture.select(false, &mut request(None));
assert_eq!(
fixture.active_connections(&backend),
1,
"selection reserves"
);
let mut deltas = vec![BackendDelta {
slot: backend.slot(),
change: BackendChange::DialFailed(fixture.now),
}];
fixture.registry.apply_all(&mut deltas, &fixture.backends);
assert!(deltas.is_empty(), "the drain empties the ledger");
let released = fixture.backend(&backend);
assert_eq!(
released.active_connections, 0,
"a failed dial must release the reservation selection took"
);
assert_eq!(released.failures, 1, "a failed dial must count one failure");
assert_eq!(
released.retry_policy.can_try(fixture.now),
Some(RetryAction::WAIT),
"a failed dial must arm the backoff at the instant it failed"
);
}
#[test]
fn a_cookie_pins_only_a_frontend_that_sticks() {
let mut unpinned = fixture();
let mut context = request(Some(STICKY_B));
let backend = unpinned.select(false, &mut context);
assert_eq!(
&*backend.backend_id, "backend-a",
"a frontend that does not stick must ignore the client's cookie"
);
assert_eq!(
context.sticky_session, None,
"a frontend that does not stick must not answer with a cookie"
);
let mut pinned = fixture();
let mut context = request(Some(STICKY_B));
let backend = pinned.select(true, &mut context);
assert_eq!(
&*backend.backend_id, "backend-b",
"a sticky frontend must follow the cookie to its backend"
);
assert_eq!(context.sticky_session.as_deref(), Some(STICKY_B));
assert_eq!(
pinned.active_connections(&backend),
1,
"a sticky selection reserves like any other"
);
let mut fallback = fixture();
let mut context = request(None);
let backend = fallback.select(true, &mut context);
assert_eq!(
context.sticky_session.as_deref(),
Some(&*backend.backend_id),
"a sticky frontend with no cookie answers with the chosen \
backend's id when it has no sticky id of its own"
);
}
fn hrw_fixture() -> Fixture {
fixture_with(LoadBalancingAlgorithms::Hrw)
}
fn keyed_request(cookie: Option<&str>, key: u64) -> HttpContext {
let mut context = request(cookie);
context.affinity_key = Some(key);
context
}
#[test]
fn a_keyed_request_stays_on_one_backend_under_hrw() {
let mut fixture = hrw_fixture();
for key in 0..16u64 {
let first = fixture.select(false, &mut keyed_request(None, key));
for _ in 0..3 {
let again = fixture.select(false, &mut keyed_request(None, key));
assert_eq!(
again.backend_id, first.backend_id,
"key {key} must stay on one backend under HRW"
);
}
}
}
#[test]
fn a_sticky_cookie_wins_over_the_affinity_key() {
let mut fixture = hrw_fixture();
let key = (0..64u64)
.find(|&key| {
&*fixture
.select(false, &mut keyed_request(None, key))
.backend_id
== "backend-a"
})
.expect("HRW sends some of 64 keys to each of two backends");
let mut context = keyed_request(Some(STICKY_B), key);
let backend = fixture.select(true, &mut context);
assert_eq!(
&*backend.backend_id, "backend-b",
"the cookie must pin the request whatever its key says"
);
let mut context = keyed_request(None, key);
let backend = fixture.select(true, &mut context);
assert_eq!(
&*backend.backend_id, "backend-a",
"without a cookie a sticky frontend follows the key"
);
}
#[test]
fn stamping_a_dialled_backend_allocates_nothing() {
use std::hint::black_box;
use super::{Affinity, BackendSelector, SelectedBackend};
use crate::{backends::BackendError, test_allocations::allocations};
const DIALS: usize = 64;
struct StubSelector {
backend: BackendId,
}
impl BackendSelector for StubSelector {
fn select(
&mut self,
_cluster_id: &str,
_affinity: Affinity<'_>,
_key: Option<u64>,
) -> Result<SelectedBackend, BackendError> {
Ok(SelectedBackend {
backend: self.backend.clone(),
sticky_session: None,
})
}
}
let mut selector = StubSelector {
backend: BackendId::new(
0,
std::rc::Rc::from("stamped-backend"),
SocketAddr::from(([127, 0, 0, 1], 9503)),
),
};
let mut router = Router::new(Duration::from_secs(10), Duration::from_secs(10));
let mut context = request(None);
drop(
router
.backend_from_request(CLUSTER, false, &mut context, &mut selector)
.expect("the stub selector always selects"),
);
let mut allocated = 0;
for _ in 0..DIALS {
let before = allocations();
let selected = router.backend_from_request(
black_box(CLUSTER),
false,
black_box(&mut context),
&mut selector,
);
allocated += allocations() - before;
drop(selected.expect("the stub selector always selects"));
}
assert_eq!(context.backend_id.as_deref(), Some("stamped-backend"));
assert_eq!(
allocated, 0,
"{DIALS} dials made {allocated} heap allocations stamping the \
backend onto the request, expected none"
);
}
}
#[cfg(test)]
mod affinity_key_tests {
use std::net::{IpAddr, Ipv4Addr, SocketAddr};
use sozu_command::proto::command::{Cluster, LoadBalancingAlgorithms};
use super::{StickyCookie, affinity_key};
use crate::{
load_balancing::{affinity_key_from_ip, affinity_key_from_value},
pool::Pool,
protocol::mux::GenericHttpStream,
};
const NO_STICKY: StickyCookie<'static> = StickyCookie {
name: "SOZUBALANCEID",
value: None,
};
const CLIENT: SocketAddr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(192, 0, 2, 10)), 40_000);
fn parsed(pool: &mut Pool, request: &[u8]) -> GenericHttpStream {
let mut kawa: GenericHttpStream = kawa::Kawa::new(
kawa::Kind::Request,
kawa::Buffer::new(
pool.checkout()
.expect("the test pool must hand out a buffer"),
),
);
kawa.storage.space()[..request.len()].copy_from_slice(request);
kawa.storage.fill(request.len());
kawa::h1::parse(&mut kawa, &mut kawa::h1::NoCallbacks);
assert!(kawa.is_main_phase(), "premise: the request must parse");
kawa
}
fn cluster(
policy: LoadBalancingAlgorithms,
header: Option<&str>,
cookie: Option<&str>,
) -> Cluster {
Cluster {
cluster_id: "affinity".to_owned(),
load_balancing: policy as i32,
affinity_header: header.map(ToOwned::to_owned),
affinity_cookie: cookie.map(ToOwned::to_owned),
..Default::default()
}
}
const REQUEST: &[u8] = b"GET / HTTP/1.1\r\nHost: example.com\r\nx-tenant: acme\r\n\
Cookie: theme=dark; tenant=globex\r\nX-Empty: \r\n\r\n";
#[test]
fn only_hrw_maglev_and_sharded_clusters_derive_a_key() {
let mut pool = Pool::with_capacity(1, 1, 4096);
let front = parsed(&mut pool, REQUEST);
for policy in [
LoadBalancingAlgorithms::RoundRobin,
LoadBalancingAlgorithms::Random,
LoadBalancingAlgorithms::LeastLoaded,
LoadBalancingAlgorithms::PowerOfTwo,
] {
assert_eq!(
affinity_key(
&cluster(policy, Some("X-Tenant"), None),
&front,
Some(CLIENT),
NO_STICKY
),
None,
"{policy:?} does not read a key"
);
}
let sharded_round_robin = Cluster {
shard_percent: Some(25),
..cluster(LoadBalancingAlgorithms::RoundRobin, None, None)
};
assert_eq!(
affinity_key(&sharded_round_robin, &front, Some(CLIENT), NO_STICKY),
Some(affinity_key_from_ip(CLIENT.ip())),
"a sharded cluster derives a key whatever its policy"
);
for policy in [
LoadBalancingAlgorithms::Hrw,
LoadBalancingAlgorithms::Maglev,
] {
assert_eq!(
affinity_key(
&cluster(policy, None, None),
&front,
Some(CLIENT),
NO_STICKY
),
Some(affinity_key_from_ip(CLIENT.ip())),
"{policy:?} keys on the source IP by default"
);
}
}
#[test]
fn a_configured_header_keys_the_request_and_the_source_ip_is_the_fallback() {
let mut pool = Pool::with_capacity(1, 1, 4096);
let front = parsed(&mut pool, REQUEST);
let hrw = LoadBalancingAlgorithms::Hrw;
assert_eq!(
affinity_key(
&cluster(hrw, Some("X-Tenant"), None),
&front,
Some(CLIENT),
NO_STICKY
),
Some(affinity_key_from_value(b"acme"))
);
assert_eq!(
affinity_key(
&cluster(hrw, Some("X-Absent"), None),
&front,
Some(CLIENT),
NO_STICKY
),
Some(affinity_key_from_ip(CLIENT.ip()))
);
assert_eq!(
affinity_key(
&cluster(hrw, Some("X-Empty"), None),
&front,
Some(CLIENT),
NO_STICKY
),
Some(affinity_key_from_ip(CLIENT.ip())),
"an empty value keys nothing and falls back to the source IP"
);
assert_eq!(
affinity_key(
&cluster(hrw, Some("X-Absent"), None),
&front,
None,
NO_STICKY
),
None,
"no header and no source address: no key, and the policy round-robins"
);
}
#[test]
fn a_configured_cookie_keys_the_request() {
let mut pool = Pool::with_capacity(1, 1, 4096);
let front = parsed(&mut pool, REQUEST);
let maglev = LoadBalancingAlgorithms::Maglev;
assert_eq!(
affinity_key(
&cluster(maglev, None, Some("tenant")),
&front,
Some(CLIENT),
NO_STICKY
),
Some(affinity_key_from_value(b"globex"))
);
assert_eq!(
affinity_key(
&cluster(maglev, None, Some("Tenant")),
&front,
Some(CLIENT),
NO_STICKY
),
Some(affinity_key_from_ip(CLIENT.ip())),
"a cookie name differing in case is another cookie"
);
}
#[test]
fn deriving_an_affinity_key_allocates_nothing() {
use std::hint::black_box;
use crate::test_allocations::allocations;
const DERIVATIONS: usize = 64;
let mut pool = Pool::with_capacity(1, 1, 4096);
let front = parsed(&mut pool, REQUEST);
let clusters = [
cluster(LoadBalancingAlgorithms::Hrw, None, None),
cluster(LoadBalancingAlgorithms::Hrw, Some("X-Tenant"), None),
cluster(LoadBalancingAlgorithms::Maglev, None, Some("tenant")),
cluster(LoadBalancingAlgorithms::Maglev, Some("X-Absent"), None),
];
let mut allocated = 0;
for cluster in &clusters {
for _ in 0..DERIVATIONS {
let before = allocations();
let key = affinity_key(
black_box(cluster),
black_box(&front),
Some(CLIENT),
NO_STICKY,
);
allocated += allocations() - before;
assert!(key.is_some());
}
}
assert_eq!(
allocated,
0,
"{} derivations made {allocated} heap allocations, expected none",
DERIVATIONS * clusters.len()
);
}
#[test]
fn an_affinity_cookie_named_like_the_sticky_cookie_reads_its_captured_value() {
let mut pool = Pool::with_capacity(1, 1, 4096);
let front = parsed(&mut pool, REQUEST);
let hrw = cluster(LoadBalancingAlgorithms::Hrw, None, Some("SOZUBALANCEID"));
assert_eq!(
affinity_key(
&hrw,
&front,
Some(CLIENT),
StickyCookie {
name: "SOZUBALANCEID",
value: Some("backend-7"),
},
),
Some(affinity_key_from_value(b"backend-7"))
);
assert_eq!(
affinity_key(&hrw, &front, Some(CLIENT), NO_STICKY),
Some(affinity_key_from_ip(CLIENT.ip())),
"no sticky cookie on the request: the source IP keys it"
);
}
}