cipherstash-client 0.34.1-alpha.6

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(())
/// # }
/// ```
///
/// # Tuning for bulk ingest vs interactive workloads
///
/// The defaults — `request_timeout` 10 s, `max_keys_per_req` 500,
/// `max_concurrent_reqs` 5 — are tuned for the interactive case (proxy
/// query-rewrite, single-document encrypt/decrypt) where fast failure on
/// network problems matters more than absolute throughput.
///
/// Bulk ingest (large backfills, batch ETL, our `cipherstash/benches`
/// suite) wants the opposite trade-off: longer per-request budgets to
/// absorb p99 latency spikes, and persistent connection pools so cold
/// TLS handshakes don't dominate. Recommended starting point:
///
/// ```no_run
/// # fn main() -> Result<(), Box<dyn std::error::Error>> {
/// use cipherstash_client::zerokms::ZeroKMSBuilder;
///
/// let zerokms = ZeroKMSBuilder::auto()?
///     // Connect phase fast-fail: if you can't open TCP+TLS in 5 s the
///     // network is broken; no point waiting longer.
///     .with_connect_timeout(5)
///     // Generous request budget — server-side processing time scales
///     // with `keys.len()`, and bulk callers care more about completing
///     // the batch than about failing fast.
///     .with_request_timeout(60)
///     // Keep warm TLS connections alive across idle gaps between
///     // batches; default is 90 s which is short for some pipelines.
///     .with_pool_idle_timeout(600)
///     // Reduce per-request server work when responses are slow;
///     // smaller batches lower the chance any single request times out.
///     .with_max_keys_per_req(100)
///     // Increase parallelism to compensate for the smaller batches.
///     .with_max_concurrent_reqs(20)
///     .build()?;
/// # Ok(())
/// # }
/// ```
///
/// Importantly: **build once, hold an `Arc` for the process lifetime**.
/// Recreating the client (or the [`ScopedCipher`](crate::encryption::ScopedCipher))
/// on every batch discards the warm reqwest connection pool and forces
/// every encrypt to pay the cold-start cost (DNS + TCP + TLS handshake).
/// Auth tokens auto-refresh through `stack_auth::AutoRefresh` underneath
/// the same client — no manual rotation needed.
pub struct ZeroKMSBuilder<C, ClientKeyState = ()> {
    credentials: C,
    request_timeout: Option<u64>,
    connect_timeout: Option<u64>,
    pool_idle_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,
            connect_timeout: None,
            pool_idle_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,
            connect_timeout: None,
            pool_idle_timeout: None,
            max_keys_per_req: DEFAULT_KEYS_PER_REQ,
            max_concurrent_reqs: DEFAULT_CONCURRENT_REQS,
            client_key: (),
            base_url_override: None,
        }
    }

    /// Set the **total request timeout** in seconds — covers connect +
    /// TLS handshake + body send + body receive together.
    ///
    /// If not set, defaults to 10 seconds. Bulk-ingest callers typically
    /// want this larger (30–60 s) — see the "Tuning for bulk ingest"
    /// section on [`ZeroKMSBuilder`].
    pub fn with_request_timeout(mut self, timeout_secs: u64) -> Self {
        self.request_timeout = Some(timeout_secs);
        self
    }

    /// Set the **connect timeout** in seconds — bound on TCP connect +
    /// TLS handshake only, separate from the total request timeout.
    ///
    /// If not set, reqwest falls back to the OS-level connect timeout
    /// (~75 s on most platforms). Setting this small (e.g. 5 s) gives
    /// fast-fail on broken networks without forcing the total request
    /// timeout to absorb both connect *and* response time.
    pub fn with_connect_timeout(mut self, timeout_secs: u64) -> Self {
        self.connect_timeout = Some(timeout_secs);
        self
    }

    /// Set the **pool idle timeout** in seconds — how long the underlying
    /// reqwest client keeps idle keep-alive connections in its pool
    /// before closing them.
    ///
    /// If not set, reqwest's default of 90 s applies. Long-lived
    /// processes (bulk ingest, daemons) benefit from raising this so
    /// warm TLS connections survive idle gaps between batches.
    pub fn with_pool_idle_timeout(mut self, timeout_secs: u64) -> Self {
        self.pool_idle_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,
            connect_timeout: self.connect_timeout,
            pool_idle_timeout: self.pool_idle_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,
            connect_timeout: self.connect_timeout,
            pool_idle_timeout: self.pool_idle_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);
        }
        if let Some(connect_timeout) = self.connect_timeout {
            connection_opts = connection_opts.with_connect_timeout(connect_timeout);
        }
        if let Some(pool_idle_timeout) = self.pool_idle_timeout {
            connection_opts = connection_opts.with_pool_idle_timeout(pool_idle_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)?)
    }
}