1use std::fmt;
35use std::path::PathBuf;
36use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
37use std::sync::{Arc, Mutex};
38use std::time::{Duration, SystemTime, UNIX_EPOCH};
39
40use async_trait::async_trait;
41use tokio::sync::{broadcast, oneshot, Notify};
42
43use crate::downloads::destination::{prepare_download_directory, resolve_download_directory};
44use crate::downloads::options::{CaptureOptions, DownloadArtifact, DownloadEvent, DownloadOptions};
45use crate::downloads::sources::{
46 prepare_staging_directory, set_download_behavior, DownloadFailure, DownloadSink, DownloadStart,
47 SourceHandle,
48};
49use crate::downloads::store::{
50 save_download, DownloadConflict, DownloadNamer, DownloadSource, DownloadValidator, SaveRequest,
51};
52use crate::downloads::watcher::{attach_filesystem_watcher, DEFAULT_POLL_INTERVAL};
53use crate::downloads::DownloadError;
54use crate::fingerprint::CdpTransport;
55
56pub const DEFAULT_CAPTURE_TIMEOUT: Duration = Duration::from_millis(30_000);
58
59const EVENT_BUFFER: usize = 256;
61
62static DOWNLOAD_COUNTER: AtomicU64 = AtomicU64::new(0);
64
65struct Record {
67 artifact: DownloadArtifact,
68 filename: Option<DownloadNamer>,
69 validate: Option<DownloadValidator>,
70 conflict: Option<DownloadConflict>,
71}
72
73struct Waiter {
75 token: u64,
76 armed_at: usize,
77 options: CaptureOptions,
78 settle: Option<oneshot::Sender<DownloadArtifact>>,
79}
80
81#[derive(Default)]
83struct State {
84 records: Vec<Record>,
85 by_handle: std::collections::HashMap<String, usize>,
86 by_id: std::collections::HashMap<String, usize>,
87 published: std::collections::HashSet<String>,
88 waiters: Vec<Waiter>,
89 next_token: u64,
90}
91
92pub struct DownloadManager {
94 pub directory: PathBuf,
96 pub persist: bool,
98 conflict: DownloadConflict,
99 filename: Option<DownloadNamer>,
100 validate: Option<DownloadValidator>,
101 poll_interval: Duration,
102 state: Mutex<State>,
103 events: broadcast::Sender<DownloadArtifact>,
104 in_flight: AtomicUsize,
105 settled: Notify,
106 sources: Mutex<Vec<SourceHandle>>,
107}
108
109impl fmt::Debug for DownloadManager {
110 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
111 f.debug_struct("DownloadManager")
112 .field("directory", &self.directory)
113 .field("persist", &self.persist)
114 .field("conflict", &self.conflict)
115 .finish_non_exhaustive()
116 }
117}
118
119impl DownloadManager {
120 pub fn create(options: DownloadOptions) -> Result<Arc<Self>, DownloadError> {
140 let root =
141 prepare_download_directory(&resolve_download_directory(options.directory.as_deref())?)?;
142
143 Ok(Arc::new(Self {
144 directory: root,
145 persist: options.persist,
146 conflict: options.conflict,
147 filename: options.filename,
148 validate: options.validate,
149 poll_interval: options.poll_interval.unwrap_or(DEFAULT_POLL_INTERVAL),
150 state: Mutex::new(State::default()),
151 events: broadcast::channel(EVENT_BUFFER).0,
152 in_flight: AtomicUsize::new(0),
153 settled: Notify::new(),
154 sources: Mutex::new(Vec::new()),
155 }))
156 }
157
158 pub async fn attach(
175 self: &Arc<Self>,
176 transport: &dyn CdpTransport,
177 ) -> Result<(), DownloadError> {
178 let staging = prepare_staging_directory(&self.directory)?;
179 set_download_behavior(transport, &staging)
180 .await
181 .map_err(|error| DownloadError::Transport {
182 reason: error.to_string(),
183 })?;
184
185 self.attach_filesystem().await
186 }
187
188 pub async fn attach_filesystem(self: &Arc<Self>) -> Result<(), DownloadError> {
191 let handle = attach_filesystem_watcher(
192 &self.directory,
193 Arc::clone(self) as Arc<dyn DownloadSink>,
194 self.poll_interval,
195 )?;
196 self.sources.lock().expect("manager lock").push(handle);
197 Ok(())
198 }
199
200 pub fn list(&self) -> Vec<DownloadArtifact> {
206 self.state
207 .lock()
208 .expect("manager lock")
209 .records
210 .iter()
211 .map(|record| record.artifact.clone())
212 .collect()
213 }
214
215 pub fn subscribe(&self) -> broadcast::Receiver<DownloadArtifact> {
221 self.events.subscribe()
222 }
223
224 pub async fn idle(&self) {
226 loop {
227 let notified = self.settled.notified();
228 tokio::pin!(notified);
229 notified.as_mut().enable();
233 if self.in_flight.load(Ordering::SeqCst) == 0 {
234 return;
235 }
236 notified.await;
237 }
238 }
239
240 pub async fn capture<F>(
263 &self,
264 options: CaptureOptions,
265 action: F,
266 ) -> Result<DownloadArtifact, DownloadError>
267 where
268 F: std::future::Future<Output = anyhow::Result<()>>,
269 {
270 let timeout = options.timeout.unwrap_or(DEFAULT_CAPTURE_TIMEOUT);
271 let (settle, settled) = oneshot::channel();
272 let token = {
273 let mut state = self.state.lock().expect("manager lock");
274 let token = state.next_token;
275 state.next_token += 1;
276 let armed_at = state.records.len();
277 state.waiters.push(Waiter {
278 token,
279 armed_at,
280 options,
281 settle: Some(settle),
282 });
283 token
284 };
285
286 if let Err(error) = action.await {
287 self.discard(token);
288 return Err(DownloadError::ActionFailed {
289 reason: error.to_string(),
290 });
291 }
292
293 let artifact = match tokio::time::timeout(timeout, settled).await {
294 Ok(Ok(artifact)) => artifact,
295 _ => {
296 self.discard(token);
297 return Err(DownloadError::CaptureTimeout {
298 timeout_ms: timeout.as_millis() as u64,
299 });
300 }
301 };
302 self.discard(token);
303
304 if artifact.state != DownloadEvent::Completed {
305 return Err(DownloadError::DownloadFailed {
306 id: artifact.id,
307 state: artifact.state.to_string(),
308 failure: artifact.failure.unwrap_or_default(),
309 });
310 }
311 Ok(artifact)
312 }
313
314 pub async fn dispose(&self) {
319 self.idle().await;
322 let handles: Vec<SourceHandle> =
323 std::mem::take(self.sources.lock().expect("manager lock").as_mut());
324 for handle in handles {
325 (handle.detach)().await;
326 }
327 self.idle().await;
330 }
331
332 fn discard(&self, token: u64) {
334 self.state
335 .lock()
336 .expect("manager lock")
337 .waiters
338 .retain(|waiter| waiter.token != token);
339 }
340
341 fn save_request(&self, id: &str, source: DownloadSource) -> Option<SaveRequest> {
346 let state = self.state.lock().expect("manager lock");
347 let record = state.by_id.get(id).map(|index| &state.records[*index])?;
348 if record.artifact.state != DownloadEvent::Started {
349 return None;
350 }
351
352 Some(SaveRequest {
353 root: self.directory.clone(),
354 source,
355 suggested_filename: record
356 .artifact
357 .suggested_filename
358 .clone()
359 .unwrap_or_default(),
360 mime_type: record.artifact.mime_type.clone(),
361 conflict: record.conflict.unwrap_or(self.conflict),
362 filename: record.filename.clone().or_else(|| self.filename.clone()),
363 validate: record.validate.clone().or_else(|| self.validate.clone()),
364 })
365 }
366
367 fn publish(&self, id: &str, state: DownloadEvent, update: impl FnOnce(&mut DownloadArtifact)) {
373 let artifact = {
374 let mut locked = self.state.lock().expect("manager lock");
375 if !locked.published.insert(id.to_string()) {
376 return;
377 }
378 let Some(index) = locked.by_id.get(id).copied() else {
379 return;
380 };
381 let record = &mut locked.records[index];
382 record.artifact.state = state;
383 update(&mut record.artifact);
384 let artifact = record.artifact.clone();
385
386 for waiter in std::mem::take(&mut locked.waiters) {
390 if index >= waiter.armed_at {
391 if let Some(settle) = waiter.settle {
392 let _ = settle.send(artifact.clone());
393 }
394 } else {
395 locked.waiters.push(waiter);
396 }
397 }
398 artifact
399 };
400 let _ = self.events.send(artifact);
401 }
402
403 fn begin_work(&self) {
405 self.in_flight.fetch_add(1, Ordering::SeqCst);
406 }
407
408 fn end_work(&self) {
410 if self.in_flight.fetch_sub(1, Ordering::SeqCst) == 1 {
411 self.settled.notify_waiters();
412 }
413 }
414}
415
416#[async_trait]
417impl DownloadSink for DownloadManager {
418 fn started(&self, start: DownloadStart) -> String {
419 let mut state = self.state.lock().expect("manager lock");
420 if let Some(index) = state.by_handle.get(&start.engine_handle).copied() {
427 if state.records[index].artifact.state == DownloadEvent::Started {
428 return state.records[index].artifact.id.clone();
429 }
430 }
431
432 let id = format!(
433 "dl-{:06}",
434 DOWNLOAD_COUNTER.fetch_add(1, Ordering::SeqCst) + 1
435 );
436 let artifact = DownloadArtifact {
437 id: id.clone(),
438 url: start.url,
439 suggested_filename: start.suggested_filename,
440 state: DownloadEvent::Started,
441 started_at: iso8601(SystemTime::now()),
442 completed_at: None,
443 path: None,
444 mime_type: start.mime_type,
445 bytes: None,
446 checksum: None,
447 failure: None,
448 };
449
450 let index = state.records.len();
453 let claimed = state
454 .waiters
455 .iter()
456 .find(|waiter| index >= waiter.armed_at)
457 .map(|waiter| waiter.options.clone());
458
459 state.records.push(Record {
460 artifact: artifact.clone(),
461 filename: claimed
462 .as_ref()
463 .and_then(|options| options.filename.clone()),
464 validate: claimed
465 .as_ref()
466 .and_then(|options| options.validate.clone()),
467 conflict: claimed.as_ref().and_then(|options| options.conflict),
468 });
469 state.by_handle.insert(start.engine_handle, index);
470 state.by_id.insert(id.clone(), index);
471 drop(state);
472
473 let _ = self.events.send(artifact);
474 id
475 }
476
477 async fn finished(&self, id: String, source: DownloadSource) {
478 let Some(request) = self.save_request(&id, source) else {
479 return;
480 };
481
482 self.begin_work();
483 let saved = tokio::task::spawn_blocking(move || save_download(request))
484 .await
485 .unwrap_or_else(|error| {
486 Err(DownloadError::ActionFailed {
487 reason: error.to_string(),
488 })
489 });
490
491 match saved {
492 Ok(saved) => self.publish(&id, DownloadEvent::Completed, |artifact| {
493 artifact.path = Some(saved.path);
494 artifact.bytes = Some(saved.bytes);
495 artifact.checksum = Some(saved.checksum);
496 artifact.completed_at = Some(iso8601(SystemTime::now()));
497 }),
498 Err(error) => self.publish(&id, DownloadEvent::Failed, |artifact| {
502 artifact.failure = Some(error.to_string());
503 artifact.completed_at = Some(iso8601(SystemTime::now()));
504 }),
505 }
506 self.end_work();
507 }
508
509 fn failed(&self, id: String, kind: DownloadFailure, reason: String) {
510 let state = match kind {
511 DownloadFailure::Cancelled => DownloadEvent::Cancelled,
512 DownloadFailure::Failed => DownloadEvent::Failed,
513 };
514 self.publish(&id, state, |artifact| {
515 artifact.failure = Some(reason);
516 artifact.completed_at = Some(iso8601(SystemTime::now()));
517 });
518 }
519}
520
521fn iso8601(time: SystemTime) -> String {
536 let since_epoch = time.duration_since(UNIX_EPOCH).unwrap_or_default();
537 let seconds = since_epoch.as_secs() as i64;
538 let (year, month, day) = civil_from_days(seconds.div_euclid(86_400));
539 let seconds_of_day = seconds.rem_euclid(86_400);
540
541 format!(
542 "{year:04}-{month:02}-{day:02}T{:02}:{:02}:{:02}.{:03}Z",
543 seconds_of_day / 3600,
544 (seconds_of_day % 3600) / 60,
545 seconds_of_day % 60,
546 since_epoch.subsec_millis()
547 )
548}
549
550fn civil_from_days(days: i64) -> (i64, i64, i64) {
555 let shifted = days + 719_468;
558 let era = shifted.div_euclid(146_097);
559 let day_of_era = shifted.rem_euclid(146_097);
560 let year_of_era =
561 (day_of_era - day_of_era / 1460 + day_of_era / 36_524 - day_of_era / 146_096) / 365;
562 let year = year_of_era + era * 400;
563 let day_of_year = day_of_era - (365 * year_of_era + year_of_era / 4 - year_of_era / 100);
564 let month_position = (5 * day_of_year + 2) / 153;
565 let day = day_of_year - (153 * month_position + 2) / 5 + 1;
566 let month = if month_position < 10 {
567 month_position + 3
568 } else {
569 month_position - 9
570 };
571
572 (if month <= 2 { year + 1 } else { year }, month, day)
573}
574
575#[cfg(test)]
576mod tests {
577 use super::*;
578 use crate::downloads::store::DownloadCandidate;
579 use crate::downloads::test_support::TempDir;
580
581 fn manager(temp: &TempDir, options: DownloadOptions) -> Arc<DownloadManager> {
583 DownloadManager::create(DownloadOptions {
584 directory: Some(temp.path().join("downloads").to_string_lossy().into_owned()),
585 ..options
586 })
587 .expect("manager")
588 }
589
590 fn staged(manager: &DownloadManager, name: &str, contents: &[u8]) -> DownloadSource {
592 let staging = prepare_staging_directory(&manager.directory).expect("staging");
593 let path = staging.join(name);
594 std::fs::write(&path, contents).expect("staged file is writable");
595 DownloadSource::staged(path)
596 }
597
598 async fn deliver(manager: &Arc<DownloadManager>, name: &str, contents: &[u8]) -> String {
600 let source = staged(manager, name, contents);
601 let id = manager.started(DownloadStart {
602 engine_handle: source.path.to_string_lossy().into_owned(),
603 suggested_filename: Some(name.to_string()),
604 ..DownloadStart::default()
605 });
606 manager.finished(id.clone(), source).await;
607 id
608 }
609
610 #[tokio::test]
611 async fn saves_a_download_where_the_caller_can_find_it_later() {
612 let temp = TempDir::new("bc-manager-save");
613 let manager = manager(&temp, DownloadOptions::default());
614
615 deliver(&manager, "report.pdf", b"%PDF-1.7 body").await;
616
617 let artifacts = manager.list();
618 assert_eq!(artifacts.len(), 1);
619 assert_eq!(artifacts[0].state, DownloadEvent::Completed);
620 let path = artifacts[0].path.clone().expect("a completed download");
621 assert_eq!(path, manager.directory.join("report.pdf"));
622 assert_eq!(std::fs::read(&path).unwrap(), b"%PDF-1.7 body");
623 assert_eq!(artifacts[0].bytes, Some(13));
624 }
625
626 #[tokio::test]
627 async fn keeps_the_file_after_everything_that_made_it_is_gone() {
628 let temp = TempDir::new("bc-manager-persist");
629 let manager = manager(&temp, DownloadOptions::default());
630
631 deliver(&manager, "keep.txt", b"still here").await;
632 let path = manager.list()[0]
633 .path
634 .clone()
635 .expect("a completed download");
636 manager.dispose().await;
637 drop(manager);
638
639 assert_eq!(std::fs::read(&path).unwrap(), b"still here");
641 }
642
643 #[tokio::test]
644 async fn reports_one_download_once_however_many_sources_saw_it() {
645 let temp = TempDir::new("bc-manager-dedupe");
646 let manager = manager(&temp, DownloadOptions::default());
647 let mut events = manager.subscribe();
648
649 let source = staged(&manager, "once.txt", b"body");
650 let handle = source.path.to_string_lossy().into_owned();
651 let first = manager.started(DownloadStart {
652 engine_handle: handle.clone(),
653 suggested_filename: Some("once.txt".to_string()),
654 ..DownloadStart::default()
655 });
656 let second = manager.started(DownloadStart {
657 engine_handle: handle,
658 suggested_filename: Some("once.txt".to_string()),
659 ..DownloadStart::default()
660 });
661 assert_eq!(first, second, "the same file became two downloads");
662
663 manager.finished(first.clone(), source.clone()).await;
664 manager.finished(first, source).await;
665
666 assert_eq!(manager.list().len(), 1);
667 let completed: Vec<DownloadArtifact> = std::iter::from_fn(|| events.try_recv().ok())
668 .filter(|artifact| artifact.state == DownloadEvent::Completed)
669 .collect();
670 assert_eq!(completed.len(), 1, "one file was reported as two");
671 assert!(!manager.directory.join("once (2).txt").exists());
673 }
674
675 #[tokio::test]
676 async fn names_a_download_that_finishes_inside_the_action() {
677 let temp = TempDir::new("bc-manager-race");
678 let manager = manager(&temp, DownloadOptions::default());
679
680 let during = Arc::clone(&manager);
683 let artifact = manager
684 .capture(CaptureOptions::named("chosen.pdf"), async move {
685 deliver(&during, "7f1c9ab2", b"%PDF-1.7").await;
686 Ok(())
687 })
688 .await
689 .expect("capture");
690
691 assert_eq!(
692 artifact.path,
693 Some(manager.directory.join("chosen.pdf")),
694 "the caller's name lost the race with the download"
695 );
696 }
697
698 #[tokio::test]
699 async fn saves_a_bare_uuid_under_a_name_a_person_can_open() {
700 let temp = TempDir::new("bc-manager-uuid");
701 let manager = manager(&temp, DownloadOptions::default());
702
703 let artifact = manager
704 .capture(
705 CaptureOptions::named("statement.pdf").validated_by(Arc::new(
706 |candidate: &DownloadCandidate| {
707 let head = std::fs::read(&candidate.path)?;
708 Ok(head.starts_with(b"%PDF-"))
709 },
710 )),
711 {
712 let manager = Arc::clone(&manager);
713 async move {
714 deliver(&manager, "7f1c9ab2-4d3e", b"%PDF-1.7 statement").await;
715 Ok(())
716 }
717 },
718 )
719 .await
720 .expect("capture");
721
722 assert_eq!(artifact.path, Some(manager.directory.join("statement.pdf")));
723 }
724
725 #[tokio::test]
726 async fn refuses_to_call_a_rejected_download_a_success() {
727 let temp = TempDir::new("bc-manager-reject");
728 let manager = manager(&temp, DownloadOptions::default());
729
730 let error = manager
731 .capture(
732 CaptureOptions::named("invoice.pdf")
733 .validated_by(Arc::new(|_candidate: &DownloadCandidate| Ok(false))),
734 {
735 let manager = Arc::clone(&manager);
736 async move {
737 deliver(&manager, "invoice.pdf", b"<html>login</html>").await;
738 Ok(())
739 }
740 },
741 )
742 .await
743 .unwrap_err();
744
745 assert!(
746 error.to_string().contains("failed"),
747 "unexpected message: {error}"
748 );
749 assert!(
750 !manager.directory.join("invoice.pdf").exists(),
751 "a rejected download was published anyway"
752 );
753 }
754
755 #[tokio::test]
756 async fn refuses_to_call_a_cancelled_download_a_success() {
757 let temp = TempDir::new("bc-manager-cancel");
758 let manager = manager(&temp, DownloadOptions::default());
759
760 let during = Arc::clone(&manager);
761 let error = manager
762 .capture(CaptureOptions::default(), async move {
763 let id = during.started(DownloadStart {
764 engine_handle: "handle-1".to_string(),
765 suggested_filename: Some("big.zip".to_string()),
766 ..DownloadStart::default()
767 });
768 during.failed(id, DownloadFailure::Cancelled, "Canceled".to_string());
769 Ok(())
770 })
771 .await
772 .unwrap_err();
773
774 assert!(
775 error.to_string().contains("cancelled: Canceled"),
776 "unexpected message: {error}"
777 );
778 }
779
780 #[tokio::test]
781 async fn gives_up_when_no_download_arrives() {
782 let temp = TempDir::new("bc-manager-timeout");
783 let manager = manager(&temp, DownloadOptions::default());
784
785 let error = manager
786 .capture(
787 CaptureOptions::default().within(Duration::from_millis(20)),
788 async { Ok(()) },
789 )
790 .await
791 .unwrap_err();
792
793 assert!(
794 error.to_string().contains("no download completed within"),
795 "unexpected message: {error}"
796 );
797 }
798
799 #[tokio::test]
800 async fn reports_the_actions_own_failure_rather_than_a_timeout() {
801 let temp = TempDir::new("bc-manager-action");
802 let manager = manager(&temp, DownloadOptions::default());
803
804 let error = manager
805 .capture(CaptureOptions::default(), async {
806 Err(anyhow::anyhow!("the link was not there"))
807 })
808 .await
809 .unwrap_err();
810
811 assert!(
812 error.to_string().contains("the link was not there"),
813 "unexpected message: {error}"
814 );
815 }
816
817 #[tokio::test]
818 async fn numbers_a_second_download_of_the_same_name() {
819 let temp = TempDir::new("bc-manager-collision");
820 let manager = manager(&temp, DownloadOptions::default());
821
822 deliver(&manager, "report.pdf", b"first").await;
823 deliver(&manager, "report.pdf", b"second").await;
824
825 assert_eq!(
826 std::fs::read(manager.directory.join("report.pdf")).unwrap(),
827 b"first"
828 );
829 assert_eq!(
830 std::fs::read(manager.directory.join("report (2).pdf")).unwrap(),
831 b"second"
832 );
833 }
834
835 #[tokio::test]
836 async fn waits_for_a_download_still_being_written() {
837 let temp = TempDir::new("bc-manager-idle");
838 let manager = manager(&temp, DownloadOptions::default());
839
840 let source = staged(&manager, "slow.bin", &vec![7u8; 512 * 1024]);
841 let id = manager.started(DownloadStart {
842 engine_handle: source.path.to_string_lossy().into_owned(),
843 suggested_filename: Some("slow.bin".to_string()),
844 ..DownloadStart::default()
845 });
846 let saving = tokio::spawn({
847 let manager = Arc::clone(&manager);
848 async move { manager.finished(id, source).await }
849 });
850
851 manager.idle().await;
852 saving.await.expect("the save task");
853 assert!(manager.directory.join("slow.bin").exists());
856 }
857
858 #[test]
859 fn writes_a_timestamp_the_other_languages_can_read() {
860 assert_eq!(iso8601(UNIX_EPOCH), "1970-01-01T00:00:00.000Z");
861 assert_eq!(
862 iso8601(UNIX_EPOCH + Duration::from_millis(1_700_000_000_123)),
863 "2023-11-14T22:13:20.123Z"
864 );
865 assert_eq!(
867 iso8601(UNIX_EPOCH + Duration::from_secs(1_709_164_800)),
868 "2024-02-29T00:00:00.000Z"
869 );
870 }
871
872 #[test]
873 fn names_the_states_the_way_the_other_languages_do() {
874 assert_eq!(DownloadEvent::Started.to_string(), "started");
875 assert_eq!(DownloadEvent::Completed.to_string(), "completed");
876 assert_eq!(DownloadEvent::Failed.to_string(), "failed");
877 assert_eq!(DownloadEvent::Cancelled.to_string(), "cancelled");
878 }
879}