1use std::io::{self, Write};
32
33use digest::Digest;
34use rand::seq::SliceRandom;
35
36#[derive(Debug, thiserror::Error)]
38pub enum Error {
39 #[error("unsupported algorithm: {0}")]
41 UnsupportedAlgorithm(String),
42
43 #[error("source_urls is required")]
45 MissingSourceUrls,
46
47 #[error("hash mismatch: expected {expected}, got {actual}")]
49 HashMismatch { expected: String, actual: String },
50}
51
52pub fn normalize_algo(name: &str) -> String {
57 name.chars()
58 .filter_map(|c| match c {
59 'A'..='Z' => Some(c.to_ascii_lowercase()),
60 'a'..='z' | '0'..='9' => Some(c),
61 _ => None,
62 })
63 .collect()
64}
65
66pub fn is_supported(algo: &str) -> bool {
68 matches!(normalize_algo(algo).as_str(), "sha1" | "sha256" | "sha512")
69}
70
71pub fn parse_fetchurl_server(value: &str) -> Vec<String> {
73 let value = value.trim();
74 if value.is_empty() {
75 return Vec::new();
76 }
77 if !value.starts_with('"') {
78 return vec![value.to_string()];
79 }
80 parse_sfv_string_list(value)
81}
82
83pub fn encode_source_urls(urls: &[impl AsRef<str>]) -> String {
85 let strs: Vec<&str> = urls.iter().map(|s| s.as_ref()).collect();
86 encode_sfv_string_list(&strs)
87}
88
89fn encode_sfv_string_list(strings: &[&str]) -> String {
92 strings
93 .iter()
94 .map(|s| {
95 let escaped = s.replace('\\', "\\\\").replace('"', "\\\"");
96 format!("\"{escaped}\"")
97 })
98 .collect::<Vec<_>>()
99 .join(", ")
100}
101
102fn parse_sfv_string_list(input: &str) -> Vec<String> {
103 let mut results = Vec::new();
104 let bytes = input.as_bytes();
105 let mut i = 0;
106
107 while i < bytes.len() {
108 while i < bytes.len() && matches!(bytes[i], b' ' | b'\t') {
110 i += 1;
111 }
112 if i >= bytes.len() {
113 break;
114 }
115
116 if bytes[i] != b'"' {
118 while i < bytes.len() && bytes[i] != b',' {
120 i += 1;
121 }
122 if i < bytes.len() {
123 i += 1;
124 }
125 continue;
126 }
127 i += 1;
128
129 let mut s = String::new();
131 while i < bytes.len() {
132 match bytes[i] {
133 b'\\' if i + 1 < bytes.len() => {
134 s.push(bytes[i + 1] as char);
135 i += 2;
136 }
137 b'"' => {
138 i += 1;
139 break;
140 }
141 c => {
142 s.push(c as char);
143 i += 1;
144 }
145 }
146 }
147 results.push(s);
148
149 while i < bytes.len() && bytes[i] != b',' {
151 i += 1;
152 }
153 if i < bytes.len() {
154 i += 1;
155 }
156 }
157
158 results
159}
160
161#[derive(Clone, Debug)]
165pub struct FetchAttempt {
166 url: String,
167 headers: Vec<(String, String)>,
168}
169
170impl FetchAttempt {
171 pub fn url(&self) -> &str {
173 &self.url
174 }
175
176 pub fn headers(&self) -> &[(String, String)] {
178 &self.headers
179 }
180}
181
182pub struct FetchSession {
193 attempts: Vec<FetchAttempt>,
194 current: usize,
195 algo: String,
196 hash: String,
197 done: bool,
198 success: bool,
199}
200
201impl FetchSession {
202 pub fn new(algo: &str, hash: &str, source_urls: &[impl AsRef<str>]) -> Result<Self, Error> {
208 if source_urls.is_empty() {
209 return Err(Error::MissingSourceUrls);
210 }
211
212 let algo = normalize_algo(algo);
213 if !is_supported(&algo) {
214 return Err(Error::UnsupportedAlgorithm(algo));
215 }
216
217 let servers_env = std::env::var("FETCHURL_SERVER").unwrap_or_default();
218 let servers = parse_fetchurl_server(&servers_env);
219
220 let source_header = if !source_urls.is_empty() {
221 Some(encode_source_urls(source_urls))
222 } else {
223 None
224 };
225
226 let mut attempts = Vec::new();
227
228 for server in servers {
230 let base = server.trim_end_matches('/');
231 let url = format!("{base}/{algo}/{hash}");
232 let mut headers = Vec::new();
233 if let Some(ref val) = source_header {
234 headers.push(("X-Source-Urls".to_string(), val.clone()));
235 }
236 attempts.push(FetchAttempt { url, headers });
237 }
238
239 let mut direct: Vec<String> = source_urls.iter().map(|s| s.as_ref().to_string()).collect();
241 direct.shuffle(&mut rand::thread_rng());
242 for url in direct {
243 attempts.push(FetchAttempt {
244 url,
245 headers: Vec::new(),
246 });
247 }
248
249 Ok(FetchSession {
250 attempts,
251 current: 0,
252 algo,
253 hash: hash.to_string(),
254 done: false,
255 success: false,
256 })
257 }
258
259 pub fn next_attempt(&mut self) -> Option<FetchAttempt> {
268 if self.done || self.current >= self.attempts.len() {
269 return None;
270 }
271 let attempt = self.attempts[self.current].clone();
272 self.current += 1;
273 Some(attempt)
274 }
275
276 pub fn report_success(&mut self) {
278 self.done = true;
279 self.success = true;
280 }
281
282 pub fn report_partial(&mut self) {
285 self.done = true;
286 }
287
288 pub fn succeeded(&self) -> bool {
290 self.success
291 }
292
293 pub fn verifier<W: Write>(&self, writer: W) -> HashVerifier<W> {
298 HashVerifier::new(&self.algo, &self.hash, writer)
299 }
300}
301
302enum HasherInner {
305 Sha1(sha1::Sha1),
306 Sha256(sha2::Sha256),
307 Sha512(sha2::Sha512),
308}
309
310impl HasherInner {
311 fn new(algo: &str) -> Option<Self> {
312 match algo {
313 "sha1" => Some(HasherInner::Sha1(sha1::Sha1::new())),
314 "sha256" => Some(HasherInner::Sha256(sha2::Sha256::new())),
315 "sha512" => Some(HasherInner::Sha512(sha2::Sha512::new())),
316 _ => None,
317 }
318 }
319
320 fn update(&mut self, data: &[u8]) {
321 match self {
322 HasherInner::Sha1(h) => h.update(data),
323 HasherInner::Sha256(h) => h.update(data),
324 HasherInner::Sha512(h) => h.update(data),
325 }
326 }
327
328 fn finalize(self) -> Vec<u8> {
329 match self {
330 HasherInner::Sha1(h) => h.finalize().to_vec(),
331 HasherInner::Sha256(h) => h.finalize().to_vec(),
332 HasherInner::Sha512(h) => h.finalize().to_vec(),
333 }
334 }
335}
336
337fn to_hex(bytes: &[u8]) -> String {
338 bytes.iter().map(|b| format!("{b:02x}")).collect()
339}
340
341pub struct HashVerifier<W: Write> {
346 inner: W,
347 hasher: HasherInner,
348 expected_hash: String,
349 bytes_written: u64,
350}
351
352impl<W: Write> HashVerifier<W> {
353 fn new(algo: &str, expected_hash: &str, inner: W) -> Self {
354 let hasher = HasherInner::new(algo).expect("HashVerifier created with validated algo");
355 HashVerifier {
356 inner,
357 hasher,
358 expected_hash: expected_hash.to_string(),
359 bytes_written: 0,
360 }
361 }
362
363 pub fn bytes_written(&self) -> u64 {
365 self.bytes_written
366 }
367
368 pub fn finish(self) -> Result<W, Error> {
372 let actual = to_hex(&self.hasher.finalize());
373 if actual == self.expected_hash {
374 Ok(self.inner)
375 } else {
376 Err(Error::HashMismatch {
377 expected: self.expected_hash,
378 actual,
379 })
380 }
381 }
382}
383
384impl<W: Write> Write for HashVerifier<W> {
385 fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
386 let n = self.inner.write(buf)?;
387 self.hasher.update(&buf[..n]);
388 self.bytes_written += n as u64;
389 Ok(n)
390 }
391
392 fn flush(&mut self) -> io::Result<()> {
393 self.inner.flush()
394 }
395}
396
397#[cfg(test)]
398mod tests {
399 use super::*;
400
401 fn sha256_hex(data: &[u8]) -> String {
402 to_hex(&sha2::Sha256::digest(data))
403 }
404
405 #[test]
406 fn test_normalize_algo() {
407 assert_eq!(normalize_algo("SHA-256"), "sha256");
408 assert_eq!(normalize_algo("sha256"), "sha256");
409 assert_eq!(normalize_algo("SHA512"), "sha512");
410 assert_eq!(normalize_algo("md5"), "md5");
411 }
412
413 #[test]
414 fn test_is_supported() {
415 assert!(is_supported("sha256"));
416 assert!(is_supported("SHA-256"));
417 assert!(is_supported("sha1"));
418 assert!(is_supported("sha512"));
419 assert!(!is_supported("md5"));
420 }
421
422 #[test]
423 fn test_sfv_encode() {
424 assert_eq!(
425 encode_sfv_string_list(&["https://a.com", "https://b.com"]),
426 r#""https://a.com", "https://b.com""#
427 );
428 }
429
430 #[test]
431 fn test_sfv_parse() {
432 let parsed = parse_sfv_string_list(r#""https://a.com", "https://b.com""#);
433 assert_eq!(parsed, vec!["https://a.com", "https://b.com"]);
434 }
435
436 #[test]
437 fn test_sfv_roundtrip() {
438 let urls = &[
439 "https://cdn.example.com/file.tar.gz",
440 "https://mirror.org/archive.tgz",
441 ];
442 let encoded = encode_sfv_string_list(urls);
443 let decoded = parse_sfv_string_list(&encoded);
444 assert_eq!(decoded, urls);
445 }
446
447 #[test]
448 fn test_sfv_parse_with_parameters() {
449 let parsed = parse_sfv_string_list(r#""https://a.com";q=0.9, "https://b.com""#);
451 assert_eq!(parsed, vec!["https://a.com", "https://b.com"]);
452 }
453
454 #[test]
455 fn test_hash_verifier_success() {
456 let data = b"hello world";
457 let hash = sha256_hex(data);
458
459 let mut output = Vec::new();
460 {
461 let mut verifier = HashVerifier::new("sha256", &hash, &mut output);
462 verifier.write_all(data).unwrap();
463 assert_eq!(verifier.bytes_written(), data.len() as u64);
464 verifier.finish().unwrap();
465 }
466 assert_eq!(output, data);
467 }
468
469 #[test]
470 fn test_hash_verifier_mismatch() {
471 let data = b"hello world";
472 let wrong_hash = sha256_hex(b"wrong");
473
474 let mut output = Vec::new();
475 let mut verifier = HashVerifier::new("sha256", &wrong_hash, &mut output);
476 verifier.write_all(data).unwrap();
477 let err = verifier.finish().unwrap_err();
478 assert!(matches!(err, Error::HashMismatch { .. }));
479 }
480
481 #[test]
482 fn test_session_unsupported_algo() {
483 let err = FetchSession::new("md5", "abc", &["http://src"]);
484 assert!(matches!(err, Err(Error::UnsupportedAlgorithm(_))));
485 }
486
487 #[test]
488 fn test_session_missing_source_urls() {
489 let err = FetchSession::new("sha256", "abc", &[] as &[&str]);
490 assert!(matches!(err, Err(Error::MissingSourceUrls)));
491 }
492
493 #[test]
494 fn test_session_attempt_ordering() {
495 let hash = sha256_hex(b"test");
496 unsafe {
497 std::env::set_var(
498 "FETCHURL_SERVER",
499 "\"http://cache1/api/fetchurl\", \"http://cache2/api/fetchurl\"",
500 );
501 }
502 let mut session = FetchSession::new("sha256", &hash, &["http://src1"]).unwrap();
503
504 let a1 = session.next_attempt().unwrap();
506 assert!(a1.url().starts_with("http://cache1/api/fetchurl/sha256/"));
507 assert!(!a1.headers().is_empty());
508
509 let a2 = session.next_attempt().unwrap();
510 assert!(a2.url().starts_with("http://cache2/api/fetchurl/sha256/"));
511
512 let a3 = session.next_attempt().unwrap();
514 assert_eq!(a3.url(), "http://src1");
515 assert!(a3.headers().is_empty());
516
517 assert!(session.next_attempt().is_none());
519 assert!(!session.succeeded());
520 }
521
522 #[test]
523 fn test_session_success_stops() {
524 let hash = sha256_hex(b"test");
525 unsafe {
526 std::env::set_var("FETCHURL_SERVER", "\"http://cache/api/fetchurl\"");
527 }
528 let mut session = FetchSession::new("sha256", &hash, &["http://src"]).unwrap();
529
530 let _ = session.next_attempt().unwrap();
531 session.report_success();
532 assert!(session.succeeded());
533 assert!(session.next_attempt().is_none());
534 }
535
536 #[test]
537 fn test_session_partial_stops() {
538 let hash = sha256_hex(b"test");
539 unsafe {
540 std::env::set_var("FETCHURL_SERVER", "\"http://cache/api/fetchurl\"");
541 }
542 let mut session = FetchSession::new("sha256", &hash, &["http://src"]).unwrap();
543
544 let _ = session.next_attempt().unwrap();
545 session.report_partial();
546 assert!(!session.succeeded());
547 assert!(session.next_attempt().is_none());
548 }
549
550 #[test]
551 fn test_session_server_has_source_header() {
552 let hash = sha256_hex(b"test");
553 unsafe {
554 std::env::set_var("FETCHURL_SERVER", "\"http://cache/api/fetchurl\"");
555 }
556 let mut session =
557 FetchSession::new("sha256", &hash, &["http://src1", "http://src2"]).unwrap();
558
559 let attempt = session.next_attempt().unwrap();
560 let source_header = attempt
561 .headers()
562 .iter()
563 .find(|(k, _)| k == "X-Source-Urls")
564 .map(|(_, v)| v.clone())
565 .unwrap();
566
567 let parsed = parse_sfv_string_list(&source_header);
568 assert!(parsed.contains(&"http://src1".to_string()));
569 assert!(parsed.contains(&"http://src2".to_string()));
570 }
571}