cipherstash-client 0.34.1-alpha.1

use super::TokenExpiry;
use miette::Diagnostic;
use std::{
    fs,
    path::{Path, PathBuf},
};
use thiserror::Error;
use tracing::instrument;

#[derive(Diagnostic, Error, Debug)]
pub enum SetTokenError {
    #[error("IOWriteError: {0}")]
    IOWriteError(#[from] std::io::Error),

    #[error("JsonError: {0}")]
    JsonError(#[from] serde_json::Error),
}

#[derive(Error, Debug, Diagnostic)]
pub enum ClearTokenError {
    #[error("IOWriteError: {0}")]
    IoError(#[from] std::io::Error),
}

#[derive(Debug)]
pub struct TokenStore<TK: for<'a> TokenExpiry<'a>> {
    token_path: PathBuf,
    cached_token: Option<TK>,
}

impl<TK: for<'a> TokenExpiry<'a>> TokenStore<TK> {
    pub fn new(token_path: &Path) -> Self {
        let cache_dir = token_path.parent().unwrap_or_else(|| {
            panic!(
                "Token path must have a parent directory: {}",
                token_path.display()
            )
        });
        if !fs::exists(cache_dir).expect("Cannot access token directory") {
            fs::create_dir_all(cache_dir).expect("Failed to create token directory");
        }
        let cached_token: Option<TK> = std::fs::read_to_string(token_path)
            .ok()
            .and_then(|x| serde_json::from_str(&x).ok());

        let token_path = token_path.to_owned();

        Self {
            token_path,
            cached_token,
        }
    }

    /// Returns a cached token if available.
    ///
    /// Note: The in-memory path only returns non-expired tokens, while the disk
    /// path returns tokens regardless of expiry. This asymmetry is intentional —
    /// callers like `UserCredentials::get_token()` need the expired token to pass
    /// its refresh token to the IdP's token-exchange endpoint.
    #[instrument(level = "debug", skip(self), fields(path = ?self.token_path))]
    pub fn get(&mut self) -> Option<TK> {
        // reads token from in-memory cache
        if let Some(token) = &self.cached_token {
            if !token.is_expired() {
                return Some(token.clone());
            }
        }

        // reads token from disk, as it might have been updated outside the scope of this struct
        let token_from_disk: Option<TK> = std::fs::read_to_string(&self.token_path)
            .ok()
            .and_then(|x| serde_json::from_str(&x).ok());

        if let Some(token) = &token_from_disk {
            self.cached_token = Some(token.clone());
            return Some(token.clone());
        }

        tracing::debug!("No valid token found in cache or disk");

        None
    }

    /// Caches the token in the store
    #[instrument(level = "debug", skip(self, token), fields(path = ?self.token_path))]
    pub fn set(&mut self, token: &TK) -> Result<(), SetTokenError> {
        // Replaces the in-memory cache
        self.cached_token = Some(token.clone());

        // Saves the token to disk
        let json_string = serde_json::to_string_pretty(token)?;
        std::fs::write(&self.token_path, json_string)?;

        Ok(())
    }

    pub fn clear(&mut self) -> Result<(), ClearTokenError> {
        self.cached_token = None;

        if self.token_path.exists() {
            std::fs::remove_file(&self.token_path)?;
        }

        Ok(())
    }
}