use base64::{engine::general_purpose, Engine as _}; use hmac::{Hmac, Mac}; use rand::{distributions::Alphanumeric, Rng}; use sha2::{Digest, Sha256}; use std::path::Path; use tokio::fs::File; use tokio::io::AsyncReadExt; use tokio_util::sync::CancellationToken; type HmacSha256 = Hmac; /// Generates random app token pub fn generate_app_token() -> String { rand::thread_rng() .sample_iter(&Alphanumeric) .take(32) .map(char::from) .collect() } /// Generate API key using HMAC-SHA256 pub fn generate_api_key(model_id: String, api_secret: String) -> Result { let mut mac = HmacSha256::new_from_slice(api_secret.as_bytes()) .map_err(|e| format!("Invalid key length: {}", e))?; mac.update(model_id.as_bytes()); let result = mac.finalize(); let code_bytes = result.into_bytes(); let hash = general_purpose::STANDARD.encode(code_bytes); Ok(hash) } /// Compute SHA256 hash of a file with cancellation support by chunking the file pub async fn compute_file_sha256_with_cancellation( file_path: &Path, cancel_token: &CancellationToken, ) -> Result { // Check for cancellation before starting if cancel_token.is_cancelled() { return Err("Hash computation cancelled".to_string()); } let mut file = File::open(file_path) .await .map_err(|e| format!("Failed to open file for hashing: {}", e))?; let mut hasher = Sha256::new(); let mut buffer = vec![0u8; 64 * 1024]; // 64KB chunks let mut total_read = 0u64; loop { // Check for cancellation every chunk (every 64KB) if cancel_token.is_cancelled() { return Err("Hash computation cancelled".to_string()); } let bytes_read = file .read(&mut buffer) .await .map_err(|e| format!("Failed to read file for hashing: {}", e))?; if bytes_read == 0 { break; // EOF } hasher.update(&buffer[..bytes_read]); total_read += bytes_read as u64; // Log progress for very large files (every 100MB) if total_read % (100 * 1024 * 1024) == 0 { #[cfg(feature = "logging")] log::debug!("Hash progress: {} MB processed", total_read / (1024 * 1024)); } } // Final cancellation check if cancel_token.is_cancelled() { return Err("Hash computation cancelled".to_string()); } let hash_bytes = hasher.finalize(); let hash_hex = format!("{:x}", hash_bytes); #[cfg(feature = "logging")] log::debug!("Hash computation completed for {} bytes", total_read); Ok(hash_hex) } #[cfg(test)] mod tests { use super::*; use tokio_util::sync::CancellationToken; #[test] fn test_generate_app_token() { let token1 = generate_app_token(); let token2 = generate_app_token(); // Should be 32 characters long assert_eq!(token1.len(), 32); assert_eq!(token2.len(), 32); // Should be different each time assert_ne!(token1, token2); // Should only contain alphanumeric characters assert!(token1.chars().all(|c| c.is_alphanumeric())); assert!(token2.chars().all(|c| c.is_alphanumeric())); } #[test] fn test_generate_api_key() { let model_id = "test-model".to_string(); let api_secret = "test-secret".to_string(); let key1 = generate_api_key(model_id.clone(), api_secret.clone()).unwrap(); let key2 = generate_api_key(model_id.clone(), api_secret.clone()).unwrap(); // Should generate same key for same inputs assert_eq!(key1, key2); // Should be base64 encoded (and thus contain base64 characters) assert!(key1.chars().all(|c| c.is_alphanumeric() || c == '+' || c == '/' || c == '=')); // Different model_id should produce different key let different_key = generate_api_key("different-model".to_string(), api_secret).unwrap(); assert_ne!(key1, different_key); } #[test] fn test_generate_api_key_empty_inputs() { let result = generate_api_key("".to_string(), "secret".to_string()); assert!(result.is_ok()); // Should still work with empty model_id let result = generate_api_key("model".to_string(), "".to_string()); assert!(result.is_ok()); // Should still work with empty secret } #[tokio::test] async fn test_compute_file_sha256_with_cancellation() { use std::io::Write; use tempfile::NamedTempFile; // Create a temporary file with known content let mut temp_file = NamedTempFile::new().unwrap(); let test_content = b"Hello, World!"; temp_file.write_all(test_content).unwrap(); temp_file.flush().unwrap(); let token = CancellationToken::new(); // Compute hash of the file let hash = compute_file_sha256_with_cancellation(temp_file.path(), &token).await.unwrap(); // Verify it's a valid hex string assert_eq!(hash.len(), 64); // SHA256 is 256 bits = 64 hex chars assert!(hash.chars().all(|c| c.is_ascii_hexdigit())); // Verify it matches expected SHA256 of "Hello, World!" let expected = "dffd6021bb2bd5b0af676290809ec3a53191dd81c7f70a4b28688a362182986f"; assert_eq!(hash, expected); } #[tokio::test] async fn test_compute_file_sha256_cancellation() { use std::io::Write; use tempfile::NamedTempFile; // Create a temporary file let mut temp_file = NamedTempFile::new().unwrap(); temp_file.write_all(b"test content").unwrap(); temp_file.flush().unwrap(); let token = CancellationToken::new(); token.cancel(); // Cancel immediately // Should return cancellation error let result = compute_file_sha256_with_cancellation(temp_file.path(), &token).await; assert!(result.is_err()); assert!(result.unwrap_err().contains("cancelled")); } #[tokio::test] async fn test_compute_file_sha256_nonexistent_file() { let token = CancellationToken::new(); let nonexistent_path = Path::new("/nonexistent/file.txt"); let result = compute_file_sha256_with_cancellation(nonexistent_path, &token).await; assert!(result.is_err()); assert!(result.unwrap_err().contains("Failed to open file for hashing")); } }