use ts_control::StableNodeId;
use ts_derp::RegionId;
use crate::status::NetcheckReport;
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ExitNodeCandidate {
pub stable_id: StableNodeId,
pub name: String,
pub derp_region: Option<RegionId>,
pub online: Option<bool>,
pub advertises_exit_route: bool,
pub has_suggest_cap: bool,
}
impl ExitNodeCandidate {
fn is_eligible(&self) -> bool {
self.online == Some(true) && self.has_suggest_cap && self.advertises_exit_route
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ExitNodeSuggestion {
pub id: StableNodeId,
pub name: String,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SuggestExitNodeError {
NoPreferredDerp,
}
impl core::fmt::Display for SuggestExitNodeError {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
match self {
Self::NoPreferredDerp => write!(f, "no preferred DERP, try again later"),
}
}
}
impl core::error::Error for SuggestExitNodeError {}
pub type SelectRegion<'a> = dyn Fn(&[RegionId]) -> RegionId + 'a;
pub type SelectNode<'a> =
dyn Fn(&[ExitNodeCandidate], Option<&StableNodeId>) -> ExitNodeCandidate + 'a;
pub(crate) fn suggest_exit_node(
report: &NetcheckReport,
candidates: &[ExitNodeCandidate],
prev_suggestion: Option<&StableNodeId>,
select_region: &SelectRegion<'_>,
select_node: &SelectNode<'_>,
) -> Result<Option<ExitNodeSuggestion>, SuggestExitNodeError> {
match report.preferred_derp {
None | Some(0) => return Err(SuggestExitNodeError::NoPreferredDerp),
Some(_) => {}
}
let eligible: Vec<&ExitNodeCandidate> = candidates.iter().filter(|c| c.is_eligible()).collect();
match eligible.as_slice() {
[] => return Ok(None),
[only] => {
return Ok(Some(ExitNodeSuggestion {
id: only.stable_id.clone(),
name: only.name.clone(),
}));
}
_ => {}
}
let mut by_region: std::collections::BTreeMap<RegionId, Vec<ExitNodeCandidate>> =
std::collections::BTreeMap::new();
let mut region_less: Vec<ExitNodeCandidate> = Vec::new();
for c in eligible {
match c.derp_region {
Some(region) => by_region.entry(region).or_default().push(c.clone()),
None => region_less.push(c.clone()),
}
}
if !by_region.is_empty() {
let regions: Vec<RegionId> = by_region.keys().copied().collect();
let min_region = match min_latency_derp_region(®ions, report) {
Some(region) => region,
None => select_region(®ions),
};
let region_candidates = by_region
.get(&min_region)
.expect("selected region must be a candidate region");
let chosen = select_node(region_candidates, prev_suggestion);
return Ok(Some(ExitNodeSuggestion {
id: chosen.stable_id,
name: chosen.name,
}));
}
let chosen = select_node(®ion_less, prev_suggestion);
Ok(Some(ExitNodeSuggestion {
id: chosen.stable_id,
name: chosen.name,
}))
}
fn min_latency_derp_region(regions: &[RegionId], report: &NetcheckReport) -> Option<RegionId> {
let latency_of = |region: RegionId| -> Option<core::time::Duration> {
report
.region_latencies
.iter()
.find(|rl| rl.region_id == region.0.get())
.map(|rl| rl.latency)
};
let max_duration = core::time::Duration::MAX;
let min = regions.iter().copied().min_by(|&i, &j| {
let il = latency_of(i).unwrap_or(max_duration);
let jl = latency_of(j).unwrap_or(max_duration);
il.cmp(&jl).then_with(|| i.0.get().cmp(&j.0.get()))
})?;
match latency_of(min) {
Some(latency) if !latency.is_zero() => Some(min),
_ => None,
}
}
pub(crate) fn random_region(regions: &[RegionId]) -> RegionId {
regions[rand::random_range(0..regions.len())]
}
pub(crate) fn random_node(
nodes: &[ExitNodeCandidate],
prefer: Option<&StableNodeId>,
) -> ExitNodeCandidate {
if let Some(prefer) = prefer.filter(|p| !p.0.is_empty())
&& let Some(found) = nodes.iter().find(|n| &n.stable_id == prefer)
{
return found.clone();
}
nodes[rand::random_range(0..nodes.len())].clone()
}
pub(crate) fn next_sticky(
prev: Option<StableNodeId>,
outcome: &Result<Option<ExitNodeSuggestion>, SuggestExitNodeError>,
) -> Option<StableNodeId> {
match outcome {
Ok(maybe) => maybe.as_ref().map(|s| s.id.clone()),
Err(_) => prev,
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::status::RegionLatency;
fn region(id: u32) -> RegionId {
RegionId(core::num::NonZeroU32::new(id).unwrap())
}
fn report(preferred: Option<u32>, latencies: &[(u32, u64)]) -> NetcheckReport {
NetcheckReport {
preferred_derp: preferred,
region_latencies: latencies
.iter()
.map(|&(region_id, ms)| RegionLatency {
region_id,
latency: core::time::Duration::from_millis(ms),
})
.collect(),
}
}
fn candidate(id: u32, derp: Option<u32>) -> ExitNodeCandidate {
ExitNodeCandidate {
stable_id: StableNodeId(format!("stable{id}")),
name: format!("peer{id}"),
derp_region: derp.map(region),
online: Some(true),
advertises_exit_route: true,
has_suggest_cap: true,
}
}
fn pick_region(want: Vec<RegionId>, use_region: RegionId) -> impl Fn(&[RegionId]) -> RegionId {
move |got: &[RegionId]| {
let mut got_sorted = got.to_vec();
got_sorted.sort();
let mut want_sorted = want.clone();
want_sorted.sort();
assert_eq!(got_sorted, want_sorted, "candidate regions mismatch");
assert!(want.contains(&use_region), "use_region must be in want");
use_region
}
}
fn pick_node(
want: Vec<&'static str>,
want_last: Option<&'static str>,
use_id: &'static str,
) -> impl Fn(&[ExitNodeCandidate], Option<&StableNodeId>) -> ExitNodeCandidate {
move |got: &[ExitNodeCandidate], last: Option<&StableNodeId>| {
let via_random = random_node(got, last);
assert!(
got.iter().any(|c| c.stable_id == via_random.stable_id),
"random_node returned a non-member"
);
let got_ids: Vec<String> = got.iter().map(|c| c.stable_id.0.clone()).collect();
let mut got_sorted = got_ids.clone();
got_sorted.sort();
let mut want_sorted: Vec<String> = want.iter().map(|s| s.to_string()).collect();
want_sorted.sort();
assert_eq!(got_sorted, want_sorted, "candidate nodes mismatch");
let last_str = last.map(|s| s.0.as_str());
assert_eq!(last_str, want_last, "last (prev suggestion) mismatch");
got.iter()
.find(|c| c.stable_id.0 == use_id)
.cloned()
.expect("use_id must be among candidates")
}
}
fn unused_region() -> impl Fn(&[RegionId]) -> RegionId {
|_: &[RegionId]| panic!("select_region must not be called on this path")
}
fn unused_node() -> impl Fn(&[ExitNodeCandidate], Option<&StableNodeId>) -> ExitNodeCandidate {
|_: &[ExitNodeCandidate], _: Option<&StableNodeId>| {
panic!("select_node must not be called on this path")
}
}
#[test]
fn no_preferred_derp_errors() {
let r = report(None, &[(1, 10)]);
let cands = [candidate(1, Some(1)), candidate(2, Some(2))];
let err = suggest_exit_node(&r, &cands, None, &unused_region(), &unused_node())
.expect_err("no preferred DERP must error");
assert_eq!(err, SuggestExitNodeError::NoPreferredDerp);
let r0 = report(Some(0), &[(1, 10)]);
assert_eq!(
suggest_exit_node(&r0, &cands, None, &unused_region(), &unused_node()),
Err(SuggestExitNodeError::NoPreferredDerp)
);
}
#[test]
fn no_candidates_returns_none() {
let r = report(Some(1), &[(1, 10)]);
assert_eq!(
suggest_exit_node(&r, &[], None, &unused_region(), &unused_node()),
Ok(None)
);
}
#[test]
fn single_candidate_returned_directly() {
let r = report(Some(1), &[(1, 10)]);
let cands = [candidate(7, Some(2))];
let got = suggest_exit_node(&r, &cands, None, &unused_region(), &unused_node())
.expect("ok")
.expect("some");
assert_eq!(got.id, StableNodeId("stable7".into()));
assert_eq!(got.name, "peer7");
}
#[test]
fn two_regions_lower_latency_wins() {
let r = report(Some(1), &[(1, 10), (2, 20), (3, 30)]);
let cands = [candidate(2, Some(1)), candidate(4, Some(3))];
let select_node = pick_node(vec!["stable2"], None, "stable2");
let got = suggest_exit_node(&r, &cands, None, &unused_region(), &select_node)
.expect("ok")
.expect("some");
assert_eq!(got.id, StableNodeId("stable2".into()));
assert_eq!(got.name, "peer2");
}
#[test]
fn two_candidates_same_region_select_node_picks() {
let r = report(Some(1), &[(1, 10), (2, 20), (3, 30)]);
let cands = [candidate(1, Some(1)), candidate(2, Some(1))];
let select_node = pick_node(vec!["stable1", "stable2"], None, "stable1");
let got = suggest_exit_node(&r, &cands, None, &unused_region(), &select_node)
.expect("ok")
.expect("some");
assert_eq!(got.id, StableNodeId("stable1".into()));
assert_eq!(got.name, "peer1");
}
#[test]
fn prev_suggestion_sticky_when_present() {
let r = report(Some(1), &[(1, 10), (2, 20), (3, 30)]);
let cands = [candidate(1, Some(1)), candidate(2, Some(1))];
let prev = StableNodeId("stable2".into());
let select_node = pick_node(vec!["stable1", "stable2"], Some("stable2"), "stable2");
let got = suggest_exit_node(&r, &cands, Some(&prev), &unused_region(), &select_node)
.expect("ok")
.expect("some");
assert_eq!(got.id, StableNodeId("stable2".into()));
assert_eq!(got.name, "peer2");
}
#[test]
fn better_region_beats_stale_prev_suggestion() {
let r = report(Some(1), &[(1, 10), (2, 20), (3, 30)]);
let cands = [candidate(2, Some(1)), candidate(3, Some(3))];
let prev = StableNodeId("stable3".into());
let select_node = pick_node(vec!["stable2"], Some("stable3"), "stable2");
let got = suggest_exit_node(&r, &cands, Some(&prev), &unused_region(), &select_node)
.expect("ok")
.expect("some");
assert_eq!(got.id, StableNodeId("stable2".into()));
}
#[test]
fn equal_latency_lower_region_id_wins() {
let r = report(Some(1), &[(1, 10), (2, 20), (3, 10)]);
let cands = [candidate(1, Some(1)), candidate(3, Some(3))];
let select_node = pick_node(vec!["stable1"], None, "stable1");
let got = suggest_exit_node(&r, &cands, None, &unused_region(), &select_node)
.expect("ok")
.expect("some");
assert_eq!(got.id, StableNodeId("stable1".into()));
assert_eq!(got.name, "peer1");
}
#[test]
fn no_usable_latency_falls_back_to_select_region() {
let r = report(Some(1), &[(1, 0), (2, 0), (3, 0)]);
let cands = [candidate(2, Some(1)), candidate(4, Some(3))];
let select_region = pick_region(vec![region(1), region(3)], region(1));
let select_node = pick_node(vec!["stable2"], None, "stable2");
let got = suggest_exit_node(&r, &cands, None, &select_region, &select_node)
.expect("ok")
.expect("some");
assert_eq!(got.id, StableNodeId("stable2".into()));
assert_eq!(got.name, "peer2");
}
#[test]
fn missing_latency_loses_to_measured_region() {
let r = report(Some(3), &[(3, 10)]); let cands = [candidate(1, Some(1)), candidate(3, Some(3))];
let select_node = pick_node(vec!["stable3"], None, "stable3");
let got = suggest_exit_node(&r, &cands, None, &unused_region(), &select_node)
.expect("ok")
.expect("some");
assert_eq!(got.id, StableNodeId("stable3".into()));
}
#[test]
fn predicate_excludes_missing_suggest_cap() {
let r = report(Some(1), &[(1, 10)]);
let mut no_cap = candidate(1, Some(1));
no_cap.has_suggest_cap = false;
let cands = [no_cap, candidate(2, Some(2))];
let got = suggest_exit_node(&r, &cands, None, &unused_region(), &unused_node())
.expect("ok")
.expect("some");
assert_eq!(got.id, StableNodeId("stable2".into()));
}
#[test]
fn predicate_excludes_no_exit_route() {
let r = report(Some(1), &[(1, 10)]);
let mut no_route = candidate(1, Some(1));
no_route.advertises_exit_route = false;
let cands = [no_route, candidate(2, Some(2))];
let got = suggest_exit_node(&r, &cands, None, &unused_region(), &unused_node())
.expect("ok")
.expect("some");
assert_eq!(got.id, StableNodeId("stable2".into()));
}
#[test]
fn predicate_excludes_offline_and_unknown() {
let r = report(Some(1), &[(1, 10)]);
let mut offline = candidate(1, Some(1));
offline.online = Some(false);
let mut unknown = candidate(3, Some(3));
unknown.online = None;
let cands = [offline, unknown, candidate(2, Some(2))];
let got = suggest_exit_node(&r, &cands, None, &unused_region(), &unused_node())
.expect("ok")
.expect("some");
assert_eq!(got.id, StableNodeId("stable2".into()));
let r2 = report(Some(1), &[(1, 10)]);
let mut lone_offline = candidate(9, Some(1));
lone_offline.online = Some(false);
assert_eq!(
suggest_exit_node(&r2, &[lone_offline], None, &unused_region(), &unused_node()),
Ok(None)
);
}
#[test]
fn all_region_less_falls_back_to_select_node() {
let r = report(Some(1), &[(1, 10)]);
let cands = [candidate(5, None), candidate(6, None)];
let select_node = pick_node(vec!["stable5", "stable6"], None, "stable5");
let got = suggest_exit_node(&r, &cands, None, &unused_region(), &select_node)
.expect("ok")
.expect("some");
assert_eq!(got.id, StableNodeId("stable5".into()));
assert_eq!(got.name, "peer5");
}
#[test]
fn region_less_skipped_when_derp_homed_exists() {
let r = report(Some(1), &[(1, 10)]);
let cands = [candidate(6, None), candidate(2, Some(1))];
let select_node = pick_node(vec!["stable2"], None, "stable2");
let got = suggest_exit_node(&r, &cands, None, &unused_region(), &select_node)
.expect("ok")
.expect("some");
assert_eq!(got.id, StableNodeId("stable2".into()));
}
#[test]
fn random_node_prefers_then_falls_back() {
let cands = [candidate(1, Some(1)), candidate(2, Some(1))];
let prefer = StableNodeId("stable2".into());
assert_eq!(random_node(&cands, Some(&prefer)).stable_id, prefer);
let absent = StableNodeId("stableX".into());
let got = random_node(&cands, Some(&absent));
assert!(cands.iter().any(|c| c.stable_id == got.stable_id));
let got2 = random_node(&cands, None);
assert!(cands.iter().any(|c| c.stable_id == got2.stable_id));
}
#[test]
fn min_latency_region_semantics() {
let r = report(Some(1), &[(1, 30), (2, 10), (3, 20)]);
assert_eq!(
min_latency_derp_region(&[region(1), region(2), region(3)], &r),
Some(region(2))
);
let req = report(Some(1), &[(1, 10), (2, 10)]);
assert_eq!(
min_latency_derp_region(&[region(1), region(2)], &req),
Some(region(1))
);
let rz = report(Some(1), &[(1, 0), (2, 0)]);
assert_eq!(min_latency_derp_region(&[region(1), region(2)], &rz), None);
let rm = report(Some(1), &[(1, 10)]);
assert_eq!(min_latency_derp_region(&[region(5)], &rm), None);
}
#[test]
fn next_sticky_matches_go_last_suggested() {
let sugg = ExitNodeSuggestion {
id: StableNodeId("stable2".to_owned()),
name: "peer2".to_owned(),
};
let prev = || Some(StableNodeId("stable1".to_owned()));
assert_eq!(
next_sticky(prev(), &Ok(Some(sugg.clone()))),
Some(StableNodeId("stable2".to_owned()))
);
assert_eq!(
next_sticky(None, &Ok(Some(sugg))),
Some(StableNodeId("stable2".to_owned()))
);
assert_eq!(next_sticky(prev(), &Ok(None)), None);
assert_eq!(
next_sticky(prev(), &Err(SuggestExitNodeError::NoPreferredDerp)),
prev()
);
assert_eq!(
next_sticky(None, &Err(SuggestExitNodeError::NoPreferredDerp)),
None
);
}
#[test]
fn random_node_ignores_empty_prefer_id() {
let only = candidate(7, Some(1));
let empty = StableNodeId(String::new());
let picked = random_node(std::slice::from_ref(&only), Some(&empty));
assert_eq!(picked.stable_id, only.stable_id);
}
}