cipherstash-client 0.34.1-alpha.1

use crate::zerokms::{
    key_provider::{KeyProvider, KeyProviderError},
    vitur_client::{connection::HttpConnectionOpts, DEFAULT_CONCURRENT_REQS, DEFAULT_KEYS_PER_REQ},
    ClientKey, ZeroKMS, ZeroKMSClientOpts, ZeroKMSWithClientKey,
};
use stack_auth::AuthStrategy;
use thiserror::Error;
use url::Url;

/// Error type for [`ZeroKMSBuilder`] operations.
#[derive(Debug, Error, miette::Diagnostic)]
pub enum ZeroKMSBuilderError {
    /// Failed to initialize the ZeroKMS client
    #[error("Failed to initialize client: {0}")]
    ClientInit(#[from] crate::zerokms::Error),

    /// Authentication strategy failed to initialize
    #[error("Auth strategy error: {0}")]
    Auth(#[from] stack_auth::AuthError),

    /// Key provider failed to load a client key
    #[error("Key provider error: {0}")]
    KeyProvider(#[from] KeyProviderError),
}

/// A builder for creating [`ZeroKMS`] and [`ZeroKMSWithClientKey`] clients.
///
/// The ZeroKMS endpoint is resolved in this order:
/// 1. Explicit URL via [`with_base_url`](Self::with_base_url)
/// 2. `CS_ZEROKMS_HOST` (or legacy `CS_VITUR_HOST`) environment variable
/// 3. Automatically from the token's `services` claim
///
/// # Examples
///
/// ## Basic usage with AutoStrategy
///
/// ```no_run
/// use cipherstash_client::zerokms::ZeroKMSBuilder;
/// use stack_auth::AutoStrategy;
///
/// # fn example() -> Result<(), Box<dyn std::error::Error>> {
/// let strategy = AutoStrategy::detect()?;
/// let zerokms = ZeroKMSBuilder::new(strategy)
///     .with_request_timeout(30)
///     .with_max_concurrent_reqs(20)
///     .build()?;
/// # Ok(())
/// # }
/// ```
///
/// ## With client key for encryption/decryption
///
/// ```no_run
/// use cipherstash_client::zerokms::{ZeroKMSBuilder, ClientKey};
/// use stack_auth::AutoStrategy;
/// use uuid::Uuid;
///
/// # fn example() -> Result<(), Box<dyn std::error::Error>> {
/// let strategy = AutoStrategy::detect()?;
/// let client_id = Uuid::parse_str("550e8400-e29b-41d4-a716-446655440000")?;
/// let client_key = ClientKey::from_hex_v1(client_id, "a4627031...")?;
///
/// let zerokms = ZeroKMSBuilder::new(strategy)
///     .with_client_key(client_key)
///     .build()?;
/// # Ok(())
/// # }
/// ```
pub struct ZeroKMSBuilder<C, ClientKeyState = ()> {
    credentials: C,
    request_timeout: Option<u64>,
    max_keys_per_req: usize,
    max_concurrent_reqs: usize,
    client_key: ClientKeyState,
    base_url_override: Option<Url>,
}

impl ZeroKMSBuilder<stack_auth::AutoStrategy, ()> {
    /// Create a [`ZeroKMSBuilder`] that automatically detects credentials from the environment.
    ///
    /// This is the recommended entry point for most use cases. It uses [`AutoStrategy`](stack_auth::AutoStrategy)
    /// to detect available credentials (access key or OAuth token) and resolves the ZeroKMS
    /// endpoint from the token's `services` claim.
    ///
    /// # Examples
    ///
    /// ```no_run
    /// use cipherstash_client::zerokms::ZeroKMSBuilder;
    ///
    /// # fn example() -> Result<(), Box<dyn std::error::Error>> {
    /// let zerokms = ZeroKMSBuilder::auto()?.build()?;
    /// # Ok(())
    /// # }
    /// ```
    pub fn auto() -> Result<Self, ZeroKMSBuilderError> {
        let strategy = stack_auth::AutoStrategy::detect()?;
        Ok(Self {
            credentials: strategy,
            request_timeout: None,
            max_keys_per_req: DEFAULT_KEYS_PER_REQ,
            max_concurrent_reqs: DEFAULT_CONCURRENT_REQS,
            client_key: (),
            base_url_override: None,
        })
    }
}

impl<C> ZeroKMSBuilder<C, ()>
where
    C: Send + Sync + 'static,
    for<'a> &'a C: AuthStrategy,
{
    /// Create a new [`ZeroKMSBuilder`].
    ///
    /// # Arguments
    ///
    /// * `credentials` - Credentials provider for obtaining access tokens
    pub fn new(credentials: C) -> Self {
        Self {
            credentials,
            request_timeout: None,
            max_keys_per_req: DEFAULT_KEYS_PER_REQ,
            max_concurrent_reqs: DEFAULT_CONCURRENT_REQS,
            client_key: (),
            base_url_override: None,
        }
    }

    /// Set the request timeout in seconds.
    ///
    /// If not set, defaults to 10 seconds.
    pub fn with_request_timeout(mut self, timeout_secs: u64) -> Self {
        self.request_timeout = Some(timeout_secs);
        self
    }

    /// Set the maximum number of keys per request.
    ///
    /// Defaults to 500 if not set.
    pub fn with_max_keys_per_req(mut self, max_keys: usize) -> Self {
        self.max_keys_per_req = max_keys;
        self
    }

    /// Set the maximum number of concurrent requests.
    ///
    /// Defaults to 5 if not set.
    pub fn with_max_concurrent_reqs(mut self, max_concurrent: usize) -> Self {
        self.max_concurrent_reqs = max_concurrent;
        self
    }

    /// Override the base URL for the ZeroKMS service.
    ///
    /// This bypasses resolving the URL from the token's `services` claim and connects
    /// directly to the specified URL. Useful for:
    /// - Connecting to local development servers
    /// - Testing against non-production environments
    /// - Custom deployment configurations
    pub fn with_base_url(mut self, base_url: Url) -> Self {
        self.base_url_override = Some(base_url);
        self
    }

    /// Add a [`KeyProvider`] to load a client key asynchronously at build time.
    ///
    /// This transforms the builder into one that will build a [`ZeroKMSWithClientKey`]
    /// via an async [`build()`](ZeroKMSBuilder::build) call.
    ///
    /// Use this instead of [`with_client_key`](Self::with_client_key) when the key needs to be
    /// loaded from an external source (environment variables, keychain, etc.).
    ///
    /// # Example
    ///
    /// ```no_run
    /// # async fn example() -> Result<(), Box<dyn std::error::Error>> {
    /// use cipherstash_client::zerokms::{
    ///     ZeroKMSBuilder, EnvKeyProvider, FallbackKeyProvider,
    /// };
    /// use stack_profile::ProfileStore;
    /// use stack_auth::AutoStrategy;
    ///
    /// let strategy = AutoStrategy::detect()?;
    ///
    /// // Try environment variables first, fall back to the on-disk profile
    /// let zerokms = ZeroKMSBuilder::new(strategy)
    ///     .with_key_provider(FallbackKeyProvider::new(
    ///         EnvKeyProvider,
    ///         ProfileStore::default(),
    ///     ))
    ///     .build()
    ///     .await?;
    /// # Ok(())
    /// # }
    /// ```
    pub fn with_key_provider<K: KeyProvider>(
        self,
        provider: K,
    ) -> ZeroKMSBuilder<C, WithKeyProvider<K>> {
        ZeroKMSBuilder {
            credentials: self.credentials,
            request_timeout: self.request_timeout,
            max_keys_per_req: self.max_keys_per_req,
            max_concurrent_reqs: self.max_concurrent_reqs,
            client_key: WithKeyProvider(provider),
            base_url_override: self.base_url_override,
        }
    }

    /// Add a client key to enable encryption and decryption operations.
    ///
    /// This transforms the builder into one that will build a [`ZeroKMSWithClientKey`] instead of a plain [`ZeroKMS`].
    pub fn with_client_key(self, client_key: ClientKey) -> ZeroKMSBuilder<C, ClientKey> {
        ZeroKMSBuilder {
            credentials: self.credentials,
            request_timeout: self.request_timeout,
            max_keys_per_req: self.max_keys_per_req,
            max_concurrent_reqs: self.max_concurrent_reqs,
            client_key,
            base_url_override: self.base_url_override,
        }
    }

    /// Build a [`ZeroKMS`] client.
    ///
    /// The ZeroKMS endpoint is resolved from (in order): an explicit
    /// [`with_base_url`](Self::with_base_url) override, the `CS_ZEROKMS_HOST`
    /// or the token's `services` claim.
    pub fn build(self) -> Result<ZeroKMS<C>, ZeroKMSBuilderError> {
        let (opts, _) = self.build_opts();
        Ok(ZeroKMS::connect(opts)?)
    }
}

impl<C, S> ZeroKMSBuilder<C, S> {
    fn build_opts(self) -> (ZeroKMSClientOpts<C>, S) {
        let base_url = self.base_url_override.or_else(Self::base_url_from_env);
        let mut connection_opts = HttpConnectionOpts::new(base_url);
        if let Some(timeout) = self.request_timeout {
            connection_opts = connection_opts.with_request_timeout(timeout);
        }

        let opts = ZeroKMSClientOpts {
            credentials: self.credentials,
            connection_opts,
            max_keys_per_req: self.max_keys_per_req,
            max_concurrent_reqs: self.max_concurrent_reqs,
        };

        (opts, self.client_key)
    }

    /// Resolve the ZeroKMS base URL from the `CS_ZEROKMS_HOST` environment
    /// variable (or legacy `CS_VITUR_HOST`) — explicit URL override
    fn base_url_from_env() -> Option<Url> {
        use crate::config::vars::CS_ZEROKMS_HOST;

        for name in CS_ZEROKMS_HOST {
            if let Ok(value) = std::env::var(name) {
                match value.parse() {
                    Ok(url) => return Some(url),
                    Err(err) => {
                        tracing::warn!(
                            target: "zerokms",
                            %err,
                            env_var = name,
                            "Ignoring invalid URL in environment variable"
                        );
                    }
                }
            }
        }

        None
    }
}

impl<C> ZeroKMSBuilder<C, ClientKey>
where
    C: Send + Sync + 'static,
    for<'a> &'a C: AuthStrategy,
{
    /// Build a [`ZeroKMSWithClientKey`] client.
    ///
    /// This is only available after calling [`with_client_key`](ZeroKMSBuilder::with_client_key).
    pub fn build(self) -> Result<ZeroKMSWithClientKey<C>, ZeroKMSBuilderError> {
        let (opts, client_key) = self.build_opts();
        Ok(ZeroKMS::connect_with_client_key(opts, client_key)?)
    }
}

/// Newtype wrapper that marks a [`KeyProvider`] in the builder's type state.
///
/// This avoids coherence issues — [`ClientKey`] does not implement [`KeyProvider`],
/// and `WithKeyProvider` keeps the two `build()` signatures unambiguous.
pub struct WithKeyProvider<K: KeyProvider>(K);

impl<C, K> ZeroKMSBuilder<C, WithKeyProvider<K>>
where
    C: Send + Sync + 'static,
    for<'a> &'a C: AuthStrategy,
    K: KeyProvider,
{
    /// Build a [`ZeroKMSWithClientKey`] client by loading the key from the provider.
    ///
    /// This is an async method because the key provider may need to perform I/O.
    pub async fn build(self) -> Result<ZeroKMSWithClientKey<C>, ZeroKMSBuilderError> {
        let (opts, provider) = self.build_opts();
        let client_key = provider.0.client_key().await?;
        Ok(ZeroKMS::connect_with_client_key(opts, client_key)?)
    }
}