rget/
mirror.rs

1//! Mirrors and source selection (PRD §15).
2//!
3//! The rule that matters: **never splice bytes from two resources you cannot
4//! show are the same file.** Same filename on two hosts proves nothing. So a
5//! mirror is admitted only when
6//!
7//! * its size matches the primary's, **and**
8//! * its strong `ETag` matches the primary's, **or** the user gave us a
9//!   checksum, which means an end-to-end verification will catch a mismatch.
10//!
11//! Anything else is reported and skipped rather than silently mixed in.
12
13use std::sync::Mutex;
14use std::sync::atomic::{AtomicU32, Ordering};
15
16use url::Url;
17
18use crate::http::RemoteInfo;
19
20#[derive(Debug, Clone, PartialEq, Eq)]
21pub enum Admission {
22    /// The URL the user asked for; always usable.
23    Primary,
24    /// Proven equivalent by strong validators.
25    Verified,
26    /// Not proven equivalent, but a checksum will catch any mistake.
27    ChecksumGuarded,
28    Rejected(String),
29}
30
31impl Admission {
32    pub fn is_usable(&self) -> bool {
33        !matches!(self, Admission::Rejected(_))
34    }
35}
36
37/// Decide whether `candidate` may contribute bytes to the same file as
38/// `primary`.
39pub fn classify(primary: &RemoteInfo, candidate: &RemoteInfo, has_checksum: bool) -> Admission {
40    match (primary.size, candidate.size) {
41        (Some(a), Some(b)) if a != b => {
42            return Admission::Rejected(format!("size differs ({b} vs {a} bytes)"));
43        }
44        (Some(_), None) => {
45            return Admission::Rejected("mirror did not report a size".into());
46        }
47        _ => {}
48    }
49
50    if !candidate.accept_ranges && primary.accept_ranges {
51        return Admission::Rejected("mirror does not support range requests".into());
52    }
53
54    if primary.has_strong_etag() && candidate.has_strong_etag() {
55        if primary.etag == candidate.etag {
56            return Admission::Verified;
57        }
58        if has_checksum {
59            return Admission::ChecksumGuarded;
60        }
61        return Admission::Rejected(
62            "ETag differs from the primary; pass a checksum to allow it anyway".into(),
63        );
64    }
65
66    if has_checksum {
67        Admission::ChecksumGuarded
68    } else {
69        Admission::Rejected(
70            "cannot prove this is the same file (no strong ETag on both sides and no checksum given)"
71                .into(),
72        )
73    }
74}
75
76pub struct Source {
77    pub url: Url,
78    pub admission: Admission,
79    /// This source's own `If-Range` validator. Validators are per-resource, not
80    /// per-download: sending the primary's ETag to a mirror would make the
81    /// mirror answer with a full body, which we would read as "the file
82    /// changed". Each source therefore carries the validator it issued.
83    pub validator: Option<String>,
84    failures: AtomicU32,
85}
86
87impl Source {
88    pub fn new(url: Url, admission: Admission, validator: Option<String>) -> Self {
89        Self {
90            url,
91            admission,
92            validator,
93            failures: AtomicU32::new(0),
94        }
95    }
96
97    pub fn failures(&self) -> u32 {
98        self.failures.load(Ordering::Relaxed)
99    }
100}
101
102/// The set of usable sources for one download, with least-failures selection so
103/// a flaky mirror drains itself out of rotation without being banned outright.
104pub struct SourceSet {
105    sources: Vec<Source>,
106    cursor: Mutex<usize>,
107}
108
109impl SourceSet {
110    /// Keeps only usable sources. The primary is always first.
111    pub fn new(sources: Vec<Source>) -> Self {
112        let sources: Vec<Source> = sources
113            .into_iter()
114            .filter(|s| s.admission.is_usable())
115            .collect();
116        Self {
117            sources,
118            cursor: Mutex::new(0),
119        }
120    }
121
122    pub fn single(url: Url, validator: Option<String>) -> Self {
123        Self::new(vec![Source::new(url, Admission::Primary, validator)])
124    }
125
126    pub fn len(&self) -> usize {
127        self.sources.len()
128    }
129
130    pub fn is_empty(&self) -> bool {
131        self.sources.is_empty()
132    }
133
134    pub fn urls(&self) -> Vec<String> {
135        self.sources.iter().map(|s| s.url.to_string()).collect()
136    }
137
138    /// Pick a source: fewest failures wins, ties broken round-robin so several
139    /// workers starting at once spread across mirrors.
140    pub fn pick(&self) -> (usize, &Source) {
141        let min = self.sources.iter().map(|s| s.failures()).min().unwrap_or(0);
142        let candidates: Vec<usize> = self
143            .sources
144            .iter()
145            .enumerate()
146            .filter(|(_, s)| s.failures() == min)
147            .map(|(i, _)| i)
148            .collect();
149        let mut cursor = self.cursor.lock().unwrap_or_else(|e| e.into_inner());
150        *cursor = cursor.wrapping_add(1);
151        let idx = candidates[*cursor % candidates.len()];
152        (idx, &self.sources[idx])
153    }
154
155    pub fn penalise(&self, idx: usize) {
156        if let Some(s) = self.sources.get(idx) {
157            s.failures.fetch_add(1, Ordering::Relaxed);
158        }
159    }
160
161    pub fn reward(&self, idx: usize) {
162        if let Some(s) = self.sources.get(idx) {
163            // Decay rather than reset: a mirror that just succeeded is not
164            // proven healthy, it is merely less suspect.
165            let _ = s
166                .failures
167                .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |v| {
168                    Some(v.saturating_sub(1))
169                });
170        }
171    }
172}
173
174#[cfg(test)]
175mod tests {
176    use super::*;
177
178    fn info(size: Option<u64>, etag: Option<&str>) -> RemoteInfo {
179        RemoteInfo {
180            final_url: Url::parse("https://a.example/f").unwrap(),
181            size,
182            accept_ranges: true,
183            etag: etag.map(String::from),
184            last_modified: None,
185            content_type: None,
186            content_disposition: None,
187            content_encoding: None,
188        }
189    }
190
191    #[test]
192    fn accepts_matching_strong_etags() {
193        let a = info(Some(100), Some("\"x\""));
194        let b = info(Some(100), Some("\"x\""));
195        assert_eq!(classify(&a, &b, false), Admission::Verified);
196    }
197
198    #[test]
199    fn rejects_size_mismatch_even_with_a_checksum() {
200        let a = info(Some(100), Some("\"x\""));
201        let b = info(Some(101), Some("\"x\""));
202        assert!(matches!(classify(&a, &b, true), Admission::Rejected(_)));
203    }
204
205    #[test]
206    fn rejects_differing_etags_without_a_checksum() {
207        let a = info(Some(100), Some("\"x\""));
208        let b = info(Some(100), Some("\"y\""));
209        assert!(matches!(classify(&a, &b, false), Admission::Rejected(_)));
210        // A checksum makes it safe to try: verification is the backstop.
211        assert_eq!(classify(&a, &b, true), Admission::ChecksumGuarded);
212    }
213
214    #[test]
215    fn rejects_unprovable_equivalence() {
216        let a = info(Some(100), None);
217        let b = info(Some(100), None);
218        assert!(matches!(classify(&a, &b, false), Admission::Rejected(_)));
219        assert_eq!(classify(&a, &b, true), Admission::ChecksumGuarded);
220    }
221
222    #[test]
223    fn rejects_weak_etags_as_proof() {
224        let a = info(Some(100), Some("W/\"x\""));
225        let b = info(Some(100), Some("W/\"x\""));
226        // Weak ETags say "semantically equivalent", not "byte-identical".
227        assert!(matches!(classify(&a, &b, false), Admission::Rejected(_)));
228    }
229
230    #[test]
231    fn rejects_mirror_without_range_support() {
232        let a = info(Some(100), Some("\"x\""));
233        let mut b = info(Some(100), Some("\"x\""));
234        b.accept_ranges = false;
235        assert!(matches!(classify(&a, &b, false), Admission::Rejected(_)));
236    }
237
238    #[test]
239    fn drops_rejected_sources() {
240        let set = SourceSet::new(vec![
241            Source::new(
242                Url::parse("https://a.example/f").unwrap(),
243                Admission::Primary,
244                None,
245            ),
246            Source::new(
247                Url::parse("https://b.example/f").unwrap(),
248                Admission::Rejected("nope".into()),
249                None,
250            ),
251        ]);
252        assert_eq!(set.len(), 1);
253    }
254
255    #[test]
256    fn selection_avoids_failing_mirrors() {
257        let set = SourceSet::new(vec![
258            Source::new(
259                Url::parse("https://a.example/f").unwrap(),
260                Admission::Primary,
261                Some("\"a\"".into()),
262            ),
263            Source::new(
264                Url::parse("https://b.example/f").unwrap(),
265                Admission::Verified,
266                Some("\"b\"".into()),
267            ),
268        ]);
269        set.penalise(0);
270        for _ in 0..6 {
271            let (idx, source) = set.pick();
272            assert_eq!(
273                idx, 1,
274                "should avoid the failing mirror, got {}",
275                source.url
276            );
277            // Each source keeps its own validator.
278            assert_eq!(source.validator.as_deref(), Some("\"b\""));
279        }
280        // Once it recovers, rotation resumes.
281        set.reward(0);
282        let picks: std::collections::HashSet<usize> = (0..8).map(|_| set.pick().0).collect();
283        assert_eq!(picks.len(), 2);
284    }
285
286    #[test]
287    fn single_source_is_always_picked() {
288        let set = SourceSet::single(Url::parse("https://a.example/f").unwrap(), None);
289        set.penalise(0);
290        assert_eq!(set.pick().0, 0);
291    }
292}
rget/mirror.rs

rget/
mirror.rs
