1
//! Declare a general purpose "document ID type" for tracking which
2
//! documents we want and which we have.
3

            
4
use std::collections::HashMap;
5
use tracing::trace;
6

            
7
use crate::storage::Store;
8
use crate::DocumentText;
9
use tor_dirclient::request;
10
#[cfg(feature = "routerdesc")]
11
use tor_netdoc::doc::routerdesc::RdDigest;
12
use tor_netdoc::doc::{authcert::AuthCertKeyIds, microdesc::MdDigest, netstatus::ConsensusFlavor};
13

            
14
/// The identity of a single document, in enough detail to load it
15
/// from storage.
16
#[derive(Clone, Copy, Debug, Hash, Eq, PartialEq)]
17
#[non_exhaustive]
18
pub enum DocId {
19
    /// A request for the most recent consensus document.
20
    LatestConsensus {
21
        /// The flavor of consensus to request.
22
        flavor: ConsensusFlavor,
23
        /// Rules for loading this consensus from the cache.
24
        cache_usage: CacheUsage,
25
    },
26
    /// A request for an authority certificate, by the SHA1 digests of
27
    /// its identity key and signing key.
28
    AuthCert(AuthCertKeyIds),
29
    /// A request for a single microdescriptor, by SHA256 digest.
30
    Microdesc(MdDigest),
31
    /// A request for the router descriptor of a public relay, by SHA1
32
    /// digest.
33
    #[cfg(feature = "routerdesc")]
34
    RouterDesc(RdDigest),
35
}
36

            
37
/// The underlying type of a DocId.
38
///
39
/// Documents with the same type can be grouped into the same query; others
40
/// cannot.
41
#[derive(Clone, Debug, Eq, PartialEq, Hash, Ord, PartialOrd)]
42
#[non_exhaustive]
43
pub(crate) enum DocType {
44
    /// A consensus document
45
    Consensus(ConsensusFlavor),
46
    /// An authority certificate
47
    AuthCert,
48
    /// A microdescriptor
49
    Microdesc,
50
    /// A router descriptor.
51
    #[cfg(feature = "routerdesc")]
52
    RouterDesc,
53
}
54

            
55
impl DocId {
56
    /// Return the associated doctype of this DocId.
57
5598
    pub(crate) fn doctype(&self) -> DocType {
58
        use DocId::*;
59
        use DocType as T;
60
5598
        match self {
61
8
            LatestConsensus { flavor: f, .. } => T::Consensus(*f),
62
518
            AuthCert(_) => T::AuthCert,
63
2556
            Microdesc(_) => T::Microdesc,
64
            #[cfg(feature = "routerdesc")]
65
2516
            RouterDesc(_) => T::RouterDesc,
66
        }
67
5598
    }
68
}
69

            
70
/// A request for a specific kind of directory resource that a DirMgr can
71
/// request.
72
#[derive(Clone, Debug)]
73
pub(crate) enum ClientRequest {
74
    /// Request for a consensus
75
    Consensus(request::ConsensusRequest),
76
    /// Request for one or more authority certificates
77
    AuthCert(request::AuthCertRequest),
78
    /// Request for one or more microdescriptors
79
    Microdescs(request::MicrodescRequest),
80
    /// Request for one or more router descriptors
81
    #[cfg(feature = "routerdesc")]
82
    RouterDescs(request::RouterDescRequest),
83
}
84

            
85
impl ClientRequest {
86
    /// Turn a ClientRequest into a Requestable.
87
    pub(crate) fn as_requestable(&self) -> &(dyn request::Requestable + Send + Sync) {
88
        use ClientRequest::*;
89
        match self {
90
            Consensus(a) => a,
91
            AuthCert(a) => a,
92
            Microdescs(a) => a,
93
            #[cfg(feature = "routerdesc")]
94
            RouterDescs(a) => a,
95
        }
96
    }
97
}
98

            
99
/// Description of how to start out a given bootstrap attempt.
100
#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash)]
101
pub enum CacheUsage {
102
    /// The bootstrap attempt will only use the cache.  Therefore, don't
103
    /// load a pending consensus from the cache, since we won't be able
104
    /// to find enough information to make it usable.
105
    CacheOnly,
106
    /// The bootstrap attempt is willing to download information or to
107
    /// use the cache.  Therefore, we want the latest cached
108
    /// consensus, whether it is pending or not.
109
    CacheOkay,
110
    /// The bootstrap attempt is trying to fetch a new consensus. Therefore,
111
    /// we don't want a consensus from the cache.
112
    MustDownload,
113
}
114

            
115
impl CacheUsage {
116
    /// Turn this CacheUsage into a pending field for use with
117
    /// SqliteStorage.
118
8
    pub(crate) fn pending_requirement(&self) -> Option<bool> {
119
8
        match self {
120
2
            CacheUsage::CacheOnly => Some(false),
121
6
            _ => None,
122
        }
123
8
    }
124
}
125

            
126
/// A group of DocIds that can be downloaded or loaded from the database
127
/// together.
128
///
129
/// TODO: Perhaps this should be the same as ClientRequest?
130
#[derive(Clone, Debug, Eq, PartialEq)]
131
pub(crate) enum DocQuery {
132
    /// A request for the latest consensus
133
    LatestConsensus {
134
        /// A desired flavor of consensus
135
        flavor: ConsensusFlavor,
136
        /// Whether we can or must use the cache
137
        cache_usage: CacheUsage,
138
    },
139
    /// A request for authority certificates
140
    AuthCert(Vec<AuthCertKeyIds>),
141
    /// A request for microdescriptors
142
    Microdesc(Vec<MdDigest>),
143
    /// A request for router descriptors
144
    #[cfg(feature = "routerdesc")]
145
    RouterDesc(Vec<RdDigest>),
146
}
147

            
148
impl DocQuery {
149
    /// Construct an "empty" docquery from the given DocId
150
48
    pub(crate) fn empty_from_docid(id: &DocId) -> Self {
151
48
        match *id {
152
            DocId::LatestConsensus {
153
8
                flavor,
154
8
                cache_usage,
155
8
            } => Self::LatestConsensus {
156
8
                flavor,
157
8
                cache_usage,
158
8
            },
159
6
            DocId::AuthCert(_) => Self::AuthCert(Vec::new()),
160
28
            DocId::Microdesc(_) => Self::Microdesc(Vec::new()),
161
            #[cfg(feature = "routerdesc")]
162
6
            DocId::RouterDesc(_) => Self::RouterDesc(Vec::new()),
163
        }
164
48
    }
165

            
166
    /// Add `id` to this query, if possible.
167
5598
    fn push(&mut self, id: DocId) {
168
5598
        match (self, id) {
169
8
            (Self::LatestConsensus { .. }, DocId::LatestConsensus { .. }) => {}
170
516
            (Self::AuthCert(ids), DocId::AuthCert(id)) => ids.push(id),
171
2560
            (Self::Microdesc(ids), DocId::Microdesc(id)) => ids.push(id),
172
            #[cfg(feature = "routerdesc")]
173
2514
            (Self::RouterDesc(ids), DocId::RouterDesc(id)) => ids.push(id),
174
            (_, _) => panic!(),
175
        }
176
5598
    }
177

            
178
    /// If this query contains too many documents to download with a single
179
    /// request, divide it up.
180
22
    pub(crate) fn split_for_download(self) -> Vec<Self> {
181
        use DocQuery::*;
182
        /// How many objects can be put in a single HTTP GET line?
183
        const N: usize = 500;
184
22
        match self {
185
6
            LatestConsensus { .. } => vec![self],
186
4
            AuthCert(mut v) => {
187
4
                v.sort_unstable();
188
14
                v[..].chunks(N).map(|s| AuthCert(s.to_vec())).collect()
189
            }
190
8
            Microdesc(mut v) => {
191
8
                v.sort_unstable();
192
26
                v[..].chunks(N).map(|s| Microdesc(s.to_vec())).collect()
193
            }
194
            #[cfg(feature = "routerdesc")]
195
4
            RouterDesc(mut v) => {
196
4
                v.sort_unstable();
197
12
                v[..].chunks(N).map(|s| RouterDesc(s.to_vec())).collect()
198
            }
199
        }
200
22
    }
201

            
202
    /// Load documents specified by this `DocQuery` from the store, if they can be found.
203
    ///
204
    /// # Note
205
    ///
206
    /// This function may not return all documents that the query asked for. If this happens, no
207
    /// error will be returned. It is the caller's responsibility to handle this case.
208
24
    pub(crate) fn load_from_store_into(
209
24
        &self,
210
24
        result: &mut HashMap<DocId, DocumentText>,
211
24
        store: &dyn Store,
212
24
    ) -> crate::Result<()> {
213
        use DocQuery::*;
214
24
        match self {
215
            LatestConsensus {
216
2
                flavor,
217
2
                cache_usage,
218
2
            } => {
219
2
                if *cache_usage == CacheUsage::MustDownload {
220
                    // Do nothing: we don't want a cached consensus.
221
                    trace!("MustDownload is set; not checking for cached consensus.");
222
2
                } else if let Some(c) =
223
2
                    store.latest_consensus(*flavor, cache_usage.pending_requirement())?
224
                {
225
2
                    trace!("Found a reasonable consensus in the cache");
226
2
                    let id = DocId::LatestConsensus {
227
2
                        flavor: *flavor,
228
2
                        cache_usage: *cache_usage,
229
2
                    };
230
2
                    result.insert(id, c.into());
231
                }
232
            }
233
2
            AuthCert(ids) => result.extend(
234
2
                store
235
2
                    .authcerts(ids)?
236
2
                    .into_iter()
237
3
                    .map(|(id, c)| (DocId::AuthCert(id), DocumentText::from_string(c))),
238
2
            ),
239
18
            Microdesc(digests) => {
240
18
                result.extend(
241
18
                    store
242
18
                        .microdescs(digests)?
243
18
                        .into_iter()
244
41
                        .map(|(id, md)| (DocId::Microdesc(id), DocumentText::from_string(md))),
245
18
                );
246
18
            }
247
            #[cfg(feature = "routerdesc")]
248
2
            RouterDesc(digests) => result.extend(
249
2
                store
250
2
                    .routerdescs(digests)?
251
2
                    .into_iter()
252
3
                    .map(|(id, rd)| (DocId::RouterDesc(id), DocumentText::from_string(rd))),
253
2
            ),
254
        }
255
24
        Ok(())
256
24
    }
257
}
258

            
259
impl From<DocId> for DocQuery {
260
8
    fn from(d: DocId) -> DocQuery {
261
8
        let mut result = DocQuery::empty_from_docid(&d);
262
8
        result.push(d);
263
8
        result
264
8
    }
265
}
266

            
267
/// Given a list of DocId, split them up into queries, by type.
268
30
pub(crate) fn partition_by_type<T>(collection: T) -> HashMap<DocType, DocQuery>
269
30
where
270
30
    T: IntoIterator<Item = DocId>,
271
30
{
272
30
    let mut result = HashMap::new();
273
5590
    for item in collection.into_iter() {
274
5590
        let tp = item.doctype();
275
5590
        result
276
5590
            .entry(tp)
277
5590
            .or_insert_with(|| DocQuery::empty_from_docid(&item))
278
5590
            .push(item);
279
5590
    }
280
30
    result
281
30
}
282

            
283
#[cfg(test)]
284
mod test {
285
    // @@ begin test lint list maintained by maint/add_warning @@
286
    #![allow(clippy::bool_assert_comparison)]
287
    #![allow(clippy::clone_on_copy)]
288
    #![allow(clippy::dbg_macro)]
289
    #![allow(clippy::mixed_attributes_style)]
290
    #![allow(clippy::print_stderr)]
291
    #![allow(clippy::print_stdout)]
292
    #![allow(clippy::single_char_pattern)]
293
    #![allow(clippy::unwrap_used)]
294
    #![allow(clippy::unchecked_duration_subtraction)]
295
    #![allow(clippy::useless_vec)]
296
    #![allow(clippy::needless_pass_by_value)]
297
    //! <!-- @@ end test lint list maintained by maint/add_warning @@ -->
298
    use super::*;
299
    use tor_basic_utils::test_rng::testing_rng;
300

            
301
    #[test]
302
    fn doctype() {
303
        assert_eq!(
304
            DocId::LatestConsensus {
305
                flavor: ConsensusFlavor::Microdesc,
306
                cache_usage: CacheUsage::CacheOkay,
307
            }
308
            .doctype(),
309
            DocType::Consensus(ConsensusFlavor::Microdesc)
310
        );
311

            
312
        let auth_id = AuthCertKeyIds {
313
            id_fingerprint: [10; 20].into(),
314
            sk_fingerprint: [12; 20].into(),
315
        };
316
        assert_eq!(DocId::AuthCert(auth_id).doctype(), DocType::AuthCert);
317

            
318
        assert_eq!(DocId::Microdesc([22; 32]).doctype(), DocType::Microdesc);
319
        #[cfg(feature = "routerdesc")]
320
        assert_eq!(DocId::RouterDesc([42; 20]).doctype(), DocType::RouterDesc);
321
    }
322

            
323
    #[test]
324
    fn partition_ids() {
325
        let mut ids = Vec::new();
326
        for byte in 0..=255 {
327
            ids.push(DocId::Microdesc([byte; 32]));
328
            #[cfg(feature = "routerdesc")]
329
            ids.push(DocId::RouterDesc([byte; 20]));
330
            ids.push(DocId::AuthCert(AuthCertKeyIds {
331
                id_fingerprint: [byte; 20].into(),
332
                sk_fingerprint: [33; 20].into(),
333
            }));
334
        }
335
        let consensus_q = DocId::LatestConsensus {
336
            flavor: ConsensusFlavor::Microdesc,
337
            cache_usage: CacheUsage::CacheOkay,
338
        };
339
        ids.push(consensus_q);
340

            
341
        let split = partition_by_type(ids);
342
        #[cfg(feature = "routerdesc")]
343
        assert_eq!(split.len(), 4); // 4 distinct types.
344
        #[cfg(not(feature = "routerdesc"))]
345
        assert_eq!(split.len(), 3); // 3 distinct types.
346

            
347
        let q = split
348
            .get(&DocType::Consensus(ConsensusFlavor::Microdesc))
349
            .unwrap();
350
        assert!(matches!(q, DocQuery::LatestConsensus { .. }));
351

            
352
        let q = split.get(&DocType::Microdesc).unwrap();
353
        assert!(matches!(q, DocQuery::Microdesc(v) if v.len() == 256));
354

            
355
        #[cfg(feature = "routerdesc")]
356
        {
357
            let q = split.get(&DocType::RouterDesc).unwrap();
358
            assert!(matches!(q, DocQuery::RouterDesc(v) if v.len() == 256));
359
        }
360
        let q = split.get(&DocType::AuthCert).unwrap();
361
        assert!(matches!(q, DocQuery::AuthCert(v) if v.len() == 256));
362
    }
363

            
364
    #[test]
365
    fn split_into_chunks() {
366
        use std::collections::HashSet;
367
        //use itertools::Itertools;
368
        use rand::Rng;
369

            
370
        // Construct a big query.
371
        let mut rng = testing_rng();
372
        let ids: HashSet<MdDigest> = (0..3400).map(|_| rng.random()).collect();
373

            
374
        // Test microdescs.
375
        let split = DocQuery::Microdesc(ids.clone().into_iter().collect()).split_for_download();
376
        assert_eq!(split.len(), 7);
377
        let mut found_ids = HashSet::new();
378
        for q in split {
379
            match q {
380
                DocQuery::Microdesc(ids) => ids.into_iter().for_each(|id| {
381
                    found_ids.insert(id);
382
                }),
383
                _ => panic!("Wrong type."),
384
            }
385
        }
386
        assert_eq!(found_ids.len(), 3400);
387
        assert_eq!(found_ids, ids);
388

            
389
        // Test routerdescs.
390
        #[cfg(feature = "routerdesc")]
391
        {
392
            let ids: HashSet<RdDigest> = (0..1001).map(|_| rng.random()).collect();
393
            let split =
394
                DocQuery::RouterDesc(ids.clone().into_iter().collect()).split_for_download();
395
            assert_eq!(split.len(), 3);
396
            let mut found_ids = HashSet::new();
397
            for q in split {
398
                match q {
399
                    DocQuery::RouterDesc(ids) => ids.into_iter().for_each(|id| {
400
                        found_ids.insert(id);
401
                    }),
402
                    _ => panic!("Wrong type."),
403
                }
404
            }
405
            assert_eq!(found_ids.len(), 1001);
406
            assert_eq!(&found_ids, &ids);
407
        }
408

            
409
        // Test authcerts.
410
        let ids: HashSet<AuthCertKeyIds> = (0..2500)
411
            .map(|_| {
412
                let id_fingerprint = rng.random::<[u8; 20]>().into();
413
                let sk_fingerprint = rng.random::<[u8; 20]>().into();
414
                AuthCertKeyIds {
415
                    id_fingerprint,
416
                    sk_fingerprint,
417
                }
418
            })
419
            .collect();
420
        let split = DocQuery::AuthCert(ids.clone().into_iter().collect()).split_for_download();
421
        assert_eq!(split.len(), 5);
422
        let mut found_ids = HashSet::new();
423
        for q in split {
424
            match q {
425
                DocQuery::AuthCert(ids) => ids.into_iter().for_each(|id| {
426
                    found_ids.insert(id);
427
                }),
428
                _ => panic!("Wrong type."),
429
            }
430
        }
431
        assert_eq!(found_ids.len(), 2500);
432
        assert_eq!(&found_ids, &ids);
433

            
434
        // Consensus is trivial?
435
        let query = DocQuery::LatestConsensus {
436
            flavor: ConsensusFlavor::Microdesc,
437
            cache_usage: CacheUsage::CacheOkay,
438
        };
439
        let split = query.clone().split_for_download();
440
        assert_eq!(split, vec![query]);
441
    }
442

            
443
    #[test]
444
    fn into_query() {
445
        let q: DocQuery = DocId::Microdesc([99; 32]).into();
446
        assert_eq!(q, DocQuery::Microdesc(vec![[99; 32]]));
447
    }
448

            
449
    #[test]
450
    fn pending_requirement() {
451
        // If we want to keep all of our activity within the cache,
452
        // we must request a non-pending consensus from the cache.
453
        assert_eq!(CacheUsage::CacheOnly.pending_requirement(), Some(false));
454
        // Otherwise, any cached consensus, pending or not, will meet
455
        // our needs.
456
        assert_eq!(CacheUsage::CacheOkay.pending_requirement(), None);
457
        assert_eq!(CacheUsage::MustDownload.pending_requirement(), None);
458
    }
459
}