use anyhow::{Context, bail};
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use std::collections::BTreeMap;
use std::fs;
use std::path::{Path, PathBuf};
use url::Url;
pub const DEFAULT_MAX_DOWNLOAD_RECORDS: usize = 500;
pub const MAX_ACTIVE_DOWNLOADS: usize = 16;
pub const MAX_DOWNLOADS_FILE_BYTES: u64 = 2 * 1024 * 1024;
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
pub struct DownloadId(pub u64);
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub enum DownloadStatus {
InProgress,
Paused,
Completed,
Cancelled,
Failed(String),
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum DownloadCommand {
Open,
ShowInFolder,
Cancel,
RemoveFromList,
DeleteLocalFile,
CopyLink,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DownloadRecord {
pub id: DownloadId,
pub url: Url,
pub file_path: PathBuf,
pub mime_type: Option<String>,
pub total_bytes: Option<u64>,
pub downloaded_bytes: u64,
pub started_at: DateTime<Utc>,
pub finished_at: Option<DateTime<Utc>>,
pub status: DownloadStatus,
}
impl DownloadRecord {
pub fn progress(&self) -> Option<f64> {
let total = self.total_bytes?;
if total == 0 {
return Some(0.0);
}
Some((self.downloaded_bytes as f64 / total as f64).clamp(0.0, 1.0))
}
pub fn file_name(&self) -> String {
self.file_path
.file_name()
.and_then(|name| name.to_str())
.filter(|name| !name.trim().is_empty())
.map(ToString::to_string)
.unwrap_or_else(|| "download".to_string())
}
pub fn status_label(&self) -> String {
match &self.status {
DownloadStatus::InProgress => match self.progress() {
Some(progress) => format!("Downloading {:.0}%", progress * 100.0),
None => format!("Downloading {}", format_bytes(self.downloaded_bytes)),
},
DownloadStatus::Paused => "Paused".to_string(),
DownloadStatus::Completed => {
format!("Complete - {}", format_bytes(self.downloaded_bytes))
}
DownloadStatus::Cancelled => "Cancelled".to_string(),
DownloadStatus::Failed(error) => format!("Failed - {error}"),
}
}
pub fn command_enabled(&self, command: DownloadCommand) -> bool {
match command {
DownloadCommand::Open | DownloadCommand::ShowInFolder => {
self.status == DownloadStatus::Completed
}
DownloadCommand::Cancel => self.status == DownloadStatus::InProgress,
DownloadCommand::RemoveFromList => self.status != DownloadStatus::InProgress,
DownloadCommand::DeleteLocalFile => matches!(
self.status,
DownloadStatus::Completed | DownloadStatus::Cancelled | DownloadStatus::Failed(_)
),
DownloadCommand::CopyLink => true,
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DownloadManager {
records: BTreeMap<DownloadId, DownloadRecord>,
next_id: u64,
#[serde(default = "default_max_download_records")]
max_records: usize,
}
impl Default for DownloadManager {
fn default() -> Self {
Self {
records: BTreeMap::new(),
next_id: 0,
max_records: DEFAULT_MAX_DOWNLOAD_RECORDS,
}
}
}
impl DownloadManager {
pub fn start(
&mut self,
url: Url,
file_path: PathBuf,
mime_type: Option<String>,
total_bytes: Option<u64>,
) -> anyhow::Result<DownloadId> {
if self.active_count() >= MAX_ACTIVE_DOWNLOADS {
bail!("download active limit reached ({MAX_ACTIVE_DOWNLOADS})");
}
let id = self.allocate_id()?;
self.records.insert(
id,
DownloadRecord {
id,
url,
file_path,
mime_type,
total_bytes,
downloaded_bytes: 0,
started_at: Utc::now(),
finished_at: None,
status: DownloadStatus::InProgress,
},
);
self.enforce_limits();
Ok(id)
}
pub fn enforce_limits(&mut self) {
self.max_records = self.max_records.clamp(1, DEFAULT_MAX_DOWNLOAD_RECORDS);
let max_records = self.max_records;
while self.records.len() > max_records {
let Some(id) = self
.records
.iter()
.find(|(_, record)| record.status != DownloadStatus::InProgress)
.map(|(id, _)| *id)
.or_else(|| self.records.keys().next().copied())
else {
break;
};
self.records.remove(&id);
}
self.normalize_next_id();
}
pub fn mark_in_progress_interrupted(&mut self) {
let now = Utc::now();
for record in self.records.values_mut() {
if record.status == DownloadStatus::InProgress {
record.status = DownloadStatus::Failed("interrupted".to_string());
record.finished_at = Some(now);
}
}
}
pub fn max_records(&self) -> usize {
self.max_records.max(1)
}
pub fn active_count(&self) -> usize {
self.records
.values()
.filter(|record| record.status == DownloadStatus::InProgress)
.count()
}
pub fn update_progress(&mut self, id: DownloadId, downloaded_bytes: u64) -> bool {
let Some(record) = self.records.get_mut(&id) else {
return false;
};
record.downloaded_bytes = downloaded_bytes;
true
}
pub fn update_destination(&mut self, id: DownloadId, file_path: PathBuf) -> bool {
let Some(record) = self.records.get_mut(&id) else {
return false;
};
record.file_path = file_path;
true
}
pub fn update_metadata(
&mut self,
id: DownloadId,
mime_type: Option<String>,
total_bytes: Option<u64>,
) -> bool {
let Some(record) = self.records.get_mut(&id) else {
return false;
};
if mime_type.is_some() {
record.mime_type = mime_type;
}
if total_bytes.is_some() {
record.total_bytes = total_bytes;
}
true
}
pub fn pause(&mut self, id: DownloadId) -> bool {
self.set_status(id, DownloadStatus::Paused)
}
pub fn resume(&mut self, id: DownloadId) -> bool {
if self
.records
.get(&id)
.is_some_and(|record| record.status == DownloadStatus::InProgress)
{
return true;
}
if self.active_count() >= MAX_ACTIVE_DOWNLOADS {
return false;
}
self.set_status(id, DownloadStatus::InProgress)
}
pub fn cancel(&mut self, id: DownloadId) -> bool {
self.set_finished_status(id, DownloadStatus::Cancelled)
}
pub fn finish(&mut self, id: DownloadId) -> bool {
self.set_finished_status(id, DownloadStatus::Completed)
}
pub fn fail(&mut self, id: DownloadId, error: impl Into<String>) -> bool {
self.set_finished_status(id, DownloadStatus::Failed(error.into()))
}
pub fn remove(&mut self, id: DownloadId) -> Option<DownloadRecord> {
self.records.remove(&id)
}
pub fn get(&self, id: DownloadId) -> Option<&DownloadRecord> {
self.records.get(&id)
}
pub fn records(&self) -> impl Iterator<Item = &DownloadRecord> {
self.records.values()
}
fn set_status(&mut self, id: DownloadId, status: DownloadStatus) -> bool {
let Some(record) = self.records.get_mut(&id) else {
return false;
};
record.status = status;
true
}
fn set_finished_status(&mut self, id: DownloadId, status: DownloadStatus) -> bool {
let Some(record) = self.records.get_mut(&id) else {
return false;
};
record.status = status;
record.finished_at = Some(Utc::now());
true
}
fn allocate_id(&mut self) -> anyhow::Result<DownloadId> {
self.normalize_next_id();
let id = DownloadId(self.next_id.max(1));
self.next_id = id.0.checked_add(1).context("download id limit reached")?;
Ok(id)
}
fn normalize_next_id(&mut self) {
let minimum_next = self
.records
.keys()
.map(|id| id.0)
.max()
.and_then(|id| id.checked_add(1))
.unwrap_or(1);
if self.next_id == 0 || self.next_id == u64::MAX || self.next_id < minimum_next {
self.next_id = minimum_next;
}
}
}
fn default_max_download_records() -> usize {
DEFAULT_MAX_DOWNLOAD_RECORDS
}
pub fn format_bytes(bytes: u64) -> String {
const UNITS: [&str; 5] = ["B", "KB", "MB", "GB", "TB"];
let mut value = bytes as f64;
let mut unit = 0;
while value >= 1024.0 && unit + 1 < UNITS.len() {
value /= 1024.0;
unit += 1;
}
if unit == 0 {
format!("{bytes} B")
} else {
format!("{value:.1} {}", UNITS[unit])
}
}
#[derive(Debug, Clone)]
pub struct DownloadStore {
path: PathBuf,
}
impl DownloadStore {
pub fn new(path: PathBuf) -> Self {
Self { path }
}
pub fn path(&self) -> &Path {
&self.path
}
pub fn load(&self) -> anyhow::Result<DownloadManager> {
if !self.path.exists() {
return Ok(DownloadManager::default());
}
let data = read_bounded_to_string(&self.path, MAX_DOWNLOADS_FILE_BYTES)?;
let mut manager: DownloadManager = serde_json::from_str(&data)?;
manager.mark_in_progress_interrupted();
manager.enforce_limits();
Ok(manager)
}
pub fn save(&self, manager: &DownloadManager) -> anyhow::Result<()> {
let mut manager = manager.clone();
manager.enforce_limits();
atomic_write_json(&self.path, &manager)
}
}
fn read_bounded_to_string(path: &Path, max_bytes: u64) -> anyhow::Result<String> {
let size = fs::metadata(path)
.with_context(|| format!("failed to inspect downloads {}", path.display()))?
.len();
if size > max_bytes {
bail!(
"downloads file {} is too large: {size} bytes exceeds {max_bytes}",
path.display()
);
}
fs::read_to_string(path).with_context(|| format!("failed to read downloads {}", path.display()))
}
fn atomic_write_json<T: Serialize>(path: &Path, value: &T) -> anyhow::Result<()> {
if let Some(parent) = path.parent() {
fs::create_dir_all(parent)?;
}
let tmp = path.with_extension("tmp");
fs::write(&tmp, serde_json::to_vec(value)?)?;
fs::rename(tmp, path)?;
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn tracks_download_lifecycle() {
let mut manager = DownloadManager::default();
let id = manager
.start(
Url::parse("https://example.com/file.zip").unwrap(),
PathBuf::from("file.zip"),
Some("application/zip".to_string()),
Some(100),
)
.unwrap();
assert!(manager.update_progress(id, 40));
let record = manager.records().next().unwrap();
assert_eq!(record.progress(), Some(0.4));
assert!(manager.finish(id));
let record = manager.records().next().unwrap();
assert_eq!(record.status, DownloadStatus::Completed);
assert!(record.finished_at.is_some());
}
#[test]
fn exposes_download_commands_by_status() {
let mut manager = DownloadManager::default();
let id = manager
.start(
Url::parse("https://example.com/file.zip").unwrap(),
PathBuf::from("file.zip"),
None,
Some(100),
)
.unwrap();
let record = manager.get(id).unwrap();
assert!(record.command_enabled(DownloadCommand::Cancel));
assert!(!record.command_enabled(DownloadCommand::Open));
manager.finish(id);
let record = manager.get(id).unwrap();
assert!(record.command_enabled(DownloadCommand::Open));
assert!(record.command_enabled(DownloadCommand::RemoveFromList));
assert!(!record.command_enabled(DownloadCommand::Cancel));
}
#[test]
fn prunes_finished_records_to_configured_limit() {
let mut manager = DownloadManager {
max_records: 2,
..DownloadManager::default()
};
for index in 0..4 {
let id = manager
.start(
Url::parse(&format!("https://example.com/{index}.zip")).unwrap(),
PathBuf::from(format!("{index}.zip")),
None,
None,
)
.unwrap();
manager.finish(id);
manager.enforce_limits();
}
assert_eq!(manager.records().count(), 2);
assert!(manager.get(DownloadId(1)).is_none());
assert!(manager.get(DownloadId(4)).is_some());
}
#[test]
fn prunes_stale_in_progress_records_to_configured_limit() {
let mut manager = DownloadManager {
max_records: 2,
..DownloadManager::default()
};
for index in 0..4 {
manager
.start(
Url::parse(&format!("https://example.com/{index}.zip")).unwrap(),
PathBuf::from(format!("{index}.zip")),
None,
None,
)
.unwrap();
}
assert_eq!(manager.records().count(), 2);
assert!(manager.get(DownloadId(1)).is_none());
assert!(manager.get(DownloadId(4)).is_some());
}
#[test]
fn clamps_tampered_download_record_limit() {
let mut manager = DownloadManager {
max_records: DEFAULT_MAX_DOWNLOAD_RECORDS + 10_000,
..DownloadManager::default()
};
manager.enforce_limits();
assert_eq!(manager.max_records(), DEFAULT_MAX_DOWNLOAD_RECORDS);
manager.max_records = 0;
manager.enforce_limits();
assert_eq!(manager.max_records(), 1);
}
#[test]
fn load_marks_stale_in_progress_downloads_interrupted() {
let directory = tempfile::tempdir().unwrap();
let store = DownloadStore::new(directory.path().join("downloads.json"));
let mut manager = DownloadManager::default();
let id = manager
.start(
Url::parse("https://example.com/file.zip").unwrap(),
PathBuf::from("file.zip"),
None,
None,
)
.unwrap();
store.save(&manager).unwrap();
let loaded = store.load().unwrap();
let record = loaded.get(id).unwrap();
assert_eq!(
record.status,
DownloadStatus::Failed("interrupted".to_string())
);
assert!(record.finished_at.is_some());
}
#[test]
fn rejects_new_downloads_over_active_cap() {
let mut manager = DownloadManager::default();
for index in 0..MAX_ACTIVE_DOWNLOADS {
manager
.start(
Url::parse(&format!("https://example.com/{index}.zip")).unwrap(),
PathBuf::from(format!("{index}.zip")),
None,
None,
)
.unwrap();
}
assert_eq!(manager.active_count(), MAX_ACTIVE_DOWNLOADS);
assert!(
manager
.start(
Url::parse("https://example.com/overflow.zip").unwrap(),
PathBuf::from("overflow.zip"),
None,
None,
)
.is_err()
);
}
#[test]
fn normalizes_tampered_next_download_id() {
let mut manager = DownloadManager {
next_id: u64::MAX,
..DownloadManager::default()
};
let id = manager
.start(
Url::parse("https://example.com/file.zip").unwrap(),
PathBuf::from("file.zip"),
None,
None,
)
.unwrap();
assert_eq!(id, DownloadId(1));
}
#[test]
fn load_rejects_oversized_download_file_before_parse() {
let directory = tempfile::tempdir().unwrap();
let path = directory.path().join("downloads.json");
fs::write(&path, vec![b' '; (MAX_DOWNLOADS_FILE_BYTES + 1) as usize]).unwrap();
let store = DownloadStore::new(path);
assert!(store.load().is_err());
}
}