1
//! Logic for selecting relays from a network directory,
2
//! and reporting the outcome of such a selection.
3

            
4
use crate::{LowLevelRelayPredicate, RelayExclusion, RelayRestriction, RelayUsage};
5
use tor_basic_utils::iter::FilterCount;
6
use tor_netdir::{NetDir, Relay, WeightRole};
7

            
8
use std::fmt;
9

            
10
/// Description of the requirements that a relay must implement in order to be selected.
11
///
12
/// This object is used to pick a [`Relay`] from a [`NetDir`], or to ensure that a
13
/// previously selected `Relay` still meets its requirements.
14
///
15
/// The requirements on a relay can be _strict_ or _flexible_.
16
/// If any restriction is flexible, and relay selection fails at first,
17
/// we _relax_ the `RelaySelector` by removing that restriction,
18
/// and trying again,
19
/// before we give up completely.
20
#[derive(Clone, Debug)]
21
pub struct RelaySelector<'a> {
22
    /// A usage that the relay must support.
23
    ///
24
    /// Invariant: This is a RelayUsage.
25
    usage: Restr<'a>,
26

            
27
    /// An excludion that the relay must obey.
28
    ///
29
    /// Invariant: This a RelayExclusion.
30
    exclusion: Restr<'a>,
31

            
32
    /// Other restrictions that a Relay must obey in order to be selected.
33
    other_restrictions: Vec<Restr<'a>>,
34
}
35

            
36
/// A single restriction, along with a flag about whether it's strict.
37
#[derive(Clone, Debug)]
38
struct Restr<'a> {
39
    /// The underlying restriction.
40
    restriction: RelayRestriction<'a>,
41
    /// Is the restriction strict or flexible?
42
    strict: bool,
43
}
44

            
45
impl<'a> Restr<'a> {
46
    /// Try relaxing this restriction.
47
    ///
48
    /// (If this can't be relaxed, just return a copy of it.)
49
1108
    fn maybe_relax(&self) -> Self {
50
1108
        if self.strict {
51
554
            self.clone()
52
        } else {
53
554
            Self {
54
554
                restriction: self.restriction.relax(),
55
554
                // The new restriction is always strict, since we don't want to
56
554
                // relax it any further.
57
554
                strict: true,
58
554
            }
59
        }
60
1108
    }
61
}
62

            
63
/// Information about how a given selection was generated.
64
///
65
/// Records the specifics of how many relays were excluded by each
66
/// requirement,
67
/// whether we had to relax the selector, and so on.
68
///
69
/// The caller should typically decide whether an error or warning is necessary,
70
/// and if so use this to generate a formattable report about what went wrong.
71
#[derive(Debug, Clone)]
72
pub struct SelectionInfo<'a> {
73
    /// Outcome of our first attempt to pick a relay.
74
    first_try: FilterCounts,
75

            
76
    /// Present if we tried again with a relaxed version of our
77
    /// flexible members.
78
    relaxed_try: Option<FilterCounts>,
79

            
80
    /// True if we eventually succeeded in picking a relay.
81
    succeeded: bool,
82

            
83
    /// The `RelaySelector` that was used.
84
    ///
85
    /// Used to produce information about which restriction is which.
86
    in_selection: &'a RelaySelector<'a>,
87
}
88

            
89
impl<'a> SelectionInfo<'a> {
90
    /// Return true if we eventually picked at least one relay.
91
    ///
92
    /// (We report success on `pick_n_relays` if we returned a nonzero
93
    /// number of relays, even if it is smaller than the requested number.)
94
200
    pub fn success(&self) -> bool {
95
200
        self.succeeded
96
200
    }
97

            
98
    /// Return true if picked at least one relay,
99
    /// but only after relaxing our initial selector.
100
200
    pub fn result_is_relaxed_success(&self) -> bool {
101
200
        self.relaxed_try.is_some() && self.succeeded
102
200
    }
103
}
104

            
105
impl<'a> fmt::Display for SelectionInfo<'a> {
106
834
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
107
834
        match (self.succeeded, &self.relaxed_try) {
108
2
            (true, None) => write!(f, "Success: {}", FcDisp(&self.first_try, self.in_selection))?,
109
830
            (false, None) => write!(f, "Failed: {}", FcDisp(&self.first_try, self.in_selection))?,
110
2
            (true, Some(retry)) => write!(
111
2
                f,
112
                "Failed at first, then succeeded. At first, {}. After relaxing requirements, {}",
113
2
                FcDisp(&self.first_try, self.in_selection),
114
2
                FcDisp(retry, self.in_selection)
115
            )?,
116
            (false, Some(retry)) => write!(
117
                f,
118
                "Failed even after relaxing requirement. At first, {}. After relaxing requirements, {}",
119
                FcDisp(&self.first_try, self.in_selection),
120
                FcDisp(retry, self.in_selection)
121
            )?,
122
        };
123
834
        Ok(())
124
834
    }
125
}
126

            
127
/// A list of [`FilterCount`], associated with a [`RelaySelector`].
128
#[derive(Debug, Clone)]
129
struct FilterCounts {
130
    /// The [`FilterCount`] created by each restriction.
131
    ///
132
    /// This `Vec` has the same length as the list of restrictions; its items
133
    /// refer to them one by one.
134
    ///
135
    /// Because restrictions are applied as a set of filters, each successive
136
    /// count will only include the relays not excluded by the previous filters.
137
    counts: Vec<FilterCount>,
138
}
139

            
140
impl FilterCounts {
141
    /// Create a new empty `FilterCounts`.
142
1169440
    fn new(selector: &RelaySelector) -> Self {
143
1169440
        let counts = vec![FilterCount::default(); selector.n_restrictions()];
144
1169440
        FilterCounts { counts }
145
1169440
    }
146
}
147

            
148
/// Helper to display filter counts
149
struct FcDisp<'a>(&'a FilterCounts, &'a RelaySelector<'a>);
150
impl<'a> fmt::Display for FcDisp<'a> {
151
836
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
152
836
        let counts = &self.0.counts;
153
836
        let restrictions = self.1.all_restrictions();
154
836
        write!(f, "rejected ")?;
155
836
        let mut first = true;
156
836
        let mut found_any_rejected = false;
157
1856
        for (c, r) in counts.iter().zip(restrictions) {
158
1856
            if c.n_rejected == 0 {
159
1016
                continue;
160
840
            }
161
840
            if let Some(desc) = r.restriction.rejection_description() {
162
840
                if first {
163
836
                    first = false;
164
836
                } else {
165
4
                    write!(f, "; ")?;
166
                }
167
840
                write!(f, "{} as {}", c.display_frac_rejected(), desc)?;
168
840
                found_any_rejected = true;
169
            } else {
170
                debug_assert_eq!(c.n_rejected, 0);
171
            }
172
        }
173
836
        if !found_any_rejected {
174
            write!(f, "none")?;
175
836
        }
176
836
        Ok(())
177
836
    }
178
}
179

            
180
impl<'a> RelaySelector<'a> {
181
    /// Create a new RelaySelector to pick relays with a given
182
    /// [`RelayUsage`] and [`RelayExclusion`].
183
    ///
184
    /// Both arguments are required, since every caller should consider them explicitly.
185
    ///
186
    /// The provided usage and exclusion are strict by default.
187
    ///
188
    // TODO: Possibly have this take a struct with named pieces instead, when we
189
    // get a third thing that we want everybody to think about.
190
1171356
    pub fn new(usage: RelayUsage, exclusion: RelayExclusion<'a>) -> Self {
191
1171356
        Self {
192
1171356
            usage: Restr {
193
1171356
                restriction: RelayRestriction::for_usage(usage),
194
1171356
                strict: true,
195
1171356
            },
196
1171356
            exclusion: Restr {
197
1171356
                restriction: exclusion.into(),
198
1171356
                strict: true,
199
1171356
            },
200
1171356
            other_restrictions: vec![],
201
1171356
        }
202
1171356
    }
203

            
204
    /// Mark the originally provided `RelayUsage` as flexible.
205
828
    pub fn mark_usage_flexible(&mut self) {
206
828
        self.usage.strict = false;
207
828
    }
208

            
209
    /// Mark the originally provided `RelayExclusion` as flexible.
210
2
    pub fn mark_exclusion_flexible(&mut self) {
211
2
        self.exclusion.strict = false;
212
2
    }
213

            
214
    /// Add a new _strict_ [`RelayRestriction`] to this selector.
215
3404
    pub fn push_restriction(&mut self, restriction: RelayRestriction<'a>) {
216
3404
        self.push_inner(restriction, true);
217
3404
    }
218

            
219
    /// Add a new _flexible_ [`RelayRestriction`] to this selector.
220
    pub fn push_flexible_restriction(&mut self, restriction: RelayRestriction<'a>) {
221
        self.push_inner(restriction, false);
222
    }
223

            
224
    /// Helper to implement adding a new restriction.
225
3404
    fn push_inner(&mut self, restriction: RelayRestriction<'a>, strict: bool) {
226
3404
        self.other_restrictions.push(Restr {
227
3404
            restriction,
228
3404
            strict,
229
3404
        });
230
3404
    }
231

            
232
    /// Return the usage for this selector.
233
1737308
    pub fn usage(&self) -> &RelayUsage {
234
        // See invariants for explanation of why these `expects` are safe.
235
1737308
        self.usage
236
1737308
            .restriction
237
1737308
            .as_usage()
238
1737308
            .expect("Usage not a usage!?")
239
1737308
    }
240

            
241
    /// Return the [`WeightRole`] to use when randomly picking relays according
242
    /// to this selector.
243
1169438
    fn weight_role(&self) -> WeightRole {
244
1169438
        self.usage().selection_weight_role()
245
1169438
    }
246

            
247
    /// Return true if `relay` is one that this selector would pick.
248
    pub fn permits_relay(&self, relay: &tor_netdir::Relay<'_>) -> bool {
249
        self.low_level_predicate_permits_relay(relay)
250
    }
251

            
252
    /// Return an iterator that yields each restriction from this selector,
253
    /// including the usage and exclusion.
254
45896128
    fn all_restrictions(&self) -> impl Iterator<Item = &Restr<'a>> {
255
        use std::iter::once;
256
45896128
        once(&self.usage)
257
45896128
            .chain(once(&self.exclusion))
258
45896128
            .chain(self.other_restrictions.iter())
259
45896128
    }
260

            
261
    /// Return the number of restrictions in this selector,
262
    /// including the usage and exclusion.
263
47038816
    fn n_restrictions(&self) -> usize {
264
47038816
        self.other_restrictions.len() + 2
265
47038816
    }
266

            
267
    /// Try to pick a random relay from `netdir`,
268
    /// according to the rules of this selector.
269
49578
    pub fn select_relay<'s, 'd, R: rand::Rng>(
270
49578
        &'s self,
271
49578
        rng: &mut R,
272
49578
        netdir: &'d NetDir,
273
49578
    ) -> (Option<Relay<'d>>, SelectionInfo<'s>) {
274
49578
        with_possible_relaxation(
275
49578
            self,
276
49604
            |selector| {
277
49604
                let role = selector.weight_role();
278
49604
                let mut fc = FilterCounts::new(selector);
279
1967080
                let relay = netdir.pick_relay(rng, role, |r| selector.relay_usable(r, &mut fc));
280
49604
                (relay, fc)
281
49604
            },
282
            Option::is_some,
283
        )
284
49578
    }
285

            
286
    /// Try to pick `n_relays` distinct random relay from `netdir`,
287
    /// according to the rules of this selector.
288
27285
    pub fn select_n_relays<'s, 'd, R: rand::Rng>(
289
27285
        &'s self,
290
27285
        rng: &mut R,
291
27285
        n_relays: usize,
292
27285
        netdir: &'d NetDir,
293
27285
    ) -> (Vec<Relay<'d>>, SelectionInfo<'s>) {
294
27285
        with_possible_relaxation(
295
27285
            self,
296
27285
            |selector| {
297
27285
                let role = selector.weight_role();
298
27285
                let mut fc = FilterCounts::new(selector);
299
27285
                let relays = netdir
300
587306
                    .pick_n_relays(rng, n_relays, role, |r| selector.relay_usable(r, &mut fc));
301
27285
                (relays, fc)
302
27285
            },
303
27285
            |relays| !relays.is_empty(),
304
        )
305
27285
    }
306

            
307
    /// Check whether a given relay `r` obeys the restrictions of this selector,
308
    /// updating `fc` according to which restrictions (if any) accepted or
309
    /// rejected it.
310
    ///
311
    /// Requires that `fc` has the same length as self.restrictions.
312
    ///
313
    /// This differs from `<Self as RelayPredicate>::permits_relay` in taking
314
    /// `fc` as an argument.
315
45869376
    fn relay_usable(&self, r: &Relay<'_>, fc: &mut FilterCounts) -> bool {
316
45869376
        debug_assert_eq!(self.n_restrictions(), fc.counts.len());
317

            
318
45869376
        self.all_restrictions()
319
45869376
            .zip(fc.counts.iter_mut())
320
78410208
            .all(|(restr, restr_count)| {
321
77408992
                restr_count.count(restr.restriction.low_level_predicate_permits_relay(r))
322
77408992
            })
323
45869376
    }
324

            
325
    /// Return true if this selector has any flexible restrictions.
326
15372
    fn can_relax(&self) -> bool {
327
30990
        self.all_restrictions().any(|restr| !restr.strict)
328
15372
    }
329

            
330
    /// Return a new selector created by relaxing every flexible restriction in
331
    /// this selector.
332
554
    fn relax(&self) -> Self {
333
554
        let new_selector = RelaySelector {
334
554
            usage: self.usage.maybe_relax(),
335
554
            exclusion: self.exclusion.maybe_relax(),
336
554
            other_restrictions: self
337
554
                .other_restrictions
338
554
                .iter()
339
554
                .map(Restr::maybe_relax)
340
554
                .collect(),
341
554
        };
342
554
        debug_assert!(!new_selector.can_relax());
343
554
        new_selector
344
554
    }
345
}
346

            
347
impl<'a> LowLevelRelayPredicate for RelaySelector<'a> {
348
10544
    fn low_level_predicate_permits_relay(&self, relay: &tor_netdir::Relay<'_>) -> bool {
349
10544
        self.all_restrictions()
350
29318
            .all(|r| r.restriction.low_level_predicate_permits_relay(relay))
351
10544
    }
352
}
353

            
354
/// Re-run relay selection, relaxing our selector as necessary.
355
///
356
/// This is a helper to implement our relay selection logic.
357
/// We try to run `select` to find one or more random relays
358
/// conforming to `selector`.
359
/// If `ok` says that the result is good (by returning true),
360
/// we return that result.
361
/// Otherwise, we try to _relax_ the selector (if possible),
362
/// and try again.
363
/// If the selector can't be relaxed any further,
364
/// we return the original (not-ok) result.
365
//
366
// TODO: Later, we might want to relax our restrictions one by one,
367
// rather than all at once.
368
76863
fn with_possible_relaxation<'a, SEL, OK, T>(
369
76863
    selector: &'a RelaySelector,
370
76863
    mut select: SEL,
371
76863
    ok: OK,
372
76863
) -> (T, SelectionInfo<'a>)
373
76863
where
374
76863
    SEL: FnMut(&RelaySelector) -> (T, FilterCounts),
375
76863
    OK: Fn(&T) -> bool,
376
{
377
76863
    let (outcome, count_strict) = select(selector);
378
76863
    let succeeded = ok(&outcome);
379
76863
    if succeeded || !selector.can_relax() {
380
76837
        let info = SelectionInfo {
381
76837
            first_try: count_strict,
382
76837
            relaxed_try: None,
383
76837
            succeeded,
384
76837
            in_selection: selector,
385
76837
        };
386
76837
        return (outcome, info);
387
26
    }
388
26
    let relaxed_selector = selector.relax();
389
26
    let (relaxed_outcome, count_relaxed) = select(&relaxed_selector);
390
26
    let info = SelectionInfo {
391
26
        first_try: count_strict,
392
26
        relaxed_try: Some(count_relaxed),
393
26
        succeeded: ok(&relaxed_outcome),
394
26
        in_selection: selector,
395
26
    };
396
26
    (relaxed_outcome, info)
397
76863
}
398

            
399
#[cfg(test)]
400
mod test {
401
    // @@ begin test lint list maintained by maint/add_warning @@
402
    #![allow(clippy::bool_assert_comparison)]
403
    #![allow(clippy::clone_on_copy)]
404
    #![allow(clippy::dbg_macro)]
405
    #![allow(clippy::mixed_attributes_style)]
406
    #![allow(clippy::print_stderr)]
407
    #![allow(clippy::print_stdout)]
408
    #![allow(clippy::single_char_pattern)]
409
    #![allow(clippy::unwrap_used)]
410
    #![allow(clippy::unchecked_time_subtraction)]
411
    #![allow(clippy::useless_vec)]
412
    #![allow(clippy::needless_pass_by_value)]
413
    #![allow(clippy::string_slice)] // See arti#2571
414
    //! <!-- @@ end test lint list maintained by maint/add_warning @@ -->
415

            
416
    use std::collections::HashSet;
417

            
418
    use tor_basic_utils::test_rng::testing_rng;
419
    use tor_linkspec::{HasRelayIds, RelayId};
420
    use tor_netdir::{FamilyRules, Relay, SubnetConfig};
421

            
422
    use super::*;
423
    use crate::{
424
        RelaySelectionConfig, TargetPort,
425
        testing::{cfg, split_netdir, testnet},
426
    };
427

            
428
    #[test]
429
    fn selector_as_predicate() {
430
        let nd = testnet();
431
        let id_4 = "$0404040404040404040404040404040404040404".parse().unwrap();
432
        let usage = RelayUsage::middle_relay(None);
433
        let exclusion = RelayExclusion::exclude_identities([id_4].into_iter().collect());
434
        let sel = RelaySelector::new(usage.clone(), exclusion.clone());
435

            
436
        let (yes, no) = split_netdir(&nd, &sel);
437
        let p = |r: &Relay<'_>| {
438
            usage.low_level_predicate_permits_relay(r)
439
                && exclusion.low_level_predicate_permits_relay(r)
440
        };
441
        assert!(yes.iter().all(p));
442
        assert!(no.iter().all(|r| !p(r)));
443
    }
444

            
445
    #[test]
446
    fn selector_as_filter() {
447
        let nd = testnet();
448
        let id_4 = "$0404040404040404040404040404040404040404".parse().unwrap();
449
        let usage = RelayUsage::middle_relay(None);
450
        let exclusion = RelayExclusion::exclude_identities([id_4].into_iter().collect());
451
        let sel = RelaySelector::new(usage.clone(), exclusion.clone());
452
        let mut fc = FilterCounts::new(&sel);
453

            
454
        let (yes, _no) = split_netdir(&nd, &sel);
455
        let filtered: Vec<_> = nd
456
            .relays()
457
            .filter(|r| sel.relay_usable(r, &mut fc))
458
            .collect();
459
        assert_eq!(yes.len(), filtered.len());
460

            
461
        let k1: HashSet<_> = yes.iter().map(|r| r.rsa_identity().unwrap()).collect();
462
        let k2: HashSet<_> = filtered.iter().map(|r| r.rsa_identity().unwrap()).collect();
463
        assert_eq!(k1, k2);
464

            
465
        // 6 relays are rejected for not being suitable as a general-purpose middle relay
466
        // (no Fast flag or no stable flag)
467
        assert_eq!(fc.counts[0].n_rejected, 12);
468
        // 1 additional relay is rejected for having id_4.
469
        assert_eq!(fc.counts[1].n_rejected, 1);
470
        // The remainder are accepted.
471
        assert_eq!(fc.counts[1].n_accepted, yes.len());
472
    }
473

            
474
    #[test]
475
    fn selector_pick_random() {
476
        let nd = testnet();
477
        let id_4 = "$0404040404040404040404040404040404040404".parse().unwrap();
478
        let usage = RelayUsage::middle_relay(None);
479
        let exclusion = RelayExclusion::exclude_identities([id_4].into_iter().collect());
480
        let sel = RelaySelector::new(usage.clone(), exclusion.clone());
481

            
482
        let (yes, _no) = split_netdir(&nd, &sel);
483
        let k_yes: HashSet<_> = yes.iter().map(|r| r.rsa_identity().unwrap()).collect();
484
        let p = |r: Relay<'_>| k_yes.contains(r.rsa_identity().unwrap());
485

            
486
        let mut rng = testing_rng();
487
        for _ in 0..50 {
488
            // Select one relay; make sure it is ok.
489
            let (r_rand, si) = sel.select_relay(&mut rng, &nd);
490
            assert!(si.success());
491
            assert!(!si.result_is_relaxed_success());
492
            assert!(p(r_rand.unwrap()));
493

            
494
            // Select 20 random relays; make sure they are distinct and ok.
495
            let (rs_rand, si) = sel.select_n_relays(&mut rng, 20, &nd);
496
            assert_eq!(rs_rand.len(), 20);
497
            assert!(si.success());
498
            assert!(!si.result_is_relaxed_success());
499
            assert!(rs_rand.iter().cloned().all(p));
500
            let k_got: HashSet<_> = rs_rand.iter().map(|r| r.rsa_identity().unwrap()).collect();
501
            assert_eq!(k_got.len(), 20);
502
        }
503
    }
504

            
505
    #[test]
506
    fn selector_report() {
507
        let nd = testnet();
508
        let id_4 = "$0404040404040404040404040404040404040404".parse().unwrap();
509
        let usage = RelayUsage::middle_relay(None);
510
        let exclusion = RelayExclusion::exclude_identities([id_4].into_iter().collect());
511
        let sel = RelaySelector::new(usage.clone(), exclusion.clone());
512

            
513
        let mut rng = testing_rng();
514
        let (_, si) = sel.select_relay(&mut rng, &nd);
515
        assert_eq!(
516
            si.to_string(),
517
            "Success: rejected 12/40 as not usable as middle relay; 1/28 as already selected"
518
        );
519

            
520
        // Now try failing.
521
        // (The test network doesn't have ipv6 support.)
522
        let unreachable_port = TargetPort::ipv6(80);
523
        let sel = RelaySelector::new(
524
            RelayUsage::exit_to_all_ports(&cfg(), vec![unreachable_port]),
525
            exclusion.clone(),
526
        );
527
        let (r_none, si) = sel.select_relay(&mut rng, &nd);
528
        assert!(r_none.is_none());
529
        assert_eq!(
530
            si.to_string(),
531
            "Failed: rejected 40/40 as not exiting to desired ports"
532
        );
533
    }
534

            
535
    #[test]
536
    fn relax() {
537
        let all_families = FamilyRules::all_family_info();
538

            
539
        let nd = testnet();
540
        let id_4: RelayId = "$0404040404040404040404040404040404040404".parse().unwrap();
541
        let r4 = nd.by_id(&id_4).unwrap();
542
        let usage = RelayUsage::middle_relay(None);
543
        let very_silly_cfg = RelaySelectionConfig {
544
            long_lived_ports: cfg().long_lived_ports,
545
            // This should exclude everyone.
546
            subnet_config: SubnetConfig::new(1, 1),
547
        };
548
        let exclude_relays = vec![r4];
549
        let exclude_everyone = RelayExclusion::exclude_relays_in_same_family(
550
            &very_silly_cfg,
551
            exclude_relays,
552
            all_families,
553
        );
554

            
555
        let mut sel = RelaySelector::new(usage.clone(), exclude_everyone.clone());
556
        let mut rng = testing_rng();
557
        let (r_none, _) = sel.select_relay(&mut rng, &nd);
558
        assert!(r_none.is_none());
559

            
560
        sel.mark_exclusion_flexible();
561
        let (r_some, si) = sel.select_relay(&mut rng, &nd);
562
        assert!(r_some.is_some());
563
        assert_eq!(
564
            si.to_string(),
565
            "Failed at first, then succeeded. At first, rejected 12/40 as not usable as middle relay; \
566
                                    28/28 as in same family as already selected. \
567
                                    After relaxing requirements, rejected 12/40 as not usable as middle relay"
568
        );
569
    }
570
}