cipherstash-client 0.12.5

use aes_gcm_siv::aead::{Aead, Payload};
use aes_gcm_siv::{Aes256GcmSiv, Key as AesKey, KeyInit, Nonce};
use async_mutex::Mutex;
use key::DataKeyWithTag;
use log::trace;
use rand::SeedableRng;
use rand_chacha::ChaChaRng;
use serde::{Deserialize, Serialize};
use std::borrow::Cow;
use std::collections::HashMap;
use uuid::Uuid;
use zerokms_protocol::cipherstash_config::{DatasetConfig, DatasetConfigWithIndexRootKey};
use zerokms_protocol::{
    CreateClientRequest, CreateClientResponse, CreateDatasetRequest, Dataset, DatasetClient,
    DisableDatasetRequest, EnableDatasetRequest, GenerateKeyRequest, GenerateKeySpec, GeneratedKey,
    GrantDatasetRequest, ListClientRequest, ListDatasetRequest, LoadConfigRequest,
    ModifyDatasetRequest, RetrieveKeyRequest, RetrieveKeySpec, RetrievedKey, RevokeClientRequest,
    RevokeClientResponse, RevokeDatasetRequest, SaveConfigRequest, ViturConnection,
};
pub mod connection;
pub mod errors;
mod futures;
pub mod key;
mod retry;
pub use connection::HttpConnection;
pub use errors::*;
use futures::map_async_chunked;
pub use key::{ClientKey, DataKey};

pub use recipher::key::{GenRandom, Iv};

#[cfg(test)]
pub mod test_connection;

const INDEX_ROOT_KEY_DESCRIPTOR: &str = "dataset-config-index-root-key";

/// Options for configuring certain behaviours of the [`Client`]
pub struct ClientOpts {
    /// The maximum number of key specs that should be in each generate or retrieve request to
    /// Vitur. Having too large a number of specs per request can panic by exceeding reqwest's max
    /// body size.
    pub max_keys_per_req: usize,
    /// The maximum number of requests that will be spun up per call to `generate_data_keys` or
    /// `retrieve_data_keys`. Having too large a number of concurrent requests can result in
    /// dropped connections which will fail the calls.
    pub max_concurrent_reqs: usize,
}

impl Default for ClientOpts {
    fn default() -> Self {
        Self {
            max_keys_per_req: 500,
            max_concurrent_reqs: 5,
        }
    }
}

pub struct Client<C> {
    rand: Mutex<ChaChaRng>,
    connection: C,
    opts: ClientOpts,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EncryptedRecord {
    pub iv: Iv,
    pub ciphertext: Vec<u8>,
    pub tag: Vec<u8>,
    pub descriptor: String,
    #[serde(default)]
    pub dataset_id: Option<Uuid>,
}

impl EncryptedRecord {
    pub fn to_vec(&self) -> Result<Vec<u8>, DecryptError> {
        serde_cbor::to_vec(&self).map_err(DecryptError::FormatError)
    }

    pub fn to_hex(&self) -> Result<String, DecryptError> {
        self.to_vec().map(hex::encode)
    }

    pub fn from_slice(bytes: &[u8]) -> Result<Self, DecryptError> {
        serde_cbor::from_slice(bytes).map_err(DecryptError::FormatError)
    }

    pub fn from_hex(hexstr: impl AsRef<[u8]>) -> Result<Self, DecryptError> {
        let bytes = hex::decode(hexstr.as_ref()).map_err(DecryptError::DecodingError)?;

        Self::from_slice(&bytes)
    }
}

#[derive(Debug, Clone)]
pub struct EncryptPayload<'a> {
    pub msg: &'a [u8],
    pub descriptor: &'a str,
}

impl<'a> EncryptPayload<'a> {
    /// Create a new [`EncryptPayload`] with the given message.
    pub fn new(msg: &'a [u8]) -> Self {
        Self {
            msg,
            descriptor: "",
        }
    }

    /// Create a new [`EncryptPayload`] with the given message and descriptor.
    /// The descriptor is used as the additional authenticated data (AAD) for the encryption.
    pub fn new_with_descriptor(msg: &'a [u8], descriptor: &'a str) -> Self {
        Self { msg, descriptor }
    }
}

/// The requirements for generating a data key from Vitur.
pub struct GenerateKeyPayload<'a> {
    pub descriptor: &'a str,
}

/// The requirements for retrieving a data key from Vitur.
pub struct RetrieveKeyPayload<'a> {
    pub iv: Iv,
    pub descriptor: &'a str,
    pub tag: &'a [u8],
}

impl Client<HttpConnection> {
    pub fn init(host: String) -> Self {
        Self::with_opts(host, Default::default())
    }

    pub fn with_opts(host: String, opts: ClientOpts) -> Self {
        Self::from_connection_with_opts(HttpConnection::init(host), opts)
    }
}

impl<C: ViturConnection + Send + Sync> Client<C> {
    pub fn from_connection_with_opts(connection: C, opts: ClientOpts) -> Self {
        Self {
            connection,
            rand: Mutex::new(ChaChaRng::from_entropy()),
            opts,
        }
    }

    pub async fn create_dataset(
        &self,
        name: &str,
        description: &str,
        access_token: &str,
    ) -> Result<Dataset, CreateDatasetError> {
        let req = CreateDatasetRequest {
            name: name.into(),
            description: description.into(),
        };

        let response = self.connection.send(req, access_token).await?;

        Ok(response)
    }

    /// Grant a Client access to a Dataset
    pub async fn grant_dataset(
        &self,
        client_id: Uuid,
        dataset_id: Uuid,
        access_token: &str,
    ) -> Result<(), GrantDatasetError> {
        let req = GrantDatasetRequest {
            client_id,
            dataset_id,
        };

        self.connection.send(req, access_token).await?;

        Ok(())
    }

    /// Revoke a Client's access to a Dataset
    pub async fn revoke_dataset(
        &self,
        client_id: Uuid,
        dataset_id: Uuid,
        access_token: &str,
    ) -> Result<(), RevokeDatasetError> {
        let req = RevokeDatasetRequest {
            client_id,
            dataset_id,
        };

        self.connection.send(req, access_token).await?;

        Ok(())
    }

    pub async fn list_datasets(
        &self,
        access_token: &str,
        show_disabled: bool,
    ) -> Result<Vec<Dataset>, ListDatasetError> {
        let response = self
            .connection
            .send(ListDatasetRequest { show_disabled }, access_token)
            .await?;
        Ok(response)
    }

    pub async fn enable_dataset(
        &self,
        dataset_id: Uuid,
        access_token: &str,
    ) -> Result<(), EnableDatasetError> {
        self.connection
            .send(EnableDatasetRequest { dataset_id }, access_token)
            .await?;

        Ok(())
    }

    pub async fn disable_dataset(
        &self,
        dataset_id: Uuid,
        access_token: &str,
    ) -> Result<(), DisableDatasetError> {
        self.connection
            .send(DisableDatasetRequest { dataset_id }, access_token)
            .await?;

        Ok(())
    }

    pub async fn modify_dataset(
        &self,
        dataset_id: Uuid,
        name: Option<&str>,
        description: Option<&str>,
        access_token: &str,
    ) -> Result<(), ModifyDatasetError> {
        self.connection
            .send(
                ModifyDatasetRequest {
                    dataset_id,
                    name: name.map(Cow::Borrowed),
                    description: description.map(Cow::Borrowed),
                },
                access_token,
            )
            .await?;

        Ok(())
    }

    pub async fn create_client(
        &self,
        name: &str,
        description: &str,
        dataset_id: Uuid,
        access_token: &str,
    ) -> Result<CreateClientResponse, CreateClientError> {
        let req = CreateClientRequest {
            name: name.into(),
            description: description.into(),
            dataset_id,
        };

        let response = self.connection.send(req, access_token).await?;

        Ok(response)
    }

    pub async fn list_clients(
        &self,
        access_token: &str,
    ) -> Result<Vec<DatasetClient>, ListClientError> {
        let response = self
            .connection
            .send(ListClientRequest, access_token)
            .await?;
        Ok(response)
    }

    // TODO: Rename to delete_client
    pub async fn revoke_client(
        &self,
        client_id: Uuid,
        access_token: &str,
    ) -> Result<RevokeClientResponse, RevokeClientError> {
        let req = RevokeClientRequest {
            client_id,
            dataset_id: None,
        };

        let response = self.connection.send(req, access_token).await?;

        Ok(response)
    }

    /// Generate multiple data keys for an iterator of [`RetrieveKeyPayload`]
    pub async fn retrieve_keys(
        &self,
        keys: impl IntoIterator<Item = RetrieveKeyPayload<'_>>,
        key: &ClientKey,
        dataset_id: Option<Uuid>,
        access_token: &str,
    ) -> Result<Vec<DataKey>, RetrieveKeyError> {
        let ClientOpts {
            max_keys_per_req,
            max_concurrent_reqs: max_parallel_reqs,
        } = self.opts;

        trace!(target: "vitur_client::retrieve_keys", "preparing payloads");

        let keys = keys
            .into_iter()
            .map(
                |RetrieveKeyPayload {
                     iv,
                     descriptor,
                     tag,
                 }| RetrieveKeySpec {
                    iv,
                    descriptor: descriptor.into(),
                    tag: tag.into(),

                    // Tag version will be ignored in future version of Vitur
                    tag_version: 0,
                },
            )
            .collect::<Vec<_>>();

        trace!(target: "vitur_client::retrieve_keys", "sending requests with {max_parallel_reqs} parallel requests and {max_keys_per_req} keys per request");

        // map_async_chunked will split the retrieve key requests up into chunks and send them to
        // Vitur concurrently. The number of concurrent requests and size of the chunks are passed through from
        // ClientOpts.
        let result = map_async_chunked(
            &keys,
            |keys| async {
                let req = RetrieveKeyRequest {
                    keys: keys.into(),
                    dataset_id,
                    client_id: key.key_id
                };

                trace!(target: "vitur_client::retrieve_keys", "sending request with {} keys", keys.len());

                self.connection
                    .send(req, access_token)
                    .await
                    .map_err(RetrieveKeyError::RequestFailed)
                    .and_then(|res| {
                        // This should never happen with Vitur but check just to be sure.
                        if res.keys.len() != keys.len() {
                            return Err(RetrieveKeyError::InvalidNumberOfKeys {
                                expected: keys.len(),
                                received: res.keys.len(),
                            });
                        }

                        trace!(target: "vitur_client::retrieve_keys", "retrieved keys - creating data keys");

                        Ok(keys
                            .iter()
                            .zip(res.keys)
                            .map(
                                |(RetrieveKeySpec { iv, .. }, RetrievedKey { key_material })| {
                                    DataKey::from_key_material(key, *iv, &key_material)
                                },
                            )
                            .collect())
                    })
            },
            max_keys_per_req,
            max_parallel_reqs,
        )
        .await;

        match &result {
            Err(x) => {
                trace!(target: "vitur_client::retrieve_keys", "failed with error: {x}");
            }
            Ok(x) => {
                trace!(target: "vitur_client::retrieve_keys", "successfully generated {} keys", x.len());
            }
        }

        result
    }

    /// Generate multiple data keys for an iterator of [`GenerateKeyPayload`]
    pub async fn generate_keys(
        &self,
        keys: impl IntoIterator<Item = GenerateKeyPayload<'_>>,
        key: &ClientKey,
        dataset_id: Option<Uuid>,
        access_token: &str,
    ) -> Result<Vec<DataKeyWithTag>, GenerateKeyError> {
        let ClientOpts {
            max_keys_per_req,
            max_concurrent_reqs: max_parallel_reqs,
        } = self.opts;

        // Use a block here so that the mutex is released as soon as the IV is generated.
        // This stops the mutex being held while the requests to Vitur are being made.
        let keys = {
            trace!(target: "vitur_client::generate_keys", "waiting for rand lock");
            let mut guard = self.rand.lock().await;
            trace!(target: "vitur_client::generate_keys", "got rand lock");

            keys.into_iter()
                .map(|GenerateKeyPayload { descriptor }| {
                    GenRandom::gen_random(&mut *guard)
                        .map(|iv: Iv| GenerateKeySpec {
                            iv,
                            descriptor: descriptor.into(),
                        })
                        .map_err(GenerateKeyError::GenerateIv)
                })
                .collect::<Result<Vec<_>, _>>()?
        };

        trace!(target: "vitur_client::generate_keys", "generated {} key payloads", keys.len());
        trace!(target: "vitur_client::generate_keys", "sending requests with {max_parallel_reqs} parallel requests and {max_keys_per_req} keys per request");

        // map_async_chunked will split the generate key requests up into chunks and send them to
        // Vitur concurrently. The number of concurrent requests and size of the chunks are passed through from
        // ClientOpts.
        let result = map_async_chunked(
            &keys,
            |keys| async {
                let req = GenerateKeyRequest {
                    keys: keys.into(),
                    dataset_id,
                    client_id: key.key_id
                };

                trace!(target: "vitur_client::generate_keys", "sending request with {} keys", keys.len());

                self.connection
                    .send(req, access_token)
                    .await
                    .map_err(GenerateKeyError::RequestFailed)
                    .and_then(|res| {
                        // This should never happen with Vitur but check just to be sure.
                        if res.keys.len() != keys.len() {
                            return Err(GenerateKeyError::InvalidNumberOfKeys {
                                expected: keys.len(),
                                received: res.keys.len(),
                            });
                        }

                        trace!(target: "vitur_client::generate_keys", "sending request with {} keys", keys.len());

                        Ok(keys
                            .iter()
                            .zip(res.keys)
                            .map(
                                |(
                                    GenerateKeySpec { iv, .. },
                                    GeneratedKey { key_material, tag },
                                )| {
                                    DataKeyWithTag::from_key_material(key, *iv, &key_material, tag)
                                },
                            )
                            .collect())
                    })
            },
            max_keys_per_req,
            max_parallel_reqs,
        )
        .await;

        match &result {
            Err(x) => {
                trace!(target: "vitur_client::generate_keys", "failed with error: {x}");
            }
            Ok(x) => {
                trace!(target: "vitur_client::generate_keys", "successfully generated {} keys", x.len());
            }
        }

        result
    }

    /// Encrypt multiple message and descriptor pairs using randomly generated data keys.
    pub async fn encrypt(
        &self,
        payloads: impl IntoIterator<Item = EncryptPayload<'_>>,
        key: &ClientKey,
        dataset_id: Option<Uuid>,
        access_token: &str,
    ) -> Result<Vec<EncryptedRecord>, EncryptError> {
        let payloads = payloads.into_iter().collect::<Vec<_>>();
        let mut output = Vec::with_capacity(payloads.len());

        trace!(target: "vitur_client::encrypt", "generating {} keys", payloads.len());

        let keys = self
            .generate_keys(
                payloads
                    .iter()
                    .map(|EncryptPayload { descriptor, .. }| GenerateKeyPayload { descriptor }),
                key,
                dataset_id,
                access_token,
            )
            .await?;

        trace!(target: "vitur_client::encrypt", "generated {} keys - encrypting records", keys.len());

        for (EncryptPayload { msg, descriptor }, DataKeyWithTag { key, tag }) in
            payloads.into_iter().zip(keys)
        {
            let DataKey { iv, key } = key;
            let key = AesKey::<Aes256GcmSiv>::from_slice(&key);

            // We're using the first 12 bytes of the IV as the nonce for AES-GCM-SIV.
            // This isn't ideal but it's not possible at the moment to use a 16 byte nonce.
            let nonce = Nonce::from_slice(&iv[..12]);

            let cipher = Aes256GcmSiv::new(key);

            let ciphertext = cipher
                .encrypt(
                    nonce,
                    Payload {
                        msg,
                        aad: descriptor.as_bytes(),
                    },
                )
                .map_err(EncryptError::FailedToEncrypt)?;

            output.push(EncryptedRecord {
                iv,
                ciphertext,
                tag,
                descriptor: descriptor.to_string(),
                dataset_id,
            })
        }

        trace!(target: "vitur_client::encrypt", "success - encrypted {} records", output.len());

        Ok(output)
    }

    /// Encrypt a single message with a particular descriptor using a randomly generated data key.
    pub async fn encrypt_single(
        &self,
        payload: EncryptPayload<'_>,
        key: &ClientKey,
        dataset_id: Option<Uuid>,
        access_token: &str,
    ) -> Result<EncryptedRecord, EncryptError> {
        let mut vec = self
            .encrypt([payload], key, dataset_id, access_token)
            .await?;
        debug_assert_eq!(vec.len(), 1);
        Ok(vec.remove(0))
    }

    /// Decrypt multiple [`EncryptedRecord`]s using their Vitur data keys and descriptors.
    pub async fn decrypt(
        &self,
        payloads: impl IntoIterator<Item = EncryptedRecord>,
        key: &ClientKey,
        access_token: &str,
    ) -> Result<Vec<Vec<u8>>, DecryptError> {
        let payloads = payloads.into_iter().collect::<Vec<_>>();
        let mut output: Vec<Option<Vec<u8>>> = vec![None; payloads.len()];

        trace!(target: "vitur_client::decrypt", "retrieving keys");

        let mut grouped_payloads = HashMap::<Option<Uuid>, Vec<(usize, EncryptedRecord)>>::new();

        for (index, record) in payloads.into_iter().enumerate() {
            grouped_payloads
                .entry(record.dataset_id)
                .or_default()
                .push((index, record))
        }

        for (dataset_id, records) in grouped_payloads.into_iter() {
            let keys = self
                .retrieve_keys(
                    records.iter().map(|(_, record)| RetrieveKeyPayload {
                        descriptor: &record.descriptor,
                        iv: record.iv,
                        tag: &record.tag,
                    }),
                    key,
                    dataset_id,
                    access_token,
                )
                .await
                .inspect_err(|_| {
                    trace!(target: "vitur_client::decrypt", "failed to retrieve keys");
                })?;

            trace!(target: "vitur_client::decrypt", "retrieved keys - decrypting records");

            let keys_len = keys.len();

            for (record, DataKey { key, .. }) in records.into_iter().zip(keys) {
                let (
                    index,
                    EncryptedRecord {
                        iv,
                        ciphertext,
                        descriptor,
                        ..
                    },
                ) = record;

                let key = AesKey::<Aes256GcmSiv>::from_slice(&key);

                // We're using the first 12 bytes of the IV as the nonce for AES-GCM-SIV.
                // This isn't ideal but it's not possible at the moment to use a 16 byte nonce.
                let nonce = Nonce::from_slice(&iv[..12]);

                let cipher = Aes256GcmSiv::new(key);

                let plaintext = cipher
                    .decrypt(
                        nonce,
                        Payload {
                            msg: &ciphertext,
                            aad: descriptor.as_bytes(),
                        },
                    )
                    .map_err(DecryptError::FailedToDecrypt)?;

                output[index] = Some(plaintext);
            }

            trace!(target: "vitur_client::decrypt", "decrypted {keys_len} records");
        }

        output
            .into_iter()
            .map(|record| {
                record.ok_or_else(|| {
                    DecryptError::Internal(
                        "Record was None but everything should be Some(_) as this point"
                            .to_string(),
                    )
                })
            })
            .collect()
    }

    /// Decrypt a single [`EncryptedRecord`] using its Vitur data key and descriptor.
    pub async fn decrypt_single(
        &self,
        payload: EncryptedRecord,
        key: &ClientKey,
        access_token: &str,
    ) -> Result<Vec<u8>, DecryptError> {
        let mut vec = self.decrypt([payload], key, access_token).await?;
        debug_assert_eq!(vec.len(), 1);
        Ok(vec.remove(0))
    }

    pub async fn save_config(
        &self,
        dataset_config: DatasetConfig,
        key: &ClientKey,
        access_token: &str,
        dataset_id: Option<Uuid>,
    ) -> Result<DatasetConfigWithIndexRootKey, SaveConfigError> {
        // Look for an existing dataset_config and attempts to reuse the index_root_key
        // This allows uploading the config multiple times while preserving the existing index_root_key
        // If an existing dataset_config is found, grab the encrypted key, and decrypts it
        // Otherwise, generate a new index_root_key and encrypts it
        // TODO: Should distinguish between an api error vs dataset_config not exist
        let (index_root_key, encrypted_index_root_key): ([u8; 32], EncryptedRecord) = {
            let load_config_req = LoadConfigRequest {
                client_id: key.key_id,
                dataset_id,
            };

            match self.connection.send(load_config_req, access_token).await {
                Ok(response) => {
                    let encrypted_index_root_key =
                        EncryptedRecord::from_slice(&response.encrypted_index_root_key)
                            .map_err(SaveConfigError::DeserializeEncryptedRootKey)?;

                    let index_root_key: [u8; 32] = self
                        .decrypt_single(encrypted_index_root_key.clone(), key, access_token)
                        .await
                        .map_err(SaveConfigError::DecryptRootKey)?
                        .try_into()
                        .map_err(|e: Vec<u8>| SaveConfigError::InvalidIndexRootKeySize(e.len()))?;

                    (index_root_key, encrypted_index_root_key)
                }
                Err(_e) => {
                    let index_root_key: [u8; 32] = {
                        trace!(target: "vitur_client::save_config", "waiting for rand lock");
                        let mut guard = self.rand.lock().await;
                        trace!(target: "vitur_client::save_config", "got rand lock");
                        GenRandom::gen_random(&mut *guard)
                            .map_err(SaveConfigError::CreateRootKey)?
                    };

                    let encrypted_index_root_key = self
                        .encrypt_single(
                            EncryptPayload {
                                msg: &index_root_key,
                                descriptor: INDEX_ROOT_KEY_DESCRIPTOR,
                            },
                            key,
                            dataset_id,
                            access_token,
                        )
                        .await
                        .map_err(SaveConfigError::EncryptRootKey)?;

                    (index_root_key, encrypted_index_root_key)
                }
            }
        };

        let req = SaveConfigRequest {
            client_id: key.key_id,
            encrypted_index_root_key: encrypted_index_root_key
                .to_vec()
                .map_err(SaveConfigError::SerializeEncryptedRootKey)?,
            dataset_config: Cow::Borrowed(&dataset_config),
            dataset_id,
        };

        self.connection.send(req, access_token).await?;

        Ok(DatasetConfigWithIndexRootKey {
            config: dataset_config,
            index_root_key,
        })
    }

    pub async fn load_config(
        &self,
        key: &ClientKey,
        access_token: &str,
        dataset_id: Option<Uuid>,
    ) -> Result<DatasetConfigWithIndexRootKey, LoadConfigError> {
        let req = LoadConfigRequest {
            client_id: key.key_id,
            dataset_id,
        };

        let response = self.connection.send(req, access_token).await?;

        let record = EncryptedRecord::from_slice(&response.encrypted_index_root_key)
            .map_err(LoadConfigError::DeserializeEncryptedRootKey)?;

        let index_root_key = self
            .decrypt_single(record, key, access_token)
            .await?
            .try_into()
            .map_err(|e: Vec<u8>| LoadConfigError::InvalidIndexRootKeySize(e.len()))?;

        Ok(DatasetConfigWithIndexRootKey {
            index_root_key,
            config: response.dataset_config,
        })
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::sync::{Arc, Mutex};
    use test_connection::*;
    use thiserror::Error;
    use uuid::uuid;
    use zerokms_protocol::*;

    fn client_key() -> ClientKey {
        ClientKey::from_bytes(
            uuid!("00000000-0000-0000-0000-000000000000"),
            &hex::decode("a4627031a16b7065726d75746174696f6e90010b0a00020d090c07080f060503040e6770325f66726f6da16b7065726d75746174696f6e900d03080c040b0f02060e05070100090a6570325f746fa16b7065726d75746174696f6e900b0a09010e0c020003050f0d07060804627033a16b7065726d75746174696f6e98211819130d11100315120c0917080704181c1818181f0e181e1820181a181b060116000f05181d140a0b02").expect("Failed to decode key material")).expect("Failed to create client key")
    }

    fn build_client(
        callback: impl FnOnce(TestConnectionBuilder) -> TestConnectionBuilder,
    ) -> Client<TestConnection> {
        Client::from_connection_with_opts(
            callback(TestConnectionBuilder::new()).build(),
            Default::default(),
        )
    }

    #[derive(Error, Debug)]
    #[error("{0}")]
    struct TestConnectionError(String);

    #[tokio::test]
    async fn test_save_config_root_key_random() {
        let encrypted_index_root_key = Arc::new(Mutex::new(Vec::<u8>::new()));

        let client = build_client(|builder| {
            builder
                .add_failed_response::<LoadConfigRequest>(ViturRequestError::response(
                    "Server returned failure response",
                    TestConnectionError("Status: 404, Body: Not Found".into()),
                ))
                .add_success_response::<GenerateKeyRequest>(GenerateKeyResponse {
                    keys: vec![GeneratedKey {
                        key_material: vec![0; 528].into(),
                        tag: vec![],
                    }],
                })
                .add_success_response::<RetrieveKeyRequest>(RetrieveKeyResponse {
                    keys: vec![RetrievedKey {
                        key_material: vec![0; 528].into(),
                    }],
                })
                .add_success_response::<SaveConfigRequest>(SaveConfigResponse {})
                .add_effect({
                    let encrypted_index_root_key = encrypted_index_root_key.clone();
                    move |mut body: SaveConfigRequest| {
                        // Fish out the generated key from the request
                        std::mem::swap(
                            &mut *encrypted_index_root_key.lock().unwrap(),
                            &mut body.encrypted_index_root_key,
                        );
                    }
                })
        });

        let key = client_key();

        client
            .save_config(DatasetConfig::init(), &key, "token", None)
            .await
            .unwrap();

        let record =
            EncryptedRecord::from_slice(&encrypted_index_root_key.lock().unwrap()).unwrap();

        let decrypted_root_key = client
            .decrypt_single(record, &key, "token")
            .await
            .expect("Failed to decrypt root key");

        assert_eq!(decrypted_root_key.len(), 32);
    }

    #[tokio::test]
    async fn test_save_config_with_existing_key() {
        let key = client_key();
        let existing_index_root_key = [2_u8; 32];

        let key_client = build_client(|builder| {
            builder.add_success_response::<GenerateKeyRequest>(GenerateKeyResponse {
                keys: vec![GeneratedKey {
                    key_material: vec![0; 528].into(),
                    tag: vec![],
                }],
            })
        });

        let encrypted_existing_index_root_key = key_client
            .encrypt_single(
                EncryptPayload {
                    descriptor: "dataset-config-index-root-key",
                    msg: &existing_index_root_key,
                },
                &key,
                None,
                "token",
            )
            .await
            .expect("Failed to encrypt root key")
            .to_vec()
            .unwrap();

        let encrypted_index_root_key = Arc::new(Mutex::new(Vec::<u8>::new()));

        let client = build_client(|builder| {
            builder
                .add_success_response::<LoadConfigRequest>(LoadConfigResponse {
                    encrypted_index_root_key: encrypted_existing_index_root_key.clone(),
                    dataset_config: DatasetConfig::init(),
                })
                .add_success_response::<RetrieveKeyRequest>(RetrieveKeyResponse {
                    keys: vec![RetrievedKey {
                        key_material: vec![0; 528].into(),
                    }],
                })
                .add_success_response::<GenerateKeyRequest>(GenerateKeyResponse {
                    keys: vec![GeneratedKey {
                        key_material: vec![0; 528].into(),
                        tag: vec![],
                    }],
                })
                .add_success_response::<RetrieveKeyRequest>(RetrieveKeyResponse {
                    keys: vec![RetrievedKey {
                        key_material: vec![0; 528].into(),
                    }],
                })
                .add_success_response::<SaveConfigRequest>(SaveConfigResponse {})
                .add_effect({
                    let encrypted_index_root_key = encrypted_index_root_key.clone();
                    move |mut body: SaveConfigRequest| {
                        // Fish out the generated key from the request
                        std::mem::swap(
                            &mut *encrypted_index_root_key.lock().unwrap(),
                            &mut body.encrypted_index_root_key,
                        );
                    }
                })
        });

        let key = client_key();

        client
            .save_config(DatasetConfig::init(), &key, "token", None)
            .await
            .unwrap();

        let record =
            EncryptedRecord::from_slice(&encrypted_index_root_key.lock().unwrap()).unwrap();

        let decrypted_root_key = client
            .decrypt_single(record, &key, "token")
            .await
            .expect("Failed to decrypt root key");

        // Check the decrypted root key is the same as the provided one
        assert_eq!(decrypted_root_key, vec![2; 32]);

        // Check the encrypted_existing_index_root_key was not changed
        assert_eq!(
            encrypted_index_root_key.lock().unwrap().to_vec(),
            encrypted_existing_index_root_key
        );
    }

    #[tokio::test]
    async fn test_load_config() {
        let key = client_key();

        let key_client = build_client(|builder| {
            builder.add_success_response::<GenerateKeyRequest>(GenerateKeyResponse {
                keys: vec![GeneratedKey {
                    key_material: vec![0; 528].into(),
                    tag: vec![],
                }],
            })
        });

        let encrypted_index_root_key = key_client
            .encrypt_single(
                EncryptPayload {
                    descriptor: "dataset-config-index-root-key",
                    msg: &[3; 32],
                },
                &key,
                None,
                "token",
            )
            .await
            .expect("Failed to encrypt root key")
            .to_vec()
            .unwrap();

        let client = build_client(move |builder| {
            builder
                .add_success_response::<RetrieveKeyRequest>(RetrieveKeyResponse {
                    keys: vec![RetrievedKey {
                        key_material: vec![0; 528].into(),
                    }],
                })
                .add_success_response::<LoadConfigRequest>(LoadConfigResponse {
                    dataset_config: DatasetConfig::init(),
                    encrypted_index_root_key,
                })
        });

        let config = client
            .load_config(&key, "token", None)
            .await
            .expect("Failed to load config");

        assert_eq!(config.index_root_key, [3; 32]);
    }

    #[tokio::test]
    async fn test_load_config_key_too_small() {
        let key = client_key();

        let key_client = build_client(|builder| {
            builder.add_success_response::<GenerateKeyRequest>(GenerateKeyResponse {
                keys: vec![GeneratedKey {
                    key_material: vec![0; 528].into(),
                    tag: vec![],
                }],
            })
        });

        let encrypted_index_root_key = key_client
            .encrypt_single(
                EncryptPayload {
                    descriptor: "dataset-config-index-root-key",
                    msg: &[3; 16],
                },
                &key,
                None,
                "token",
            )
            .await
            .expect("Failed to encrypt root key")
            .to_vec()
            .unwrap();

        let client = build_client(move |builder| {
            builder
                .add_success_response::<RetrieveKeyRequest>(RetrieveKeyResponse {
                    keys: vec![RetrievedKey {
                        key_material: vec![0; 528].into(),
                    }],
                })
                .add_success_response::<LoadConfigRequest>(LoadConfigResponse {
                    dataset_config: DatasetConfig::init(),
                    encrypted_index_root_key,
                })
        });

        let err = client
            .load_config(&key, "token", None)
            .await
            .expect_err("Expected loading config to fail");

        assert_eq!(err.to_string(), "Invalid index root key length: 16");
    }
}