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