use std::{
collections::HashMap,
sync::{Arc, RwLock},
};
use nfrp_common::config::ProxyConfig;
#[derive(Clone, Copy, PartialEq, Eq, Debug)]
pub enum ProxyOrigin {
Config,
Dynamic,
}
#[derive(Clone)]
struct ProxyEntry {
cfg: ProxyConfig,
origin: ProxyOrigin,
}
#[derive(Clone, Default)]
pub struct ProxyTable {
inner: Arc<RwLock<HashMap<String, ProxyEntry>>>,
}
impl ProxyTable {
pub fn from_iter<I: IntoIterator<Item = ProxyConfig>>(items: I) -> Self {
let map = items
.into_iter()
.map(|p| {
(
p.name.clone(),
ProxyEntry {
cfg: p,
origin: ProxyOrigin::Config,
},
)
})
.collect::<HashMap<_, _>>();
Self {
inner: Arc::new(RwLock::new(map)),
}
}
pub fn load_dynamic<I: IntoIterator<Item = ProxyConfig>>(&self, items: I) {
let mut g = match self.inner.write() {
Ok(g) => g,
Err(e) => e.into_inner(),
};
for p in items {
g.entry(p.name.clone()).or_insert(ProxyEntry {
cfg: p,
origin: ProxyOrigin::Dynamic,
});
}
}
pub fn get(&self, name: &str) -> Option<Arc<ProxyConfig>> {
self.read().get(name).map(|e| Arc::new(e.cfg.clone()))
}
pub fn origin_of(&self, name: &str) -> Option<ProxyOrigin> {
self.read().get(name).map(|e| e.origin)
}
fn read(&self) -> std::sync::RwLockReadGuard<'_, HashMap<String, ProxyEntry>> {
self.inner.read().unwrap_or_else(|e| e.into_inner())
}
pub fn names(&self) -> Vec<String> {
self.read().keys().cloned().collect()
}
pub fn list(&self) -> Vec<ProxyConfig> {
let mut v: Vec<ProxyConfig> = self.read().values().map(|e| e.cfg.clone()).collect();
v.sort_by(|a, b| a.name.cmp(&b.name));
v
}
pub fn insert_dynamic(&self, p: ProxyConfig) -> Result<Option<ProxyConfig>, String> {
let mut g = match self.inner.write() {
Ok(g) => g,
Err(e) => e.into_inner(),
};
if let Some(existing) = g.get(&p.name) {
if existing.origin == ProxyOrigin::Config {
return Err(format!(
"代理 [{}] 来自本机配置文件,远程下发不能覆盖它;\
请在本机改配置文件(或换个名字)",
p.name
));
}
}
Ok(g.insert(
p.name.clone(),
ProxyEntry {
cfg: p,
origin: ProxyOrigin::Dynamic,
},
)
.map(|e| e.cfg))
}
pub fn remove_dynamic(&self, name: &str) -> Result<Option<ProxyConfig>, String> {
let mut g = match self.inner.write() {
Ok(g) => g,
Err(e) => e.into_inner(),
};
if let Some(existing) = g.get(name) {
if existing.origin == ProxyOrigin::Config {
return Err(format!(
"代理 [{name}] 来自本机配置文件,远程下发不能删除它;\
请在本机改配置文件"
));
}
}
Ok(g.remove(name).map(|e| e.cfg))
}
pub fn insert(&self, p: ProxyConfig) -> Option<ProxyConfig> {
let mut g = match self.inner.write() {
Ok(g) => g,
Err(e) => e.into_inner(),
};
let origin = g
.get(&p.name)
.map(|e| e.origin)
.unwrap_or(ProxyOrigin::Dynamic);
g.insert(p.name.clone(), ProxyEntry { cfg: p, origin })
.map(|e| e.cfg)
}
pub fn remove(&self, name: &str) -> Option<ProxyConfig> {
let mut g = match self.inner.write() {
Ok(g) => g,
Err(e) => e.into_inner(),
};
g.remove(name).map(|e| e.cfg)
}
pub fn len(&self) -> usize {
self.read().len()
}
}
#[cfg(test)]
mod tests {
use super::*;
use nfrp_common::config::ProxyConfig;
fn p(name: &str) -> ProxyConfig {
ProxyConfig {
name: name.to_string(),
..Default::default()
}
}
#[test]
fn 句柄共享同一份表() {
let t = ProxyTable::from_iter([p("a"), p("b")]);
let t2 = t.clone();
assert_eq!(t.len(), 2);
t2.insert(p("c"));
assert_eq!(t.len(), 3, "clone 出来的句柄应当看到同一份表");
assert!(t.get("c").is_some());
}
#[test]
fn 按原始名增删() {
let t = ProxyTable::from_iter([p("web")]);
assert_eq!(t.get("web").unwrap().name, "web");
assert!(t.get("nope").is_none());
let old = t.insert(p("web"));
assert!(old.is_some(), "覆盖同名代理应当返回旧值");
assert_eq!(t.len(), 1, "覆盖不该把条目数变多");
assert!(t.remove("web").is_some());
assert_eq!(t.len(), 0);
assert!(t.remove("web").is_none(), "删两次不该 panic");
}
#[test]
fn 远程删除不得动配置里的代理() {
let t = ProxyTable::from_iter([p("victim-ssh")]);
assert_eq!(t.origin_of("victim-ssh"), Some(ProxyOrigin::Config));
let e = t
.remove_dynamic("victim-ssh")
.expect_err("删配置代理必须被拒");
assert!(e.contains("配置文件"), "错误要说清原因:{e}");
assert!(t.get("victim-ssh").is_some(), "被拒之后那条代理必须还在");
}
#[test]
fn 远程新增不得覆盖配置里的代理() {
let mut cfg = p("my-web");
cfg.remote_port = 18202;
let t = ProxyTable::from_iter([cfg.clone()]);
let mut evil = p("my-web");
evil.local_addr = "127.0.0.1:9".into();
let e = t.insert_dynamic(evil).expect_err("覆盖配置代理必须被拒");
assert!(e.contains("配置文件"), "错误要说清原因:{e}");
let got = t.get("my-web").expect("配置代理必须还在");
assert_eq!(got.remote_port, 18202, "配置代理的端口不该被改掉");
}
#[test]
fn 动态代理之间照常增删() {
let t = ProxyTable::from_iter([p("cfg-one")]);
t.insert_dynamic(p("dyn-one")).expect("新增动态代理");
assert_eq!(t.origin_of("dyn-one"), Some(ProxyOrigin::Dynamic));
t.insert_dynamic(p("dyn-one"))
.expect("覆盖动态代理应当允许");
assert!(t.remove_dynamic("dyn-one").expect("删动态代理").is_some());
assert!(t.get("dyn-one").is_none());
assert!(t.get("cfg-one").is_some());
}
#[test]
fn store_恢复的代理算动态() {
let t = ProxyTable::from_iter([p("from-config")]);
t.load_dynamic([p("from-store")]);
assert_eq!(t.origin_of("from-config"), Some(ProxyOrigin::Config));
assert_eq!(t.origin_of("from-store"), Some(ProxyOrigin::Dynamic));
assert!(t.remove_dynamic("from-store").expect("应可删").is_some());
}
#[test]
fn 锁损坏时不panic() {
let t = ProxyTable::from_iter([p("a")]);
let t2 = t.clone();
let _ = std::thread::spawn(move || {
let _g = t2.inner.write().unwrap();
panic!("故意把锁写脏");
})
.join();
t.insert(p("b"));
assert!(t.get("b").is_some(), "poison 之后仍要能读");
assert!(t.names().contains(&"b".to_string()));
assert_eq!(t.len(), 2);
}
}