use futures_util::StreamExt; use log::{debug, info, warn}; use sha2::{Digest, Sha256}; use std::io::Read; use std::path::PathBuf; use tokio::fs; use tokio::fs::File; use tokio::io::AsyncWriteExt; use url::Url; const CHECKSUM_FAILED_MSG: &str = "Downloaded file failed checksum verification"; /// Represents an asset that can be downloaded from a URL with checksum verification. /// /// This struct facilitates secure downloading of files from remote URLs by /// verifying the integrity of the downloaded content using SHA-256 checksums. /// It handles downloading the file, saving it to disk, and verifying the checksum matches /// the expected value. /// /// # Examples /// /// ```compile_fail /// # use url::Url; /// # use std::path::PathBuf; /// /// # async fn example() -> Result<(), String> { /// let asset = DownloadableAsset { /// url: Url::parse("https://example.com/file.zip").unwrap(), /// file_name: "file.zip".to_string(), /// checksum: "a1b2c3d4e5f6...".to_string(), /// }; /// /// let download_dir = PathBuf::from("/tmp/downloads"); /// let file_path = asset.download_to_path(download_dir).await?; /// # Ok(()) /// # } /// ``` #[derive(Debug)] pub(crate) struct DownloadableAsset { pub(crate) url: Url, pub(crate) file_name: String, pub(crate) checksum: String, } impl DownloadableAsset { fn verify_checksum(&self, file: PathBuf) -> bool { if !file.exists() { warn!("File does not exist: {:?}", file); return false; } let mut file = match std::fs::File::open(&file) { Ok(file) => file, Err(e) => { warn!("Failed to open file for checksum verification: {:?}", e); return false; } }; let mut hasher = Sha256::new(); let mut buffer = [0; 1024 * 1024]; // 1MB buffer loop { let bytes_read = match file.read(&mut buffer) { Ok(0) => break, Ok(n) => n, Err(e) => { warn!("Error reading file for checksum: {:?}", e); return false; } }; hasher.update(&buffer[..bytes_read]); } let result = hasher.finalize(); let calculated_hash = format!("{:x}", result); debug!("Expected checksum: {}", self.checksum); debug!("Calculated checksum: {}", calculated_hash); calculated_hash == self.checksum } /// Downloads the asset to the specified directory, verifying its checksum. /// /// This function will: /// 1. Create the target directory if it doesn't exist /// 2. Check if the file already exists with the correct checksum /// 3. If not, download the file from the URL /// 4. Verify the downloaded file's checksum matches the expected value /// /// # Arguments /// /// * `folder` - The directory path where the file should be saved /// /// # Returns /// /// * `Ok(PathBuf)` - The path to the downloaded file on success /// * `Err(String)` - A descriptive error message if the download or verification fails /// /// # Errors /// /// This function will return an error if: /// - The network request fails /// - The server responds with a non-success status code /// - Writing to disk fails /// - The checksum verification fails pub(crate) async fn download_to_path(&self, folder: PathBuf) -> Result { if !folder.exists() { fs::create_dir_all(&folder) .await .expect("Failed to create download directory"); } let target_file_path = folder.join(&self.file_name); debug!("Downloading to path: {:?}", target_file_path); if self.verify_checksum(target_file_path.clone()) { debug!("File already exists with correct checksum, skipping download"); return Ok(target_file_path); } debug!("Downloading from URL: {}", self.url); let client = reqwest::Client::new(); let response = client .get(self.url.clone()) .send() .await .map_err(|e| format!("Failed to download file: {e}"))?; if !response.status().is_success() { return Err(format!( "Failed to download file, status: {}", response.status() )); } let mut file = File::create(&target_file_path) .await .expect("Failed to create target file"); let mut stream = response.bytes_stream(); while let Some(chunk_result) = stream.next().await { let chunk = chunk_result.expect("Error while downloading file"); file.write_all(&chunk) .await .expect("Failed to write data to file"); } file.flush().await.expect("Failed to flush file"); drop(file); if !self.verify_checksum(target_file_path.clone()) { return Err(CHECKSUM_FAILED_MSG.to_string()); } info!( "File downloaded and verified successfully: {}", target_file_path.to_string_lossy() ); Ok(target_file_path) } } #[cfg(test)] mod tests { use super::*; use httptest::{ matchers::{self, request}, responders, Expectation, Server, }; const BASE_TEST_PATH: &str = "/tmp/harmony-test-k3d-download"; const TEST_CONTENT: &str = "This is a test file."; const TEST_CONTENT_HASH: &str = "f29bc64a9d3732b4b9035125fdb3285f5b6455778edca72414671e0ca3b2e0de"; fn setup_test() -> (PathBuf, Server) { let _ = env_logger::builder().try_init(); // Create unique test directory let test_id = std::time::SystemTime::now() .duration_since(std::time::UNIX_EPOCH) .unwrap() .as_millis(); let download_path = format!("{}/test_{}", BASE_TEST_PATH, test_id); std::fs::create_dir_all(&download_path).unwrap(); (PathBuf::from(download_path), Server::run()) } #[tokio::test] async fn test_download_to_path_success() { let (folder, server) = setup_test(); server.expect( Expectation::matching(request::method_path("GET", "/test.txt")) .respond_with(responders::status_code(200).body(TEST_CONTENT)), ); let asset = DownloadableAsset { url: Url::parse(&server.url("/test.txt").to_string()).unwrap(), file_name: "test.txt".to_string(), checksum: TEST_CONTENT_HASH.to_string(), }; let result = asset .download_to_path(folder.join("success")) .await .unwrap(); let downloaded_content = std::fs::read_to_string(result).unwrap(); assert_eq!(downloaded_content, TEST_CONTENT); } #[tokio::test] async fn test_download_to_path_already_exists() { let (folder, server) = setup_test(); server.expect( Expectation::matching(matchers::any()) .times(0) .respond_with(responders::status_code(200).body(TEST_CONTENT)), ); let asset = DownloadableAsset { url: Url::parse(&server.url("/test.txt").to_string()).unwrap(), file_name: "test.txt".to_string(), checksum: TEST_CONTENT_HASH.to_string(), }; let target_file_path = folder.join(&asset.file_name); std::fs::write(&target_file_path, TEST_CONTENT).unwrap(); let result = asset.download_to_path(folder).await.unwrap(); let content = std::fs::read_to_string(result).unwrap(); assert_eq!(content, TEST_CONTENT); } #[tokio::test] async fn test_download_to_path_server_error() { let (folder, server) = setup_test(); server.expect( Expectation::matching(matchers::any()).respond_with(responders::status_code(404)), ); let asset = DownloadableAsset { url: Url::parse(&server.url("/test.txt").to_string()).unwrap(), file_name: "test.txt".to_string(), checksum: TEST_CONTENT_HASH.to_string(), }; let result = asset.download_to_path(folder.join("error")).await; assert!(result.is_err()); assert!(result.unwrap_err().contains("status: 404")); } #[tokio::test] async fn test_download_to_path_checksum_failure() { let (folder, server) = setup_test(); let invalid_content = "This is NOT the expected content"; server.expect( Expectation::matching(matchers::any()) .respond_with(responders::status_code(200).body(invalid_content)), ); let asset = DownloadableAsset { url: Url::parse(&server.url("/test.txt").to_string()).unwrap(), file_name: "test.txt".to_string(), checksum: TEST_CONTENT_HASH.to_string(), }; let join_handle = tokio::spawn(async move { asset.download_to_path(folder.join("failure")).await }); assert_eq!( join_handle.await.unwrap().err().unwrap(), CHECKSUM_FAILED_MSG ); } #[tokio::test] async fn test_download_with_specific_path_matcher() { let (folder, server) = setup_test(); server.expect( Expectation::matching(matchers::request::path("/specific/path.txt")) .respond_with(responders::status_code(200).body(TEST_CONTENT)), ); let asset = DownloadableAsset { url: Url::parse(&server.url("/specific/path.txt").to_string()).unwrap(), file_name: "path.txt".to_string(), checksum: TEST_CONTENT_HASH.to_string(), }; let result = asset.download_to_path(folder).await.unwrap(); let downloaded_content = std::fs::read_to_string(result).unwrap(); assert_eq!(downloaded_content, TEST_CONTENT); } }