shared-context-engineering 0.3.1

//! Encryption key management backed by an env-secret override or the OS credential store.
//!
//! Provides a single entry point to get-or-create a 64-character hex
//! encryption key sourced from `SCE_AUTH_DB_ENCRYPTION_KEY` when set, or
//! stored in the platform-native credential store (macOS Keychain,
//! Linux Secret Service via zbus, Windows Credential Store) otherwise.
//!
//! On first use for a given database name (when the database file does
//! not yet exist), a random 32-byte key is generated, hex-encoded,
//! persisted in the credential store, and returned. Subsequent calls
//! read the key from the credential store.

use std::path::Path;
use std::sync::{
    atomic::{AtomicBool, Ordering},
    OnceLock,
};
use std::thread;

use anyhow::{Context, Result};
use keyring_core::Entry;
use sha2::{Digest, Sha256};

/// Environment variable that overrides keyring-backed auth DB encryption.
pub const AUTH_DB_ENCRYPTION_KEY_ENV: &str = "SCE_AUTH_DB_ENCRYPTION_KEY";

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
enum CredentialStorePlatform {
    #[cfg(target_os = "linux")]
    Linux,
    #[cfg(target_os = "macos")]
    Macos,
    #[cfg(target_os = "windows")]
    Windows,
    #[cfg(not(any(target_os = "linux", target_os = "macos", target_os = "windows")))]
    Unsupported,
}

/// Guards the one-time registration of the platform-native credential store.
///
/// `OnceLock::get_or_try_init` is still unstable in the effective stable
/// toolchain, so a small atomic in-progress guard serializes fallible
/// initialization. The `OnceLock` is set only after success; on panic or
/// error, the in-progress guard is released and the next caller can retry.
static DEFAULT_STORE: OnceLock<bool> = OnceLock::new();
static DEFAULT_STORE_INITIALIZING: AtomicBool = AtomicBool::new(false);

struct DefaultStoreInitGuard;

impl Drop for DefaultStoreInitGuard {
    fn drop(&mut self) {
        DEFAULT_STORE_INITIALIZING.store(false, Ordering::Release);
    }
}

fn ensure_default_store() -> Result<()> {
    if DEFAULT_STORE.get().is_some() {
        return Ok(());
    }

    loop {
        if DEFAULT_STORE.get().is_some() {
            return Ok(());
        }

        if DEFAULT_STORE_INITIALIZING
            .compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire)
            .is_ok()
        {
            let _guard = DefaultStoreInitGuard;
            register_default_store()?;
            let _ = DEFAULT_STORE.set(true);
            return Ok(());
        }

        thread::yield_now();
    }
}

fn register_default_store() -> Result<()> {
    #[cfg(target_os = "linux")]
    {
        keyring_core::set_default_store(
            zbus_secret_service_keyring_store::Store::new().with_context(|| {
                credential_store_unavailable_message(CredentialStorePlatform::Linux)
            })?,
        );
    }
    #[cfg(target_os = "macos")]
    {
        keyring_core::set_default_store(
            apple_native_keyring_store::keychain::Store::new().with_context(|| {
                credential_store_unavailable_message(CredentialStorePlatform::Macos)
            })?,
        );
    }
    #[cfg(target_os = "windows")]
    {
        keyring_core::set_default_store(windows_native_keyring_store::Store::new().with_context(
            || credential_store_unavailable_message(CredentialStorePlatform::Windows),
        )?);
    }
    #[cfg(not(any(target_os = "linux", target_os = "macos", target_os = "windows")))]
    {
        anyhow::bail!(credential_store_unavailable_message(
            CredentialStorePlatform::Unsupported
        ));
    }

    Ok(())
}

/// Retrieve or create the encryption key for the named database.
///
/// The key is stored in the platform-native credential store under the
/// service name `"sce"` with the database name as the user/account
/// identifier.
///
/// # Arguments
/// * `db_path` — Canonical path to the database file. Used to decide
///   whether this is a first-use (generate new key) or a subsequent
///   access (read existing key).
/// * `db_name` — Logical database name used as the credential store
///   username (e.g. `"auth_db"`, `"agent_trace_db"`).
///
/// # Returns
/// A 64-character lowercase hex string that is the encryption key.
///
/// # Errors
/// - Returns an error if the credential store cannot be initialised on
///   the current platform.
/// - Returns an error if the database file exists but the keyring entry
///   is missing (e.g. keyring was cleared on Linux).
/// - Returns an error if key generation or credential store I/O fails.
pub fn get_or_create_encryption_key(db_path: &Path, db_name: &str) -> Result<String> {
    get_or_create_encryption_key_with_keyring_resolver(
        db_path,
        db_name,
        get_or_create_keyring_encryption_key,
    )
}

fn get_or_create_encryption_key_with_keyring_resolver<F>(
    db_path: &Path,
    db_name: &str,
    keyring_resolver: F,
) -> Result<String>
where
    F: FnOnce(&Path, &str) -> Result<String>,
{
    if let Some(env_key) = env_encryption_key()? {
        return Ok(env_key);
    }

    keyring_resolver(db_path, db_name)
}

fn get_or_create_keyring_encryption_key(db_path: &Path, db_name: &str) -> Result<String> {
    ensure_default_store()?;

    let entry = Entry::new("sce", db_name).with_context(|| {
        format!(
            "failed to create keyring entry for service 'sce' / user '{db_name}'. \
             {}",
            credential_store_unavailable_message(current_credential_store_platform())
        )
    })?;

    // Try to retrieve an existing password from the credential store.
    if let Ok(password) = entry.get_password() {
        return Ok(password);
    }

    // No existing key was found. If the database file does not exist,
    // this is a first-use scenario: generate and store a new key.
    if !db_path.exists() {
        let hex_key = generate_key();

        entry.set_password(&hex_key).with_context(|| {
            format!(
                "failed to store encryption key for '{db_name}' in credential store. {}",
                credential_store_unavailable_message(current_credential_store_platform())
            )
        })?;

        return Ok(hex_key);
    }

    // The database file exists but the keyring entry is missing.
    // This can happen when credentials were removed from the OS keyring.
    // Provide clear remediation.
    anyhow::bail!(
        "encryption key for '{db_name}' not found in credential store.\n\
         The database file exists at '{}' but no matching credential was found \
         for service 'sce' / user '{db_name}'.\n\
         This can happen when the {} was cleared or unavailable when the key \
         was first stored. {}\n\
         If the database file is also stale or you no longer need its data, \
         delete it and the key will be regenerated automatically on next use.",
        db_path.display(),
        credential_store_name(current_credential_store_platform()),
        credential_store_unavailable_message(current_credential_store_platform())
    );
}

fn current_credential_store_platform() -> CredentialStorePlatform {
    #[cfg(target_os = "linux")]
    {
        CredentialStorePlatform::Linux
    }
    #[cfg(target_os = "macos")]
    {
        CredentialStorePlatform::Macos
    }
    #[cfg(target_os = "windows")]
    {
        CredentialStorePlatform::Windows
    }
    #[cfg(not(any(target_os = "linux", target_os = "macos", target_os = "windows")))]
    {
        CredentialStorePlatform::Unsupported
    }
}

fn credential_store_name(platform: CredentialStorePlatform) -> &'static str {
    match platform {
        #[cfg(target_os = "linux")]
        CredentialStorePlatform::Linux => "Linux system keyring/Secret Service",
        #[cfg(target_os = "macos")]
        CredentialStorePlatform::Macos => "macOS Keychain",
        #[cfg(target_os = "windows")]
        CredentialStorePlatform::Windows => "Windows Credential Store",
        #[cfg(not(any(target_os = "linux", target_os = "macos", target_os = "windows")))]
        CredentialStorePlatform::Unsupported => "OS credential store",
    }
}

fn credential_store_unavailable_message(platform: CredentialStorePlatform) -> String {
    match platform {
        #[cfg(not(any(target_os = "linux", target_os = "macos", target_os = "windows")))]
        CredentialStorePlatform::Unsupported => format!(
            "No supported OS credential store is available for auth DB encryption key management. \
             For headless or CI use, set {AUTH_DB_ENCRYPTION_KEY_ENV} to provide the auth DB \
             encryption secret."
        ),
        _ => format!(
            "Could not access the {}. For headless or CI use, set {AUTH_DB_ENCRYPTION_KEY_ENV} \
             to provide the auth DB encryption secret.",
            credential_store_name(platform)
        ),
    }
}

fn env_encryption_key() -> Result<Option<String>> {
    match std::env::var(AUTH_DB_ENCRYPTION_KEY_ENV) {
        Ok(secret) => Ok(Some(derive_env_encryption_key(&secret)?)),
        Err(std::env::VarError::NotPresent) => Ok(None),
        Err(std::env::VarError::NotUnicode(_)) => anyhow::bail!(
            "{AUTH_DB_ENCRYPTION_KEY_ENV} must contain valid UTF-8 secret text. \
             Try: set {AUTH_DB_ENCRYPTION_KEY_ENV} to a non-empty auth DB encryption secret."
        ),
    }
}

fn derive_env_encryption_key(secret: &str) -> Result<String> {
    if secret.trim().is_empty() {
        anyhow::bail!(
            "{AUTH_DB_ENCRYPTION_KEY_ENV} is set but empty. \
             Try: set {AUTH_DB_ENCRYPTION_KEY_ENV} to a non-empty auth DB encryption secret."
        );
    }

    Ok(hex_encode(&Sha256::digest(secret.as_bytes())))
}

/// Generate a 64-character lowercase hex key from 32 random bytes.
fn generate_key() -> String {
    use rand::Rng;

    let key: [u8; 32] = rand::thread_rng().gen();
    hex_encode(&key)
}

/// Hex-encode a byte slice into a lowercase hex string.
fn hex_encode(bytes: &[u8]) -> String {
    use std::fmt::Write;

    let mut hex = String::with_capacity(bytes.len() * 2);
    for &b in bytes {
        let _ = write!(hex, "{b:02x}");
    }
    hex
}

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

    #[test]
    fn test_hex_encode_empty() {
        assert_eq!(hex_encode(&[]), "");
    }

    #[test]
    fn test_hex_encode_32_bytes() {
        let bytes = [0xdeu8; 32];
        let expected = "de".repeat(32);
        assert_eq!(hex_encode(&bytes), expected);
    }

    #[test]
    fn test_hex_encode_mixed() {
        let bytes = [0x00u8, 0x01, 0xfe, 0xff, 0xab, 0xcd];
        assert_eq!(hex_encode(&bytes), "0001feffabcd");
    }

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

    #[test]
    fn test_generate_key_lowercase() {
        let key = generate_key();
        assert!(key
            .chars()
            .all(|c| c.is_ascii_lowercase() || c.is_ascii_digit()));
    }

    #[test]
    fn test_generate_key_is_random() {
        let key1 = generate_key();
        let key2 = generate_key();
        assert_ne!(key1, key2);
    }
}