use crate::browser::GithubRelease; use crate::events; use crate::profile::manager::ProfileManager; use reqwest::Client; use serde::{Deserialize, Serialize}; use std::path::PathBuf; use std::sync::atomic::{AtomicBool, Ordering}; use tokio::fs; use tokio::io::AsyncWriteExt; const MMDB_REPO: &str = "P3TERX/GeoLite.mmdb"; #[derive(Debug, Clone, Serialize, Deserialize)] pub struct GeoIPDownloadProgress { pub stage: String, // "downloading", "extracting", "completed" pub percentage: f64, pub message: String, // Extra fields to mirror browser download progress payload pub downloaded_bytes: Option, pub total_bytes: Option, pub speed_bytes_per_sec: Option, pub eta_seconds: Option, } static DOWNLOAD_IN_PROGRESS: AtomicBool = AtomicBool::new(false); pub struct GeoIPDownloader { client: Client, } impl GeoIPDownloader { fn new() -> Self { Self { client: Client::new(), } } pub fn instance() -> &'static GeoIPDownloader { &GEOIP_DOWNLOADER } /// Create a new downloader with custom client (for testing) #[cfg(test)] pub fn new_with_client(client: Client) -> Self { Self { client } } fn get_cache_dir() -> PathBuf { crate::app_dirs::cache_dir() } pub fn get_mmdb_file_path() -> Result> { Ok(Self::get_cache_dir().join("GeoLite2-City.mmdb")) } pub fn is_geoip_database_available() -> bool { if let Ok(mmdb_path) = Self::get_mmdb_file_path() { mmdb_path.exists() } else { false } } fn get_timestamp_path() -> PathBuf { crate::app_dirs::cache_dir().join("geoip_last_download") } fn is_geoip_stale() -> bool { let timestamp_path = Self::get_timestamp_path(); let Ok(content) = std::fs::read_to_string(×tamp_path) else { return true; }; let Ok(timestamp) = content.trim().parse::() else { return true; }; let now = std::time::SystemTime::now() .duration_since(std::time::UNIX_EPOCH) .unwrap_or_default() .as_secs(); const SEVEN_DAYS: u64 = 7 * 24 * 60 * 60; now.saturating_sub(timestamp) > SEVEN_DAYS } /// Check if GeoIP database is missing or stale for Wayfern fingerprint geo. pub fn check_missing_geoip_database( &self, ) -> Result> { let profiles = ProfileManager::instance() .list_profiles() .map_err(|e| format!("Failed to list profiles: {e}"))?; let needs_geoip = profiles.iter().any(|profile| profile.browser == "wayfern"); if needs_geoip { return Ok(!Self::is_geoip_database_available() || Self::is_geoip_stale()); } Ok(false) } fn find_city_mmdb_asset(&self, release: &GithubRelease) -> Option { for asset in &release.assets { if asset.name.ends_with("-City.mmdb") { return Some(asset.browser_download_url.clone()); } } None } pub async fn download_geoip_database( &self, _app_handle: &tauri::AppHandle, ) -> Result<(), Box> { if DOWNLOAD_IN_PROGRESS .compare_exchange(false, true, Ordering::SeqCst, Ordering::SeqCst) .is_err() { log::info!("GeoIP database download already in progress, skipping"); return Ok(()); } let result = self.download_geoip_database_inner(_app_handle).await; DOWNLOAD_IN_PROGRESS.store(false, Ordering::SeqCst); result } async fn download_geoip_database_inner( &self, _app_handle: &tauri::AppHandle, ) -> Result<(), Box> { // Emit initial progress let _ = events::emit( "geoip-download-progress", GeoIPDownloadProgress { stage: "downloading".to_string(), percentage: 0.0, message: "Starting GeoIP database download".to_string(), downloaded_bytes: Some(0), total_bytes: None, speed_bytes_per_sec: Some(0.0), eta_seconds: None, }, ); #[cfg(feature = "e2e")] let fixture_url = std::env::var("DONUT_E2E_GEOIP_DOWNLOAD_URL") .ok() .filter(|url| !url.is_empty()); #[cfg(not(feature = "e2e"))] let fixture_url: Option = None; let download_url = if let Some(url) = fixture_url { url } else { let releases = self.fetch_geoip_releases().await?; let latest_release = releases.first().ok_or("No GeoIP database releases found")?; self .find_city_mmdb_asset(latest_release) .ok_or("No compatible GeoIP database asset found")? }; // Create cache directory let cache_dir = Self::get_cache_dir(); fs::create_dir_all(&cache_dir).await?; let mmdb_path = Self::get_mmdb_file_path()?; // Always download to a temp file first, then atomically rename. // This prevents corruption if the app is closed mid-download. let temp_path = mmdb_path.with_extension("mmdb.downloading"); // Remove any leftover temp file from a previous interrupted download let _ = fs::remove_file(&temp_path).await; // Download the file let response = self.client.get(&download_url).send().await?; if !response.status().is_success() { return Err( format!( "Failed to download GeoIP database: HTTP {}", response.status() ) .into(), ); } let total_size = response.content_length().unwrap_or(0); let mut downloaded: u64 = 0; let mut file = fs::File::create(&temp_path).await?; let mut stream = response.bytes_stream(); use futures_util::StreamExt; use std::time::Instant; let start_time = Instant::now(); let mut last_update = Instant::now(); while let Some(chunk) = stream.next().await { let chunk = chunk?; downloaded += chunk.len() as u64; file.write_all(&chunk).await?; let now = Instant::now(); if now.duration_since(last_update).as_millis() >= 100 { let elapsed = start_time.elapsed().as_secs_f64(); let speed = if elapsed > 0.0 { downloaded as f64 / elapsed } else { 0.0 }; let percentage = if total_size > 0 { (downloaded as f64 / total_size as f64) * 100.0 } else { 0.0 }; let eta = if speed > 0.0 && total_size > 0 { Some((total_size.saturating_sub(downloaded)) as f64 / speed) } else { None }; let _ = events::emit( "geoip-download-progress", GeoIPDownloadProgress { stage: "downloading".to_string(), percentage, message: format!("Downloaded {downloaded} / {total_size} bytes"), downloaded_bytes: Some(downloaded), total_bytes: Some(total_size), speed_bytes_per_sec: Some(speed), eta_seconds: eta, }, ); last_update = now; } } file.flush().await?; drop(file); // Atomically replace the old database with the new one fs::rename(&temp_path, &mmdb_path).await?; // Write download timestamp let timestamp_path = Self::get_timestamp_path(); if let Some(parent) = timestamp_path.parent() { let _ = fs::create_dir_all(parent).await; } let now = std::time::SystemTime::now() .duration_since(std::time::UNIX_EPOCH) .unwrap_or_default() .as_secs(); let _ = fs::write(×tamp_path, now.to_string()).await; // Emit completion let _ = events::emit( "geoip-download-progress", GeoIPDownloadProgress { stage: "completed".to_string(), percentage: 100.0, message: "GeoIP database download completed".to_string(), downloaded_bytes: Some(downloaded), total_bytes: Some(total_size), speed_bytes_per_sec: Some(0.0), eta_seconds: Some(0.0), }, ); Ok(()) } async fn fetch_geoip_releases( &self, ) -> Result, Box> { let url = format!("https://api.github.com/repos/{MMDB_REPO}/releases"); let response = self .client .get(&url) .header("User-Agent", "Mozilla/5.0 (compatible; donutbrowser)") .send() .await?; if !response.status().is_success() { return Err(format!("Failed to fetch releases: HTTP {}", response.status()).into()); } let releases: Vec = response.json().await?; Ok(releases) } } #[tauri::command] pub fn check_missing_geoip_database() -> Result { let geoip_downloader = GeoIPDownloader::instance(); geoip_downloader .check_missing_geoip_database() .map_err(|e| format!("Failed to check missing GeoIP database: {e}")) } // Global singleton instance lazy_static::lazy_static! { static ref GEOIP_DOWNLOADER: GeoIPDownloader = GeoIPDownloader::new(); } #[cfg(test)] mod tests { use super::*; use crate::browser::GithubRelease; use wiremock::matchers::{method, path}; use wiremock::{Mock, MockServer, ResponseTemplate}; fn create_mock_release() -> GithubRelease { GithubRelease { tag_name: "v1.0.0".to_string(), name: "Test Release".to_string(), body: Some("Test release body".to_string()), published_at: "2023-01-01T00:00:00Z".to_string(), created_at: Some("2023-01-01T00:00:00Z".to_string()), html_url: Some("https://example.com/release".to_string()), tarball_url: Some("https://example.com/tarball".to_string()), zipball_url: Some("https://example.com/zipball".to_string()), draft: false, prerelease: false, is_nightly: false, id: Some(1), node_id: Some("test_node_id".to_string()), target_commitish: None, assets: vec![crate::browser::GithubAsset { id: Some(1), node_id: Some("test_asset_node_id".to_string()), name: "GeoLite2-City.mmdb".to_string(), label: None, content_type: Some("application/octet-stream".to_string()), state: Some("uploaded".to_string()), size: 1024, download_count: Some(0), created_at: Some("2023-01-01T00:00:00Z".to_string()), updated_at: Some("2023-01-01T00:00:00Z".to_string()), browser_download_url: "https://example.com/GeoLite2-City.mmdb".to_string(), }], } } #[tokio::test] async fn test_fetch_geoip_releases_success() { let mock_server = MockServer::start().await; let releases = vec![create_mock_release()]; Mock::given(method("GET")) .and(path(format!("/repos/{MMDB_REPO}/releases"))) .respond_with(ResponseTemplate::new(200).set_body_json(&releases)) .mount(&mock_server) .await; let client = Client::builder() .build() .expect("Failed to create HTTP client"); let downloader = GeoIPDownloader::new_with_client(client); // Override the URL for testing let url = format!("{}/repos/{}/releases", mock_server.uri(), MMDB_REPO); let response = downloader .client .get(&url) .header("User-Agent", "Mozilla/5.0 (compatible; donutbrowser)") .send() .await .expect("Request should succeed"); assert!(response.status().is_success()); let fetched_releases: Vec = response.json().await.expect("Should parse JSON"); assert_eq!(fetched_releases.len(), 1); assert_eq!(fetched_releases[0].tag_name, "v1.0.0"); } #[tokio::test] async fn test_find_city_mmdb_asset() { let downloader = GeoIPDownloader::new(); let release = create_mock_release(); let asset_url = downloader.find_city_mmdb_asset(&release); assert!(asset_url.is_some()); assert_eq!(asset_url.unwrap(), "https://example.com/GeoLite2-City.mmdb"); } #[tokio::test] async fn test_find_city_mmdb_asset_not_found() { let downloader = GeoIPDownloader::new(); let mut release = create_mock_release(); release.assets[0].name = "wrong-file.txt".to_string(); let asset_url = downloader.find_city_mmdb_asset(&release); assert!(asset_url.is_none()); } #[test] fn test_get_cache_dir() { let cache_dir = GeoIPDownloader::get_cache_dir(); assert!(!cache_dir.as_os_str().is_empty()); } #[test] fn test_get_mmdb_file_path() { let mmdb_path = GeoIPDownloader::get_mmdb_file_path(); assert!(mmdb_path.is_ok()); let path = mmdb_path.unwrap(); assert!(path.to_string_lossy().ends_with("GeoLite2-City.mmdb")); } #[test] fn test_is_geoip_stale() { let tmp = tempfile::TempDir::new().unwrap(); let _guard = crate::app_dirs::set_test_cache_dir(tmp.path().to_path_buf()); // No timestamp file → stale assert!(GeoIPDownloader::is_geoip_stale()); let timestamp_path = GeoIPDownloader::get_timestamp_path(); std::fs::create_dir_all(timestamp_path.parent().unwrap()).unwrap(); // Recent timestamp → not stale let now = std::time::SystemTime::now() .duration_since(std::time::UNIX_EPOCH) .unwrap() .as_secs(); std::fs::write(×tamp_path, now.to_string()).unwrap(); assert!(!GeoIPDownloader::is_geoip_stale()); // 8 days ago → stale let eight_days_ago = now - 8 * 24 * 60 * 60; std::fs::write(×tamp_path, eight_days_ago.to_string()).unwrap(); assert!(GeoIPDownloader::is_geoip_stale()); } #[test] fn test_is_geoip_database_available() { // Test that the function works correctly regardless of file system state. // This test only verifies the function doesn't panic and returns a boolean. // We cannot reliably compare is_available to file existence due to potential // race conditions with other tests that modify environment variables. let _is_available = GeoIPDownloader::is_geoip_database_available(); // Verify the function logic by checking if the path resolution works let mmdb_path_result = GeoIPDownloader::get_mmdb_file_path(); assert!( mmdb_path_result.is_ok(), "Should be able to get MMDB file path" ); let mmdb_path = mmdb_path_result.unwrap(); assert!( mmdb_path.to_string_lossy().ends_with("GeoLite2-City.mmdb"), "Path should end with expected filename" ); } }