tor_circmgr/
preemptive.rs

1//! Tools for determining what circuits to preemptively build.
2
3use crate::{PathConfig, PreemptiveCircuitConfig, TargetPort, TargetTunnelUsage};
4use std::collections::HashMap;
5use std::sync::Arc;
6use tracing::warn;
7use web_time_compat::{Instant, InstantExt};
8
9/// Predicts what circuits might be used in future based on past activity, and suggests
10/// circuits to preemptively build as a result.
11pub(crate) struct PreemptiveCircuitPredictor {
12    /// A map of every exit port we've observed being used (or `None` if we observed an exit being
13    /// used to resolve DNS names instead of building a stream), to the last time we encountered
14    /// such usage.
15    // TODO(nickm): Let's have a mechanism for cleaning this out from time to time.
16    usages: HashMap<Option<TargetPort>, Instant>,
17
18    /// Configuration for this predictor.
19    config: tor_config::MutCfg<PreemptiveCircuitConfig>,
20}
21
22impl PreemptiveCircuitPredictor {
23    /// Create a new predictor, starting out with a set of ports we think are likely to be used.
24    pub(crate) fn new(config: PreemptiveCircuitConfig) -> Self {
25        let mut usages = HashMap::new();
26        for port in &config.initial_predicted_ports {
27            // TODO(nickm) should this be IPv6? Should we have a way to configure IPv6 initial ports?
28            usages.insert(Some(TargetPort::ipv4(*port)), Instant::get());
29        }
30
31        // We want to build circuits for resolving DNS, too.
32        usages.insert(None, Instant::get());
33
34        Self {
35            usages,
36            config: config.into(),
37        }
38    }
39
40    /// Return the configuration for this PreemptiveCircuitPredictor.
41    pub(crate) fn config(&self) -> Arc<PreemptiveCircuitConfig> {
42        self.config.get()
43    }
44
45    /// Replace the current configuration for this PreemptiveCircuitPredictor
46    /// with `new_config`.
47    pub(crate) fn set_config(&self, mut new_config: PreemptiveCircuitConfig) {
48        self.config.map_and_replace(|cfg| {
49            // Force this to stay the same, since it can't meaningfully be changed.
50            new_config
51                .initial_predicted_ports
52                .clone_from(&cfg.initial_predicted_ports);
53            new_config
54        });
55    }
56
57    /// Make some predictions for what circuits should be built.
58    pub(crate) fn predict(&self, path_config: &PathConfig) -> Vec<TargetTunnelUsage> {
59        let config = self.config();
60        let now = Instant::get();
61        let circs = config.min_exit_circs_for_port;
62        self.usages
63            .iter()
64            .filter(|&(_, &time)| {
65                time.checked_add(config.prediction_lifetime)
66                    .map(|t| t > now)
67                    .unwrap_or_else(|| {
68                        // FIXME(eta): this is going to be a bit noisy if it triggers, but that's better
69                        //             than panicking or silently doing the wrong thing?
70                        warn!("failed to represent preemptive circuit prediction lifetime as an Instant");
71                        false
72                    })
73            })
74            .map(|(&port, _)| {
75                let require_stability = port.is_some_and(|p| path_config.long_lived_ports.contains(&p.port));
76                TargetTunnelUsage::Preemptive {
77                    port, circs, require_stability,
78                }
79            })
80            .collect()
81    }
82
83    /// Note the use of a new port at the provided `time`.
84    ///
85    /// # Limitations
86    ///
87    /// This function assumes that the `time` values it receives are
88    /// monotonically increasing.
89    pub(crate) fn note_usage(&mut self, port: Option<TargetPort>, time: Instant) {
90        self.usages.insert(port, time);
91    }
92}
93
94#[cfg(test)]
95mod test {
96    // @@ begin test lint list maintained by maint/add_warning @@
97    #![allow(clippy::bool_assert_comparison)]
98    #![allow(clippy::clone_on_copy)]
99    #![allow(clippy::dbg_macro)]
100    #![allow(clippy::mixed_attributes_style)]
101    #![allow(clippy::print_stderr)]
102    #![allow(clippy::print_stdout)]
103    #![allow(clippy::single_char_pattern)]
104    #![allow(clippy::unwrap_used)]
105    #![allow(clippy::unchecked_time_subtraction)]
106    #![allow(clippy::useless_vec)]
107    #![allow(clippy::needless_pass_by_value)]
108    #![allow(clippy::string_slice)] // See arti#2571
109    //! <!-- @@ end test lint list maintained by maint/add_warning @@ -->
110    use crate::{
111        PathConfig, PreemptiveCircuitConfig, PreemptiveCircuitPredictor, TargetPort,
112        TargetTunnelUsage,
113    };
114    use web_time_compat::{Duration, Instant, InstantExt};
115
116    use crate::isolation::test::{IsolationTokenEq, assert_isoleq};
117
118    #[test]
119    fn predicts_starting_ports() {
120        let path_config = PathConfig::default();
121        let mut cfg = PreemptiveCircuitConfig::builder();
122        cfg.set_initial_predicted_ports(vec![]);
123        cfg.prediction_lifetime(Duration::from_secs(2));
124        let predictor = PreemptiveCircuitPredictor::new(cfg.build().unwrap());
125
126        assert_isoleq!(
127            predictor.predict(&path_config),
128            vec![TargetTunnelUsage::Preemptive {
129                port: None,
130                circs: 2,
131                require_stability: false,
132            }]
133        );
134
135        let mut cfg = PreemptiveCircuitConfig::builder();
136        cfg.set_initial_predicted_ports(vec![80]);
137        cfg.prediction_lifetime(Duration::from_secs(2));
138        let predictor = PreemptiveCircuitPredictor::new(cfg.build().unwrap());
139
140        let results = predictor.predict(&path_config);
141        assert_eq!(results.len(), 2);
142        assert!(
143            results
144                .iter()
145                .any(|r| r.isol_eq(&TargetTunnelUsage::Preemptive {
146                    port: None,
147                    circs: 2,
148                    require_stability: false,
149                }))
150        );
151        assert!(
152            results
153                .iter()
154                .any(|r| r.isol_eq(&TargetTunnelUsage::Preemptive {
155                    port: Some(TargetPort::ipv4(80)),
156                    circs: 2,
157                    require_stability: false,
158                }))
159        );
160    }
161
162    #[test]
163    fn predicts_used_ports() {
164        let path_config = PathConfig::default();
165        let mut cfg = PreemptiveCircuitConfig::builder();
166        cfg.set_initial_predicted_ports(vec![]);
167        cfg.prediction_lifetime(Duration::from_secs(2));
168        let mut predictor = PreemptiveCircuitPredictor::new(cfg.build().unwrap());
169
170        assert_isoleq!(
171            predictor.predict(&path_config),
172            vec![TargetTunnelUsage::Preemptive {
173                port: None,
174                circs: 2,
175                require_stability: false,
176            }]
177        );
178
179        predictor.note_usage(Some(TargetPort::ipv4(1234)), Instant::get());
180
181        let results = predictor.predict(&path_config);
182        assert_eq!(results.len(), 2);
183        assert!(
184            results
185                .iter()
186                .any(|r| r.isol_eq(&TargetTunnelUsage::Preemptive {
187                    port: None,
188                    circs: 2,
189                    require_stability: false,
190                }))
191        );
192        assert!(
193            results
194                .iter()
195                .any(|r| r.isol_eq(&TargetTunnelUsage::Preemptive {
196                    port: Some(TargetPort::ipv4(1234)),
197                    circs: 2,
198                    require_stability: false,
199                }))
200        );
201    }
202
203    #[test]
204    fn does_not_predict_old_ports() {
205        let path_config = PathConfig::default();
206        let mut cfg = PreemptiveCircuitConfig::builder();
207        cfg.set_initial_predicted_ports(vec![]);
208        cfg.prediction_lifetime(Duration::from_secs(2));
209        let mut predictor = PreemptiveCircuitPredictor::new(cfg.build().unwrap());
210        let now = Instant::get();
211        let three_seconds_ago = now - Duration::from_secs(2 + 1);
212
213        predictor.note_usage(Some(TargetPort::ipv4(2345)), three_seconds_ago);
214
215        assert_isoleq!(
216            predictor.predict(&path_config),
217            vec![TargetTunnelUsage::Preemptive {
218                port: None,
219                circs: 2,
220                require_stability: false,
221            }]
222        );
223    }
224}
tor_circmgr/preemptive.rs

tor_circmgr/
preemptive.rs
