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

            
4
use std::{borrow::Borrow, collections::HashMap};
5

            
6
use tor_dirclient::request;
7
#[cfg(feature = "routerdesc")]
8
use tor_netdoc::doc::routerdesc::RdDigest;
9
use tor_netdoc::doc::{authcert::AuthCertKeyIds, microdesc::MdDigest, netstatus::ConsensusFlavor};
10

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

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

            
52
impl DocId {
53
    /// Return the associated doctype of this DocId.
54
795
    pub(crate) fn doctype(&self) -> DocType {
55
795
        use DocId::*;
56
795
        use DocType as T;
57
795
        match self {
58
2
            LatestConsensus { flavor: f, .. } => T::Consensus(*f),
59
258
            AuthCert(_) => T::AuthCert,
60
277
            Microdesc(_) => T::Microdesc,
61
            #[cfg(feature = "routerdesc")]
62
258
            RouterDesc(_) => T::RouterDesc,
63
        }
64
795
    }
65
}
66

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

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

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

            
112
impl CacheUsage {
113
    /// Turn this CacheUsage into a pending field for use with
114
    /// SqliteStorage.
115
4
    pub(crate) fn pending_requirement(&self) -> Option<bool> {
116
4
        match self {
117
1
            CacheUsage::CacheOnly => Some(false),
118
3
            _ => None,
119
        }
120
4
    }
121
}
122

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

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

            
163
    /// Add `id` to this query, if possible.
164
797
    fn push(&mut self, id: DocId) {
165
797
        match (self, id) {
166
3
            (Self::LatestConsensus { .. }, DocId::LatestConsensus { .. }) => {}
167
257
            (Self::AuthCert(ids), DocId::AuthCert(id)) => ids.push(id),
168
280
            (Self::Microdesc(ids), DocId::Microdesc(id)) => ids.push(id),
169
            #[cfg(feature = "routerdesc")]
170
257
            (Self::RouterDesc(ids), DocId::RouterDesc(id)) => ids.push(id),
171
            (_, _) => panic!(),
172
        }
173
797
    }
174

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

            
200
impl From<DocId> for DocQuery {
201
6
    fn from(d: DocId) -> DocQuery {
202
6
        let mut result = DocQuery::empty_from_docid(&d);
203
6
        result.push(d);
204
6
        result
205
6
    }
206
}
207

            
208
/// Given a list of DocId, split them up into queries, by type.
209
9
pub(crate) fn partition_by_type<T>(collection: T) -> HashMap<DocType, DocQuery>
210
9
where
211
9
    T: IntoIterator<Item = DocId>,
212
9
{
213
9
    let mut result = HashMap::new();
214
791
    for item in collection.into_iter() {
215
791
        let b = item.borrow();
216
791
        let tp = b.doctype();
217
791
        result
218
791
            .entry(tp)
219
791
            .or_insert_with(|| DocQuery::empty_from_docid(b))
220
791
            .push(item);
221
791
    }
222
9
    result
223
9
}
224

            
225
#[cfg(test)]
226
mod test {
227
    #![allow(clippy::unwrap_used)]
228
    use super::*;
229

            
230
    #[test]
231
    fn doctype() {
232
        assert_eq!(
233
            DocId::LatestConsensus {
234
                flavor: ConsensusFlavor::Microdesc,
235
                cache_usage: CacheUsage::CacheOkay,
236
            }
237
            .doctype(),
238
            DocType::Consensus(ConsensusFlavor::Microdesc)
239
        );
240

            
241
        let auth_id = AuthCertKeyIds {
242
            id_fingerprint: [10; 20].into(),
243
            sk_fingerprint: [12; 20].into(),
244
        };
245
        assert_eq!(DocId::AuthCert(auth_id).doctype(), DocType::AuthCert);
246

            
247
        assert_eq!(DocId::Microdesc([22; 32]).doctype(), DocType::Microdesc);
248
        #[cfg(feature = "routerdesc")]
249
        assert_eq!(DocId::RouterDesc([42; 20]).doctype(), DocType::RouterDesc);
250
    }
251

            
252
    #[test]
253
    fn partition_ids() {
254
        let mut ids = Vec::new();
255
        for byte in 0..=255 {
256
            ids.push(DocId::Microdesc([byte; 32]));
257
            #[cfg(feature = "routerdesc")]
258
            ids.push(DocId::RouterDesc([byte; 20]));
259
            ids.push(DocId::AuthCert(AuthCertKeyIds {
260
                id_fingerprint: [byte; 20].into(),
261
                sk_fingerprint: [33; 20].into(),
262
            }));
263
        }
264
        let consensus_q = DocId::LatestConsensus {
265
            flavor: ConsensusFlavor::Microdesc,
266
            cache_usage: CacheUsage::CacheOkay,
267
        };
268
        ids.push(consensus_q);
269

            
270
        let split = partition_by_type(ids);
271
        #[cfg(feature = "routerdesc")]
272
        assert_eq!(split.len(), 4); // 4 distinct types.
273
        #[cfg(not(feature = "routerdesc"))]
274
        assert_eq!(split.len(), 3); // 3 distinct types.
275

            
276
        let q = split
277
            .get(&DocType::Consensus(ConsensusFlavor::Microdesc))
278
            .unwrap();
279
        assert!(matches!(q, DocQuery::LatestConsensus { .. }));
280

            
281
        let q = split.get(&DocType::Microdesc).unwrap();
282
        assert!(matches!(q, DocQuery::Microdesc(v) if v.len() == 256));
283

            
284
        #[cfg(feature = "routerdesc")]
285
        {
286
            let q = split.get(&DocType::RouterDesc).unwrap();
287
            assert!(matches!(q, DocQuery::RouterDesc(v) if v.len() == 256));
288
        }
289
        let q = split.get(&DocType::AuthCert).unwrap();
290
        assert!(matches!(q, DocQuery::AuthCert(v) if v.len() == 256));
291
    }
292

            
293
    #[test]
294
    fn split_into_chunks() {
295
        use std::collections::HashSet;
296
        //use itertools::Itertools;
297
        use rand::Rng;
298

            
299
        // Construct a big query.
300
        let mut rng = rand::thread_rng();
301
        let ids: HashSet<MdDigest> = (0..3400).into_iter().map(|_| rng.gen()).collect();
302

            
303
        // Test microdescs.
304
        let split = DocQuery::Microdesc(ids.clone().into_iter().collect()).split_for_download();
305
        assert_eq!(split.len(), 7);
306
        let mut found_ids = HashSet::new();
307
        for q in split {
308
            match q {
309
                DocQuery::Microdesc(ids) => ids.into_iter().for_each(|id| {
310
                    found_ids.insert(id);
311
                }),
312
                _ => panic!("Wrong type."),
313
            }
314
        }
315
        assert_eq!(found_ids.len(), 3400);
316
        assert_eq!(found_ids, ids);
317

            
318
        // Test routerdescs.
319
        #[cfg(feature = "routerdesc")]
320
        {
321
            let ids: HashSet<RdDigest> = (0..1001).into_iter().map(|_| rng.gen()).collect();
322
            let split =
323
                DocQuery::RouterDesc(ids.clone().into_iter().collect()).split_for_download();
324
            assert_eq!(split.len(), 3);
325
            let mut found_ids = HashSet::new();
326
            for q in split {
327
                match q {
328
                    DocQuery::RouterDesc(ids) => ids.into_iter().for_each(|id| {
329
                        found_ids.insert(id);
330
                    }),
331
                    _ => panic!("Wrong type."),
332
                }
333
            }
334
            assert_eq!(found_ids.len(), 1001);
335
            assert_eq!(&found_ids, &ids);
336
        }
337

            
338
        // Test authcerts.
339
        let ids: HashSet<AuthCertKeyIds> = (0..2500)
340
            .into_iter()
341
            .map(|_| {
342
                let id_fingerprint = rng.gen::<[u8; 20]>().into();
343
                let sk_fingerprint = rng.gen::<[u8; 20]>().into();
344
                AuthCertKeyIds {
345
                    id_fingerprint,
346
                    sk_fingerprint,
347
                }
348
            })
349
            .collect();
350
        let split = DocQuery::AuthCert(ids.clone().into_iter().collect()).split_for_download();
351
        assert_eq!(split.len(), 5);
352
        let mut found_ids = HashSet::new();
353
        for q in split {
354
            match q {
355
                DocQuery::AuthCert(ids) => ids.into_iter().for_each(|id| {
356
                    found_ids.insert(id);
357
                }),
358
                _ => panic!("Wrong type."),
359
            }
360
        }
361
        assert_eq!(found_ids.len(), 2500);
362
        assert_eq!(&found_ids, &ids);
363

            
364
        // Consensus is trivial?
365
        let query = DocQuery::LatestConsensus {
366
            flavor: ConsensusFlavor::Microdesc,
367
            cache_usage: CacheUsage::CacheOkay,
368
        };
369
        let split = query.clone().split_for_download();
370
        assert_eq!(split, vec![query]);
371
    }
372

            
373
    #[test]
374
    fn into_query() {
375
        let q: DocQuery = DocId::Microdesc([99; 32]).into();
376
        assert_eq!(q, DocQuery::Microdesc(vec![[99; 32]]));
377
    }
378

            
379
    #[test]
380
    fn pending_requirement() {
381
        // If we want to keep all of our activity within the cache,
382
        // we must request a non-pending consensus from the cache.
383
        assert_eq!(CacheUsage::CacheOnly.pending_requirement(), Some(false));
384
        // Otherwise, any cached consensus, pending or not, will meet
385
        // our needs.
386
        assert_eq!(CacheUsage::CacheOkay.pending_requirement(), None);
387
        assert_eq!(CacheUsage::MustDownload.pending_requirement(), None);
388
    }
389
}