tryaudex_core/

rotation.rs

use std::path::PathBuf;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;
use std::time::Duration;

use crate::credentials::TempCredentials;
use crate::error::Result;

/// Configuration for automatic credential rotation.
#[derive(Debug, Clone)]
pub struct RotationConfig {
    /// Rotate credentials this many seconds before they expire.
    pub rotate_before_secs: u64,
    /// Minimum session TTL (in seconds) to enable rotation.
    /// Sessions shorter than this use static credentials.
    pub min_ttl_for_rotation_secs: u64,
}

impl Default for RotationConfig {
    fn default() -> Self {
        Self {
            rotate_before_secs: 300,        // 5 minutes before expiry
            min_ttl_for_rotation_secs: 900, // only rotate if TTL >= 15 minutes
        }
    }
}

/// A file that holds provider credentials, updated atomically.
/// Automatically cleaned up when dropped.
pub struct CredentialFile {
    path: PathBuf,
}

impl CredentialFile {
    /// Create a new credential file in a temp directory.
    pub fn new() -> Result<Self> {
        let dir = std::env::temp_dir().join("audex");
        std::fs::create_dir_all(&dir)?;
        let path = dir.join(format!("creds-{}.ini", uuid::Uuid::new_v4()));
        Ok(Self { path })
    }

    /// Path to the credential file.
    pub fn path(&self) -> &std::path::Path {
        &self.path
    }

    /// Write credentials in AWS shared credentials file format.
    /// Uses atomic write (write to tmp, then rename) to avoid partial reads.
    pub fn write_aws(&self, creds: &TempCredentials) -> Result<()> {
        let content = format!(
            "[default]\naws_access_key_id = {}\naws_secret_access_key = {}\naws_session_token = {}\n",
            creds.access_key_id, creds.secret_access_key, creds.session_token
        );
        atomic_write(&self.path, content.as_bytes())
    }

    /// Write a GCP access token to the credential file.
    pub fn write_gcp(&self, token: &str) -> Result<()> {
        atomic_write(&self.path, token.as_bytes())
    }

    /// Write an Azure access token to the credential file.
    pub fn write_azure(&self, token: &str) -> Result<()> {
        atomic_write(&self.path, token.as_bytes())
    }

    /// Remove the credential file from disk.
    pub fn cleanup(&self) {
        let _ = std::fs::remove_file(&self.path);
        let _ = std::fs::remove_file(self.path.with_extension("tmp"));
    }
}

impl Drop for CredentialFile {
    fn drop(&mut self) {
        self.cleanup();
    }
}

/// Atomic write: write to a .tmp sibling then rename over the target.
fn atomic_write(path: &std::path::Path, data: &[u8]) -> Result<()> {
    let tmp = path.with_extension("tmp");
    std::fs::write(&tmp, data)?;
    std::fs::rename(&tmp, path)?;
    Ok(())
}

/// Handle to a running credential rotation background thread.
/// Dropping the handle signals the thread to stop.
pub struct RotationHandle {
    stop: Arc<AtomicBool>,
    thread: Option<std::thread::JoinHandle<()>>,
}

impl RotationHandle {
    /// Signal the rotation thread to stop and wait for it to finish.
    pub fn stop(mut self) {
        self.stop.store(true, Ordering::Relaxed);
        if let Some(handle) = self.thread.take() {
            let _ = handle.join();
        }
    }
}

impl Drop for RotationHandle {
    fn drop(&mut self) {
        self.stop.store(true, Ordering::Relaxed);
        if let Some(handle) = self.thread.take() {
            let _ = handle.join();
        }
    }
}

/// Start a background credential rotation thread.
///
/// The `refresh_fn` closure is called whenever new credentials are needed.
/// It runs inside a dedicated tokio runtime on the background thread.
/// The rotation loop checks every 10 seconds whether it's time to rotate.
pub fn start_rotation<F>(
    cred_file: Arc<CredentialFile>,
    initial_expires_at: chrono::DateTime<chrono::Utc>,
    rotate_before: Duration,
    refresh_fn: F,
) -> RotationHandle
where
    F: Fn() -> std::result::Result<TempCredentials, String> + Send + 'static,
{
    let stop = Arc::new(AtomicBool::new(false));
    let stop_clone = stop.clone();

    let thread = std::thread::Builder::new()
        .name("audex-rotation".into())
        .spawn(move || {
            let rotate_chrono =
                chrono::Duration::from_std(rotate_before).unwrap_or(chrono::Duration::seconds(300));
            let mut next_rotation = initial_expires_at - rotate_chrono;

            loop {
                if stop_clone.load(Ordering::Relaxed) {
                    break;
                }

                let now = chrono::Utc::now();
                if now >= next_rotation {
                    tracing::info!("Credential rotation triggered");
                    match refresh_fn() {
                        Ok(new_creds) => {
                            match cred_file.write_aws(&new_creds) {
                                Ok(()) => {
                                    tracing::info!(
                                        "Credentials rotated successfully, new expiry: {}",
                                        new_creds.expires_at
                                    );
                                    next_rotation = new_creds.expires_at - rotate_chrono;
                                }
                                Err(e) => {
                                    tracing::error!("Failed to write rotated credentials: {}", e);
                                    // Retry in 30 seconds
                                    next_rotation = now + chrono::Duration::seconds(30);
                                }
                            }
                        }
                        Err(e) => {
                            tracing::error!("Failed to refresh credentials: {}", e);
                            // Retry in 30 seconds
                            next_rotation = now + chrono::Duration::seconds(30);
                        }
                    }
                }

                // Check every 10 seconds
                std::thread::sleep(Duration::from_secs(10));
            }

            tracing::debug!("Credential rotation thread stopped");
        })
        .expect("failed to spawn rotation thread");

    RotationHandle {
        stop,
        thread: Some(thread),
    }
}

/// Check if rotation should be enabled for a given TTL.
pub fn should_rotate(ttl: Duration, config: &RotationConfig) -> bool {
    ttl.as_secs() >= config.min_ttl_for_rotation_secs
}

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

    #[test]
    fn test_should_rotate_long_ttl() {
        let config = RotationConfig::default();
        assert!(should_rotate(Duration::from_secs(3600), &config)); // 1 hour
        assert!(should_rotate(Duration::from_secs(900), &config)); // exactly 15 min
    }

    #[test]
    fn test_should_not_rotate_short_ttl() {
        let config = RotationConfig::default();
        assert!(!should_rotate(Duration::from_secs(600), &config)); // 10 min
        assert!(!should_rotate(Duration::from_secs(300), &config)); // 5 min
        assert!(!should_rotate(Duration::from_secs(60), &config)); // 1 min
    }

    #[test]
    fn test_credential_file_write_aws() {
        let cred_file = CredentialFile::new().unwrap();
        let creds = TempCredentials {
            access_key_id: "AKIAIOSFODNN7EXAMPLE".to_string(),
            secret_access_key: "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY".to_string(),
            session_token: "FwoGZXtoken123".to_string(),
            expires_at: chrono::Utc::now() + chrono::Duration::hours(1),
        };
        cred_file.write_aws(&creds).unwrap();

        let content = std::fs::read_to_string(cred_file.path()).unwrap();
        assert!(content.contains("[default]"));
        assert!(content.contains("AKIAIOSFODNN7EXAMPLE"));
        assert!(content.contains("wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY"));
        assert!(content.contains("FwoGZXtoken123"));
    }

    #[test]
    fn test_credential_file_write_gcp() {
        let cred_file = CredentialFile::new().unwrap();
        cred_file.write_gcp("ya29.test-token-123").unwrap();

        let content = std::fs::read_to_string(cred_file.path()).unwrap();
        assert_eq!(content, "ya29.test-token-123");
    }

    #[test]
    fn test_credential_file_cleanup_on_drop() {
        let path;
        {
            let cred_file = CredentialFile::new().unwrap();
            cred_file.write_gcp("test").unwrap();
            path = cred_file.path().to_path_buf();
            assert!(path.exists());
        }
        // After drop, file should be cleaned up
        assert!(!path.exists());
    }

    #[test]
    fn test_rotation_handle_stop() {
        use std::sync::atomic::AtomicU32;

        let cred_file = Arc::new(CredentialFile::new().unwrap());
        let call_count = Arc::new(AtomicU32::new(0));
        let call_count_clone = call_count.clone();

        // Set expiry in the past so rotation triggers immediately
        let expires_at = chrono::Utc::now() - chrono::Duration::seconds(10);

        let initial_creds = TempCredentials {
            access_key_id: "AKIATEST".to_string(),
            secret_access_key: "secret".to_string(),
            session_token: "token".to_string(),
            expires_at: chrono::Utc::now() + chrono::Duration::hours(1),
        };

        // Write initial credentials so the file exists
        cred_file.write_aws(&initial_creds).unwrap();

        let handle = start_rotation(
            cred_file.clone(),
            expires_at,
            Duration::from_secs(60),
            move || {
                call_count_clone.fetch_add(1, Ordering::Relaxed);
                Ok(TempCredentials {
                    access_key_id: "AKIAROTATED".to_string(),
                    secret_access_key: "new-secret".to_string(),
                    session_token: "new-token".to_string(),
                    expires_at: chrono::Utc::now() + chrono::Duration::hours(1),
                })
            },
        );

        // Give the rotation thread time to fire
        std::thread::sleep(Duration::from_millis(500));
        handle.stop();

        // refresh_fn should have been called at least once
        assert!(call_count.load(Ordering::Relaxed) >= 1);

        // Credential file should have rotated credentials
        let content = std::fs::read_to_string(cred_file.path()).unwrap();
        assert!(content.contains("AKIAROTATED"));
    }

    #[test]
    fn test_rotation_config_defaults() {
        let config = RotationConfig::default();
        assert_eq!(config.rotate_before_secs, 300);
        assert_eq!(config.min_ttl_for_rotation_secs, 900);
    }

    #[test]
    fn test_atomic_write_is_atomic() {
        let dir = std::env::temp_dir().join("audex-test-atomic");
        let _ = std::fs::create_dir_all(&dir);
        let path = dir.join("test-atomic.txt");

        atomic_write(&path, b"hello world").unwrap();
        assert_eq!(std::fs::read_to_string(&path).unwrap(), "hello world");

        // Overwrite
        atomic_write(&path, b"updated").unwrap();
        assert_eq!(std::fs::read_to_string(&path).unwrap(), "updated");

        // tmp file should not exist
        assert!(!path.with_extension("tmp").exists());

        let _ = std::fs::remove_file(&path);
    }
}