tryaudex_core/

rotation.rs

use std::path::PathBuf;
use std::sync::atomic::AtomicBool;
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)?;
        // Restrict directory to owner-only to prevent other users from listing
        // credential files on shared machines.
        // R6-M58: propagate permission errors instead of silently
        // discarding them. A credential file in a world-readable temp
        // directory is a credential leak on shared machines.
        #[cfg(unix)]
        {
            use std::os::unix::fs::PermissionsExt;
            std::fs::set_permissions(&dir, std::fs::Permissions::from_mode(0o700))?;
        }
        // R6-M14: include PID in the filename so concurrent audex processes
        // owned by the same UID can be distinguished in the shared temp
        // directory (and so a stale file from a crashed process can be
        // correlated to the process that left it behind). The UUIDv4 is
        // retained to prevent collision if the same PID issues multiple
        // credential files in the same second.
        let pid = std::process::id();
        let path = dir.join(format!("creds-{}-{}.ini", pid, 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 as plain text.
    /// GCP SDKs don't support reading pre-fetched tokens from a credential file,
    /// so we write the raw token for tools that can read a token file directly.
    /// The subprocess should use `CLOUDSDK_AUTH_ACCESS_TOKEN` (set at launch)
    /// and `AUDEX_GCP_TOKEN_FILE` (for tools that re-read on rotation).
    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<()> {
        let json = serde_json::json!({
            "type": "azure_access_token",
            "token": token,
        });
        atomic_write(&self.path, json.to_string().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.
/// Sets file permissions to 0600 (owner-only) on Unix to prevent other
/// users from reading credentials on shared machines.
///
/// # Windows file permission limitation
///
/// On Windows, this function sets the hidden attribute via `attrib +H` but
/// does **not** enforce ACL-based access control. The credential file remains
/// readable by other users with access to the same filesystem path. For
/// production use on Windows, replace the `attrib` call with proper DACL
/// manipulation using the [`windows-acl`](https://crates.io/crates/windows-acl)
/// crate (e.g. `SecurityDescriptor::new` + `set_dacl`).
fn atomic_write(path: &std::path::Path, data: &[u8]) -> Result<()> {
    use std::io::Write;

    let tmp = path.with_extension("tmp");
    // R6-M21: open the file explicitly so we can fsync before rename.
    // Previously std::fs::write + rename left a window where power loss
    // could commit a rename over an empty/partial data block because the
    // rename only guarantees metadata durability — the data pages may
    // still be in the page cache. fsync the data, then rename; a
    // crash-consistent filesystem will either retain the old contents
    // or present the fully-written new contents.
    {
        let mut file = std::fs::OpenOptions::new()
            .write(true)
            .create(true)
            .truncate(true)
            .open(&tmp)?;
        file.write_all(data)?;
        file.sync_all()?;
    }
    #[cfg(unix)]
    {
        use std::os::unix::fs::PermissionsExt;
        std::fs::set_permissions(&tmp, std::fs::Permissions::from_mode(0o600))?;
    }
    #[cfg(windows)]
    {
        // On Windows, mark the file as hidden and restrict via read-only attribute.
        // Full ACL restriction would require the windows-acl crate; this is a
        // best-effort guard for shared machines.
        tracing::warn!(
            path = %tmp.display(),
            "Windows credential file permissions are not enforced: the file may be \
             readable by other users on this machine. Use the windows-acl crate for \
             production deployments to apply a restrictive DACL."
        );
        let _ = std::process::Command::new("attrib")
            .args(["+H", &tmp.to_string_lossy()])
            .status();
    }
    std::fs::rename(&tmp, path)?;
    Ok(())
}

/// Handle to a running credential rotation background thread.
/// Dropping the handle signals the thread to stop.
///
/// R6-M22: exposes a `failed` flag that the background thread sets when it
/// gives up after `MAX_RETRY_FAILURES` consecutive rotation attempts.
/// Callers should check `has_failed()` on session end so operators learn
/// that the subprocess was running against a stale credential file; the
/// previous behaviour was to silently stop the rotation thread and leave
/// the main thread unaware that scoping had effectively lapsed.
pub struct RotationHandle {
    stop: Arc<AtomicBool>,
    failed: 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, std::sync::atomic::Ordering::Release);
        if let Some(handle) = self.thread.take() {
            let _ = handle.join();
        }
    }

    /// Returns true if the background rotation thread gave up after the
    /// configured retry ceiling, meaning the credential file has not been
    /// updated since the first failure and the subprocess has been running
    /// with stale credentials.
    pub fn has_failed(&self) -> bool {
        self.failed.load(std::sync::atomic::Ordering::Acquire)
    }
}

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

/// Start a background credential rotation thread.
///
/// The `refresh_and_write` closure is called whenever new credentials are needed.
/// It must fetch new credentials, write them to disk, and return the new expiry time.
/// This design is provider-agnostic: the caller decides how to fetch and persist
/// credentials (AWS via write_aws, GCP via write_gcp, Azure via write_azure).
/// The rotation loop checks every 10 seconds whether it's time to rotate.
/// R6-M60: returns Result so a thread-spawn failure surfaces as an error
/// to the caller instead of panicking the entire process. In constrained
/// environments (cgroup thread limits, seccomp filters) spawn can fail.
pub fn start_rotation<F>(
    initial_expires_at: chrono::DateTime<chrono::Utc>,
    rotate_before: Duration,
    refresh_and_write: F,
) -> Result<RotationHandle>
where
    F: Fn() -> std::result::Result<chrono::DateTime<chrono::Utc>, String> + Send + 'static,
{
    let stop = Arc::new(AtomicBool::new(false));
    let stop_clone = stop.clone();
    let failed = Arc::new(AtomicBool::new(false));
    let failed_clone = failed.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));
            // R6-M59: use checked subtraction. If rotate_before is larger
            // than the distance from epoch to initial_expires_at (an
            // impossible-in-practice but defensible edge case), the
            // previous unchecked subtraction would wrap to a far-future
            // timestamp and skip the first rotation entirely.
            let mut next_rotation = initial_expires_at
                .checked_sub_signed(rotate_chrono)
                .unwrap_or(initial_expires_at);
            let mut consecutive_failures: u32 = 0;
            const MAX_RETRY_FAILURES: u32 = 10;
            const MAX_BACKOFF_SECS: i64 = 300; // 5 minutes

            loop {
                if stop_clone.load(std::sync::atomic::Ordering::Acquire) {
                    break;
                }

                let now = chrono::Utc::now();
                if now >= next_rotation {
                    if consecutive_failures >= MAX_RETRY_FAILURES {
                        tracing::error!(
                            "Credential rotation gave up after {} consecutive failures — \
                             subprocess is running with stale credentials",
                            consecutive_failures
                        );
                        // R6-M22: signal the main thread so it can surface
                        // the rotation failure to the operator instead of
                        // silently continuing to read the stale credential
                        // file.
                        failed_clone.store(true, std::sync::atomic::Ordering::Release);
                        break;
                    }

                    tracing::info!("Credential rotation triggered");
                    match refresh_and_write() {
                        Ok(new_expires_at) => {
                            tracing::info!(
                                "Credentials rotated successfully, new expiry: {}",
                                new_expires_at
                            );
                            next_rotation = new_expires_at - rotate_chrono;
                            consecutive_failures = 0;
                        }
                        Err(e) => {
                            consecutive_failures += 1;
                            let backoff = 30i64
                                .saturating_mul(2i64.saturating_pow(consecutive_failures - 1))
                                .min(MAX_BACKOFF_SECS);
                            tracing::error!(
                                "Failed to rotate credentials (attempt {}/{}): {}",
                                consecutive_failures,
                                MAX_RETRY_FAILURES,
                                e
                            );
                            next_rotation = now + chrono::Duration::seconds(backoff);
                        }
                    }
                }

                // R6-M57: sleep in 1-second ticks and re-check the stop
                // flag each tick. The previous 10-second sleep meant the
                // process hung for up to 10 seconds after stop() was called,
                // delaying shutdown and credential cleanup.
                for _ in 0..10 {
                    if stop_clone.load(std::sync::atomic::Ordering::Acquire) {
                        break;
                    }
                    std::thread::sleep(Duration::from_secs(1));
                }
            }

            tracing::debug!("Credential rotation thread stopped");
        })
        .map_err(|e| {
            crate::error::AvError::InvalidPolicy(format!(
                "Failed to spawn credential rotation thread: {}",
                e
            ))
        })?;

    Ok(RotationHandle {
        stop,
        failed,
        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 cred_file_clone = cred_file.clone();
        let handle = start_rotation(expires_at, Duration::from_secs(60), move || {
            call_count_clone.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
            let new_creds = 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),
            };
            let expires = new_creds.expires_at;
            cred_file_clone
                .write_aws(&new_creds)
                .map_err(|e| e.to_string())?;
            Ok(expires)
        })
        .unwrap();

        // 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(std::sync::atomic::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);
    }
}