diff --git a/Cargo.lock b/Cargo.lock index ce12778..df34e82 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2488,7 +2488,7 @@ checksum = "8b870d8c151b6f2fb93e84a13146138f05d02ed11c7e7c54f8826aaaf7c9f184" [[package]] name = "pingora" version = "0.7.0" -source = "git+https://github.com/directxman12/pingora?branch=custom%2Fmain#771c7926b9635f716ddb30a24ba1d941a77596aa" +source = "git+https://github.com/directxman12/pingora?branch=custom%2Fmain#4a9a8abc2c25267bde433bb2b8555736989c4075" dependencies = [ "pingora-cache", "pingora-core", @@ -2501,7 +2501,7 @@ dependencies = [ [[package]] name = "pingora-cache" version = "0.7.0" -source = "git+https://github.com/directxman12/pingora?branch=custom%2Fmain#771c7926b9635f716ddb30a24ba1d941a77596aa" +source = "git+https://github.com/directxman12/pingora?branch=custom%2Fmain#4a9a8abc2c25267bde433bb2b8555736989c4075" dependencies = [ "ahash", "async-trait", @@ -2537,7 +2537,7 @@ dependencies = [ [[package]] name = "pingora-core" version = "0.7.0" -source = "git+https://github.com/directxman12/pingora?branch=custom%2Fmain#771c7926b9635f716ddb30a24ba1d941a77596aa" +source = "git+https://github.com/directxman12/pingora?branch=custom%2Fmain#4a9a8abc2c25267bde433bb2b8555736989c4075" dependencies = [ "ahash", "async-trait", @@ -2589,12 +2589,12 @@ dependencies = [ [[package]] name = "pingora-error" version = "0.7.0" -source = "git+https://github.com/directxman12/pingora?branch=custom%2Fmain#771c7926b9635f716ddb30a24ba1d941a77596aa" +source = "git+https://github.com/directxman12/pingora?branch=custom%2Fmain#4a9a8abc2c25267bde433bb2b8555736989c4075" [[package]] name = "pingora-header-serde" version = "0.7.0" -source = "git+https://github.com/directxman12/pingora?branch=custom%2Fmain#771c7926b9635f716ddb30a24ba1d941a77596aa" +source = "git+https://github.com/directxman12/pingora?branch=custom%2Fmain#4a9a8abc2c25267bde433bb2b8555736989c4075" dependencies = [ "bytes", "http", @@ -2609,7 +2609,7 @@ dependencies = [ [[package]] name = "pingora-http" version = "0.7.0" -source = "git+https://github.com/directxman12/pingora?branch=custom%2Fmain#771c7926b9635f716ddb30a24ba1d941a77596aa" +source = "git+https://github.com/directxman12/pingora?branch=custom%2Fmain#4a9a8abc2c25267bde433bb2b8555736989c4075" dependencies = [ "bytes", "http", @@ -2619,7 +2619,7 @@ dependencies = [ [[package]] name = "pingora-ketama" version = "0.7.0" -source = "git+https://github.com/directxman12/pingora?branch=custom%2Fmain#771c7926b9635f716ddb30a24ba1d941a77596aa" +source = "git+https://github.com/directxman12/pingora?branch=custom%2Fmain#4a9a8abc2c25267bde433bb2b8555736989c4075" dependencies = [ "crc32fast", ] @@ -2627,7 +2627,7 @@ dependencies = [ [[package]] name = "pingora-load-balancing" version = "0.7.0" -source = "git+https://github.com/directxman12/pingora?branch=custom%2Fmain#771c7926b9635f716ddb30a24ba1d941a77596aa" +source = "git+https://github.com/directxman12/pingora?branch=custom%2Fmain#4a9a8abc2c25267bde433bb2b8555736989c4075" dependencies = [ "arc-swap", "async-trait", @@ -2648,7 +2648,7 @@ dependencies = [ [[package]] name = "pingora-lru" version = "0.7.0" -source = "git+https://github.com/directxman12/pingora?branch=custom%2Fmain#771c7926b9635f716ddb30a24ba1d941a77596aa" +source = "git+https://github.com/directxman12/pingora?branch=custom%2Fmain#4a9a8abc2c25267bde433bb2b8555736989c4075" dependencies = [ "arrayvec", "hashbrown 0.16.1", @@ -2659,7 +2659,7 @@ dependencies = [ [[package]] name = "pingora-pool" version = "0.7.0" -source = "git+https://github.com/directxman12/pingora?branch=custom%2Fmain#771c7926b9635f716ddb30a24ba1d941a77596aa" +source = "git+https://github.com/directxman12/pingora?branch=custom%2Fmain#4a9a8abc2c25267bde433bb2b8555736989c4075" dependencies = [ "crossbeam-queue", "log", @@ -2673,7 +2673,7 @@ dependencies = [ [[package]] name = "pingora-proxy" version = "0.7.0" -source = "git+https://github.com/directxman12/pingora?branch=custom%2Fmain#771c7926b9635f716ddb30a24ba1d941a77596aa" +source = "git+https://github.com/directxman12/pingora?branch=custom%2Fmain#4a9a8abc2c25267bde433bb2b8555736989c4075" dependencies = [ "async-trait", "bytes", @@ -2695,7 +2695,7 @@ dependencies = [ [[package]] name = "pingora-runtime" version = "0.7.0" -source = "git+https://github.com/directxman12/pingora?branch=custom%2Fmain#771c7926b9635f716ddb30a24ba1d941a77596aa" +source = "git+https://github.com/directxman12/pingora?branch=custom%2Fmain#4a9a8abc2c25267bde433bb2b8555736989c4075" dependencies = [ "once_cell", "rand 0.8.5", @@ -2706,7 +2706,7 @@ dependencies = [ [[package]] name = "pingora-rustls" version = "0.7.0" -source = "git+https://github.com/directxman12/pingora?branch=custom%2Fmain#771c7926b9635f716ddb30a24ba1d941a77596aa" +source = "git+https://github.com/directxman12/pingora?branch=custom%2Fmain#4a9a8abc2c25267bde433bb2b8555736989c4075" dependencies = [ "log", "no_debug", @@ -2722,7 +2722,7 @@ dependencies = [ [[package]] name = "pingora-timeout" version = "0.7.0" -source = "git+https://github.com/directxman12/pingora?branch=custom%2Fmain#771c7926b9635f716ddb30a24ba1d941a77596aa" +source = "git+https://github.com/directxman12/pingora?branch=custom%2Fmain#4a9a8abc2c25267bde433bb2b8555736989c4075" dependencies = [ "once_cell", "parking_lot", diff --git a/src/backend_selection.rs b/src/backend_selection.rs new file mode 100644 index 0000000..53e7905 --- /dev/null +++ b/src/backend_selection.rs @@ -0,0 +1,79 @@ +//! _grumbles_ ketama doesn't support uds _grumbles_ +//! +//! (also this lets us add an option to choose in the future) + +use std::sync::Arc; + +use pingora::lb::selection::BackendSelection; +use pingora::lb::selection::consistent::{KetamaConfig, KetamaHashing}; +use pingora::lb::{self, Backend}; +use pingora::prelude::*; + +pub enum SelectionChoice { + Ketama(Arc), + RoundRobin(Arc), +} + +#[allow( + dead_code, + reason = "required for trait impl, will use the in the future" +)] +pub enum SelectionConfig { + Ketama(KetamaConfig), + RoundRobin(()), +} +pub enum SelectionIter { + Ketama(::Iter), + RoundRobin(::Iter), +} +impl lb::selection::BackendIter for SelectionIter { + fn next(&mut self) -> Option<&Backend> { + match self { + SelectionIter::Ketama(k) => k.next(), + SelectionIter::RoundRobin(rr) => rr.next(), + } + } +} +impl BackendSelection for SelectionChoice { + type Iter = SelectionIter; + + type Config = SelectionConfig; + + fn build(backends: &std::collections::BTreeSet) -> Self { + // need the arc here for iter + if backends + .iter() + .any(|b| matches!(b.addr, pingora::protocols::l4::socket::SocketAddr::Unix(_))) + { + // ketama does not support unix domain sockets + Self::RoundRobin(Arc::new(RoundRobin::build(backends))) + } else { + Self::Ketama(Arc::new(KetamaHashing::build(backends))) + } + } + + fn iter(self: &std::sync::Arc, key: &[u8]) -> Self::Iter + where + Self::Iter: lb::selection::BackendIter, + { + match &**self { + SelectionChoice::Ketama(k) => SelectionIter::Ketama(k.iter(key)), + SelectionChoice::RoundRobin(rr) => SelectionIter::RoundRobin(rr.iter(key)), + } + } + + fn build_with_config( + backends: &std::collections::BTreeSet, + config: &Self::Config, + ) -> Self { + // need the arc here for iter + match config { + SelectionConfig::Ketama(ketama) => { + Self::Ketama(Arc::new(KetamaHashing::build_with_config(backends, ketama))) + } + SelectionConfig::RoundRobin(rr) => { + Self::RoundRobin(Arc::new(RoundRobin::build_with_config(backends, rr))) + } + } + } +} diff --git a/src/gateway.rs b/src/gateway.rs index 0474015..2cb4554 100644 --- a/src/gateway.rs +++ b/src/gateway.rs @@ -8,7 +8,6 @@ use async_trait::async_trait; use cookie_rs::CookieJar; use http::status::StatusCode; use http::{HeaderName, HeaderValue}; -use pingora::lb::selection::consistent::KetamaHashing; use pingora::prelude::*; use pingora::protocols::tls::CaType; use url::Url; @@ -23,7 +22,7 @@ pub mod oidc; /// per-domain information about backends and such pub struct DomainInfo { /// the load balancer to use to select backends - pub balancer: Arc>, + pub balancer: Arc>, /// whether or not we allow insecure connections from clients pub tls_mode: config::format::domain::TlsMode, /// the sni name of this domain, used to pass to backends @@ -636,6 +635,9 @@ impl ProxyHttp for AuthGateway { peer.options.ca = ca; peer } + BackendData::HttpOverUds => { + HttpPeer::new_from_sockaddr(backend.addr, false, backends.sni_name.to_string()) + } }; Ok(Box::new(peer)) } @@ -692,6 +694,7 @@ impl ProxyHttp for AuthGateway { /// for use in [`AuthGateway::upstream_peer`] #[derive(Clone)] pub enum BackendData { + HttpOverUds, HttpOnly, Tls { /// should skip we verifying the cert (useful if the certs are selfsigned and we don't want to diff --git a/src/main.rs b/src/main.rs index 3a4de07..c6252ad 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,14 +1,15 @@ use std::collections::HashMap; use color_eyre::eyre::Context as _; -use pingora::lb; -use pingora::lb::selection::consistent::KetamaHashing; +use pingora::lb::{self, Backend}; use pingora::listeners::tls::{BundleCert, CertAndKey, TlsSettings}; use pingora::prelude::*; use pingora::utils::tls::CertKey; +use self::backend_selection::SelectionChoice; use self::gateway::{AuthGateway, BackendData, DomainInfo, oidc}; +mod backend_selection; mod config; mod cookies; mod gateway; @@ -17,14 +18,14 @@ mod oauth; /// constructed load balancer, with [backend info][`BackendInfo`] to be passed to [`AuthGateway`] type BalancerInfo = ( - Vec>>, + Vec>>, HashMap, ); /// construct the load balancer and initialize the [backend info][`BackendInfo`] for the /// [`AuthGateway`] fn balancer(domains: &HashMap) -> color_eyre::Result { - use lb::{self, Backend, discovery}; + use lb::{self, discovery}; use pingora::protocols::l4::socket::SocketAddr; let mut balancers = HashMap::with_capacity(domains.len()); @@ -68,7 +69,7 @@ fn balancer(domains: &HashMap) -> color_eyre::Re })) .chain(domain.uds.iter().map(|backend| { let mut ext = lb::Extensions::new(); - ext.insert(BackendData::HttpOnly); + ext.insert(BackendData::HttpOverUds); Ok(Backend { addr: SocketAddr::Unix( std::os::unix::net::SocketAddr::from_pathname(&backend.path)