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    //! <!-- @@ end test lint list maintained by maint/add_warning @@ -->
109    use crate::{
110        PathConfig, PreemptiveCircuitConfig, PreemptiveCircuitPredictor, TargetPort,
111        TargetTunnelUsage,
112    };
113    use web_time_compat::{Duration, Instant, InstantExt};
114
115    use crate::isolation::test::{IsolationTokenEq, assert_isoleq};
116
117    #[test]
118    fn predicts_starting_ports() {
119        let path_config = PathConfig::default();
120        let mut cfg = PreemptiveCircuitConfig::builder();
121        cfg.set_initial_predicted_ports(vec![]);
122        cfg.prediction_lifetime(Duration::from_secs(2));
123        let predictor = PreemptiveCircuitPredictor::new(cfg.build().unwrap());
124
125        assert_isoleq!(
126            predictor.predict(&path_config),
127            vec![TargetTunnelUsage::Preemptive {
128                port: None,
129                circs: 2,
130                require_stability: false,
131            }]
132        );
133
134        let mut cfg = PreemptiveCircuitConfig::builder();
135        cfg.set_initial_predicted_ports(vec![80]);
136        cfg.prediction_lifetime(Duration::from_secs(2));
137        let predictor = PreemptiveCircuitPredictor::new(cfg.build().unwrap());
138
139        let results = predictor.predict(&path_config);
140        assert_eq!(results.len(), 2);
141        assert!(
142            results
143                .iter()
144                .any(|r| r.isol_eq(&TargetTunnelUsage::Preemptive {
145                    port: None,
146                    circs: 2,
147                    require_stability: false,
148                }))
149        );
150        assert!(
151            results
152                .iter()
153                .any(|r| r.isol_eq(&TargetTunnelUsage::Preemptive {
154                    port: Some(TargetPort::ipv4(80)),
155                    circs: 2,
156                    require_stability: false,
157                }))
158        );
159    }
160
161    #[test]
162    fn predicts_used_ports() {
163        let path_config = PathConfig::default();
164        let mut cfg = PreemptiveCircuitConfig::builder();
165        cfg.set_initial_predicted_ports(vec![]);
166        cfg.prediction_lifetime(Duration::from_secs(2));
167        let mut predictor = PreemptiveCircuitPredictor::new(cfg.build().unwrap());
168
169        assert_isoleq!(
170            predictor.predict(&path_config),
171            vec![TargetTunnelUsage::Preemptive {
172                port: None,
173                circs: 2,
174                require_stability: false,
175            }]
176        );
177
178        predictor.note_usage(Some(TargetPort::ipv4(1234)), Instant::get());
179
180        let results = predictor.predict(&path_config);
181        assert_eq!(results.len(), 2);
182        assert!(
183            results
184                .iter()
185                .any(|r| r.isol_eq(&TargetTunnelUsage::Preemptive {
186                    port: None,
187                    circs: 2,
188                    require_stability: false,
189                }))
190        );
191        assert!(
192            results
193                .iter()
194                .any(|r| r.isol_eq(&TargetTunnelUsage::Preemptive {
195                    port: Some(TargetPort::ipv4(1234)),
196                    circs: 2,
197                    require_stability: false,
198                }))
199        );
200    }
201
202    #[test]
203    fn does_not_predict_old_ports() {
204        let path_config = PathConfig::default();
205        let mut cfg = PreemptiveCircuitConfig::builder();
206        cfg.set_initial_predicted_ports(vec![]);
207        cfg.prediction_lifetime(Duration::from_secs(2));
208        let mut predictor = PreemptiveCircuitPredictor::new(cfg.build().unwrap());
209        let now = Instant::get();
210        let three_seconds_ago = now - Duration::from_secs(2 + 1);
211
212        predictor.note_usage(Some(TargetPort::ipv4(2345)), three_seconds_ago);
213
214        assert_isoleq!(
215            predictor.predict(&path_config),
216            vec![TargetTunnelUsage::Preemptive {
217                port: None,
218                circs: 2,
219                require_stability: false,
220            }]
221        );
222    }
223}
tor_circmgr/preemptive.rs

tor_circmgr/
preemptive.rs
