tor_circmgr/
preemptive.rs

1//! Tools for determining what circuits to preemptively build.
2
3use crate::{PathConfig, PreemptiveCircuitConfig, TargetCircUsage, TargetPort};
4use std::collections::HashMap;
5use std::sync::Arc;
6use std::time::Instant;
7use tracing::warn;
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::now());
29        }
30
31        // We want to build circuits for resolving DNS, too.
32        usages.insert(None, Instant::now());
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<TargetCircUsage> {
59        let config = self.config();
60        let now = Instant::now();
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                TargetCircUsage::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_duration_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, TargetCircUsage,
111        TargetPort,
112    };
113    use std::time::{Duration, Instant};
114
115    use crate::isolation::test::{assert_isoleq, IsolationTokenEq};
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![TargetCircUsage::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!(results
142            .iter()
143            .any(|r| r.isol_eq(&TargetCircUsage::Preemptive {
144                port: None,
145                circs: 2,
146                require_stability: false,
147            })));
148        assert!(results
149            .iter()
150            .any(|r| r.isol_eq(&TargetCircUsage::Preemptive {
151                port: Some(TargetPort::ipv4(80)),
152                circs: 2,
153                require_stability: false,
154            })));
155    }
156
157    #[test]
158    fn predicts_used_ports() {
159        let path_config = PathConfig::default();
160        let mut cfg = PreemptiveCircuitConfig::builder();
161        cfg.set_initial_predicted_ports(vec![]);
162        cfg.prediction_lifetime(Duration::from_secs(2));
163        let mut predictor = PreemptiveCircuitPredictor::new(cfg.build().unwrap());
164
165        assert_isoleq!(
166            predictor.predict(&path_config),
167            vec![TargetCircUsage::Preemptive {
168                port: None,
169                circs: 2,
170                require_stability: false,
171            }]
172        );
173
174        predictor.note_usage(Some(TargetPort::ipv4(1234)), Instant::now());
175
176        let results = predictor.predict(&path_config);
177        assert_eq!(results.len(), 2);
178        assert!(results
179            .iter()
180            .any(|r| r.isol_eq(&TargetCircUsage::Preemptive {
181                port: None,
182                circs: 2,
183                require_stability: false,
184            })));
185        assert!(results
186            .iter()
187            .any(|r| r.isol_eq(&TargetCircUsage::Preemptive {
188                port: Some(TargetPort::ipv4(1234)),
189                circs: 2,
190                require_stability: false,
191            })));
192    }
193
194    #[test]
195    fn does_not_predict_old_ports() {
196        let path_config = PathConfig::default();
197        let mut cfg = PreemptiveCircuitConfig::builder();
198        cfg.set_initial_predicted_ports(vec![]);
199        cfg.prediction_lifetime(Duration::from_secs(2));
200        let mut predictor = PreemptiveCircuitPredictor::new(cfg.build().unwrap());
201        let now = Instant::now();
202        let three_seconds_ago = now - Duration::from_secs(2 + 1);
203
204        predictor.note_usage(Some(TargetPort::ipv4(2345)), three_seconds_ago);
205
206        assert_isoleq!(
207            predictor.predict(&path_config),
208            vec![TargetCircUsage::Preemptive {
209                port: None,
210                circs: 2,
211                require_stability: false,
212            }]
213        );
214    }
215}