refactor: split sticky-cookie into a separated feature

This commit is contained in:
Jun Kurihara 2023-06-15 15:38:19 +09:00
commit d4040b9b98
No known key found for this signature in database
GPG key ID: 48ADFD173ED22B03
14 changed files with 248 additions and 186 deletions

View file

@ -12,12 +12,13 @@ publish = false
# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
[features] [features]
default = ["http3"] default = ["http3", "sticky-cookie"]
http3 = ["quinn", "h3", "h3-quinn"] http3 = ["quinn", "h3", "h3-quinn"]
sticky-cookie = ["base64", "sha2", "chrono"]
[dependencies] [dependencies]
anyhow = "1.0.71" anyhow = "1.0.71"
clap = { version = "4.3.3", features = ["std", "cargo", "wrap_help"] } clap = { version = "4.3.4", features = ["std", "cargo", "wrap_help"] }
rand = "0.8.5" rand = "0.8.5"
toml = { version = "0.7.4", default-features = false, features = ["parse"] } toml = { version = "0.7.4", default-features = false, features = ["parse"] }
rustc-hash = "1.1.0" rustc-hash = "1.1.0"
@ -51,7 +52,7 @@ hyper-rustls = { version = "0.24.0", default-features = false, features = [
] } ] }
tokio-rustls = { version = "0.24.1", features = ["early-data"] } tokio-rustls = { version = "0.24.1", features = ["early-data"] }
rustls-pemfile = "1.0.2" rustls-pemfile = "1.0.2"
rustls = { version = "0.21.1", default-features = false } rustls = { version = "0.21.2", default-features = false }
webpki = "0.22.0" webpki = "0.22.0"
# logging # logging
@ -65,14 +66,14 @@ h3 = { path = "./h3/h3/", optional = true }
# h3-quinn = { path = "./h3/h3-quinn/", optional = true } # h3-quinn = { path = "./h3/h3-quinn/", optional = true }
h3-quinn = { path = "./h3-quinn/", optional = true } # Tentative to support rustls-0.21 h3-quinn = { path = "./h3-quinn/", optional = true } # Tentative to support rustls-0.21
# cookie handling # cookie handling for sticky cookie
chrono = { version = "0.4.26", default-features = false, features = [ chrono = { version = "0.4.26", default-features = false, features = [
"unstable-locales", "unstable-locales",
"alloc", "alloc",
"clock", "clock",
] } ], optional = true }
base64 = "0.21.2" base64 = { version = "0.21.2", optional = true }
sha2 = { version = "0.10.6", default-features = false } sha2 = { version = "0.10.6", default-features = false, optional = true }
[target.'cfg(not(target_env = "msvc"))'.dependencies] [target.'cfg(not(target_env = "msvc"))'.dependencies]

View file

@ -3,6 +3,13 @@
- Improvement of path matcher - Improvement of path matcher
- More flexible option for rewriting path - More flexible option for rewriting path
- Refactoring - Refactoring
Split `backend` module into three parts
- backend(s): struct containing info, defined for each served domain with multiple paths
- upstream/upstream group: information on targeted destinations for each set of (a domain + a path)
- load-balance: load balancing mod for a domain + path
- Unit tests - Unit tests
- Options to serve custom http_error page. - Options to serve custom http_error page.
- Prometheus metrics - Prometheus metrics

2
h3

@ -1 +1 @@
Subproject commit 22da9387f19d724852b3bf1dfd7e66f0fd45cb81 Subproject commit 3ef7c1a37b635e8446322d8f8d3a68580a208ad8

View file

@ -1,14 +1,13 @@
use super::{load_balance_sticky_cookie::StickyCookieConfig, LbContext, Upstream}; #[cfg(feature = "sticky-cookie")]
use crate::{constants::STICKY_COOKIE_NAME, log::*}; pub use super::{
load_balance_sticky::{LbStickyRoundRobin, LbStickyRoundRobinBuilder},
sticky_cookie::StickyCookie,
};
use derive_builder::Builder; use derive_builder::Builder;
use rand::Rng; use rand::Rng;
use rustc_hash::FxHashMap as HashMap; use std::sync::{
use std::{
borrow::Cow,
sync::{
atomic::{AtomicUsize, Ordering}, atomic::{AtomicUsize, Ordering},
Arc, Arc,
},
}; };
/// Constants to specify a load balance option /// Constants to specify a load balance option
@ -16,6 +15,7 @@ pub(super) mod load_balance_options {
pub const FIX_TO_FIRST: &str = "none"; pub const FIX_TO_FIRST: &str = "none";
pub const ROUND_ROBIN: &str = "round_robin"; pub const ROUND_ROBIN: &str = "round_robin";
pub const RANDOM: &str = "random"; pub const RANDOM: &str = "random";
#[cfg(feature = "sticky-cookie")]
pub const STICKY_ROUND_ROBIN: &str = "sticky"; pub const STICKY_ROUND_ROBIN: &str = "sticky";
} }
@ -27,7 +27,7 @@ pub(super) struct PointerToUpstream {
pub context_lb: Option<LbContext>, pub context_lb: Option<LbContext>,
} }
/// Trait for LB /// Trait for LB
trait LbWithPointer { pub(super) trait LbWithPointer {
fn get_ptr(&self, req_info: Option<&LbContext>) -> PointerToUpstream; fn get_ptr(&self, req_info: Option<&LbContext>) -> PointerToUpstream;
} }
@ -85,123 +85,6 @@ impl LbWithPointer for LbRandom {
} }
} }
#[derive(Debug, Clone, Builder)]
/// Round Robin LB object in the sticky cookie manner
pub struct LbStickyRoundRobin {
#[builder(default)]
/// Pointer to the index of the last served upstream destination
ptr: Arc<AtomicUsize>,
#[builder(setter(custom), default)]
/// Number of upstream destinations
num_upstreams: usize,
#[builder(setter(custom))]
/// Information to build the cookie to stick clients to specific backends
pub sticky_config: StickyCookieConfig,
#[builder(setter(custom))]
/// Hashmaps:
/// - Hashmap that maps server indices to server id (string)
/// - Hashmap that maps server ids (string) to server indices, for fast reverse lookup
upstream_maps: UpstreamMap,
}
#[derive(Debug, Clone)]
pub struct UpstreamMap {
/// Hashmap that maps server indices to server id (string)
upstream_index_map: Vec<String>,
/// Hashmap that maps server ids (string) to server indices, for fast reverse lookup
upstream_id_map: HashMap<String, usize>,
}
impl LbStickyRoundRobinBuilder {
pub fn num_upstreams(&mut self, v: &usize) -> &mut Self {
self.num_upstreams = Some(*v);
self
}
pub fn sticky_config(&mut self, server_name: &str, path_opt: &Option<String>) -> &mut Self {
self.sticky_config = Some(StickyCookieConfig {
name: STICKY_COOKIE_NAME.to_string(), // TODO: config等で変更できるように
domain: server_name.to_ascii_lowercase(),
path: if let Some(v) = path_opt {
v.to_ascii_lowercase()
} else {
"/".to_string()
},
duration: 300, // TODO: config等で変更できるように
});
self
}
pub fn upstream_maps(&mut self, upstream_vec: &[Upstream]) -> &mut Self {
let upstream_index_map: Vec<String> = upstream_vec
.iter()
.enumerate()
.map(|(i, v)| v.calculate_id_with_index(i))
.collect();
let mut upstream_id_map = HashMap::default();
for (i, v) in upstream_index_map.iter().enumerate() {
upstream_id_map.insert(v.to_string(), i);
}
self.upstream_maps = Some(UpstreamMap {
upstream_index_map,
upstream_id_map,
});
self
}
}
impl<'a> LbStickyRoundRobin {
fn simple_increment_ptr(&self) -> usize {
// Get a current count of upstream served
let current_ptr = self.ptr.load(Ordering::Relaxed);
if current_ptr < self.num_upstreams - 1 {
self.ptr.fetch_add(1, Ordering::Relaxed)
} else {
// Clear the counter
self.ptr.fetch_and(0, Ordering::Relaxed)
}
}
/// This is always called only internally. So 'unwrap()' is executed.
fn get_server_id_from_index(&self, index: usize) -> String {
self.upstream_maps.upstream_index_map.get(index).unwrap().to_owned()
}
/// This function takes value passed from outside. So 'result' is used.
fn get_server_index_from_id(&self, id: impl Into<Cow<'a, str>>) -> Option<usize> {
let id_str = id.into().to_string();
self.upstream_maps.upstream_id_map.get(&id_str).map(|v| v.to_owned())
}
}
impl LbWithPointer for LbStickyRoundRobin {
fn get_ptr(&self, req_info: Option<&LbContext>) -> PointerToUpstream {
// If given context is None or invalid (not contained), get_ptr() is invoked to increment the pointer.
// Otherwise, get the server index indicated by the server_id inside the cookie
let ptr = match req_info {
None => {
debug!("No sticky cookie");
self.simple_increment_ptr()
}
Some(context) => {
let server_id = &context.sticky_cookie.value.value;
if let Some(server_index) = self.get_server_index_from_id(server_id) {
debug!("Valid sticky cookie: id={}, index={}", server_id, server_index);
server_index
} else {
debug!("Invalid sticky cookie: id={}", server_id);
self.simple_increment_ptr()
}
}
};
// Get the server id from the ptr.
// TODO: This should be simplified and optimized if ptr is not changed (id value exists in cookie).
let upstream_id = self.get_server_id_from_index(ptr);
let new_cookie = self.sticky_config.build_sticky_cookie(upstream_id).unwrap();
let new_context = Some(LbContext {
sticky_cookie: new_cookie,
});
PointerToUpstream {
ptr,
context_lb: new_context,
}
}
}
#[derive(Debug, Clone)] #[derive(Debug, Clone)]
/// Load Balancing Option /// Load Balancing Option
pub enum LoadBalance { pub enum LoadBalance {
@ -211,6 +94,7 @@ pub enum LoadBalance {
Random(LbRandom), Random(LbRandom),
/// Simple round robin without session persistance /// Simple round robin without session persistance
RoundRobin(LbRoundRobin), RoundRobin(LbRoundRobin),
#[cfg(feature = "sticky-cookie")]
/// Round robin with session persistance using cookie /// Round robin with session persistance using cookie
StickyRoundRobin(LbStickyRoundRobin), StickyRoundRobin(LbStickyRoundRobin),
} }
@ -222,7 +106,7 @@ impl Default for LoadBalance {
impl LoadBalance { impl LoadBalance {
/// Get the index of the upstream serving the incoming request /// Get the index of the upstream serving the incoming request
pub(super) fn get_context(&self, context_to_lb: &Option<LbContext>) -> PointerToUpstream { pub(super) fn get_context(&self, _context_to_lb: &Option<LbContext>) -> PointerToUpstream {
match self { match self {
LoadBalance::FixToFirst => PointerToUpstream { LoadBalance::FixToFirst => PointerToUpstream {
ptr: 0usize, ptr: 0usize,
@ -230,10 +114,22 @@ impl LoadBalance {
}, },
LoadBalance::RoundRobin(ptr) => ptr.get_ptr(None), LoadBalance::RoundRobin(ptr) => ptr.get_ptr(None),
LoadBalance::Random(ptr) => ptr.get_ptr(None), LoadBalance::Random(ptr) => ptr.get_ptr(None),
#[cfg(feature = "sticky-cookie")]
LoadBalance::StickyRoundRobin(ptr) => { LoadBalance::StickyRoundRobin(ptr) => {
// Generate new context if sticky round robin is enabled. // Generate new context if sticky round robin is enabled.
ptr.get_ptr(context_to_lb.as_ref()) ptr.get_ptr(_context_to_lb.as_ref())
} }
} }
} }
} }
#[derive(Debug, Clone)]
/// Struct to handle the sticky cookie string,
/// - passed from Rp module (http handler) to LB module, manipulated from req, only StickyCookieValue exists.
/// - passed from LB module to Rp module (http handler), will be inserted into res, StickyCookieValue and Info exist.
pub struct LbContext {
#[cfg(feature = "sticky-cookie")]
pub sticky_cookie: StickyCookie,
#[cfg(not(feature = "sticky-cookie"))]
pub sticky_cookie: (),
}

View file

@ -0,0 +1,132 @@
use super::{
load_balance::{LbContext, LbWithPointer, PointerToUpstream},
sticky_cookie::StickyCookieConfig,
Upstream,
};
use crate::{constants::STICKY_COOKIE_NAME, log::*};
use derive_builder::Builder;
use rustc_hash::FxHashMap as HashMap;
use std::{
borrow::Cow,
sync::{
atomic::{AtomicUsize, Ordering},
Arc,
},
};
#[derive(Debug, Clone, Builder)]
/// Round Robin LB object in the sticky cookie manner
pub struct LbStickyRoundRobin {
#[builder(default)]
/// Pointer to the index of the last served upstream destination
ptr: Arc<AtomicUsize>,
#[builder(setter(custom), default)]
/// Number of upstream destinations
num_upstreams: usize,
#[builder(setter(custom))]
/// Information to build the cookie to stick clients to specific backends
pub sticky_config: StickyCookieConfig,
#[builder(setter(custom))]
/// Hashmaps:
/// - Hashmap that maps server indices to server id (string)
/// - Hashmap that maps server ids (string) to server indices, for fast reverse lookup
upstream_maps: UpstreamMap,
}
#[derive(Debug, Clone)]
pub struct UpstreamMap {
/// Hashmap that maps server indices to server id (string)
upstream_index_map: Vec<String>,
/// Hashmap that maps server ids (string) to server indices, for fast reverse lookup
upstream_id_map: HashMap<String, usize>,
}
impl LbStickyRoundRobinBuilder {
pub fn num_upstreams(&mut self, v: &usize) -> &mut Self {
self.num_upstreams = Some(*v);
self
}
pub fn sticky_config(&mut self, server_name: &str, path_opt: &Option<String>) -> &mut Self {
self.sticky_config = Some(StickyCookieConfig {
name: STICKY_COOKIE_NAME.to_string(), // TODO: config等で変更できるように
domain: server_name.to_ascii_lowercase(),
path: if let Some(v) = path_opt {
v.to_ascii_lowercase()
} else {
"/".to_string()
},
duration: 300, // TODO: config等で変更できるように
});
self
}
pub fn upstream_maps(&mut self, upstream_vec: &[Upstream]) -> &mut Self {
let upstream_index_map: Vec<String> = upstream_vec
.iter()
.enumerate()
.map(|(i, v)| v.calculate_id_with_index(i))
.collect();
let mut upstream_id_map = HashMap::default();
for (i, v) in upstream_index_map.iter().enumerate() {
upstream_id_map.insert(v.to_string(), i);
}
self.upstream_maps = Some(UpstreamMap {
upstream_index_map,
upstream_id_map,
});
self
}
}
impl<'a> LbStickyRoundRobin {
fn simple_increment_ptr(&self) -> usize {
// Get a current count of upstream served
let current_ptr = self.ptr.load(Ordering::Relaxed);
if current_ptr < self.num_upstreams - 1 {
self.ptr.fetch_add(1, Ordering::Relaxed)
} else {
// Clear the counter
self.ptr.fetch_and(0, Ordering::Relaxed)
}
}
/// This is always called only internally. So 'unwrap()' is executed.
fn get_server_id_from_index(&self, index: usize) -> String {
self.upstream_maps.upstream_index_map.get(index).unwrap().to_owned()
}
/// This function takes value passed from outside. So 'result' is used.
fn get_server_index_from_id(&self, id: impl Into<Cow<'a, str>>) -> Option<usize> {
let id_str = id.into().to_string();
self.upstream_maps.upstream_id_map.get(&id_str).map(|v| v.to_owned())
}
}
impl LbWithPointer for LbStickyRoundRobin {
fn get_ptr(&self, req_info: Option<&LbContext>) -> PointerToUpstream {
// If given context is None or invalid (not contained), get_ptr() is invoked to increment the pointer.
// Otherwise, get the server index indicated by the server_id inside the cookie
let ptr = match req_info {
None => {
debug!("No sticky cookie");
self.simple_increment_ptr()
}
Some(context) => {
let server_id = &context.sticky_cookie.value.value;
if let Some(server_index) = self.get_server_index_from_id(server_id) {
debug!("Valid sticky cookie: id={}, index={}", server_id, server_index);
server_index
} else {
debug!("Invalid sticky cookie: id={}", server_id);
self.simple_increment_ptr()
}
}
};
// Get the server id from the ptr.
// TODO: This should be simplified and optimized if ptr is not changed (id value exists in cookie).
let upstream_id = self.get_server_id_from_index(ptr);
let new_cookie = self.sticky_config.build_sticky_cookie(upstream_id).unwrap();
let new_context = Some(LbContext {
sticky_cookie: new_cookie,
});
PointerToUpstream {
ptr,
context_lb: new_context,
}
}
}

View file

@ -1,11 +1,15 @@
mod load_balance; mod load_balance;
mod load_balance_sticky_cookie; #[cfg(feature = "sticky-cookie")]
mod load_balance_sticky;
#[cfg(feature = "sticky-cookie")]
mod sticky_cookie;
mod upstream; mod upstream;
mod upstream_opts; mod upstream_opts;
#[cfg(feature = "sticky-cookie")]
pub use self::sticky_cookie::{StickyCookie, StickyCookieValue};
pub use self::{ pub use self::{
load_balance::LoadBalance, load_balance::{LbContext, LoadBalance},
load_balance_sticky_cookie::{LbContext, StickyCookie, StickyCookieBuilder, StickyCookieValue},
upstream::{ReverseProxy, Upstream, UpstreamGroup, UpstreamGroupBuilder}, upstream::{ReverseProxy, Upstream, UpstreamGroup, UpstreamGroupBuilder},
upstream_opts::UpstreamOption, upstream_opts::UpstreamOption,
}; };
@ -270,15 +274,22 @@ impl Backends {
let mut server_config_local = if client_ca_roots_local.is_empty() { let mut server_config_local = if client_ca_roots_local.is_empty() {
// with no client auth, enable http1.1 -- 3 // with no client auth, enable http1.1 -- 3
#[cfg(not(feature = "http3"))]
{
ServerConfig::builder()
.with_safe_defaults()
.with_no_client_auth()
.with_cert_resolver(Arc::new(resolver_local))
}
#[cfg(feature = "http3")]
{
let mut sc = ServerConfig::builder() let mut sc = ServerConfig::builder()
.with_safe_defaults() .with_safe_defaults()
.with_no_client_auth() .with_no_client_auth()
.with_cert_resolver(Arc::new(resolver_local)); .with_cert_resolver(Arc::new(resolver_local));
#[cfg(feature = "http3")]
{
sc.alpn_protocols = vec![b"h3".to_vec(), b"hq-29".to_vec()]; // TODO: remove hq-29 later? sc.alpn_protocols = vec![b"h3".to_vec(), b"hq-29".to_vec()]; // TODO: remove hq-29 later?
}
sc sc
}
} else { } else {
// with client auth, enable only http1.1 and 2 // with client auth, enable only http1.1 and 2
// let client_certs_verifier = rustls::server::AllowAnyAnonymousOrAuthenticatedClient::new(client_ca_roots); // let client_certs_verifier = rustls::server::AllowAnyAnonymousOrAuthenticatedClient::new(client_ca_roots);
@ -320,7 +331,7 @@ impl Backends {
} }
#[cfg(not(feature = "http3"))] #[cfg(not(feature = "http3"))]
{ {
server_config.alpn_protocols = vec![b"h2".to_vec(), b"http/1.1".to_vec()]; server_crypto_global.alpn_protocols = vec![b"h2".to_vec(), b"http/1.1".to_vec()];
} }
Ok(ServerCrypto { Ok(ServerCrypto {

View file

@ -4,14 +4,6 @@ use crate::error::*;
use chrono::{TimeZone, Utc}; use chrono::{TimeZone, Utc};
use derive_builder::Builder; use derive_builder::Builder;
#[derive(Debug, Clone)]
/// Struct to handle the sticky cookie string,
/// - passed from Rp module (http handler) to LB module, manipulated from req, only StickyCookieValue exists.
/// - passed from LB module to Rp module (http handler), will be inserted into res, StickyCookieValue and Info exist.
pub struct LbContext {
pub sticky_cookie: StickyCookie,
}
#[derive(Debug, Clone, Builder)] #[derive(Debug, Clone, Builder)]
/// Cookie value only, used for COOKIE in req /// Cookie value only, used for COOKIE in req
pub struct StickyCookieValue { pub struct StickyCookieValue {

View file

@ -1,14 +1,13 @@
use super::{ #[cfg(feature = "sticky-cookie")]
load_balance::{ use super::load_balance::LbStickyRoundRobinBuilder;
load_balance_options as lb_opts, LbRandomBuilder, LbRoundRobinBuilder, LbStickyRoundRobinBuilder, LoadBalance, use super::load_balance::{load_balance_options as lb_opts, LbRandomBuilder, LbRoundRobinBuilder, LoadBalance};
}, use super::{BytesName, LbContext, PathNameBytesExp, UpstreamOption};
load_balance_sticky_cookie::LbContext,
BytesName, PathNameBytesExp, UpstreamOption,
};
use crate::log::*; use crate::log::*;
#[cfg(feature = "sticky-cookie")]
use base64::{engine::general_purpose, Engine as _}; use base64::{engine::general_purpose, Engine as _};
use derive_builder::Builder; use derive_builder::Builder;
use rustc_hash::{FxHashMap as HashMap, FxHashSet as HashSet}; use rustc_hash::{FxHashMap as HashMap, FxHashSet as HashSet};
#[cfg(feature = "sticky-cookie")]
use sha2::{Digest, Sha256}; use sha2::{Digest, Sha256};
use std::borrow::Cow; use std::borrow::Cow;
#[derive(Debug, Clone)] #[derive(Debug, Clone)]
@ -58,6 +57,7 @@ pub struct Upstream {
pub uri: hyper::Uri, pub uri: hyper::Uri,
} }
impl Upstream { impl Upstream {
#[cfg(feature = "sticky-cookie")]
/// Hashing uri with index to avoid collision /// Hashing uri with index to avoid collision
pub fn calculate_id_with_index(&self, index: usize) -> String { pub fn calculate_id_with_index(&self, index: usize) -> String {
let mut hasher = Sha256::new(); let mut hasher = Sha256::new();
@ -114,8 +114,8 @@ impl UpstreamGroupBuilder {
v: &Option<String>, v: &Option<String>,
// upstream_num: &usize, // upstream_num: &usize,
upstream_vec: &Vec<Upstream>, upstream_vec: &Vec<Upstream>,
server_name: &str, _server_name: &str,
path_opt: &Option<String>, _path_opt: &Option<String>,
) -> &mut Self { ) -> &mut Self {
let upstream_num = &upstream_vec.len(); let upstream_num = &upstream_vec.len();
let lb = if let Some(x) = v { let lb = if let Some(x) = v {
@ -128,10 +128,11 @@ impl UpstreamGroupBuilder {
.build() .build()
.unwrap(), .unwrap(),
), ),
#[cfg(feature = "sticky-cookie")]
lb_opts::STICKY_ROUND_ROBIN => LoadBalance::StickyRoundRobin( lb_opts::STICKY_ROUND_ROBIN => LoadBalance::StickyRoundRobin(
LbStickyRoundRobinBuilder::default() LbStickyRoundRobinBuilder::default()
.num_upstreams(upstream_num) .num_upstreams(upstream_num)
.sticky_config(server_name, path_opt) .sticky_config(_server_name, _path_opt)
.upstream_maps(upstream_vec) // TODO: .upstream_maps(upstream_vec) // TODO:
.build() .build()
.unwrap(), .unwrap(),
@ -180,7 +181,10 @@ impl UpstreamGroup {
#[cfg(test)] #[cfg(test)]
mod test { mod test {
#[allow(unused)]
use super::*; use super::*;
#[cfg(feature = "sticky-cookie")]
#[test] #[test]
fn calc_id_works() { fn calc_id_works() {
let uri = "https://www.rust-lang.org".parse::<hyper::Uri>().unwrap(); let uri = "https://www.rust-lang.org".parse::<hyper::Uri>().unwrap();

View file

@ -25,5 +25,6 @@ pub mod H3 {
pub const MAX_IDLE_TIMEOUT: u64 = 10; // secs pub const MAX_IDLE_TIMEOUT: u64 = 10; // secs
} }
// For load-balancing with sticky cookie #[cfg(feature = "sticky-cookie")]
/// For load-balancing with sticky cookie
pub const STICKY_COOKIE_NAME: &str = "rpxy_srv_id"; pub const STICKY_COOKIE_NAME: &str = "rpxy_srv_id";

View file

@ -22,15 +22,18 @@ pub enum RpxyError {
#[error("TCP/UDP Proxy Layer Error: {0}")] #[error("TCP/UDP Proxy Layer Error: {0}")]
Proxy(String), Proxy(String),
#[allow(unused)]
#[error("LoadBalance Layer Error: {0}")] #[error("LoadBalance Layer Error: {0}")]
LoadBalance(String), LoadBalance(String),
#[error("I/O Error")] #[error("I/O Error")]
Io(#[from] io::Error), Io(#[from] io::Error),
#[cfg(feature = "http3")]
#[error("Quic Connection Error")] #[error("Quic Connection Error")]
QuicConn(#[from] quinn::ConnectionError), QuicConn(#[from] quinn::ConnectionError),
#[cfg(feature = "http3")]
#[error("H3 Error")] #[error("H3 Error")]
H3(#[from] h3::Error), H3(#[from] h3::Error),

View file

@ -1,7 +1,7 @@
// Highly motivated by https://github.com/felipenoris/hyper-reverse-proxy // Highly motivated by https://github.com/felipenoris/hyper-reverse-proxy
use super::{utils_headers::*, utils_request::*, utils_synth_response::*, HandlerContext}; use super::{utils_headers::*, utils_request::*, utils_synth_response::*, HandlerContext};
use crate::{ use crate::{
backend::{Backend, LoadBalance, UpstreamGroup}, backend::{Backend, UpstreamGroup},
error::*, error::*,
globals::Globals, globals::Globals,
log::*, log::*,
@ -91,7 +91,7 @@ where
let request_upgraded = req.extensions_mut().remove::<hyper::upgrade::OnUpgrade>(); let request_upgraded = req.extensions_mut().remove::<hyper::upgrade::OnUpgrade>();
// Build request from destination information // Build request from destination information
let context = match self.generate_request_forwarded( let _context = match self.generate_request_forwarded(
&client_addr, &client_addr,
&listen_addr, &listen_addr,
&mut req, &mut req,
@ -127,7 +127,8 @@ where
}; };
// Process reverse proxy context generated during the forwarding request generation. // Process reverse proxy context generated during the forwarding request generation.
if let Some(context_from_lb) = context.context_lb { #[cfg(feature = "sticky-cookie")]
if let Some(context_from_lb) = _context.context_lb {
let res_headers = res_backend.headers_mut(); let res_headers = res_backend.headers_mut();
if let Err(e) = set_sticky_cookie_lb_context(res_headers, &context_from_lb) { if let Err(e) = set_sticky_cookie_lb_context(res_headers, &context_from_lb) {
error!("Failed to append context to the response given from backend: {}", e); error!("Failed to append context to the response given from backend: {}", e);
@ -279,15 +280,24 @@ where
///////////////////////////////////////////// /////////////////////////////////////////////
// Fix unique upstream destination since there could be multiple ones. // Fix unique upstream destination since there could be multiple ones.
let context_to_lb = if let LoadBalance::StickyRoundRobin(lb) = &upstream_group.lb { #[cfg(feature = "sticky-cookie")]
let (upstream_chosen_opt, context_from_lb) = {
let context_to_lb = if let crate::backend::LoadBalance::StickyRoundRobin(lb) = &upstream_group.lb {
takeout_sticky_cookie_lb_context(req.headers_mut(), &lb.sticky_config.name)? takeout_sticky_cookie_lb_context(req.headers_mut(), &lb.sticky_config.name)?
} else { } else {
None None
}; };
let (upstream_chosen_opt, context_from_lb) = upstream_group.get(&context_to_lb); upstream_group.get(&context_to_lb)
};
#[cfg(not(feature = "sticky-cookie"))]
let (upstream_chosen_opt, _) = upstream_group.get(&None);
let upstream_chosen = upstream_chosen_opt.ok_or_else(|| anyhow!("Failed to get upstream"))?; let upstream_chosen = upstream_chosen_opt.ok_or_else(|| anyhow!("Failed to get upstream"))?;
let context = HandlerContext { let context = HandlerContext {
#[cfg(feature = "sticky-cookie")]
context_lb: context_from_lb, context_lb: context_from_lb,
#[cfg(not(feature = "sticky-cookie"))]
context_lb: None,
}; };
///////////////////////////////////////////// /////////////////////////////////////////////

View file

@ -3,11 +3,15 @@ mod utils_headers;
mod utils_request; mod utils_request;
mod utils_synth_response; mod utils_synth_response;
#[cfg(feature = "sticky-cookie")]
use crate::backend::LbContext;
pub use handler_main::{HttpMessageHandler, HttpMessageHandlerBuilder, HttpMessageHandlerBuilderError}; pub use handler_main::{HttpMessageHandler, HttpMessageHandlerBuilder, HttpMessageHandlerBuilderError};
use crate::backend::LbContext; #[allow(dead_code)]
#[derive(Debug)] #[derive(Debug)]
struct HandlerContext { struct HandlerContext {
#[cfg(feature = "sticky-cookie")]
context_lb: Option<LbContext>, context_lb: Option<LbContext>,
#[cfg(not(feature = "sticky-cookie"))]
context_lb: Option<()>,
} }

View file

@ -1,9 +1,8 @@
use crate::{ #[cfg(feature = "sticky-cookie")]
backend::{LbContext, StickyCookie, StickyCookieValue, UpstreamGroup, UpstreamOption}, use crate::backend::{LbContext, StickyCookie, StickyCookieValue};
error::*, use crate::backend::{UpstreamGroup, UpstreamOption};
log::*,
utils::*, use crate::{error::*, log::*, utils::*};
};
use bytes::BufMut; use bytes::BufMut;
use hyper::{ use hyper::{
header::{self, HeaderMap, HeaderName, HeaderValue}, header::{self, HeaderMap, HeaderName, HeaderValue},
@ -14,6 +13,7 @@ use std::net::SocketAddr;
//////////////////////////////////////////////////// ////////////////////////////////////////////////////
// Functions to manipulate headers // Functions to manipulate headers
#[cfg(feature = "sticky-cookie")]
/// Take sticky cookie header value from request header, /// Take sticky cookie header value from request header,
/// and returns LbContext to be forwarded to LB if exist and if needed. /// and returns LbContext to be forwarded to LB if exist and if needed.
/// Removing sticky cookie is needed and it must not be passed to the upstream. /// Removing sticky cookie is needed and it must not be passed to the upstream.
@ -55,6 +55,7 @@ pub(super) fn takeout_sticky_cookie_lb_context(
} }
} }
#[cfg(feature = "sticky-cookie")]
/// Set-Cookie if LB Sticky is enabled and if cookie is newly created/updated. /// Set-Cookie if LB Sticky is enabled and if cookie is newly created/updated.
/// Set-Cookie response header could be in multiple lines. /// Set-Cookie response header could be in multiple lines.
/// https://developer.mozilla.org/ja/docs/Web/HTTP/Headers/Set-Cookie /// https://developer.mozilla.org/ja/docs/Web/HTTP/Headers/Set-Cookie

View file

@ -6,10 +6,10 @@ use crate::{
log::*, log::*,
utils::BytesName, utils::BytesName,
}; };
#[cfg(feature = "http3")]
use hyper::{client::connect::Connect, server::conn::Http}; use hyper::{client::connect::Connect, server::conn::Http};
#[cfg(feature = "http3")] #[cfg(feature = "http3")]
use quinn::{crypto::rustls::HandshakeData, Endpoint, ServerConfig as QuicServerConfig, TransportConfig}; use quinn::{crypto::rustls::HandshakeData, Endpoint, ServerConfig as QuicServerConfig, TransportConfig};
#[cfg(feature = "http3")]
use rustls::ServerConfig; use rustls::ServerConfig;
use std::sync::Arc; use std::sync::Arc;
use tokio::{ use tokio::{
@ -196,14 +196,14 @@ where
let (tx, rx) = watch::channel::<Option<Arc<ServerCrypto>>>(None); let (tx, rx) = watch::channel::<Option<Arc<ServerCrypto>>>(None);
#[cfg(not(feature = "http3"))] #[cfg(not(feature = "http3"))]
{ {
select! { tokio::select! {
_= self.cert_service(tx).fuse() => { _= self.cert_service(tx) => {
error!("Cert service for TLS exited"); error!("Cert service for TLS exited");
}, },
_ = self.listener_service(server, rx).fuse() => { _ = self.listener_service(server, rx) => {
error!("TCP proxy service for TLS exited"); error!("TCP proxy service for TLS exited");
}, },
complete => { else => {
error!("Something went wrong"); error!("Something went wrong");
return Ok(()) return Ok(())
} }