cipherstash-client 0.12.5

use std::{path::Path, time::Duration};

use async_mutex::Mutex as AsyncMutex;
use async_trait::async_trait;
use log::debug;
use miette::Diagnostic;
use serde_json::json;
use thiserror::Error;
use url::Url;

use super::service_token::ServiceToken;
use crate::credentials::{
    token_store::TokenStore, user_credentials::UserToken, AutoRefreshable, ClearTokenError,
    Credentials, GetTokenError, TokenExpiry,
};
use crate::reqwest_client::create_client;
use crate::user_agent::get_user_agent;

pub struct ServiceUserCredentials<C: Credentials<Token = UserToken>> {
    user_credentials: C,
    cts_base_url: Url,
    workspace_id: String,
    token_store: AsyncMutex<TokenStore<ServiceToken>>,
    client: reqwest_middleware::ClientWithMiddleware,
}

#[derive(Diagnostic, Error, Debug)]
pub enum AcquireTokenError {
    #[error("Failed to acquire console token: {0}")]
    GetTokenError(#[from] GetTokenError),

    #[error("Failed to acquire token: {0}")]
    RequestFailed(Box<dyn std::error::Error + Send + Sync>),

    #[error("Failed to parse json response: {0}")]
    BadResponse(Box<dyn std::error::Error + Sync + Send>),
}

impl<C: Credentials<Token = UserToken>> ServiceUserCredentials<C> {
    pub fn new(
        token_path: &Path,
        user_credentials: C,
        cts_base_url: &Url,
        workspace_id: &str,
    ) -> Self {
        Self {
            user_credentials,
            cts_base_url: cts_base_url.to_owned(),
            workspace_id: workspace_id.to_string(),
            token_store: AsyncMutex::new(TokenStore::new(token_path)),
            client: create_client(),
        }
    }

    async fn federate_token(&self) -> Result<ServiceToken, AcquireTokenError> {
        debug!(target: "service_user_credentials", "Exchanging Access Token with CTS");
        let url = self.cts_base_url.join("/api/federate").unwrap();

        let user_token = self.user_credentials.get_token().await?;

        let token: ServiceToken = self
            .client
            .post(url)
            .json(&json!({
                "accessToken": user_token.access_token(),
                "workspaceId": self.workspace_id.clone(),
            }))
            .header("authorization", user_token.as_header())
            .header("user-agent", get_user_agent())
            .send()
            .await
            .map_err(|e| AcquireTokenError::RequestFailed(Box::new(e)))?
            .error_for_status()
            .map_err(|e| AcquireTokenError::RequestFailed(Box::new(e)))?
            .json()
            .await
            .map_err(|e| AcquireTokenError::BadResponse(Box::new(e)))?;

        debug!(target: "service_user_credentials",
            "Access Token Acquired - expiry(epoch seconds): {}",
            &token.expiry
        );
        Ok(token)
    }
}

#[async_trait]
impl<C: Credentials<Token = UserToken>> Credentials for ServiceUserCredentials<C> {
    type Token = ServiceToken;

    async fn get_token(&self) -> Result<Self::Token, GetTokenError> {
        debug!(target: "service_user_credentials", "getting token (waiting for lock)");

        let mut token_store = self.token_store.lock().await;

        debug!(target: "service_user_credentials", "getting token (got lock)");

        // Check to see if we can get a new token from the cache
        if let Some(cached_token) = token_store.get() {
            debug!(target: "service_user_credentials", "found cached token");

            // If the token hasn't expired yet just return with it immediately
            if !cached_token.is_expired() {
                debug!(target: "service_user_credentials", "using cached token");
                return Ok(cached_token);
            }

            debug!(target: "service_user_credentials", "cached token is expired");
        }

        debug!(target: "service_user_credentials", "fetching new token");

        // The cached token has expired so try and get a new one from CTS.
        let new_token = self
            .federate_token()
            .await
            .map_err(|e| GetTokenError::AcquireNewTokenFailed(Box::new(e)))?;

        // Saves the token to disk
        token_store
            .set(&new_token)
            .map_err(|e| GetTokenError::PersistTokenError(Box::new(e)))?;

        Ok(new_token)
    }

    async fn clear_token(&self) -> Result<(), ClearTokenError> {
        let mut token_store = self.token_store.lock().await;
        token_store
            .clear()
            .map_err(|e| ClearTokenError(Box::new(e)))
    }
}

#[async_trait]
impl<C: Credentials<Token = UserToken>> AutoRefreshable for ServiceUserCredentials<C> {
    async fn refresh(&self) -> Duration {
        let token = {
            // Drop the guard early to allow other tasks get the current cached token if still valid
            let mut token_store = self.token_store.lock().await;
            token_store.get()
        };

        // Check to see if we have a token from the cache or disk
        if let Some(cached_token) = token {
            debug!(target: "service_user_credentials", "Found token on disk");

            // If the token is still new, we do an early return
            if !cached_token.should_refresh() {
                debug!(target: "service_user_credentials", "Access token is still new");
                return cached_token.refresh_interval();
            }
        }

        // The cached token is either missing, or needs a refresh, so try acquiring from the user credentials
        debug!(target: "service_user_credentials", "Access token is missing or close to expiry, refreshing");
        if let Ok(new_token) = self.federate_token().await {
            let mut token_store = self.token_store.lock().await;
            if token_store.set(&new_token).is_ok() {
                debug!(target: "service_user_credentials", "Access token refreshed and saved to disk");
                return new_token.refresh_interval();
            }
        }

        Self::Token::min_refresh_interval()
    }
}