use hkdf::Hkdf;
use hmac::{Hmac, Mac};
use sha2::Sha256;
use super::error::AuditError;
use super::key_store::SigningKeyStore;
use super::record::{self, PerQueryAudit};
type HmacSha256 = Hmac<Sha256>;
const HKDF_INFO: &[u8] = b"jammi-audit-search-v1";
pub fn ensure_master_key_present(store: &dyn SigningKeyStore) -> Result<(), AuditError> {
store.master_key().map(|_| ())
}
pub fn derive_tenant_secret(master: &[u8; 32], tenant_id: &str) -> Result<[u8; 32], AuditError> {
let hk = Hkdf::<Sha256>::new(Some(tenant_id.as_bytes()), master);
let mut secret = [0u8; 32];
hk.expand(HKDF_INFO, &mut secret)
.map_err(|e| AuditError::MasterKey(format!("hkdf expand failed: {e}")))?;
Ok(secret)
}
pub fn hmac_sign(canonical: &[u8], secret: &[u8; 32]) -> String {
let mut mac = <HmacSha256 as Mac>::new_from_slice(secret).expect("HMAC accepts any key length");
mac.update(canonical);
hex::encode(mac.finalize().into_bytes())
}
pub fn sign_record(
record: &mut PerQueryAudit,
store: &dyn SigningKeyStore,
) -> Result<(), AuditError> {
let tenant = record
.tenant_id
.clone()
.ok_or(AuditError::NoTenantBinding)?;
let master = store.master_key()?;
let secret = derive_tenant_secret(&master, &tenant)?;
let canonical = record::canonical_serialize(record)?;
record.signature = hmac_sign(&canonical, &secret);
Ok(())
}
pub fn verify(record: &PerQueryAudit, secret: &[u8; 32]) -> Result<(), AuditError> {
let canonical = record::canonical_serialize(record)?;
let expected = hmac_sign(&canonical, secret);
if !constant_time_eq(expected.as_bytes(), record.signature.as_bytes()) {
return Err(AuditError::SignatureMismatch(record.query_id));
}
Ok(())
}
pub fn verify_with_store(
record: &PerQueryAudit,
store: &dyn SigningKeyStore,
) -> Result<(), AuditError> {
let tenant = record
.tenant_id
.clone()
.ok_or(AuditError::NoTenantBinding)?;
let master = store.master_key()?;
let secret = derive_tenant_secret(&master, &tenant)?;
verify(record, &secret)
}
fn constant_time_eq(a: &[u8], b: &[u8]) -> bool {
if a.len() != b.len() {
return false;
}
let mut diff = 0u8;
for (x, y) in a.iter().zip(b.iter()) {
diff |= x ^ y;
}
diff == 0
}
#[cfg(test)]
mod tests {
use super::*;
use crate::audit::{EnvSigningKeyStore, MASTER_KEY_ENV};
use std::sync::Mutex;
use uuid::Uuid;
static ENV_LOCK: Mutex<()> = Mutex::new(());
const TEST_KEY: &str = "0000000000000000000000000000000000000000000000000000000000000001";
fn scoped() -> PerQueryAudit {
let mut r = PerQueryAudit::new(
Uuid::nil(),
"m",
"v",
serde_json::json!({ "k": 1 }),
vec!["a".into()],
vec![0.5],
)
.unwrap();
r.tenant_id = Some("tenant-a".into());
r
}
#[test]
fn sign_then_verify_roundtrips() {
let _g = ENV_LOCK.lock().unwrap();
std::env::set_var(MASTER_KEY_ENV, TEST_KEY);
let store = EnvSigningKeyStore;
let mut r = scoped();
sign_record(&mut r, &store).unwrap();
assert!(!r.signature.is_empty());
verify_with_store(&r, &store).unwrap();
}
#[test]
fn tampering_breaks_signature() {
let _g = ENV_LOCK.lock().unwrap();
std::env::set_var(MASTER_KEY_ENV, TEST_KEY);
let store = EnvSigningKeyStore;
let mut r = scoped();
sign_record(&mut r, &store).unwrap();
r.model_id = "tampered".into();
assert!(matches!(
verify_with_store(&r, &store),
Err(AuditError::SignatureMismatch(_))
));
}
#[test]
fn signing_is_deterministic_across_calls() {
let _g = ENV_LOCK.lock().unwrap();
std::env::set_var(MASTER_KEY_ENV, TEST_KEY);
let master = EnvSigningKeyStore.master_key().unwrap();
let secret = derive_tenant_secret(&master, "tenant-a").unwrap();
let r = scoped();
let canonical = record::canonical_serialize(&r).unwrap();
assert_eq!(
hmac_sign(&canonical, &secret),
hmac_sign(&canonical, &secret)
);
}
#[test]
fn different_tenants_get_different_secrets() {
let _g = ENV_LOCK.lock().unwrap();
std::env::set_var(MASTER_KEY_ENV, TEST_KEY);
let master = EnvSigningKeyStore.master_key().unwrap();
assert_ne!(
derive_tenant_secret(&master, "tenant-a").unwrap(),
derive_tenant_secret(&master, "tenant-b").unwrap()
);
}
#[test]
fn missing_master_key_makes_present_check_fail() {
let _g = ENV_LOCK.lock().unwrap();
std::env::remove_var(MASTER_KEY_ENV);
assert!(matches!(
ensure_master_key_present(&EnvSigningKeyStore),
Err(AuditError::MasterKey(_))
));
}
}