cipherstash-client 0.12.5

//! 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.
//!
//! ## Key Heirarchy
//!
//! ## Workspaces, Datasets and Clients
//!
//! ZeroKMS is built around the concept of Datasets and Clients that are encapsulated within a Workspace.
//!
//! * **Workspaces** are the top-level entity that encapsulates all Datasets and Clients.
//! * **Datasets** 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 datasets.
//!
//! ## 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 Datasets 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 Dataset
//! * Grant and Revoke access to a Dataset for a client
//! * List all Datasets in a Workspace
//! * Enable and Disable a Dataset
//! * Modify a Dataset's name or description
//! * Create a new Client for a Dataset
//! * List all Clients in a Workspace
//! * Revoke (Delete) a Client
//!
//! ### Example: Create a Dataset
//!
//! ```no_run
//! use cipherstash_client::zerokms::Error;
//! use cipherstash_client::config::ZeroKMSConfig;
//!
//! #[tokio::main]
//! async fn main() -> Result<(), Error> {
//!     let zerokms = ZeroKMSConfig::builder()
//!         .with_env()
//!         .build()
//!         .expect("failed to build config")
//!         .create_client();
//!
//!     zerokms.create_dataset("my-dataset", "My dataset description").await?;
//!
//!     Ok(())
//! }
//! ```
//!
//! ## [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 Datasets 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 Dataset's configuration
//!
//! ### Example: Encrypt and Decrypt
//!
//! ```no_run
//! use cipherstash_client::zerokms::{ClientKey, EncryptPayload, Error};
//! use cipherstash_client::config::ZeroKMSConfig;
//! # // TODO: Use a mock client to run this example
//!
//! # #[tokio::main]
//! # async fn main() -> Result<(), Error> {
//! let zerokms = ZeroKMSConfig::builder()
//!   .with_env()
//!   .client_id("29f4c409-da51-496d-a1b8-cea9a9fd2660")
//!   .client_key("e27aaa358784340a822d692a99c12692cd0...")
//!   .build_with_client_key()
//!   .expect("failed to build config")
//!   .create_client();
//!
//! let message = b"Hello, World!";
//! let payload = EncryptPayload::new(message);
//! let encrypted = zerokms.encrypt_single(payload, None).await?;
//! let result = zerokms.decrypt_single(encrypted).await?;
//! assert_eq!(message, result.as_slice());
//! # Ok(())
//! # }
//! ```
//!
//! See also [ClientKey] and [EncryptPayload].
//!
//! ## Specifying Dataset ID
//!
//! When a client has access only to one dataset, ZeroKMS will automatically use that dataset for encryption and decryption.
//! However, when a client has access to multiple datasets, you need to specify the dataset ID.
//!
//! ```no_run
//! use cipherstash_client::zerokms::{ClientKey, EncryptPayload, Error};
//! use cipherstash_client::config::ZeroKMSConfig;
//! # // TODO: Use a mock client to run this example
//!
//! # #[tokio::main]
//! # async fn main() -> Result<(), Error> {
//! # let zerokms = ZeroKMSConfig::builder()
//! #  .with_env()
//! #  .client_id("29f4c409-da51-496d-a1b8-cea9a9fd2660")
//! #  .client_key("e27aaa358784340a822d692a99c12692cd0...")
//! #  .build_with_client_key()
//! #  .expect("failed to build config")
//! #  .create_client();
//! use uuid::Uuid;
//! // Dataset ID for the dataset you want to use
//! let dataset_id: Uuid = "901C968B-9AA9-453D-B024-B26739E3AC88".parse().unwrap();
//! let message = b"Hello, World!";
//! let payload = EncryptPayload::new(message);
//! let encrypted = zerokms.encrypt_single(payload, Some(dataset_id)).await?;
//! let result = zerokms.decrypt_single(encrypted).await?;
//! assert_eq!(message, result.as_slice());
//! # Ok(())
//! # }
//! ```
//!
//! ## 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.
//!
//! ### Using Descriptors
//!
//! A common use for descriptors is to identify the data being encrypted like the table, column and row ID in a database.
//!
//! ```no_run
//! # use cipherstash_client::zerokms::{ClientKey, EncryptPayload, Error};
//! # use cipherstash_client::config::ZeroKMSConfig;
//! # // TODO: Use a mock client to run this example
//! # #[tokio::main]
//! # async fn main() -> Result<(), Error> {
//! # let zerokms = ZeroKMSConfig::builder()
//! #  .with_env()
//! #  .client_id("29f4c409-da51-496d-a1b8-cea9a9fd2660")
//! #  .client_key("e27aaa358784340a822d692a99c12692cd0...")
//! #  .build_with_client_key()
//! #  .expect("failed to build config")
//! #  .create_client();
//! let message = b"Kate Smith";
//! // Descriptor for record with ID 1 in the "users" table and "name" column
//! let payload = EncryptPayload::new_with_descriptor(message, "users/name/1");
//! let encrypted = zerokms.encrypt_single(payload, None).await?;
//! let result = zerokms.decrypt_single(encrypted).await?;
//! # // FIXME: This currently doesn't work - decrypt single should take a descriptor as an argument to check
//! # // Relates to: https://linear.app/cipherstash/issue/CIP-811/prevent-decryption-of-plaintext-columns
//! # // Check that the descriptor is as we expect
//! # // assert_eq!(result.descriptor(), "users/name/1");
//! # // assert_eq!(message, result.as_slice());
//! # Ok(())
//! # }
//! ```
//!
// TODO: Link to a Glossary page for AAD above
mod errors;
mod local_log;
mod vitur_client;

use crate::credentials::{service_credentials::ServiceToken, Credentials};
use local_log::log_decryptions;
use log::debug;
use std::path::{Path, PathBuf};
use url::Url;
use uuid::Uuid;

// Re-exports
pub use errors::Error;
pub use vitur_client::{ClientKey, EncryptPayload, EncryptedRecord};
pub use zerokms_protocol::cipherstash_config::{DatasetConfig, DatasetConfigWithIndexRootKey};
pub use zerokms_protocol::{
    ClientDatasetId, CreateClientResponse, Dataset, DatasetClient, RevokeClientResponse,
};

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

pub struct ZeroKMS<C: Credentials<Token = ServiceToken>, ClientKeyState = ()> {
    client: ViturClient,
    credentials: C,
    decryption_log_path: Option<PathBuf>,
    client_key: ClientKeyState,
}

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

impl<C: Credentials<Token = ServiceToken>> ZeroKMS<C> {
    /// Create a new instance of the [`ZeroKMS`] client.
    ///
    /// In most cases it is prefered to use [`crate::config::ZeroKMSConfig::create_client`] instead of calling
    /// this manually.
    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();
        }

        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 prefered to use [`crate::config::ZeroKMSConfigWithClientKey::create_client`] instead of calling
    /// this manually.
    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();
        }

        let client = ViturClient::init(host);

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

impl<C: Credentials<Token = ServiceToken>, K> ZeroKMS<C, K> {
    pub fn log_decryptions(&self, records: &[EncryptedRecord], access_token: &str) {
        if let Some(log_path) = &self.decryption_log_path {
            // ignore log errors
            _ = log_decryptions(records, access_token, log_path);
        }
    }

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

        self.client
            .create_dataset(name, description, &access_token)
            .await
            .map_err(Error::from)
    }

    /// Grant a client with the given `client_id` access to a [Dataset] with an ID of `dataset_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 dataset must be in the same workspace.
    pub async fn grant_dataset(&self, client_id: Uuid, dataset_id: Uuid) -> Result<(), Error> {
        let access_token = self.credentials.get_token().await?.access_token();

        self.client
            .grant_dataset(client_id, dataset_id, &access_token)
            .await
            .map_err(Error::from)
    }

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

        self.client
            .revoke_dataset(client_id, dataset_id, &access_token)
            .await
            .map_err(Error::from)
    }

    /// List all [Dataset]s in ZeroKMS for the current workspace.
    pub async fn list_datasets(&self) -> Result<Vec<Dataset>, Error> {
        let access_token = self.credentials.get_token().await?.access_token();
        let show_disabled = false;

        self.client
            .list_datasets(&access_token, show_disabled)
            .await
            .map_err(Error::from)
    }

    /// Enable a [Dataset] by ID if it has been disabled.
    pub async fn enable_dataset(&self, dataset_id: Uuid) -> Result<(), Error> {
        let access_token = self.credentials.get_token().await?.access_token();

        self.client
            .enable_dataset(dataset_id, &access_token)
            .await
            .map_err(Error::from)
    }

    /// Disable a [Dataset] by ID.
    ///
    /// A disabled dataset will deny all attempts to encrypt and decrypt data.
    pub async fn disable_dataset(&self, dataset_id: Uuid) -> Result<(), Error> {
        let access_token = self.credentials.get_token().await?.access_token();

        self.client
            .disable_dataset(dataset_id, &access_token)
            .await
            .map_err(Error::from)
    }

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

        self.client
            .modify_dataset(dataset_id, name, description, &access_token)
            .await
            .map_err(Error::from)
    }

    /// Create a new client for the specified dataset.
    ///
    /// Clients are required to generate and retrieve datasets key a specified dataset. 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::revoke_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_dataset`] instead.
    pub async fn create_client(
        &self,
        name: &str,
        description: &str,
        dataset_id: Uuid,
    ) -> Result<CreateClientResponse, Error> {
        let access_token = self.credentials.get_token().await?.access_token();

        self.client
            .create_client(name, description, dataset_id, &access_token)
            .await
            .map_err(Error::from)
    }

    /// List clients for the current workspace in ZeroKMS.
    pub async fn list_clients(&self) -> Result<Vec<DatasetClient>, Error> {
        let access_token = self.credentials.get_token().await?.access_token();
        self.client
            .list_clients(&access_token)
            .await
            .map_err(Error::from)
    }

    /// Revoke a specific client by ID.
    ///
    /// Once a client is revoked 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.
    pub async fn revoke_client(&self, client_id: Uuid) -> Result<RevokeClientResponse, Error> {
        let access_token = self.credentials.get_token().await?.access_token();
        self.client
            .revoke_client(client_id, &access_token)
            .await
            .map_err(Error::from)
    }
}

impl<C: Credentials<Token = ServiceToken>> ZeroKMSWithClientKey<C> {
    /// Save a configuration file to the current dataset.
    ///
    /// The [`DatasetConfig`] is used by Proxy to store index and column encryption configuration.
    pub async fn save_dataset_config(
        &self,
        config: DatasetConfig,
    ) -> Result<DatasetConfigWithIndexRootKey, Error> {
        let access_token = self.credentials.get_token().await?.access_token();
        // TODO: Temporarily send None until we properly handle the dataset_id field
        let dataset_id = None;

        self.client
            .save_config(config, &self.client_key, &access_token, dataset_id)
            .await
            .map_err(Error::from)
    }

    /// Retrieve the [`DatasetConfig`] for the current dataset.
    pub async fn load_dataset_config(&self) -> Result<DatasetConfigWithIndexRootKey, Error> {
        let access_token = self.credentials.get_token().await?.access_token();
        let dataset_id = None;
        self.client
            .load_config(&self.client_key, &access_token, dataset_id)
            .await
            .map_err(Error::from)
    }

    /// 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<'_>>,
        dataset_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 access_token = self.credentials.get_token().await?.access_token();

        debug!(target: "zero_kms::encrypt", "got token, encrypting");
        let res = self
            .client
            .encrypt(payloads, &self.client_key, dataset_id, &access_token)
            .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<'_>,
        dataset_id: Option<Uuid>,
    ) -> Result<EncryptedRecord, Error> {
        debug!(target: "zero_kms::encrypt_single", "encrypting record - waiting for access token");
        let access_token = self.credentials.get_token().await?.access_token();

        debug!(target: "zero_kms::encrypt_single", "got token, encrypting");
        let res = self
            .client
            .encrypt_single(payload, &self.client_key, dataset_id, &access_token)
            .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.
    pub async fn decrypt(
        &self,
        payloads: impl IntoIterator<Item = EncryptedRecord>,
    ) -> Result<Vec<Vec<u8>>, Error> {
        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![]);
        }

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

        self.log_decryptions(&payloads[..], &access_token);

        debug!(target: "zero_kms::decrypt", "got token, decrypting {} records", payloads.len());
        let res = self
            .client
            .decrypt(payloads, &self.client_key, &access_token)
            .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(&self, payload: EncryptedRecord) -> Result<Vec<u8>, Error> {
        debug!(target: "zero_kms::decrypt_single", "decrypting record - waiting for access token");
        let access_token = self.credentials.get_token().await?.access_token();

        self.log_decryptions(&[payload.clone()], &access_token);

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

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