tor_netdir/
weight.rs

1//! Functions for applying the correct weights to relays when choosing
2//! a relay at random.
3//!
4//! The weight to use when picking a relay depends on several factors:
5//!
6//! - The relay's *apparent bandwidth*.  (This is ideally measured by a set of
7//!   bandwidth authorities, but if no bandwidth authorities are running (as on
8//!   a test network), we might fall back either to relays' self-declared
9//!   values, or we might treat all relays as having equal bandwidth.)
10//! - The role that we're selecting a relay to play.  (See [`WeightRole`]).
11//! - The flags that a relay has in the consensus, and their scarcity.  If a
12//!   relay provides particularly scarce functionality, we might choose not to
13//!   use it for other roles, or to use it less commonly for them.
14
15use crate::params::NetParameters;
16use crate::ConsensusRelays;
17use bitflags::bitflags;
18use tor_netdoc::doc::netstatus::{self, MdConsensus, MdConsensusRouterStatus, NetParams};
19
20/// Helper: Calculate the function we should use to find initial relay
21/// bandwidths.
22fn pick_bandwidth_fn<'a, I>(mut weights: I) -> BandwidthFn
23where
24    I: Clone + Iterator<Item = &'a netstatus::RelayWeight>,
25{
26    let has_measured = weights.clone().any(|w| w.is_measured());
27    let has_nonzero = weights.clone().any(|w| w.is_nonzero());
28    let has_nonzero_measured = weights.any(|w| w.is_measured() && w.is_nonzero());
29
30    if !has_nonzero {
31        // If every value is zero, we should just pretend everything has
32        // bandwidth == 1.
33        BandwidthFn::Uniform
34    } else if !has_measured {
35        // If there are no measured values, then we can look at unmeasured
36        // weights.
37        BandwidthFn::IncludeUnmeasured
38    } else if has_nonzero_measured {
39        // Otherwise, there are measured values; we should look at those only, if
40        // any of them is nonzero.
41        BandwidthFn::MeasuredOnly
42    } else {
43        // This is a bit of an ugly case: We have measured values, but they're
44        // all zero.  If this happens, the bandwidth authorities exist but they
45        // very confused: we should fall back to uniform weighting.
46        BandwidthFn::Uniform
47    }
48}
49
50/// Internal: how should we find the base bandwidth of each relay?  This
51/// value is global over a whole directory, and depends on the bandwidth
52/// weights in the consensus.
53#[derive(Copy, Clone, Debug, PartialEq, Eq)]
54enum BandwidthFn {
55    /// There are no weights at all in the consensus: weight every
56    /// relay as 1.
57    Uniform,
58    /// There are no measured weights in the consensus: count
59    /// unmeasured weights as the weights for relays.
60    IncludeUnmeasured,
61    /// There are measured relays in the consensus; only use those.
62    MeasuredOnly,
63}
64
65impl BandwidthFn {
66    /// Apply this function to the measured or unmeasured bandwidth
67    /// of a single relay.
68    fn apply(&self, w: &netstatus::RelayWeight) -> u32 {
69        use netstatus::RelayWeight::*;
70        use BandwidthFn::*;
71        match (self, w) {
72            (Uniform, _) => 1,
73            (IncludeUnmeasured, Unmeasured(u)) => *u,
74            (IncludeUnmeasured, Measured(m)) => *m,
75            (MeasuredOnly, Unmeasured(_)) => 0,
76            (MeasuredOnly, Measured(m)) => *m,
77            (_, _) => 0,
78        }
79    }
80}
81
82/// Possible ways to weight relays when selecting them a random.
83///
84/// Relays are weighted by a function of their bandwidth that
85/// depends on how scarce that "kind" of bandwidth is.  For
86/// example, if Exit bandwidth is rare, then Exits should be
87/// less likely to get chosen for the middle hop of a path.
88#[derive(Clone, Debug, Copy)]
89#[non_exhaustive]
90pub enum WeightRole {
91    /// Selecting a relay to use as a guard
92    Guard,
93    /// Selecting a relay to use as a middle relay in a circuit.
94    Middle,
95    /// Selecting a relay to use to deliver traffic to the internet.
96    Exit,
97    /// Selecting a relay for a one-hop BEGIN_DIR directory request.
98    BeginDir,
99    /// Selecting a relay with no additional weight beyond its bandwidth.
100    Unweighted,
101    /// Selecting a relay for use as a hidden service introduction point
102    HsIntro,
103    // Note: There is no `HsRend` role, since in practice when we want to pick a
104    // rendezvous point we use a pre-built circuit from our circuit-pool, the
105    // last hop of which was selected with the `Middle` weight.  Fortunately,
106    // the weighting rules for picking rendezvous points are the same as for
107    // picking middle relays.
108}
109
110/// Description for how to weight a single kind of relay for each WeightRole.
111#[derive(Clone, Debug, Copy)]
112struct RelayWeight {
113    /// How to weight this kind of relay when picking a guard relay.
114    as_guard: u32,
115    /// How to weight this kind of relay when picking a middle relay.
116    as_middle: u32,
117    /// How to weight this kind of relay when picking a exit relay.
118    as_exit: u32,
119    /// How to weight this kind of relay when picking a one-hop BEGIN_DIR.
120    as_dir: u32,
121}
122
123impl std::ops::Mul<u32> for RelayWeight {
124    type Output = Self;
125    fn mul(self, rhs: u32) -> Self {
126        RelayWeight {
127            as_guard: self.as_guard * rhs,
128            as_middle: self.as_middle * rhs,
129            as_exit: self.as_exit * rhs,
130            as_dir: self.as_dir * rhs,
131        }
132    }
133}
134impl std::ops::Div<u32> for RelayWeight {
135    type Output = Self;
136    fn div(self, rhs: u32) -> Self {
137        RelayWeight {
138            as_guard: self.as_guard / rhs,
139            as_middle: self.as_middle / rhs,
140            as_exit: self.as_exit / rhs,
141            as_dir: self.as_dir / rhs,
142        }
143    }
144}
145
146impl RelayWeight {
147    /// Return the largest weight that we give for this kind of relay.
148    // The unwrap() is safe because array is nonempty.
149    #[allow(clippy::unwrap_used)]
150    fn max_weight(&self) -> u32 {
151        [self.as_guard, self.as_middle, self.as_exit, self.as_dir]
152            .iter()
153            .max()
154            .copied()
155            .unwrap()
156    }
157    /// Return the weight we should give this kind of relay's
158    /// bandwidth for a given role.
159    fn for_role(&self, role: WeightRole) -> u32 {
160        match role {
161            WeightRole::Guard => self.as_guard,
162            WeightRole::Middle => self.as_middle,
163            WeightRole::Exit => self.as_exit,
164            WeightRole::BeginDir => self.as_dir,
165            WeightRole::HsIntro => self.as_middle, // TODO SPEC is this right?
166            WeightRole::Unweighted => 1,
167        }
168    }
169}
170
171bitflags! {
172    /// A kind of relay, for the purposes of selecting a relay by weight.
173    ///
174    /// Relays can have or lack the Guard flag, the Exit flag, and the
175    /// V2Dir flag. All together, this makes 8 kinds of relays.
176    #[derive(Clone, Copy, Debug, Eq, PartialEq)]
177    struct WeightKind: u8 {
178        /// Flag in weightkind for Guard relays.
179        const GUARD = 1 << 0;
180        /// Flag in weightkind for Exit relays.
181        const EXIT = 1 << 1;
182        /// Flag in weightkind for V2Dir relays.
183        const DIR = 1 << 2;
184    }
185}
186
187impl WeightKind {
188    /// Return the appropriate WeightKind for a relay.
189    fn for_rs(rs: &MdConsensusRouterStatus) -> Self {
190        let mut r = WeightKind::empty();
191        if rs.is_flagged_guard() {
192            r |= WeightKind::GUARD;
193        }
194        if rs.is_flagged_exit() {
195            r |= WeightKind::EXIT;
196        }
197        if rs.is_flagged_v2dir() {
198            r |= WeightKind::DIR;
199        }
200        r
201    }
202    /// Return the index to use for this kind of a relay within a WeightSet.
203    fn idx(self) -> usize {
204        self.bits() as usize
205    }
206}
207
208/// Information derived from a consensus to use when picking relays by
209/// weighted bandwidth.
210#[derive(Debug, Clone)]
211pub(crate) struct WeightSet {
212    /// How to find the bandwidth to use when picking a relay by weighted
213    /// bandwidth.
214    ///
215    /// (This tells us us whether to count unmeasured relays, whether
216    /// to look at bandwidths at all, etc.)
217    bandwidth_fn: BandwidthFn,
218    /// Number of bits that we need to right-shift our weighted products
219    /// so that their sum won't overflow u64::MAX.
220    //
221    // TODO: Perhaps we should use f64 to hold our weights instead,
222    // so we don't need to keep this ad-hoc fixed-point implementation?
223    // If we did so, we won't have to worry about overflows.
224    // (When we call choose_multiple_weighted, it already converts into
225    // f64 internally.  (Though choose_weighted doesn't.))
226    // Before making this change, however,
227    // we should think a little about performance and precision.
228    shift: u8,
229    /// A set of RelayWeight values, indexed by [`WeightKind::idx`], used
230    /// to weight different kinds of relays.
231    w: [RelayWeight; 8],
232}
233
234impl WeightSet {
235    /// Find the actual 64-bit weight to use for a given routerstatus when
236    /// considering it for a given role.
237    ///
238    /// NOTE: This function _does not_ consider whether the relay in question
239    /// actually matches the given role.  For example, if `role` is Guard
240    /// we don't check whether or not `rs` actually has the Guard flag.
241    pub(crate) fn weight_rs_for_role(&self, rs: &MdConsensusRouterStatus, role: WeightRole) -> u64 {
242        self.weight_bw_for_role(WeightKind::for_rs(rs), rs.weight(), role)
243    }
244
245    /// Find the 64-bit weight to report for a relay of `kind` whose weight in
246    /// the consensus is `relay_weight` when using it for `role`.
247    fn weight_bw_for_role(
248        &self,
249        kind: WeightKind,
250        relay_weight: &netstatus::RelayWeight,
251        role: WeightRole,
252    ) -> u64 {
253        let ws = &self.w[kind.idx()];
254
255        let router_bw = self.bandwidth_fn.apply(relay_weight);
256        // Note a subtlety here: we multiply the two values _before_
257        // we shift, to improve accuracy.  We know that this will be
258        // safe, since the inputs are both u32, and so cannot overflow
259        // a u64.
260        let router_weight = u64::from(router_bw) * u64::from(ws.for_role(role));
261        router_weight >> self.shift
262    }
263
264    /// Compute the correct WeightSet for a provided MdConsensus.
265    pub(crate) fn from_consensus(consensus: &MdConsensus, params: &NetParameters) -> Self {
266        let bandwidth_fn = pick_bandwidth_fn(consensus.c_relays().iter().map(|rs| rs.weight()));
267        let weight_scale = params.bw_weight_scale.into();
268
269        let total_bw = consensus
270            .c_relays()
271            .iter()
272            .map(|rs| u64::from(bandwidth_fn.apply(rs.weight())))
273            .sum();
274        let p = consensus.bandwidth_weights();
275
276        Self::from_parts(bandwidth_fn, total_bw, weight_scale, p).validate(consensus)
277    }
278
279    /// Compute the correct WeightSet given a bandwidth function, a
280    /// weight-scaling parameter, a total amount of bandwidth for all
281    /// relays in the consensus, and a set of bandwidth parameters.
282    fn from_parts(
283        bandwidth_fn: BandwidthFn,
284        total_bw: u64,
285        weight_scale: u32,
286        p: &NetParams<i32>,
287    ) -> Self {
288        /// Find a single RelayWeight, given the names that its bandwidth
289        /// parameters have. The `g` parameter is the weight as a guard, the
290        /// `m` parameter is the weight as a middle relay, the `e` parameter is
291        /// the weight as an exit, and the `d` parameter is the weight as a
292        /// directory.
293        #[allow(clippy::many_single_char_names)]
294        fn single(p: &NetParams<i32>, g: &str, m: &str, e: &str, d: &str) -> RelayWeight {
295            RelayWeight {
296                as_guard: w_param(p, g),
297                as_middle: w_param(p, m),
298                as_exit: w_param(p, e),
299                as_dir: w_param(p, d),
300            }
301        }
302
303        // Prevent division by zero in case we're called with a bogus
304        // input.  (That shouldn't be possible.)
305        let weight_scale = weight_scale.max(1);
306
307        // For non-V2Dir relays, we have names for most of their weights.
308        //
309        // (There is no Wge, since we only use Guard relays as guards.  By the
310        // same logic, Wme has no reason to exist, but according to the spec it
311        // does.)
312        let w_none = single(p, "Wgm", "Wmm", "Wem", "Wbm");
313        let w_guard = single(p, "Wgg", "Wmg", "Weg", "Wbg");
314        let w_exit = single(p, "---", "Wme", "Wee", "Wbe");
315        let w_both = single(p, "Wgd", "Wmd", "Wed", "Wbd");
316
317        // Note that the positions of the elements in this array need to
318        // match the values returned by WeightKind.as_idx().
319        let w = [
320            w_none,
321            w_guard,
322            w_exit,
323            w_both,
324            // The V2Dir values are the same as the non-V2Dir values, except
325            // each is multiplied by an additional factor.
326            //
327            // (We don't need to check for overflow here, since the
328            // authorities make sure that the inputs don't get too big.)
329            (w_none * w_param(p, "Wmb")) / weight_scale,
330            (w_guard * w_param(p, "Wgb")) / weight_scale,
331            (w_exit * w_param(p, "Web")) / weight_scale,
332            (w_both * w_param(p, "Wdb")) / weight_scale,
333        ];
334
335        // This is the largest weight value.
336        // The unwrap() is safe because `w` is nonempty.
337        #[allow(clippy::unwrap_used)]
338        let w_max = w.iter().map(RelayWeight::max_weight).max().unwrap();
339
340        // We want "shift" such that (total * w_max) >> shift <= u64::max
341        let shift = calculate_shift(total_bw, u64::from(w_max)) as u8;
342
343        WeightSet {
344            bandwidth_fn,
345            shift,
346            w,
347        }
348    }
349
350    /// Assert that we have correctly computed our shift values so that
351    /// our total weighted bws do not exceed u64::MAX.
352    fn validate(self, consensus: &MdConsensus) -> Self {
353        use WeightRole::*;
354        for role in [Guard, Middle, Exit, BeginDir, Unweighted] {
355            let _: u64 = consensus
356                .c_relays()
357                .iter()
358                .map(|rs| self.weight_rs_for_role(rs, role))
359                .fold(0_u64, |a, b| {
360                    a.checked_add(b)
361                        .expect("Incorrect relay weight calculation: total exceeded u64::MAX!")
362                });
363        }
364        self
365    }
366}
367
368/// The value to return if a weight parameter is absent.
369///
370/// (If there are no weights at all, then it's correct to set them all to 1,
371/// and just use the bandwidths.  If _some_ are present and some are absent,
372/// then the spec doesn't say what to do, but this behavior appears
373/// reasonable.)
374const DFLT_WEIGHT: i32 = 1;
375
376/// Return the weight param named 'kwd' in p.
377///
378/// Returns DFLT_WEIGHT if there is no such parameter, and 0
379/// if `kwd` is "---".
380fn w_param(p: &NetParams<i32>, kwd: &str) -> u32 {
381    if kwd == "---" {
382        0
383    } else {
384        clamp_to_pos(*p.get(kwd).unwrap_or(&DFLT_WEIGHT))
385    }
386}
387
388/// If `inp` is less than 0, return 0.  Otherwise return `inp` as a u32.
389fn clamp_to_pos(inp: i32) -> u32 {
390    // (The spec says that we might encounter negative values here, though
391    // we never actually generate them, and don't plan to generate them.)
392    if inp < 0 {
393        0
394    } else {
395        inp as u32
396    }
397}
398
399/// Compute a 'shift' value such that `(a * b) >> shift` will be contained
400/// inside 64 bits.
401fn calculate_shift(a: u64, b: u64) -> u32 {
402    let bits_for_product = log2_upper(a) + log2_upper(b);
403    if bits_for_product < 64 {
404        0
405    } else {
406        bits_for_product - 64
407    }
408}
409
410/// Return an upper bound for the log2 of n.
411///
412/// This function overestimates whenever n is a power of two, but that doesn't
413/// much matter for the uses we're giving it here.
414fn log2_upper(n: u64) -> u32 {
415    64 - n.leading_zeros()
416}
417
418#[cfg(test)]
419mod test {
420    // @@ begin test lint list maintained by maint/add_warning @@
421    #![allow(clippy::bool_assert_comparison)]
422    #![allow(clippy::clone_on_copy)]
423    #![allow(clippy::dbg_macro)]
424    #![allow(clippy::mixed_attributes_style)]
425    #![allow(clippy::print_stderr)]
426    #![allow(clippy::print_stdout)]
427    #![allow(clippy::single_char_pattern)]
428    #![allow(clippy::unwrap_used)]
429    #![allow(clippy::unchecked_duration_subtraction)]
430    #![allow(clippy::useless_vec)]
431    #![allow(clippy::needless_pass_by_value)]
432    //! <!-- @@ end test lint list maintained by maint/add_warning @@ -->
433    use super::*;
434    use netstatus::RelayWeight as RW;
435    use std::net::SocketAddr;
436    use std::time::{Duration, SystemTime};
437    use tor_basic_utils::test_rng::testing_rng;
438    use tor_netdoc::doc::netstatus::{Lifetime, RelayFlags, RouterStatusBuilder};
439
440    #[test]
441    fn t_clamp() {
442        assert_eq!(clamp_to_pos(32), 32);
443        assert_eq!(clamp_to_pos(i32::MAX), i32::MAX as u32);
444        assert_eq!(clamp_to_pos(0), 0);
445        assert_eq!(clamp_to_pos(-1), 0);
446        assert_eq!(clamp_to_pos(i32::MIN), 0);
447    }
448
449    #[test]
450    fn t_log2() {
451        assert_eq!(log2_upper(u64::MAX), 64);
452        assert_eq!(log2_upper(0), 0);
453        assert_eq!(log2_upper(1), 1);
454        assert_eq!(log2_upper(63), 6);
455        assert_eq!(log2_upper(64), 7); // a little buggy but harmless.
456    }
457
458    #[test]
459    fn t_calc_shift() {
460        assert_eq!(calculate_shift(1 << 20, 1 << 20), 0);
461        assert_eq!(calculate_shift(1 << 50, 1 << 10), 0);
462        assert_eq!(calculate_shift(1 << 32, 1 << 33), 3);
463        assert!(((1_u64 << 32) >> 3).checked_mul(1_u64 << 33).is_some());
464        assert_eq!(calculate_shift(432 << 40, 7777 << 40), 38);
465        assert!(((432_u64 << 40) >> 38)
466            .checked_mul(7777_u64 << 40)
467            .is_some());
468    }
469
470    #[test]
471    fn t_pick_bwfunc() {
472        let empty = [];
473        assert_eq!(pick_bandwidth_fn(empty.iter()), BandwidthFn::Uniform);
474
475        let all_zero = [RW::Unmeasured(0), RW::Measured(0), RW::Unmeasured(0)];
476        assert_eq!(pick_bandwidth_fn(all_zero.iter()), BandwidthFn::Uniform);
477
478        let all_unmeasured = [RW::Unmeasured(9), RW::Unmeasured(2222)];
479        assert_eq!(
480            pick_bandwidth_fn(all_unmeasured.iter()),
481            BandwidthFn::IncludeUnmeasured
482        );
483
484        let some_measured = [
485            RW::Unmeasured(10),
486            RW::Measured(7),
487            RW::Measured(4),
488            RW::Unmeasured(0),
489        ];
490        assert_eq!(
491            pick_bandwidth_fn(some_measured.iter()),
492            BandwidthFn::MeasuredOnly
493        );
494
495        // This corresponds to an open question in
496        // `pick_bandwidth_fn`, about what to do when the only nonzero
497        // weights are unmeasured.
498        let measured_all_zero = [RW::Unmeasured(10), RW::Measured(0)];
499        assert_eq!(
500            pick_bandwidth_fn(measured_all_zero.iter()),
501            BandwidthFn::Uniform
502        );
503    }
504
505    #[test]
506    fn t_apply_bwfn() {
507        use netstatus::RelayWeight::*;
508        use BandwidthFn::*;
509
510        assert_eq!(Uniform.apply(&Measured(7)), 1);
511        assert_eq!(Uniform.apply(&Unmeasured(0)), 1);
512
513        assert_eq!(IncludeUnmeasured.apply(&Measured(7)), 7);
514        assert_eq!(IncludeUnmeasured.apply(&Unmeasured(8)), 8);
515
516        assert_eq!(MeasuredOnly.apply(&Measured(9)), 9);
517        assert_eq!(MeasuredOnly.apply(&Unmeasured(10)), 0);
518    }
519
520    // From a fairly recent Tor consensus.
521    const TESTVEC_PARAMS: &str =
522        "Wbd=0 Wbe=0 Wbg=4096 Wbm=10000 Wdb=10000 Web=10000 Wed=10000 Wee=10000 Weg=10000 Wem=10000 Wgb=10000 Wgd=0 Wgg=5904 Wgm=5904 Wmb=10000 Wmd=0 Wme=0 Wmg=4096 Wmm=10000";
523
524    #[test]
525    fn t_weightset_basic() {
526        let total_bandwidth = 1_000_000_000;
527        let params = TESTVEC_PARAMS.parse().unwrap();
528        let ws = WeightSet::from_parts(BandwidthFn::MeasuredOnly, total_bandwidth, 10000, &params);
529
530        assert_eq!(ws.bandwidth_fn, BandwidthFn::MeasuredOnly);
531        assert_eq!(ws.shift, 0);
532
533        assert_eq!(ws.w[0].as_guard, 5904);
534        assert_eq!(ws.w[(WeightKind::GUARD.bits()) as usize].as_guard, 5904);
535        assert_eq!(ws.w[(WeightKind::EXIT.bits()) as usize].as_exit, 10000);
536        assert_eq!(
537            ws.w[(WeightKind::EXIT | WeightKind::GUARD).bits() as usize].as_dir,
538            0
539        );
540        assert_eq!(
541            ws.w[(WeightKind::GUARD | WeightKind::DIR).bits() as usize].as_dir,
542            4096
543        );
544        assert_eq!(
545            ws.w[(WeightKind::GUARD | WeightKind::DIR).bits() as usize].as_dir,
546            4096
547        );
548
549        assert_eq!(
550            ws.weight_bw_for_role(
551                WeightKind::GUARD | WeightKind::DIR,
552                &RW::Unmeasured(7777),
553                WeightRole::Guard
554            ),
555            0
556        );
557
558        assert_eq!(
559            ws.weight_bw_for_role(
560                WeightKind::GUARD | WeightKind::DIR,
561                &RW::Measured(7777),
562                WeightRole::Guard
563            ),
564            7777 * 5904
565        );
566
567        assert_eq!(
568            ws.weight_bw_for_role(
569                WeightKind::GUARD | WeightKind::DIR,
570                &RW::Measured(7777),
571                WeightRole::Middle
572            ),
573            7777 * 4096
574        );
575
576        assert_eq!(
577            ws.weight_bw_for_role(
578                WeightKind::GUARD | WeightKind::DIR,
579                &RW::Measured(7777),
580                WeightRole::Exit
581            ),
582            7777 * 10000
583        );
584
585        assert_eq!(
586            ws.weight_bw_for_role(
587                WeightKind::GUARD | WeightKind::DIR,
588                &RW::Measured(7777),
589                WeightRole::BeginDir
590            ),
591            7777 * 4096
592        );
593
594        assert_eq!(
595            ws.weight_bw_for_role(
596                WeightKind::GUARD | WeightKind::DIR,
597                &RW::Measured(7777),
598                WeightRole::Unweighted
599            ),
600            7777
601        );
602
603        // Now try those last few with routerstatuses.
604        let rs = rs_builder()
605            .set_flags(RelayFlags::GUARD | RelayFlags::V2DIR)
606            .weight(RW::Measured(7777))
607            .build()
608            .unwrap();
609        assert_eq!(ws.weight_rs_for_role(&rs, WeightRole::Exit), 7777 * 10000);
610        assert_eq!(
611            ws.weight_rs_for_role(&rs, WeightRole::BeginDir),
612            7777 * 4096
613        );
614        assert_eq!(ws.weight_rs_for_role(&rs, WeightRole::Unweighted), 7777);
615    }
616
617    /// Return a routerstatus builder set up to deliver a routerstatus
618    /// with most features disabled.
619    fn rs_builder() -> RouterStatusBuilder<[u8; 32]> {
620        MdConsensus::builder()
621            .rs()
622            .identity([9; 20].into())
623            .add_or_port(SocketAddr::from(([127, 0, 0, 1], 9001)))
624            .doc_digest([9; 32])
625            .protos("".parse().unwrap())
626            .clone()
627    }
628
629    #[test]
630    fn weight_flags() {
631        let rs1 = rs_builder().set_flags(RelayFlags::EXIT).build().unwrap();
632        assert_eq!(WeightKind::for_rs(&rs1), WeightKind::EXIT);
633
634        let rs1 = rs_builder().set_flags(RelayFlags::GUARD).build().unwrap();
635        assert_eq!(WeightKind::for_rs(&rs1), WeightKind::GUARD);
636
637        let rs1 = rs_builder().set_flags(RelayFlags::V2DIR).build().unwrap();
638        assert_eq!(WeightKind::for_rs(&rs1), WeightKind::DIR);
639
640        let rs1 = rs_builder().build().unwrap();
641        assert_eq!(WeightKind::for_rs(&rs1), WeightKind::empty());
642
643        let rs1 = rs_builder().set_flags(RelayFlags::all()).build().unwrap();
644        assert_eq!(
645            WeightKind::for_rs(&rs1),
646            WeightKind::EXIT | WeightKind::GUARD | WeightKind::DIR
647        );
648    }
649
650    #[test]
651    fn weightset_from_consensus() {
652        use rand::Rng;
653        let now = SystemTime::now();
654        let one_hour = Duration::new(3600, 0);
655        let mut rng = testing_rng();
656        let mut bld = MdConsensus::builder();
657        bld.consensus_method(34)
658            .lifetime(Lifetime::new(now, now + one_hour, now + 2 * one_hour).unwrap())
659            .weights(TESTVEC_PARAMS.parse().unwrap());
660
661        // We're going to add a huge amount of unmeasured bandwidth,
662        // and a reasonable amount of  measured bandwidth.
663        for _ in 0..10 {
664            rs_builder()
665                .identity(rng.random::<[u8; 20]>().into()) // random id
666                .weight(RW::Unmeasured(1_000_000))
667                .set_flags(RelayFlags::GUARD | RelayFlags::EXIT)
668                .build_into(&mut bld)
669                .unwrap();
670        }
671        for n in 0..30 {
672            rs_builder()
673                .identity(rng.random::<[u8; 20]>().into()) // random id
674                .weight(RW::Measured(1_000 * n))
675                .set_flags(RelayFlags::GUARD | RelayFlags::EXIT)
676                .build_into(&mut bld)
677                .unwrap();
678        }
679
680        let consensus = bld.testing_consensus().unwrap();
681        let params = NetParameters::default();
682        let ws = WeightSet::from_consensus(&consensus, &params);
683
684        assert_eq!(ws.bandwidth_fn, BandwidthFn::MeasuredOnly);
685        assert_eq!(ws.shift, 0);
686        assert_eq!(ws.w[0].as_guard, 5904);
687        assert_eq!(ws.w[5].as_guard, 5904);
688        assert_eq!(ws.w[5].as_middle, 4096);
689    }
690}