openlatch-client 0.1.12

The open-source security layer for AI agents — client forwarder
use std::path::Path;

use sha2::{Digest, Sha256};

use crate::error::{OlError, ERR_HMAC_KEY_UNAVAILABLE};

const KEYRING_SERVICE: &str = "ai.openlatch.client";
const KEYRING_USER: &str = "hmac/kid-01";
const KEY_LENGTH: usize = 32;

pub struct HmacKeyStore {
    openlatch_dir: std::path::PathBuf,
}

impl HmacKeyStore {
    pub fn new(openlatch_dir: &Path) -> Self {
        Self {
            openlatch_dir: openlatch_dir.to_path_buf(),
        }
    }

    pub fn load_or_create(&self) -> Result<Vec<u8>, OlError> {
        match self.load_from_keyring() {
            Ok(key) => return Ok(key),
            Err(e) => {
                tracing::debug!(error = %e, "keyring unavailable for HMAC key, trying file fallback");
            }
        }

        let fallback_path = self.openlatch_dir.join("hmac.key");
        if fallback_path.exists() {
            return self.load_from_file(&fallback_path);
        }

        let key = generate_key();

        if let Err(e) = self.store_to_keyring(&key) {
            tracing::debug!(error = %e, "cannot store HMAC key in keyring, using file fallback");
        }

        self.store_to_file(&fallback_path, &key)?;
        Ok(key)
    }

    fn load_from_keyring(&self) -> Result<Vec<u8>, OlError> {
        let entry = keyring::Entry::new(KEYRING_SERVICE, KEYRING_USER).map_err(|e| {
            OlError::new(
                ERR_HMAC_KEY_UNAVAILABLE,
                format!("Cannot create keyring entry: {e}"),
            )
        })?;

        let secret = entry.get_password().map_err(|e| {
            OlError::new(
                ERR_HMAC_KEY_UNAVAILABLE,
                format!("Cannot read HMAC key from keyring: {e}"),
            )
        })?;

        use base64::engine::general_purpose::URL_SAFE_NO_PAD;
        use base64::Engine;
        let key = URL_SAFE_NO_PAD.decode(&secret).map_err(|e| {
            OlError::new(
                ERR_HMAC_KEY_UNAVAILABLE,
                format!("HMAC key in keyring is not valid base64: {e}"),
            )
        })?;

        if key.len() != KEY_LENGTH {
            return Err(OlError::new(
                ERR_HMAC_KEY_UNAVAILABLE,
                format!(
                    "HMAC key in keyring has wrong length ({} bytes, expected {KEY_LENGTH})",
                    key.len()
                ),
            ));
        }

        Ok(key)
    }

    fn store_to_keyring(&self, key: &[u8]) -> Result<(), OlError> {
        use base64::engine::general_purpose::URL_SAFE_NO_PAD;
        use base64::Engine;

        let entry = keyring::Entry::new(KEYRING_SERVICE, KEYRING_USER).map_err(|e| {
            OlError::new(
                ERR_HMAC_KEY_UNAVAILABLE,
                format!("Cannot create keyring entry: {e}"),
            )
        })?;

        let encoded = URL_SAFE_NO_PAD.encode(key);
        entry.set_password(&encoded).map_err(|e| {
            OlError::new(
                ERR_HMAC_KEY_UNAVAILABLE,
                format!("Cannot store HMAC key in keyring: {e}"),
            )
        })?;

        Ok(())
    }

    fn load_from_file(&self, path: &Path) -> Result<Vec<u8>, OlError> {
        let content = std::fs::read(path).map_err(|e| {
            OlError::new(
                ERR_HMAC_KEY_UNAVAILABLE,
                format!("Cannot read HMAC key file: {e}"),
            )
        })?;

        if content.len() != KEY_LENGTH {
            return Err(OlError::new(
                ERR_HMAC_KEY_UNAVAILABLE,
                format!(
                    "HMAC key file has wrong length ({} bytes, expected {KEY_LENGTH})",
                    content.len()
                ),
            ));
        }

        Ok(content)
    }

    fn store_to_file(&self, path: &Path, key: &[u8]) -> Result<(), OlError> {
        if let Some(parent) = path.parent() {
            std::fs::create_dir_all(parent).map_err(|e| {
                OlError::new(
                    ERR_HMAC_KEY_UNAVAILABLE,
                    format!("Cannot create HMAC key directory: {e}"),
                )
            })?;
        }

        std::fs::write(path, key).map_err(|e| {
            OlError::new(
                ERR_HMAC_KEY_UNAVAILABLE,
                format!("Cannot write HMAC key file: {e}"),
            )
        })?;

        #[cfg(unix)]
        {
            use std::os::unix::fs::PermissionsExt;
            let perms = std::fs::Permissions::from_mode(0o600);
            std::fs::set_permissions(path, perms).map_err(|e| {
                OlError::new(
                    ERR_HMAC_KEY_UNAVAILABLE,
                    format!("Cannot set HMAC key file permissions: {e}"),
                )
            })?;
        }

        Ok(())
    }
}

fn generate_key() -> Vec<u8> {
    let a = uuid::Uuid::new_v4();
    let b = uuid::Uuid::new_v4();
    let mut key = Vec::with_capacity(KEY_LENGTH);
    key.extend_from_slice(a.as_bytes());
    key.extend_from_slice(b.as_bytes());
    key
}

pub fn key_fingerprint(key: &[u8]) -> String {
    let hash = Sha256::digest(key);
    hex::encode(&hash[..16])
}

pub(crate) mod hex {
    pub fn encode(bytes: &[u8]) -> String {
        let mut s = String::with_capacity(bytes.len() * 2);
        for b in bytes {
            s.push_str(&format!("{b:02x}"));
        }
        s
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn generated_key_is_32_bytes() {
        let key = generate_key();
        assert_eq!(key.len(), KEY_LENGTH);
    }

    #[test]
    fn key_fingerprint_is_32_hex_chars() {
        let key = vec![0x42; 32];
        let fp = key_fingerprint(&key);
        assert_eq!(fp.len(), 32);
        assert!(fp.chars().all(|c| c.is_ascii_hexdigit()));
    }

    #[test]
    fn different_keys_different_fingerprints() {
        let fp1 = key_fingerprint(&[0x42; 32]);
        let fp2 = key_fingerprint(&[0x43; 32]);
        assert_ne!(fp1, fp2);
    }

    #[test]
    fn file_fallback_roundtrip() {
        let dir = tempfile::tempdir().unwrap();
        let store = HmacKeyStore::new(dir.path());
        let path = dir.path().join("hmac.key");

        let key = generate_key();
        store.store_to_file(&path, &key).unwrap();
        let loaded = store.load_from_file(&path).unwrap();
        assert_eq!(key, loaded);
    }

    #[test]
    #[cfg(unix)]
    fn file_fallback_mode_0600() {
        use std::os::unix::fs::PermissionsExt;

        let dir = tempfile::tempdir().unwrap();
        let store = HmacKeyStore::new(dir.path());
        let path = dir.path().join("hmac.key");

        let key = generate_key();
        store.store_to_file(&path, &key).unwrap();
        let meta = std::fs::metadata(&path).unwrap();
        assert_eq!(meta.permissions().mode() & 0o777, 0o600);
    }

    #[test]
    fn rejects_wrong_length_file() {
        let dir = tempfile::tempdir().unwrap();
        let store = HmacKeyStore::new(dir.path());
        let path = dir.path().join("hmac.key");

        std::fs::write(&path, [0u8; 16]).unwrap();
        let err = store.load_from_file(&path).unwrap_err();
        assert_eq!(err.code, "OL-1900");
    }
}