#[cfg(feature = "sticky-cookie")] use super::load_balance::LbStickyRoundRobinBuilder; use super::load_balance::{load_balance_options as lb_opts, LbRandomBuilder, LbRoundRobinBuilder, LoadBalance}; use super::{BytesName, LbContext, PathNameBytesExp, UpstreamOption}; use crate::log::*; #[cfg(feature = "sticky-cookie")] use base64::{engine::general_purpose, Engine as _}; use derive_builder::Builder; use rustc_hash::{FxHashMap as HashMap, FxHashSet as HashSet}; #[cfg(feature = "sticky-cookie")] use sha2::{Digest, Sha256}; use std::borrow::Cow; #[derive(Debug, Clone)] pub struct ReverseProxy { pub upstream: HashMap, // TODO: HashMapでいいのかは疑問。max_by_keyでlongest prefix matchしてるのも無駄っぽいが。。。 } impl ReverseProxy { /// Get an appropriate upstream destination for given path string. pub fn get<'a>(&self, path_str: impl Into>) -> Option<&UpstreamGroup> { // trie使ってlongest prefix match させてもいいけどルート記述は少ないと思われるので、 // コスト的にこの程度で十分 let path_bytes = &path_str.to_path_name_vec(); let matched_upstream = self .upstream .iter() .filter(|(route_bytes, _)| { match path_bytes.starts_with(route_bytes) { true => { route_bytes.len() == 1 // route = '/', i.e., default || match path_bytes.get(route_bytes.len()) { None => true, // exact case Some(p) => p == &b'/', // sub-path case } } _ => false, } }) .max_by_key(|(route_bytes, _)| route_bytes.len()); if let Some((_path, u)) = matched_upstream { debug!( "Found upstream: {:?}", String::from_utf8(_path.0.clone()).unwrap_or_else(|_| "".to_string()) ); Some(u) } else { None } } } #[derive(Debug, Clone)] /// Upstream struct just containing uri without path pub struct Upstream { /// Base uri without specific path pub uri: hyper::Uri, } impl Upstream { #[cfg(feature = "sticky-cookie")] /// Hashing uri with index to avoid collision pub fn calculate_id_with_index(&self, index: usize) -> String { let mut hasher = Sha256::new(); let uri_string = format!("{}&index={}", self.uri.clone(), index); hasher.update(uri_string.as_bytes()); let digest = hasher.finalize(); general_purpose::URL_SAFE_NO_PAD.encode(digest) } } #[derive(Debug, Clone, Builder)] /// Struct serving multiple upstream servers for, e.g., load balancing. pub struct UpstreamGroup { #[builder(setter(custom))] /// Upstream server(s) pub upstream: Vec, #[builder(setter(custom), default)] /// Path like "/path" in [[PathNameBytesExp]] associated with the upstream server(s) pub path: PathNameBytesExp, #[builder(setter(custom), default)] /// Path in [[PathNameBytesExp]] that will be used to replace the "path" part of incoming url pub replace_path: Option, #[builder(setter(custom), default)] /// Load balancing option pub lb: LoadBalance, #[builder(setter(custom), default)] /// Activated upstream options defined in [[UpstreamOption]] pub opts: HashSet, } impl UpstreamGroupBuilder { pub fn upstream(&mut self, upstream_vec: &[Upstream]) -> &mut Self { self.upstream = Some(upstream_vec.to_vec()); self } pub fn path(&mut self, v: &Option) -> &mut Self { let path = match v { Some(p) => p.to_path_name_vec(), None => "/".to_path_name_vec(), }; self.path = Some(path); self } pub fn replace_path(&mut self, v: &Option) -> &mut Self { self.replace_path = Some( v.to_owned() .as_ref() .map_or_else(|| None, |v| Some(v.to_path_name_vec())), ); self } pub fn lb( &mut self, v: &Option, // upstream_num: &usize, upstream_vec: &Vec, _server_name: &str, _path_opt: &Option, ) -> &mut Self { let upstream_num = &upstream_vec.len(); let lb = if let Some(x) = v { match x.as_str() { lb_opts::FIX_TO_FIRST => LoadBalance::FixToFirst, lb_opts::RANDOM => LoadBalance::Random(LbRandomBuilder::default().num_upstreams(upstream_num).build().unwrap()), lb_opts::ROUND_ROBIN => LoadBalance::RoundRobin( LbRoundRobinBuilder::default() .num_upstreams(upstream_num) .build() .unwrap(), ), #[cfg(feature = "sticky-cookie")] lb_opts::STICKY_ROUND_ROBIN => LoadBalance::StickyRoundRobin( LbStickyRoundRobinBuilder::default() .num_upstreams(upstream_num) .sticky_config(_server_name, _path_opt) .upstream_maps(upstream_vec) // TODO: .build() .unwrap(), ), _ => { error!("Specified load balancing option is invalid."); LoadBalance::default() } } } else { LoadBalance::default() }; self.lb = Some(lb); self } pub fn opts(&mut self, v: &Option>) -> &mut Self { let opts = if let Some(opts) = v { opts .iter() .filter_map(|str| UpstreamOption::try_from(str.as_str()).ok()) .collect::>() } else { Default::default() }; self.opts = Some(opts); self } } impl UpstreamGroup { /// Get an enabled option of load balancing [[LoadBalance]] pub fn get(&self, context_to_lb: &Option) -> (Option<&Upstream>, Option) { let pointer_to_upstream = self.lb.get_context(context_to_lb); debug!("Upstream of index {} is chosen.", pointer_to_upstream.ptr); debug!("Context to LB (Cookie in Req): {:?}", context_to_lb); debug!( "Context from LB (Set-Cookie in Res): {:?}", pointer_to_upstream.context_lb ); ( self.upstream.get(pointer_to_upstream.ptr), pointer_to_upstream.context_lb, ) } } #[cfg(test)] mod test { #[allow(unused)] use super::*; #[cfg(feature = "sticky-cookie")] #[test] fn calc_id_works() { let uri = "https://www.rust-lang.org".parse::().unwrap(); let upstream = Upstream { uri }; assert_eq!( "eGsjoPbactQ1eUJjafYjPT3ekYZQkaqJnHdA_FMSkgM", upstream.calculate_id_with_index(0) ); assert_eq!( "tNVXFJ9eNCT2mFgKbYq35XgH5q93QZtfU8piUiiDxVA", upstream.calculate_id_with_index(1) ); } }