harmony/k3d/src/downloadable_asset.rs

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<PathBuf, String> {
        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);
    }
}