use dashmap::DashMap;
use ed25519_dalek::{Signature, Verifier, VerifyingKey};
use sha2::{Digest, Sha256};
use std::convert::TryInto;
use std::sync::OnceLock;
fn verified_cache() -> &'static DashMap<String, String> {
static CACHE: OnceLock<DashMap<String, String>> = OnceLock::new();
CACHE.get_or_init(DashMap::new)
}
#[derive(Debug, Clone)]
pub struct IntegrityViolationError {
pub message: String,
}
impl std::fmt::Display for IntegrityViolationError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "IntegrityViolation: {}", self.message)
}
}
impl std::error::Error for IntegrityViolationError {}
pub fn check_cache(module_name: &str, expected_hash: &str) -> bool {
verified_cache()
.get(module_name)
.map_or(false, |v| v.value() == expected_hash)
}
pub fn update_cache(module_name: &str, expected_hash: &str) {
verified_cache().insert(module_name.to_string(), expected_hash.to_string());
}
pub fn clear_cache() {
verified_cache().clear();
}
pub fn verify_module_integrity(source_code: &str, expected_sha256: &str) -> bool {
let normalized = source_code.replace("\r\n", "\n");
let mut hasher = Sha256::new();
hasher.update(normalized.as_bytes());
let actual_hash = format!("{:x}", hasher.finalize());
actual_hash == expected_sha256
}
pub fn verify_file_hash(file_path: &str, expected_sha256: &str) -> Result<bool, String> {
use memmap2::Mmap;
use std::fs::File;
let file = File::open(file_path).map_err(|e| format!("Failed to open file: {}", e))?;
let mmap = unsafe { Mmap::map(&file).map_err(|e| format!("Failed to mmap file: {}", e))? };
let mut normalized = Vec::with_capacity(mmap.len());
let mut i = 0;
while i < mmap.len() {
if i + 1 < mmap.len() && mmap[i] == b'\r' && mmap[i + 1] == b'\n' {
normalized.push(b'\n');
i += 2;
} else {
normalized.push(mmap[i]);
i += 1;
}
}
let mut hasher = Sha256::new();
hasher.update(&normalized);
let actual_hash = format!("{:x}", hasher.finalize());
Ok(actual_hash == expected_sha256)
}
pub fn verify_signature(
message: &str,
signature_hex: &str,
public_key_hex: &str,
) -> Result<bool, String> {
let public_bytes = hex::decode(public_key_hex).map_err(|e| e.to_string())?;
let signature_bytes = hex::decode(signature_hex).map_err(|e| e.to_string())?;
let public_array: [u8; 32] = public_bytes
.as_slice()
.try_into()
.map_err(|_| "Public key must be 32 bytes".to_string())?;
let signature_array: [u8; 64] = signature_bytes
.as_slice()
.try_into()
.map_err(|_| "Signature must be 64 bytes".to_string())?;
let verifying_key = VerifyingKey::from_bytes(&public_array)
.map_err(|e| format!("Invalid public key: {}", e))?;
let signature = Signature::from_bytes(&signature_array);
Ok(verifying_key.verify(message.as_bytes(), &signature).is_ok())
}
pub fn assert_module_integrity(
module_name: &str,
source_code: &str,
expected_hash: &str,
) -> Result<(), IntegrityViolationError> {
if check_cache(module_name, expected_hash) {
return Ok(());
}
if !verify_module_integrity(source_code, expected_hash) {
let normalized = source_code.replace("\r\n", "\n");
let mut hasher = Sha256::new();
hasher.update(normalized.as_bytes());
let actual_hash = format!("{:x}", hasher.finalize());
return Err(IntegrityViolationError {
message: format!(
"CRITICAL INTEGRITY VIOLATION: Module {} has been tampered with! \
Expected hash {}... but got {}...",
module_name,
&expected_hash[..expected_hash.len().min(8)],
&actual_hash[..actual_hash.len().min(8)],
),
});
}
update_cache(module_name, expected_hash);
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_verify_module_integrity() {
let source = "fn main() { println!(\"hello\"); }";
let mut hasher = Sha256::new();
hasher.update(source.as_bytes());
let expected = format!("{:x}", hasher.finalize());
assert!(verify_module_integrity(source, &expected));
assert!(!verify_module_integrity(source, "badhash"));
}
#[test]
fn test_verify_module_normalizes_crlf() {
let source_lf = "line1\nline2";
let source_crlf = "line1\r\nline2";
let mut hasher = Sha256::new();
hasher.update(source_lf.as_bytes());
let expected = format!("{:x}", hasher.finalize());
assert!(verify_module_integrity(source_lf, &expected));
assert!(verify_module_integrity(source_crlf, &expected));
}
#[test]
fn test_cache_operations() {
let test_key = "test_mod_ops_unique";
update_cache(test_key, "hash123");
assert!(check_cache(test_key, "hash123"));
assert!(!check_cache(test_key, "wrong_hash"));
}
#[test]
fn test_assert_module_integrity_valid() {
let test_key = "test_mod_valid_unique";
let source = "test source code";
let mut hasher = Sha256::new();
hasher.update(source.as_bytes());
let expected = format!("{:x}", hasher.finalize());
assert!(assert_module_integrity(test_key, source, &expected).is_ok());
assert!(check_cache(test_key, &expected));
}
#[test]
fn test_assert_module_integrity_tampered() {
let test_key = "test_mod_tampered_unique";
let source = "original source";
let result = assert_module_integrity(test_key, source, "badhash0000");
assert!(result.is_err());
assert!(result
.unwrap_err()
.message
.contains("CRITICAL INTEGRITY VIOLATION"));
}
#[test]
fn test_assert_module_integrity_cached() {
let unique_key = format!(
"cached_mod_{}_{}",
std::process::id(),
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap()
.as_nanos()
);
update_cache(&unique_key, "known_hash");
assert!(check_cache(&unique_key, "known_hash"));
assert!(assert_module_integrity(&unique_key, "doesn't matter", "known_hash").is_ok());
}
#[test]
fn test_hex_decode() {
assert_eq!(hex::decode("48656c6c6f").unwrap(), b"Hello");
assert!(hex::decode("abc").is_err()); assert!(hex::decode("zz").is_err()); }
}