harmony/harmony_assets/src/hash.rs

use crate::errors::AssetError;
use std::path::Path;

#[cfg(feature = "blake3")]
use blake3::Hasher as B3Hasher;
#[cfg(feature = "sha256")]
use sha2::{Digest, Sha256};

#[derive(Debug, Clone)]
pub enum ChecksumAlgo {
    BLAKE3,
    SHA256,
}

impl Default for ChecksumAlgo {
    fn default() -> Self {
        #[cfg(feature = "blake3")]
        return ChecksumAlgo::BLAKE3;
        #[cfg(not(feature = "blake3"))]
        return ChecksumAlgo::SHA256;
    }
}

impl ChecksumAlgo {
    pub fn name(&self) -> &'static str {
        match self {
            ChecksumAlgo::BLAKE3 => "blake3",
            ChecksumAlgo::SHA256 => "sha256",
        }
    }
}

impl std::str::FromStr for ChecksumAlgo {
    type Err = AssetError;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match s.to_lowercase().as_str() {
            "blake3" | "b3" => Ok(ChecksumAlgo::BLAKE3),
            "sha256" | "sha-256" => Ok(ChecksumAlgo::SHA256),
            _ => Err(AssetError::ChecksumAlgoNotAvailable(s.to_string())),
        }
    }
}

impl std::fmt::Display for ChecksumAlgo {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{}", self.name())
    }
}

pub async fn checksum_for_file<R>(reader: R, algo: ChecksumAlgo) -> Result<String, AssetError>
where
    R: tokio::io::AsyncRead + Unpin,
{
    match algo {
        #[cfg(feature = "blake3")]
        ChecksumAlgo::BLAKE3 => {
            let mut hasher = B3Hasher::new();
            let mut reader = reader;
            let mut buf = vec![0u8; 65536];
            loop {
                let n = tokio::io::AsyncReadExt::read(&mut reader, &mut buf).await?;
                if n == 0 {
                    break;
                }
                hasher.update(&buf[..n]);
            }
            Ok(hasher.finalize().to_hex().to_string())
        }
        #[cfg(not(feature = "blake3"))]
        ChecksumAlgo::BLAKE3 => Err(AssetError::ChecksumAlgoNotAvailable("blake3".to_string())),
        #[cfg(feature = "sha256")]
        ChecksumAlgo::SHA256 => {
            let mut hasher = Sha256::new();
            let mut reader = reader;
            let mut buf = vec![0u8; 65536];
            loop {
                let n = tokio::io::AsyncReadExt::read(&mut reader, &mut buf).await?;
                if n == 0 {
                    break;
                }
                hasher.update(&buf[..n]);
            }
            Ok(format!("{:x}", hasher.finalize()))
        }
        #[cfg(not(feature = "sha256"))]
        ChecksumAlgo::SHA256 => Err(AssetError::ChecksumAlgoNotAvailable("sha256".to_string())),
    }
}

pub async fn checksum_for_path(path: &Path, algo: ChecksumAlgo) -> Result<String, AssetError> {
    let file = tokio::fs::File::open(path)
        .await
        .map_err(AssetError::IoError)?;
    let reader = tokio::io::BufReader::with_capacity(65536, file);
    checksum_for_file(reader, algo).await
}

pub async fn checksum_for_path_with_progress<F>(
    path: &Path,
    algo: ChecksumAlgo,
    mut progress: F,
) -> Result<String, AssetError>
where
    F: FnMut(u64, Option<u64>) + Send,
{
    let file = tokio::fs::File::open(path)
        .await
        .map_err(AssetError::IoError)?;
    let metadata = file.metadata().await.map_err(AssetError::IoError)?;
    let total = Some(metadata.len());
    let reader = tokio::io::BufReader::with_capacity(65536, file);

    match algo {
        #[cfg(feature = "blake3")]
        ChecksumAlgo::BLAKE3 => {
            let mut hasher = B3Hasher::new();
            let mut reader = reader;
            let mut buf = vec![0u8; 65536];
            let mut read: u64 = 0;
            loop {
                let n = tokio::io::AsyncReadExt::read(&mut reader, &mut buf).await?;
                if n == 0 {
                    break;
                }
                hasher.update(&buf[..n]);
                read += n as u64;
                progress(read, total);
            }
            Ok(hasher.finalize().to_hex().to_string())
        }
        #[cfg(not(feature = "blake3"))]
        ChecksumAlgo::BLAKE3 => Err(AssetError::ChecksumAlgoNotAvailable("blake3".to_string())),
        #[cfg(feature = "sha256")]
        ChecksumAlgo::SHA256 => {
            let mut hasher = Sha256::new();
            let mut reader = reader;
            let mut buf = vec![0u8; 65536];
            let mut read: u64 = 0;
            loop {
                let n = tokio::io::AsyncReadExt::read(&mut reader, &mut buf).await?;
                if n == 0 {
                    break;
                }
                hasher.update(&buf[..n]);
                read += n as u64;
                progress(read, total);
            }
            Ok(format!("{:x}", hasher.finalize()))
        }
        #[cfg(not(feature = "sha256"))]
        ChecksumAlgo::SHA256 => Err(AssetError::ChecksumAlgoNotAvailable("sha256".to_string())),
    }
}

pub async fn verify_checksum(
    path: &Path,
    expected: &str,
    algo: ChecksumAlgo,
) -> Result<(), AssetError> {
    let actual = checksum_for_path(path, algo).await?;
    let expected_clean = expected
        .trim_start_matches("blake3:")
        .trim_start_matches("sha256:")
        .trim_start_matches("b3:")
        .trim_start_matches("sha-256:");
    if actual != expected_clean {
        return Err(AssetError::ChecksumMismatch {
            path: path.to_path_buf(),
            expected: expected_clean.to_string(),
            actual,
        });
    }
    Ok(())
}

pub fn format_checksum(checksum: &str, algo: ChecksumAlgo) -> String {
    format!("{}:{}", algo.name(), checksum)
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::io::Write;
    use tempfile::NamedTempFile;

    async fn create_temp_file(content: &[u8]) -> NamedTempFile {
        let mut file = NamedTempFile::new().unwrap();
        file.write_all(content).unwrap();
        file.flush().unwrap();
        file
    }

    #[tokio::test]
    async fn test_checksum_blake3() {
        let file = create_temp_file(b"hello world").await;
        let checksum = checksum_for_path(file.path(), ChecksumAlgo::BLAKE3)
            .await
            .unwrap();
        assert_eq!(
            checksum,
            "d74981efa70a0c880b8d8c1985d075dbcbf679b99a5f9914e5aaf96b831a9e24"
        );
    }

    #[tokio::test]
    async fn test_verify_checksum_success() {
        let file = create_temp_file(b"hello world").await;
        let checksum = checksum_for_path(file.path(), ChecksumAlgo::BLAKE3)
            .await
            .unwrap();
        let result = verify_checksum(file.path(), &checksum, ChecksumAlgo::BLAKE3).await;
        assert!(result.is_ok());
    }

    #[tokio::test]
    async fn test_verify_checksum_failure() {
        let file = create_temp_file(b"hello world").await;
        let result = verify_checksum(
            file.path(),
            "blake3:0000000000000000000000000000000000000000000000000000000000000000",
            ChecksumAlgo::BLAKE3,
        )
        .await;
        assert!(matches!(result, Err(AssetError::ChecksumMismatch { .. })));
    }

    #[tokio::test]
    async fn test_checksum_with_prefix() {
        let file = create_temp_file(b"hello world").await;
        let checksum = checksum_for_path(file.path(), ChecksumAlgo::BLAKE3)
            .await
            .unwrap();
        let formatted = format_checksum(&checksum, ChecksumAlgo::BLAKE3);
        assert!(formatted.starts_with("blake3:"));
    }
}