1#![allow(clippy::unwrap_used)]
14
15use crate::{MdDigest, MdReceiver, PartialNetDir};
16use std::iter;
17use std::net::SocketAddr;
18use std::time::{Duration, SystemTime};
19#[cfg(feature = "geoip")]
20use tor_geoip::GeoipDb;
21use tor_netdoc::doc::microdesc::{Microdesc, MicrodescBuilder};
22use tor_netdoc::doc::netstatus::{ConsensusBuilder, MdConsensus, MdConsensusRouterStatus};
23use tor_netdoc::doc::netstatus::{Lifetime, RelayFlags, RelayWeight, RouterStatusBuilder};
24
25pub use tor_netdoc::{BuildError, BuildResult};
26
27#[derive(Debug, Clone)]
29#[non_exhaustive]
30pub struct NodeBuilders {
31 pub rs: RouterStatusBuilder<MdDigest>,
35
36 pub md: MicrodescBuilder,
41
42 pub omit_md: bool,
44
45 pub omit_rs: bool,
47}
48
49pub fn simple_net_func(
51 _idx: usize,
52 _nb: &mut NodeBuilders,
53 _bld: &mut ConsensusBuilder<MdConsensusRouterStatus>,
54) {
55}
56
57pub fn construct_netdir() -> PartialNetDir {
59 construct_custom_netdir(simple_net_func).expect("failed to build default testing netdir")
60}
61
62pub fn construct_custom_netdir_with_params<F, P, PK>(
65 func: F,
66 params: P,
67 lifetime: Option<Lifetime>,
68) -> BuildResult<PartialNetDir>
69where
70 F: FnMut(usize, &mut NodeBuilders, &mut ConsensusBuilder<MdConsensusRouterStatus>),
71 P: IntoIterator<Item = (PK, i32)>,
72 PK: Into<String>,
73{
74 construct_custom_netdir_with_params_inner(
75 func,
76 params,
77 lifetime,
78 #[cfg(feature = "geoip")]
79 None,
80 )
81}
82
83fn construct_custom_netdir_with_params_inner<F, P, PK>(
86 func: F,
87 params: P,
88 lifetime: Option<Lifetime>,
89 #[cfg(feature = "geoip")] geoip_db: Option<&GeoipDb>,
90) -> BuildResult<PartialNetDir>
91where
92 F: FnMut(usize, &mut NodeBuilders, &mut ConsensusBuilder<MdConsensusRouterStatus>),
93 P: IntoIterator<Item = (PK, i32)>,
94 PK: Into<String>,
95{
96 let (consensus, microdescs) = construct_custom_network(func, lifetime)?;
97 #[cfg(feature = "geoip")]
98 let mut dir = if let Some(db) = geoip_db {
99 PartialNetDir::new_with_geoip(consensus, Some(¶ms.into_iter().collect()), db)
100 } else {
101 PartialNetDir::new(consensus, Some(¶ms.into_iter().collect()))
102 };
103 #[cfg(not(feature = "geoip"))]
104 let mut dir = PartialNetDir::new(consensus, Some(¶ms.into_iter().collect()));
105 for md in microdescs {
106 dir.add_microdesc(md);
107 }
108
109 Ok(dir)
110}
111
112pub fn construct_custom_netdir<F>(func: F) -> BuildResult<PartialNetDir>
114where
115 F: FnMut(usize, &mut NodeBuilders, &mut ConsensusBuilder<MdConsensusRouterStatus>),
116{
117 construct_custom_netdir_with_params(func, iter::empty::<(&str, _)>(), None)
118}
119
120#[cfg(feature = "geoip")]
121pub fn construct_custom_netdir_with_geoip<F>(func: F, db: &GeoipDb) -> BuildResult<PartialNetDir>
123where
124 F: FnMut(usize, &mut NodeBuilders, &mut ConsensusBuilder<MdConsensusRouterStatus>),
125{
126 construct_custom_netdir_with_params_inner(func, iter::empty::<(&str, _)>(), None, Some(db))
127}
128
129pub fn construct_network() -> BuildResult<(MdConsensus, Vec<Microdesc>)> {
132 construct_custom_network(simple_net_func, None)
133}
134
135pub fn construct_custom_network<F>(
189 mut func: F,
190 lifetime: Option<Lifetime>,
191) -> BuildResult<(MdConsensus, Vec<Microdesc>)>
192where
193 F: FnMut(usize, &mut NodeBuilders, &mut ConsensusBuilder<MdConsensusRouterStatus>),
194{
195 let f = RelayFlags::RUNNING
196 | RelayFlags::VALID
197 | RelayFlags::V2DIR
198 | RelayFlags::FAST
199 | RelayFlags::STABLE;
200 let flags = [
202 f | RelayFlags::HSDIR,
203 f | RelayFlags::EXIT,
204 f | RelayFlags::GUARD,
205 f | RelayFlags::EXIT | RelayFlags::GUARD,
206 ];
207
208 let lifetime = lifetime.map(Ok).unwrap_or_else(|| {
209 let now = SystemTime::now();
210 let one_day = Duration::new(86400, 0);
211
212 Lifetime::new(now, now + one_day / 2, now + one_day)
213 })?;
214
215 let mut bld = MdConsensus::builder();
216 bld.consensus_method(34)
217 .lifetime(lifetime)
218 .param("bwweightscale", 1)
219 .weights("".parse()?);
220
221 let mut microdescs = Vec::new();
222 for idx in 0..40_u8 {
223 let flags = flags[(idx / 10) as usize];
227 let policy = if flags.contains(RelayFlags::EXIT) {
228 if idx % 2 == 1 {
229 "accept 80,443"
230 } else {
231 "accept 1-65535"
232 }
233 } else {
234 "reject 1-65535"
235 };
236 let fam_id = [idx ^ 1; 20];
238 let family = hex::encode(fam_id);
239
240 let mut md_builder = Microdesc::builder();
241 md_builder
242 .ntor_key((*b"----nothing in dirmgr uses this-").into())
243 .ed25519_id([idx; 32].into())
244 .family(family.parse().unwrap())
245 .parse_ipv4_policy(policy)
246 .unwrap();
247 let protocols = if idx % 2 == 0 {
248 "DirCache=2".parse().unwrap()
250 } else {
251 "".parse().unwrap()
252 };
253 let weight = RelayWeight::Measured(1000 * u32::from(idx % 10 + 1));
254 let mut rs_builder = bld.rs();
255 rs_builder
256 .identity([idx; 20].into())
257 .add_or_port(SocketAddr::from(([idx % 5, 0, 0, 3], 9001)))
258 .protos(protocols)
259 .set_flags(flags)
260 .weight(weight);
261
262 let mut node_builders = NodeBuilders {
263 rs: rs_builder,
264 md: md_builder,
265 omit_rs: false,
266 omit_md: false,
267 };
268
269 func(idx as usize, &mut node_builders, &mut bld);
270
271 let md = node_builders.md.testing_md()?;
272 let md_digest = *md.digest();
273 if !node_builders.omit_md {
274 microdescs.push(md);
275 }
276
277 if !node_builders.omit_rs {
278 node_builders
279 .rs
280 .doc_digest(md_digest)
281 .build_into(&mut bld)?;
282 }
283 }
284
285 let consensus = bld.testing_consensus()?;
286
287 Ok((consensus, microdescs))
288}
289
290#[cfg(test)]
291mod test {
292 #![allow(clippy::bool_assert_comparison)]
294 #![allow(clippy::clone_on_copy)]
295 #![allow(clippy::dbg_macro)]
296 #![allow(clippy::mixed_attributes_style)]
297 #![allow(clippy::print_stderr)]
298 #![allow(clippy::print_stdout)]
299 #![allow(clippy::single_char_pattern)]
300 #![allow(clippy::unwrap_used)]
301 #![allow(clippy::unchecked_duration_subtraction)]
302 #![allow(clippy::useless_vec)]
303 #![allow(clippy::needless_pass_by_value)]
304 use super::*;
306 #[test]
307 fn try_with_function() {
308 let mut val = 0_u32;
309 let _net = construct_custom_netdir(|_idx, _nb, _bld| {
310 val += 1;
311 });
312 assert_eq!(val, 40);
313 }
314}