cipherstash-client 0.34.1-alpha.1

//! The `zerokms` module provides a client for interacting with the ZeroKMS service.
//!
//! ZeroKMS provides key management and encryption services using CipherStash's "Vitur" protocol.
//! This client provides a high-level interface for managing keys, granting access and encrypting/decrypting data.
//!
//! ## Creating a ZeroKMS Client
//!
//! The recommended way to create a ZeroKMS client is using [`ZeroKMSBuilder`], which resolves
//! the ZeroKMS endpoint automatically from the token's `services` claim:
//!
//! ```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)
//!     .build()?;
//! # Ok(())
//! # }
//! ```
//!
//! ## Workspaces, Keysets and Clients
//!
//! ZeroKMS is built around the concept of Keysets and Clients that are encapsulated within a Workspace.
//!
//! * **Workspaces** are the top-level entity that encapsulates all Keysets and Clients.
//! * **Keysets** represent a set of data that is encrypted under the same key heirarchy.
//! * **Clients** are entities that can generate and retrieve keys for one or more keysets.
//!
//! ## Authentication
//!
//! To use the [ZeroKMS] client, a valid access token is required.
//! You don't need to provide this directly: it is handled by the client using one of 2 methods:
//!
//! * **Stash CLI**: You can use the [stash](https://cipherstash.com/docs/reference/cli#install-the-cipher-stash-cli) CLI to authenticate and store the access token
//! * **Client Access Key**: You can provide a [client access key](https://cipherstash.com/docs/how-to/creating-access-keys) to the client to authenticate
//!
//! Stash CLI is best for local development and testing, while the client access key is best for production.
//!
//! ## [ZeroKMS] Client
//!
//! The [ZeroKMS] client provides a high-level interface for managing Keysets and Clients in ZeroKMS.
//! It _does not_ provide methods to encrypt or decrypt data.
//! For that, you need to use a `ZeroKMSWithClientKey` client.
//!
//! The [ZeroKMS] client can perform the following operations:
//!
//! * Create a new Keyset
//! * Grant and Revoke access to a Keyset for a client
//! * List all Keysets in a Workspace
//! * Enable and Disable a Keyset
//! * Modify a Keyset's name or description
//! * Create a new Client for a Keyset
//! * List all Clients in a Workspace
//! * Revoke (Delete) a Client
//!
//! ## [ZeroKMSWithClientKey] Client
//!
//! The [ZeroKMSWithClientKey] is the same as [ZeroKMS] but with the ability to encrypt and decrypt data.
//! Use it only when you need to perform encryption and decryption operations.
//! If your client only needs to manage Keysets and Clients, use the [ZeroKMS] client.
//!
//! The [ZeroKMSWithClientKey] client can perform the following operations:
//!
//! * **Encrypt** one or many values
//! * **Decrypt** one or many values
//! * **Save** and **Load** a Keyset's configuration
//!
//! See also [ClientKey] and [EncryptPayload].
//!
//! ## Descriptors
//!
//! [ZeroKMSWithClientKey] includes the concept of a "descriptor" which is a string that describes the data being encrypted.
//! It provides a way to identify the data being encrypted and decrypted which can be used to ensure the data is being used correctly.
//!
//! Descriptors are included in the "Authenticated Associated Data" (AAD) of the encryption and decryption operations.
//! This means that the descriptor is included in the authentication tag and any tampering with the descriptor will cause the decryption to fail.
//! Descriptors are also used by ZeroKMS to track the usage of keys and data.
//!
// TODO: Link to a Glossary page for AAD above
mod builder;
mod errors;
pub mod key_provider;
mod local_log;
mod secret_key;
mod vitur_client;

use crate::zerokms::vitur_client::{connection::HttpConnectionOpts, ClientOpts};
use local_log::log_decryptions;
use log::debug;
use std::{
    borrow::Cow,
    path::{Path, PathBuf},
};
use url::Url;
use uuid::Uuid;
use vitur_client::{DecryptionTarget, EncryptionTarget};
use zeroize::{Zeroize, ZeroizeOnDrop};
use zerokms_protocol::{IdentifiedBy, UnverifiedContext};

use stack_auth::AuthStrategy;

use crate::credentials::ServiceToken;

// Re-exports
pub use builder::{WithKeyProvider, ZeroKMSBuilder, ZeroKMSBuilderError};
pub use errors::Error;
pub use key_provider::{
    EnvKeyProvider, FallbackKeyProvider, KeyProvider, KeyProviderError, StaticKeyProvider,
};
pub use secret_key::SecretKey;
pub use vitur_client::{
    encrypted_record,
    encrypted_record::{Decryptable, EncryptedRecord, WithContext},
    errors::RecordDecryptError,
    ClientKey, DataKey, DataKeyWithTag, EncryptPayload, GenerateKeyPayload, RetrieveKeyPayload,
};
pub use zerokms_protocol::cipherstash_config::{DatasetConfig, DatasetConfigWithIndexRootKey};
pub use zerokms_protocol::{
    ClientKeysetId, ClientType, Context, CreateClientResponse, CreateClientSpec,
    CreateKeysetResponse, CreatedClient, DeleteClientResponse, Keyset, KeysetClient,
    ViturKeyMaterial,
};

// Crate only re-exports
// FIXME: Don't make this public once the Session is updated to use ScopedCipher
#[doc(hidden)]
pub use vitur_client::IndexKey;

type ViturClient = vitur_client::Client<vitur_client::HttpConnection>;

#[derive(Zeroize, ZeroizeOnDrop)]
pub struct ZeroKMS<C, ClientKeyState = ()>
where
    ClientKeyState: Zeroize,
{
    #[zeroize(skip)]
    client: ViturClient,
    // TODO: This should absolutely not be skipped
    #[zeroize(skip)]
    credentials: C,
    #[zeroize(skip)]
    decryption_log_path: Option<PathBuf>,
    client_key: ClientKeyState,
}

pub type ZeroKMSWithClientKey<C> = ZeroKMS<C, ClientKey>;

pub struct ZeroKMSClientOpts<C> {
    credentials: C,
    connection_opts: HttpConnectionOpts,
    max_keys_per_req: usize,
    max_concurrent_reqs: usize,
}

impl<C> ZeroKMS<C>
where
    C: Send + Sync + 'static,
    for<'a> &'a C: AuthStrategy,
{
    pub fn connect(opts: ZeroKMSClientOpts<C>) -> Result<Self, Error> {
        let client_opts: ClientOpts<HttpConnectionOpts> = ClientOpts {
            max_keys_per_req: opts.max_keys_per_req,
            max_concurrent_reqs: opts.max_concurrent_reqs,
            connection_opts: opts.connection_opts,
        };

        let client = ViturClient::init_opts(client_opts)?;
        Ok(Self {
            client,
            credentials: opts.credentials,
            decryption_log_path: None,
            client_key: (),
        })
    }

    pub fn connect_with_client_key(
        opts: ZeroKMSClientOpts<C>,
        client_key: ClientKey,
    ) -> Result<ZeroKMSWithClientKey<C>, Error> {
        let client_opts: ClientOpts<HttpConnectionOpts> = ClientOpts {
            max_keys_per_req: opts.max_keys_per_req,
            max_concurrent_reqs: opts.max_concurrent_reqs,
            connection_opts: opts.connection_opts,
        };

        let client = ViturClient::init_opts(client_opts)?;
        Ok(ZeroKMSWithClientKey {
            client,
            credentials: opts.credentials,
            decryption_log_path: None,
            client_key,
        })
    }

    /// Create a new instance of the [`ZeroKMS`] client.
    ///
    /// In most cases it is preferred to use [`ZeroKMSBuilder`] instead of calling this manually.
    #[deprecated(note = "replaced by connect", since = "0.32.2")]
    pub fn new(base_url: &Url, credentials: C, decryption_log_path: Option<&Path>) -> Self {
        let mut host = base_url.to_string();
        if host.ends_with('/') {
            host.pop();
        }

        // Allowed because this method is also deprecated
        #[allow(deprecated)]
        let client = ViturClient::init(host);
        Self {
            client,
            credentials,
            decryption_log_path: decryption_log_path.map(|p| p.to_path_buf()),
            client_key: (),
        }
    }

    /// Create a new instance of the [`ZeroKMS`] client with a [`ClientKey`].
    ///
    /// In most cases it is preferred to use [`ZeroKMSBuilder`] instead of calling this manually.
    #[deprecated(note = "replaced by connect_with_client_key", since = "0.32.2")]
    pub fn new_with_client_key(
        base_url: &Url,
        credentials: C,
        decryption_log_path: Option<&Path>,
        client_key: ClientKey,
    ) -> ZeroKMSWithClientKey<C> {
        let mut host = base_url.to_string();
        if host.ends_with('/') {
            host.pop();
        }

        // Allowed because this method is also deprecated
        #[allow(deprecated)]
        let client = ViturClient::init(host);

        ZeroKMSWithClientKey {
            client,
            credentials,
            decryption_log_path: decryption_log_path.map(|p| p.to_path_buf()),
            client_key,
        }
    }
}

impl<C, K> ZeroKMS<C, K>
where
    C: Send + Sync + 'static,
    for<'a> &'a C: AuthStrategy,
    K: Zeroize,
{
    /// Fetch a token from the credentials provider and ensure the ZeroKMS
    /// base URL has been resolved on the connection (from the token's
    /// `services` claim). The URL is only resolved once; subsequent calls
    /// skip the resolution.
    async fn get_token(&self) -> Result<stack_auth::ServiceToken, Error> {
        let token = (&self.credentials).get_token().await?;
        if !self.client.connection().has_base_url() {
            let url = token.zerokms_url()?;
            self.client.connection().ensure_base_url(url);
        }
        Ok(token)
    }

    #[deprecated(note = "removed in favor of server-side auditing", since = "0.32.2")]
    pub fn log_decryptions<P>(&self, records: &[P], access_token: &str)
    where
        P: Decryptable,
    {
        if let Some(log_path) = &self.decryption_log_path {
            // ignore log errors
            _ = log_decryptions(records, access_token, log_path);
        }
    }

    /// Create a [Keyset] (previously known as a `Keyset`) in ZeroKMS used to encrypt data.
    /// The name and description are used to identify the keyset.
    #[deprecated(note = "replaced by create_keyset", since = "0.26.0")]
    pub async fn create_dataset(&self, name: &str, description: &str) -> Result<Keyset, Error> {
        self.create_keyset(name, description).await
    }

    /// Create a [Keyset] in ZeroKMS used to encrypt data.
    /// The name and description are used to identify the keyset.
    pub async fn create_keyset(&self, name: &str, description: &str) -> Result<Keyset, Error> {
        let token = self.get_token().await?;

        self.client
            .create_keyset(name, description, token.as_str())
            .await
            .map_err(Error::from)
    }

    /// Create a [Keyset] in ZeroKMS along with a new device client in one atomic operation.
    ///
    /// This is the preferred way to provision a new device: the keyset and client are created
    /// together so there is no window where one exists without the other. The caller receives
    /// the keyset details and client key material in a single [`CreateKeysetResponse`].
    ///
    /// If the keyset name or client name already exists, the returned error will have
    /// [`Error::is_conflict`] `== true`.
    pub async fn create_keyset_with_client(
        &self,
        name: &str,
        description: &str,
        client_spec: CreateClientSpec<'_>,
    ) -> Result<CreateKeysetResponse, Error> {
        let token = self.get_token().await?;

        self.client
            .create_keyset_with_client(name, description, client_spec, token.as_str())
            .await
            .map_err(Error::from)
    }

    /// Grant a client with the given `client_id` access to a [Keyset] with an ID of `keyset_id`.
    /// For this to work, the client must already exist and have access to at least one dayaset.
    ///
    /// If you are creating a new client, use [`Self::create_client`] instead.
    /// Note that the client and keyset must be in the same workspace.
    #[deprecated(note = "replaced by grant_keyset", since = "0.26.0")]
    pub async fn grant_dataset(&self, client_id: Uuid, keyset_id: Uuid) -> Result<(), Error> {
        self.grant_keyset(client_id, keyset_id).await
    }

    /// Grant a client with the given `client_id` access to a [Keyset] with an ID of `keyset_id`.
    /// For this to work, the client must already exist and have access to at least one dayaset.
    ///
    /// If you are creating a new client, use [`Self::create_client`] instead.
    /// Note that the client and keyset must be in the same workspace.
    pub async fn grant_keyset(&self, client_id: Uuid, keyset_id: Uuid) -> Result<(), Error> {
        let token = self.get_token().await?;

        self.client
            .grant_keyset(client_id, keyset_id, token.as_str())
            .await
            .map_err(Error::from)
    }

    /// Revoke a Client with the given `client_id` access to the [Keyset] with `keyset_id`.
    /// If the client only has access to one keyset, this is the same as deleting the client.
    #[deprecated(note = "replaced by revoke_keyset", since = "0.26.0")]
    pub async fn revoke_dataset(&self, client_id: Uuid, keyset_id: Uuid) -> Result<(), Error> {
        self.revoke_keyset(client_id, keyset_id).await
    }

    /// Revoke a Client with the given `client_id` access to the [Keyset] with `keyset_id`.
    /// If the client only has access to one keyset, this is the same as deleting the client.
    pub async fn revoke_keyset(&self, client_id: Uuid, keyset_id: Uuid) -> Result<(), Error> {
        let token = self.get_token().await?;

        self.client
            .revoke_keyset(client_id, keyset_id, token.as_str())
            .await
            .map_err(Error::from)
    }

    /// List all [Keyset]s in ZeroKMS for the current workspace.
    #[deprecated(note = "replaced by list_datasets", since = "0.26.0")]
    pub async fn list_datasets(&self, include_disabled: bool) -> Result<Vec<Keyset>, Error> {
        self.list_keysets(include_disabled).await
    }

    /// List all [Keyset]s in ZeroKMS for the current workspace.
    pub async fn list_keysets(&self, include_disabled: bool) -> Result<Vec<Keyset>, Error> {
        let token = self.get_token().await?;

        self.client
            .list_keysets(token.as_str(), include_disabled)
            .await
            .map_err(Error::from)
    }

    /// Enable a [Keyset] by ID if it has been disabled.
    #[deprecated(note = "replaced by enable_keyset", since = "0.26.0")]
    pub async fn enable_dataset(&self, keyset_id: Uuid) -> Result<(), Error> {
        self.enable_keyset(keyset_id).await
    }

    /// Enable a [Keyset] by ID if it has been disabled.
    pub async fn enable_keyset(&self, keyset_id: Uuid) -> Result<(), Error> {
        let token = self.get_token().await?;

        self.client
            .enable_keyset(keyset_id, token.as_str())
            .await
            .map_err(Error::from)
    }

    /// Disable a [Keyset] by ID.
    ///
    /// A disabled keyset will deny all attempts to encrypt and decrypt data.
    #[deprecated(note = "replaced by disable_keyset", since = "0.26.0")]
    pub async fn disable_dataset(&self, keyset_id: Uuid) -> Result<(), Error> {
        self.disable_keyset(keyset_id).await
    }

    /// Disable a [Keyset] by ID.
    ///
    /// A disabled keyset will deny all attempts to encrypt and decrypt data.
    pub async fn disable_keyset(&self, keyset_id: Uuid) -> Result<(), Error> {
        let token = self.get_token().await?;

        self.client
            .disable_keyset(keyset_id, token.as_str())
            .await
            .map_err(Error::from)
    }

    /// Modify a [Keyset] by ID by setting a new name or description.
    #[deprecated(note = "replaced by modify_keyset", since = "0.26.0")]
    pub async fn modify_dataset(
        &self,
        keyset_id: Uuid,
        name: Option<&str>,
        description: Option<&str>,
    ) -> Result<(), Error> {
        self.modify_keyset(keyset_id, name, description).await
    }

    /// Modify a [Keyset] by ID by setting a new name or description.
    pub async fn modify_keyset(
        &self,
        keyset_id: Uuid,
        name: Option<&str>,
        description: Option<&str>,
    ) -> Result<(), Error> {
        let token = self.get_token().await?;

        self.client
            .modify_keyset(keyset_id, name, description, token.as_str())
            .await
            .map_err(Error::from)
    }

    /// Create a new client for the specified keyset.
    ///
    /// Clients are required to generate and retrieve keysets key a specified keyset. Use the
    /// [`ClientKey`] returned by [`CreateClientResponse`] to create a [`ZeroKMSWithClientKey`] client that can
    /// encrypt and decrypt.
    ///
    /// This [`ClientKey`] can not be retrieved again after creating the client. So it's important
    /// to keep it somewhere safe.
    ///
    /// ## [ClientKey] compromise
    ///
    /// If you suspect that a [`ClientKey`] has been compromised, you should revoke the client and create a new one.
    /// See [`Self::delete_client`] for more information.
    ///
    /// ## Create vs Grant
    ///
    /// If you are creating a new client, use this method. If you are granting access to an existing client,
    /// use [`Self::grant_keyset`] instead.
    pub async fn create_client(
        &self,
        name: &str,
        description: &str,
        keyset_id: Option<Uuid>,
    ) -> Result<CreateClientResponse, Error> {
        let token = self.get_token().await?;

        self.client
            .create_client(name, description, keyset_id, token.as_str())
            .await
            .map_err(Error::from)
    }

    /// List clients for the current workspace in ZeroKMS.
    pub async fn list_clients(&self) -> Result<Vec<KeysetClient>, Error> {
        let token = self.get_token().await?;

        self.client
            .list_clients(token.as_str())
            .await
            .map_err(Error::from)
    }

    /// Delete client by ID.
    ///
    /// Once a client is deleted it can't be used to generate or retrieve data keys.
    /// This method nullifies the [`ClientKey`] for the client.
    /// Even if an attacker has the [`ClientKey`], they can't use it to decrypt data.
    ///
    /// To revoke access only to a specific keyset, use [`Self::revoke_keyset`] instead.
    pub async fn delete_client(&self, client_id: Uuid) -> Result<DeleteClientResponse, Error> {
        let token = self.get_token().await?;
        self.client
            .delete_client(client_id, token.as_str())
            .await
            .map_err(Error::from)
    }
}

impl<C> ZeroKMSWithClientKey<C>
where
    C: Send + Sync + 'static,
    for<'a> &'a C: AuthStrategy,
{
    /// Encrypt a stream of [`EncryptPayload`] and return them as an [`EncryptedRecord`].
    /// Note that this only works when Self is a [`ZeroKMSWithClientKey`] client.
    pub async fn encrypt(
        &self,
        payloads: impl IntoIterator<Item = EncryptPayload<'_>>,
        keyset_id: Option<Uuid>,
    ) -> Result<Vec<EncryptedRecord>, Error> {
        debug!(target: "zero_kms::encrypt", "encrypting records");
        let payloads: Vec<_> = payloads.into_iter().collect();

        if payloads.is_empty() {
            debug!(target: "zero_kms::encrypt", "no records to encrypt");
            return Ok(vec![]);
        }

        debug!(target: "zero_kms::encrypt", "waiting for access token");
        let token = self.get_token().await?;

        debug!(target: "zero_kms::encrypt", "got token, encrypting");
        let res = self
            .client
            .encrypt(payloads, &self.client_key, keyset_id, token.as_str())
            .await?;

        debug!(target: "zero_kms::encrypt", "success, encrypted {} records", res.len());
        Ok(res)
    }

    /// Encrypt a single [`EncryptPayload`].
    /// Note that this only works when Self is a [`ZeroKMSWithClientKey`] client.
    pub async fn encrypt_single(
        &self,
        payload: EncryptPayload<'_>,
        keyset_id: Option<Uuid>,
    ) -> Result<EncryptedRecord, Error> {
        debug!(target: "zero_kms::encrypt_single", "encrypting record - waiting for access token");
        let token = self.get_token().await?;

        debug!(target: "zero_kms::encrypt_single", "got token, encrypting");
        let res = self
            .client
            .encrypt_single(payload, &self.client_key, keyset_id, token.as_str())
            .await?;

        debug!(target: "zero_kms::encrypt_single", "success");
        Ok(res)
    }

    /// Decrypt a stream of [`EncryptedRecord`] and return the raw decrypted binary blob.
    /// Note that this only works when Self is a [`ZeroKMSWithClientKey`] client.
    ///
    /// This function will decrypt records from *any* keyset that the client has access to.
    pub async fn decrypt<'a, P>(
        &self,
        payloads: impl IntoIterator<Item = P>,
        keyset_id: Option<Uuid>,
        service_token: Option<Cow<'a, ServiceToken>>,
        unverified_context: Option<&'a UnverifiedContext>, // TODO: make this a Cow
    ) -> Result<Vec<Vec<u8>>, Error>
    where
        P: Decryptable,
    {
        debug!(target: "zero_kms::decrypt", "decrypting records");
        let payloads: Vec<_> = payloads.into_iter().collect();

        if payloads.is_empty() {
            debug!(target: "zero_kms::decrypt", "no records to decrypt");
            return Ok(vec![]);
        }

        let token = self.resolve_token(service_token).await?;

        debug!(target: "zero_kms::decrypt", "got token, decrypting {} records", payloads.len());
        let res = self
            .client
            .decrypt(
                payloads,
                &self.client_key,
                keyset_id,
                &token,
                unverified_context,
            )
            .await?;

        debug!(target: "zero_kms::decrypt", "success, decrypted {} records", res.len());
        Ok(res)
    }

    /// Decrypt a stream of [`EncryptedRecord`] and return the raw decrypted binary blob.
    /// Note that this only works when Self is a [`ZeroKMSWithClientKey`] client.
    ///
    /// This function will decrypt records from *any* keyset that the client has access to.
    pub async fn decrypt_fallible<'a, P>(
        &self,
        payloads: impl IntoIterator<Item = P>,
        service_token: Option<Cow<'a, ServiceToken>>,
        unverified_context: Option<Cow<'a, UnverifiedContext>>,
    ) -> Result<Vec<Result<Vec<u8>, RecordDecryptError>>, Error>
    where
        P: Decryptable,
    {
        debug!(target: "zero_kms::decrypt", "decrypting records");
        let payloads: Vec<_> = payloads.into_iter().collect();

        if payloads.is_empty() {
            debug!(target: "zero_kms::decrypt", "no records to decrypt");
            return Ok(vec![]);
        }

        let token = self.resolve_token(service_token).await?;

        debug!(target: "zero_kms::decrypt", "got token, decrypting {} records", payloads.len());
        let unverified_context = unverified_context.as_deref().map(Cow::Borrowed);
        let res = self
            .client
            .decrypt_fallible(payloads, &self.client_key, &token, unverified_context)
            .await?;

        debug!(target: "zero_kms::decrypt", "success, decrypted {} records", res.len());
        Ok(res)
    }

    /// Decrypt a single [`EncryptedRecord`].
    /// Note that this only works when Self is a [`ZeroKMSWithClientKey`] client.
    pub async fn decrypt_single<'a, P>(
        &self,
        payload: P,
        keyset_id: Option<Uuid>,
        service_token: Option<Cow<'a, ServiceToken>>,
        unverified_context: Option<&UnverifiedContext>,
    ) -> Result<Vec<u8>, Error>
    where
        P: Decryptable,
    {
        let token = self.resolve_token(service_token).await?;

        debug!(target: "zero_kms::decrypt_single", "got token, decrypting record");
        let res = self
            .client
            .decrypt_single(
                payload,
                &self.client_key,
                keyset_id,
                &token,
                unverified_context,
            )
            .await?;

        debug!(target: "zero_kms::decrypt_single", "success");
        Ok(res)
    }

    pub(crate) async fn generate_data_keys<'a>(
        &self,
        payloads: impl IntoIterator<Item = GenerateKeyPayload<'_>>,
        keyset_id: Option<Uuid>,
        service_token: Option<Cow<'a, ServiceToken>>,
        unverified_context: Option<Cow<'a, UnverifiedContext>>,
    ) -> Result<Vec<DataKeyWithTag>, Error> {
        let token = self.resolve_token(service_token).await?;

        self.client
            .generate_keys(
                payloads,
                &self.client_key,
                keyset_id,
                &token,
                unverified_context,
            )
            .await
            .map_err(Error::from)
    }

    /// Internal method for loading the index key for a keyset.
    /// If keyset_id is None, it will use the client's default keyset.
    pub(crate) async fn load_keyset_index_key(
        &self,
        keyset_id: Option<IdentifiedBy>,
    ) -> Result<(Uuid, IndexKey), Error> {
        let token = self.get_token().await?;

        let (keyset, index_key) = self
            .client
            .load_keyset(&self.client_key, keyset_id, token.as_str())
            .await
            .map_err(Error::from)?;

        debug!(target: "zero_kms::load_keyset", "loaded keyset: [{}]({})", keyset.id, keyset.name);

        Ok((keyset.id, index_key))
    }

    /// If `service_token` is `Some`, use its `.token()` string.
    /// Otherwise, fall back to `self.credentials` via [`get_token`](Self::get_token).
    ///
    /// In both cases, ensures the ZeroKMS base URL has been resolved on the
    /// underlying connection (from the token's `services` claim).
    async fn resolve_token<'a>(
        &self,
        service_token: Option<Cow<'a, ServiceToken>>,
    ) -> Result<String, Error> {
        if let Some(st) = service_token {
            debug!(target: "zero_kms::resolve_token", "using service token from args");
            let raw = st.token().to_owned();
            if !self.client.connection().has_base_url() {
                let parsed = stack_auth::ServiceToken::new(stack_auth::SecretToken::new(&raw));
                let url = parsed.zerokms_url()?;
                self.client.connection().ensure_base_url(url);
            }
            Ok(raw)
        } else {
            debug!(target: "zero_kms::resolve_token", "getting token from credentials");
            let token = self.get_token().await?;
            Ok(token.as_str().to_owned())
        }
    }
}

// TODO: Move the keyset_id to the DataKeyWithTag
#[inline]
pub fn encrypt<'p, T>(
    plaintext: T,
    key: DataKeyWithTag,
    keyset_id: Option<Uuid>,
) -> Result<EncryptedRecord, Error>
where
    EncryptPayload<'p>: From<T>,
{
    let target = EncryptionTarget::new(EncryptPayload::from(plaintext), key, keyset_id);
    vitur_client::encrypt(target).map_err(Error::from)
}

#[inline]
pub fn decrypt<D>(encrypted: D, key: DataKey) -> Result<Vec<u8>, RecordDecryptError>
where
    D: Decryptable,
{
    let target = encrypted
        .into_encrypted_record()
        .map_err(|e| Error::Unexpected(e.to_string()))
        .map(|record| DecryptionTarget::new(record, key))
        .map_err(|e| RecordDecryptError {
            reason: e.to_string(),
        })?;

    vitur_client::decrypt(target)
}