zerokms_protocol/

lib.rs

mod base64_array;
mod base64_vec;
mod error;
mod identified_by;

use cts_common::claims::{
    ClientPermission, DataKeyPermission, KeysetPermission, Permission, Scope,
};
pub use identified_by::*;

mod unverified_context;

use serde::{Deserialize, Serialize};
use std::{
    borrow::Cow,
    fmt::{self, Debug, Display, Formatter},
    ops::Deref,
};
use utoipa::ToSchema;
use uuid::Uuid;
use validator::Validate;
use zeroize::{Zeroize, ZeroizeOnDrop};

pub use cipherstash_config;
/// Re-exports
pub use error::*;

pub use crate::unverified_context::{UnverifiedContext, UnverifiedContextValue};
pub use crate::{IdentifiedBy, Name};
pub mod testing;

/// The longest `descriptor` a data-key request may carry, in bytes.
///
/// ZeroKMS derives key material over a fixed 512-byte block holding the
/// descriptor, so a longer one cannot be bound. Clients check a descriptor
/// against this before a request is built; the server is expected to refuse
/// a longer one as a bad request (BUG-309 tracks making it do so — today it
/// panics on the copy).
pub const MAX_DESCRIPTOR_LEN: usize = 512;

pub trait ViturResponse: Serialize + for<'de> Deserialize<'de> + Send {}

pub trait ViturRequest: Serialize + for<'de> Deserialize<'de> + Sized + Send {
    type Response: ViturResponse;

    const SCOPE: Scope;
    const ENDPOINT: &'static str;
}

/// The type of client to create.
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, ToSchema)]
#[serde(rename_all = "snake_case")]
pub enum ClientType {
    Device,
}

/// Specification for creating a client alongside a keyset.
#[derive(Debug, Serialize, Deserialize, Validate, ToSchema)]
pub struct CreateClientSpec<'a> {
    pub client_type: ClientType,
    /// A human-readable name for the client.
    #[validate(length(min = 1, max = 64))]
    #[schema(value_type = String, min_length = 1, max_length = 64)]
    pub name: Cow<'a, str>,
}

/// Details of a client created as part of a [CreateKeysetRequest].
#[derive(Debug, Serialize, Deserialize, ToSchema)]
pub struct CreatedClient {
    pub id: Uuid,
    /// Base64-encoded 32-byte key material for the client. Store this securely.
    #[schema(value_type = String, format = Byte)]
    pub client_key: ViturKeyMaterial,
}

/// Response to a [CreateKeysetRequest].
///
/// Contains the created keyset and optionally a client if one was requested.
#[derive(Debug, Serialize, Deserialize, ToSchema)]
pub struct CreateKeysetResponse {
    #[serde(flatten)]
    pub keyset: Keyset,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub client: Option<CreatedClient>,
}

impl ViturResponse for CreateKeysetResponse {}

fn validate_keyset_name(name: &str) -> Result<(), validator::ValidationError> {
    if name.eq_ignore_ascii_case("default") {
        let mut err = validator::ValidationError::new("reserved_name");
        err.message =
            Some("the name 'default' is reserved for the workspace default keyset".into());
        return Err(err);
    }
    if !name
        .chars()
        .all(|c| c.is_ascii_alphanumeric() || c == '_' || c == '-' || c == '/')
    {
        let mut err = validator::ValidationError::new("invalid_characters");
        err.message = Some("name must only contain: A-Z a-z 0-9 _ - /".into());
        return Err(err);
    }
    Ok(())
}

/// Request message to create a new [Keyset] with the given name and description.
///
/// Requires the `dataset:create` scope.
#[derive(Debug, Serialize, Deserialize, Validate, ToSchema)]
pub struct CreateKeysetRequest<'a> {
    /// A human-readable name for the keyset.
    /// Must be 1–64 characters using only `A-Z a-z 0-9 _ - /`. The name `default` is reserved.
    #[validate(length(min = 1, max = 64), custom(function = "validate_keyset_name"))]
    #[schema(value_type = String, min_length = 1, max_length = 64, pattern = r"^[A-Za-z0-9_\-/]+$")]
    pub name: Cow<'a, str>,
    /// A description of the keyset.
    #[validate(length(min = 1, max = 256))]
    #[schema(value_type = String, min_length = 1, max_length = 256)]
    pub description: Cow<'a, str>,
    #[validate(nested)]
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub client: Option<CreateClientSpec<'a>>,
}

impl ViturRequest for CreateKeysetRequest<'_> {
    type Response = CreateKeysetResponse;

    const ENDPOINT: &'static str = "create-keyset";
    const SCOPE: Scope = Scope::with_permission(Permission::Keyset(KeysetPermission::Create));
}

/// Request message to list all [Keyset]s.
///
/// Requires the `dataset:list` scope.
/// Response is a vector of [Keyset]s.
#[derive(Default, Debug, Serialize, Deserialize, ToSchema)]
pub struct ListKeysetRequest {
    #[serde(default)]
    pub show_disabled: bool,
}

impl ViturRequest for ListKeysetRequest {
    type Response = Vec<Keyset>;

    const ENDPOINT: &'static str = "list-keysets";
    const SCOPE: Scope = Scope::with_permission(Permission::Keyset(KeysetPermission::List));
}

/// Struct representing a keyset.
/// This is the response to a [CreateKeysetRequest] and a in a vector in the response to a [ListKeysetRequest].
#[derive(Debug, Serialize, Deserialize, PartialEq, Eq, ToSchema)]
pub struct Keyset {
    pub id: Uuid,
    pub name: String,
    pub description: String,
    pub is_disabled: bool,
    #[serde(default)]
    pub is_default: bool,
}

impl ViturResponse for Vec<Keyset> {}

/// Represents an empty response for requests that don't return any data.
#[derive(Default, Debug, Serialize, Deserialize, ToSchema)]
pub struct EmptyResponse {}

impl ViturResponse for EmptyResponse {}

/// Request message to create a new client with the given name and description.
///
/// If `keyset_id` is omitted, the workspace's default keyset is used (created if necessary).
///
/// Requires the `client:create` scope.
/// Response is a [CreateClientResponse].
#[derive(Debug, Serialize, Deserialize, ToSchema)]
pub struct CreateClientRequest<'a> {
    /// The keyset to associate the client with. Accepts a UUID or a name string.
    /// If omitted, the workspace's default keyset is used.
    #[serde(alias = "dataset_id", default, skip_serializing_if = "Option::is_none")]
    #[schema(value_type = Option<String>, example = "550e8400-e29b-41d4-a716-446655440000")]
    pub keyset_id: Option<IdentifiedBy>,
    /// A human-readable name for the client.
    #[schema(value_type = String)]
    pub name: Cow<'a, str>,
    /// A description of the client.
    #[schema(value_type = String)]
    pub description: Cow<'a, str>,
}

impl ViturRequest for CreateClientRequest<'_> {
    type Response = CreateClientResponse;

    const ENDPOINT: &'static str = "create-client";
    const SCOPE: Scope = Scope::with_permission(Permission::Client(ClientPermission::Create));
}

/// Response message to a [CreateClientRequest].
///
/// Contains the `client_id` and the `client_key`, the latter being a base64 encoded 32 byte key.
/// The `client_key` should be considered sensitive and should be stored securely.
#[derive(Debug, Serialize, Deserialize, ToSchema)]
pub struct CreateClientResponse {
    /// The unique ID of the newly created client.
    pub id: Uuid,
    /// The ID of the keyset this client is associated with.
    #[serde(rename = "dataset_id")]
    pub keyset_id: Uuid,
    /// The name of the client.
    pub name: String,
    /// The description of the client.
    pub description: String,
    /// Base64-encoded 32-byte key material for the client. Store this securely.
    #[schema(value_type = String, format = Byte)]
    pub client_key: ViturKeyMaterial,
}

impl ViturResponse for CreateClientResponse {}

/// Request message to list clients.
///
/// If `keyset_id` is provided, only clients granted access to that keyset are returned;
/// otherwise all clients in the workspace are returned.
///
/// Requires the `client:list` scope.
/// Response is a vector of [KeysetClient]s.
#[derive(Debug, Default, Serialize, Deserialize, ToSchema)]
pub struct ListClientRequest {
    /// Optional keyset filter — either a keyset UUID or a keyset name (resolved within the
    /// workspace). Only clients granted that keyset are returned; if omitted, all clients in the
    /// workspace are returned.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub keyset_id: Option<IdentifiedBy>,
}

impl ViturRequest for ListClientRequest {
    type Response = Vec<KeysetClient>;

    const ENDPOINT: &'static str = "list-clients";
    const SCOPE: Scope = Scope::with_permission(Permission::Client(ClientPermission::List));
}

/// Struct representing the keyset ids associated with a client
/// which could be a single keyset or multiple keysets.
#[derive(Debug, Deserialize, Serialize, PartialEq, Eq, ToSchema)]
#[serde(untagged)]
pub enum ClientKeysetId {
    Single(Uuid),
    Multiple(Vec<Uuid>),
}

/// A `Uuid` is comparable with `ClientKeysetId` if the `ClientKeysetId` is a `Single` variant.
impl PartialEq<Uuid> for ClientKeysetId {
    fn eq(&self, other: &Uuid) -> bool {
        if let ClientKeysetId::Single(id) = self {
            id == other
        } else {
            false
        }
    }
}

/// Response type for a [ListClientRequest].
#[derive(Debug, Serialize, Deserialize, ToSchema)]
pub struct KeysetClient {
    pub id: Uuid,
    #[serde(alias = "dataset_id")]
    pub keyset_id: ClientKeysetId,
    pub name: String,
    pub description: String,
    pub created_by: Option<String>,
}

impl ViturResponse for Vec<KeysetClient> {}

/// Request message to delete a client and all associated authority keys.
///
/// Requires the `client:revoke` scope.
/// Response is an [DeleteClientResponse].
#[derive(Debug, Serialize, Deserialize, ToSchema)]
pub struct DeleteClientRequest {
    pub client_id: Uuid,
}

impl ViturRequest for DeleteClientRequest {
    type Response = DeleteClientResponse;

    const ENDPOINT: &'static str = "delete-client";
    const SCOPE: Scope = Scope::with_permission(Permission::Client(ClientPermission::Delete));
}

#[derive(Default, Debug, Serialize, Deserialize, ToSchema)]
pub struct DeleteClientResponse {}

impl ViturResponse for DeleteClientResponse {}

/// Key material type used in [GenerateKeyRequest] and [RetrieveKeyRequest] as well as [CreateClientResponse].
#[derive(Serialize, Deserialize, Zeroize, ZeroizeOnDrop)]
pub struct ViturKeyMaterial(#[serde(with = "base64_vec")] Vec<u8>);
opaque_debug::implement!(ViturKeyMaterial);

impl From<Vec<u8>> for ViturKeyMaterial {
    fn from(inner: Vec<u8>) -> Self {
        Self(inner)
    }
}

impl Deref for ViturKeyMaterial {
    type Target = [u8];

    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq, Serialize, Deserialize, Zeroize)]
#[serde(transparent)]
pub struct KeyId(#[serde(with = "base64_array")] [u8; 16]);

impl KeyId {
    pub fn into_inner(self) -> [u8; 16] {
        self.0
    }
}

impl Display for KeyId {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        write!(f, "{}", const_hex::encode(self.0))
    }
}

impl From<[u8; 16]> for KeyId {
    fn from(inner: [u8; 16]) -> Self {
        Self(inner)
    }
}

impl AsRef<[u8; 16]> for KeyId {
    fn as_ref(&self) -> &[u8; 16] {
        &self.0
    }
}

/// Represents generated data key material which is used by the client to derive data keys with its own key material.
///
/// Returned in the response to a [GenerateKeyRequest].
#[derive(Debug, Serialize, Deserialize, ToSchema)]
pub struct GeneratedKey {
    #[schema(value_type = String, format = Byte)]
    pub key_material: ViturKeyMaterial,
    // FIXME: Use Vitamin C Equatable type
    #[serde(with = "base64_vec")]
    #[schema(value_type = String, format = Byte)]
    pub tag: Vec<u8>,
    /// The decryption policy with server-generated MAC (tag_version >= 1 only).
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub decryption_policy: Option<DecryptionPolicy>,
}

/// Response to a [GenerateKeyRequest].
#[derive(Debug, Serialize, Deserialize, ToSchema)]
pub struct GenerateKeyResponse {
    pub keys: Vec<GeneratedKey>,
}

impl ViturResponse for GenerateKeyResponse {}

/// A specification for generating a data key used in a [GenerateKeyRequest].
#[derive(Debug, Serialize, Deserialize, Clone, ToSchema)]
pub struct GenerateKeySpec<'a> {
    // FIXME: Remove ID and have the server generate it instead
    #[serde(alias = "id")]
    #[schema(value_type = String, format = Byte)]
    pub iv: KeyId,
    // TODO: Deprecate descriptor in favor of context
    #[schema(value_type = String)]
    pub descriptor: Cow<'a, str>,

    #[serde(default)]
    #[schema(value_type = Vec<Context>)]
    pub context: Cow<'a, [Context]>,

    /// Optional decryption policy for OR-style lock context.
    /// When present, tag_version=1 is used (context-free base tag + policy MAC).
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub decryption_policy: Option<DecryptionPolicy>,
}

impl<'a> GenerateKeySpec<'a> {
    pub fn new(iv: [u8; 16], descriptor: &'a str) -> Self {
        Self {
            iv: KeyId(iv),
            descriptor: Cow::from(descriptor),
            context: Default::default(),
            decryption_policy: None,
        }
    }

    pub fn new_with_context(
        iv: [u8; 16],
        descriptor: &'a str,
        context: Cow<'a, [Context]>,
    ) -> Self {
        Self {
            iv: KeyId(iv),
            descriptor: Cow::from(descriptor),
            context,
            decryption_policy: None,
        }
    }

    pub fn new_with_policy(iv: [u8; 16], descriptor: &'a str, policy: DecryptionPolicy) -> Self {
        Self {
            iv: KeyId(iv),
            descriptor: Cow::from(descriptor),
            context: Default::default(),
            decryption_policy: Some(policy),
        }
    }
}
/// An identity claim condition in a decryption policy.
///
/// When `value` is `None`, the server resolves the claim value from the caller's JWT
/// at key generation time (secure default). When `value` is `Some`, the provided value
/// is used as-is (for cross-identity use cases like admin encrypting for a user).
///
/// The resolved policy (with all values filled in) is returned in the `GeneratedKey`
/// response and stored alongside the ciphertext.
#[derive(Debug, Serialize, Deserialize, Clone, PartialEq, Eq, ToSchema)]
pub struct PolicyCondition {
    /// The JWT claim name (e.g., "sub", "actor_id").
    pub claim: String,
    /// The expected claim value. When `None`, resolved from the caller's JWT at generation time.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub value: Option<String>,
}

/// A decryption policy: flat OR of identity claim conditions.
///
/// Used with `tag_version=1`. The policy conditions are included in the tag HMAC,
/// so stripping or swapping the policy causes a tag mismatch.
/// Stored alongside the ciphertext.
#[derive(Debug, Serialize, Deserialize, Clone, PartialEq, Eq, ToSchema)]
pub struct DecryptionPolicy {
    pub conditions: Vec<PolicyCondition>,
}

/// Represents a contextual attribute for a data key which is used to "lock" the key to a specific context.
/// Context attributes are included key tag generation which is in turn used as AAD in the final encryption step in the client.
/// Context attributes should _never_ include any sensitive information.
#[derive(Debug, Serialize, Deserialize, Clone, ToSchema)]
// TODO: Use Cow?
pub enum Context {
    /// A tag that can be used to identify the key.
    Tag(String),

    /// A key-value pair that can be used to identify the key.
    /// For example, a key-value pair could be `("user_id", "1234")`.
    Value(String, String),

    /// A claim from the identity of the principal that is requesting the key.
    /// The claim value is read from the claims list after token verification and prior to key generation.
    ///
    /// For example, a claim could be `"sub"`.
    #[serde(alias = "identityClaim")]
    IdentityClaim(String),
}

impl Context {
    pub fn new_tag(tag: impl Into<String>) -> Self {
        Self::Tag(tag.into())
    }

    pub fn new_value(key: impl Into<String>, value: impl Into<String>) -> Self {
        Self::Value(key.into(), value.into())
    }

    pub fn new_identity_claim(claim: &str) -> Self {
        Self::IdentityClaim(claim.to_string())
    }
}

/// A request message to generate a data key made on behalf of a client
/// in the given keyset.
///
/// Requires the `data_key:generate` scope.
/// Response is a [GenerateKeyResponse].
///
/// See also [GenerateKeySpec].
#[derive(Debug, Serialize, Deserialize, ToSchema)]
pub struct GenerateKeyRequest<'a> {
    pub client_id: Uuid,
    #[serde(alias = "dataset_id")]
    #[schema(value_type = Option<String>, example = "550e8400-e29b-41d4-a716-446655440000")]
    pub keyset_id: Option<IdentifiedBy>,
    #[schema(value_type = Vec<GenerateKeySpec>)]
    pub keys: Cow<'a, [GenerateKeySpec<'a>]>,
    #[serde(default)]
    #[schema(value_type = Object)]
    pub unverified_context: Cow<'a, UnverifiedContext>,
}

impl ViturRequest for GenerateKeyRequest<'_> {
    type Response = GenerateKeyResponse;

    const ENDPOINT: &'static str = "generate-data-key";
    const SCOPE: Scope = Scope::with_permission(Permission::DataKey(DataKeyPermission::Generate));
}

/// Returned type from a [RetrieveKeyRequest].
#[derive(Debug, Serialize, Deserialize, ToSchema)]
pub struct RetrievedKey {
    /// Base64-encoded key material.
    #[schema(value_type = String, format = Byte)]
    pub key_material: ViturKeyMaterial,
}

/// Response to a [RetrieveKeyRequest].
/// Contains a list of [RetrievedKey]s.
#[derive(Debug, Serialize, Deserialize, ToSchema)]
pub struct RetrieveKeyResponse {
    pub keys: Vec<RetrievedKey>,
}

impl ViturResponse for RetrieveKeyResponse {}

/// A specification for retrieving a data key used in a [RetrieveKeyRequest].
#[derive(Debug, Serialize, Deserialize, Clone, ToSchema)]
pub struct RetrieveKeySpec<'a> {
    #[serde(alias = "id")]
    #[schema(value_type = String, format = Byte)]
    pub iv: KeyId,
    // TODO: Make Descriptor Optional
    #[schema(value_type = String)]
    pub descriptor: Cow<'a, str>,
    #[schema(value_type = String, format = Byte)]
    pub tag: Cow<'a, [u8]>,

    #[serde(default)]
    #[schema(value_type = Vec<Context>)]
    pub context: Cow<'a, [Context]>,

    // Since this field will be removed in the future allow older versions of Vitur to be able to
    // parse a RetrieveKeySpec that doesn't include the tag_version.
    #[serde(default)]
    pub tag_version: usize,

    /// The decryption policy with MAC (only for tag_version >= 1).
    /// Server verifies the caller's claims satisfy at least one condition.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub decryption_policy: Option<DecryptionPolicy>,
}

impl<'a> RetrieveKeySpec<'a> {
    const DEFAULT_TAG_VERSION: usize = 0;

    pub fn new(id: KeyId, tag: &'a [u8], descriptor: &'a str) -> Self {
        Self {
            iv: id,
            descriptor: Cow::from(descriptor),
            tag: Cow::from(tag),
            context: Cow::Owned(Vec::new()),
            tag_version: Self::DEFAULT_TAG_VERSION,
            decryption_policy: None,
        }
    }

    pub fn with_context(mut self, context: Cow<'a, [Context]>) -> Self {
        self.context = context;
        self
    }

    pub fn with_policy(mut self, policy: DecryptionPolicy) -> Self {
        self.decryption_policy = Some(policy);
        self.tag_version = 1;
        self
    }
}

/// Request to retrieve a data key on behalf of a client in the given keyset.
/// Requires the `data_key:retrieve` scope.
/// Response is a [RetrieveKeyResponse].
///
/// See also [RetrieveKeySpec].
#[derive(Debug, Serialize, Deserialize, ToSchema)]
pub struct RetrieveKeyRequest<'a> {
    pub client_id: Uuid,
    #[serde(alias = "dataset_id")]
    #[schema(value_type = Option<String>, example = "550e8400-e29b-41d4-a716-446655440000")]
    pub keyset_id: Option<IdentifiedBy>,
    #[schema(value_type = Vec<RetrieveKeySpec>)]
    pub keys: Cow<'a, [RetrieveKeySpec<'a>]>,
    #[serde(default)]
    #[schema(value_type = Object)]
    pub unverified_context: UnverifiedContext,
}

impl ViturRequest for RetrieveKeyRequest<'_> {
    type Response = RetrieveKeyResponse;

    const ENDPOINT: &'static str = "retrieve-data-key";
    const SCOPE: Scope = Scope::with_permission(Permission::DataKey(DataKeyPermission::Retrieve));
}

/// Request to retrieve a data key on behalf of a client in the given keyset.
/// Requires the `data_key:retrieve` scope.
/// Response is a [RetrieveKeyResponse].
///
/// See also [RetrieveKeySpec].
#[derive(Debug, Serialize, Deserialize)]
pub struct RetrieveKeyRequestFallible<'a> {
    pub client_id: Uuid,
    #[serde(alias = "dataset_id")]
    pub keyset_id: Option<IdentifiedBy>,
    pub keys: Cow<'a, [RetrieveKeySpec<'a>]>,
    #[serde(default)]
    pub unverified_context: Cow<'a, UnverifiedContext>,
}

impl ViturRequest for RetrieveKeyRequestFallible<'_> {
    type Response = RetrieveKeyResponseFallible;

    const ENDPOINT: &'static str = "retrieve-data-key-fallible";
    const SCOPE: Scope = Scope::with_permission(Permission::DataKey(DataKeyPermission::Retrieve));
}

/// Response to a [RetrieveKeyRequest] with per-key error handling
#[derive(Debug, Serialize, Deserialize, ToSchema)]
pub struct RetrieveKeyResponseFallible {
    #[schema(value_type = Vec<serde_json::Value>)]
    pub keys: Vec<Result<RetrievedKey, String>>, // TODO: Error?
}

impl ViturResponse for RetrieveKeyResponseFallible {}

/// Request message to disable a keyset.
/// Requires the `dataset:disable` scope.
/// Response is an [EmptyResponse].
#[derive(Debug, Serialize, Deserialize, ToSchema)]
pub struct DisableKeysetRequest {
    /// The keyset to disable. Accepts a UUID or a name string.
    #[serde(alias = "dataset_id")]
    #[schema(value_type = String, example = "550e8400-e29b-41d4-a716-446655440000")]
    pub keyset_id: IdentifiedBy,
}

impl ViturRequest for DisableKeysetRequest {
    type Response = EmptyResponse;

    const ENDPOINT: &'static str = "disable-keyset";
    const SCOPE: Scope = Scope::with_permission(Permission::Keyset(KeysetPermission::Disable));
}

/// Request message to enable a keyset that has was previously disabled.
/// Requires the `dataset:enable` scope.
/// Response is an [EmptyResponse].
#[derive(Debug, Serialize, Deserialize, ToSchema)]
pub struct EnableKeysetRequest {
    /// The keyset to enable. Accepts a UUID or a name string.
    #[serde(alias = "dataset_id")]
    #[schema(value_type = String, example = "550e8400-e29b-41d4-a716-446655440000")]
    pub keyset_id: IdentifiedBy,
}

impl ViturRequest for EnableKeysetRequest {
    type Response = EmptyResponse;

    const ENDPOINT: &'static str = "enable-keyset";
    const SCOPE: Scope = Scope::with_permission(Permission::Keyset(KeysetPermission::Enable));
}

/// Request message to modify a keyset with the given keyset_id.
/// `name` and `description` are optional and will be updated if provided.
///
/// Requires the `dataset:modify` scope.
/// Response is an [EmptyResponse].
#[derive(Debug, Serialize, Deserialize, ToSchema)]
pub struct ModifyKeysetRequest<'a> {
    /// The keyset to modify. Accepts a UUID or a name string.
    #[serde(alias = "dataset_id")]
    #[schema(value_type = String, example = "550e8400-e29b-41d4-a716-446655440000")]
    pub keyset_id: IdentifiedBy,
    /// Optional new name for the keyset.
    #[schema(value_type = Option<String>)]
    pub name: Option<Cow<'a, str>>,
    /// Optional new description for the keyset.
    #[schema(value_type = Option<String>)]
    pub description: Option<Cow<'a, str>>,
}

impl ViturRequest for ModifyKeysetRequest<'_> {
    type Response = EmptyResponse;

    const ENDPOINT: &'static str = "modify-keyset";
    const SCOPE: Scope = Scope::with_permission(Permission::Keyset(KeysetPermission::Modify));
}

/// Request message to grant a client access to a keyset.
/// Requires the `dataset:grant` scope.
///
/// Response is an [EmptyResponse].
#[derive(Debug, Serialize, Deserialize, ToSchema)]
pub struct GrantKeysetRequest {
    pub client_id: Uuid,
    /// The keyset to grant access to. Accepts a UUID or a name string.
    #[serde(alias = "dataset_id")]
    #[schema(value_type = String, example = "550e8400-e29b-41d4-a716-446655440000")]
    pub keyset_id: IdentifiedBy,
}

impl ViturRequest for GrantKeysetRequest {
    type Response = EmptyResponse;

    const ENDPOINT: &'static str = "grant-keyset";
    const SCOPE: Scope = Scope::with_permission(Permission::Keyset(KeysetPermission::Grant));
}

/// Request message to revoke a client's access to a keyset.
/// Requires the `dataset:revoke` scope.
/// Response is an [EmptyResponse].
#[derive(Debug, Serialize, Deserialize, ToSchema)]
pub struct RevokeKeysetRequest {
    pub client_id: Uuid,
    /// The keyset to revoke access from. Accepts a UUID or a name string.
    #[serde(alias = "dataset_id")]
    #[schema(value_type = String, example = "550e8400-e29b-41d4-a716-446655440000")]
    pub keyset_id: IdentifiedBy,
}

impl ViturRequest for RevokeKeysetRequest {
    type Response = EmptyResponse;

    const ENDPOINT: &'static str = "revoke-keyset";
    const SCOPE: Scope = Scope::with_permission(Permission::Keyset(KeysetPermission::Revoke));
}

/// Request to load a keyset on behalf of a client.
/// This is used by clients before indexing or querying data and includes
/// key material which can be derived by the client to generate encrypted index terms.
///
/// If a keyset_id is not provided the client's default keyset will be loaded.
///
/// Requires the `data_key:retrieve` scope (though this may change in the future).
/// Response is a [LoadKeysetResponse].
#[derive(Debug, Serialize, Deserialize, PartialEq, PartialOrd, ToSchema)]
pub struct LoadKeysetRequest {
    pub client_id: Uuid,
    /// The keyset to load. Accepts a UUID or a name string. If omitted, the client's default keyset is used.
    #[serde(alias = "dataset_id")]
    #[schema(value_type = Option<String>, example = "550e8400-e29b-41d4-a716-446655440000")]
    pub keyset_id: Option<IdentifiedBy>,
}

impl ViturRequest for LoadKeysetRequest {
    type Response = LoadKeysetResponse;

    const ENDPOINT: &'static str = "load-keyset";

    // NOTE: We don't currently support the ability to allow an operation
    // based on any one of several possible scopes so we'll just use `data_key:retrieve` for now.
    // This should probably be allowed for any operation that requires indexing or querying.
    const SCOPE: Scope = Scope::with_permission(Permission::DataKey(DataKeyPermission::Retrieve));
}

/// Response to a [LoadKeysetRequest].
/// The response includes the key material required to derive data keys.
/// It is analogous to a [RetrieveKeyResponse] but where the server generated the key.
#[derive(Debug, Serialize, Deserialize, ToSchema)]
pub struct LoadKeysetResponse {
    pub partial_index_key: RetrievedKey,
    #[serde(rename = "dataset")]
    pub keyset: Keyset,
}

impl ViturResponse for LoadKeysetResponse {}

#[cfg(test)]
mod test {
    use serde_json::json;
    use uuid::Uuid;

    use crate::{CreateKeysetResponse, CreatedClient, IdentifiedBy, LoadKeysetRequest, Name};

    mod create_keyset_response_serialization {
        use super::*;
        use crate::{Keyset, ViturKeyMaterial};

        #[test]
        fn without_client_is_flat_keyset() {
            let id = Uuid::new_v4();
            let response = CreateKeysetResponse {
                keyset: Keyset {
                    id,
                    name: "test-keyset".into(),
                    description: "A test keyset".into(),
                    is_disabled: false,
                    is_default: false,
                },
                client: None,
            };

            let serialized = serde_json::to_value(&response).unwrap();

            // Should be a flat object identical to the old Keyset response
            assert_eq!(
                serialized,
                json!({
                    "id": id,
                    "name": "test-keyset",
                    "description": "A test keyset",
                    "is_disabled": false,
                    "is_default": false,
                })
            );

            // Should round-trip
            let deserialized: CreateKeysetResponse = serde_json::from_value(serialized).unwrap();
            assert_eq!(deserialized.keyset.id, id);
            assert!(deserialized.client.is_none());
        }

        #[test]
        fn with_client_includes_client_field() {
            let keyset_id = Uuid::new_v4();
            let client_id = Uuid::new_v4();

            let response = CreateKeysetResponse {
                keyset: Keyset {
                    id: keyset_id,
                    name: "device-keyset".into(),
                    description: "Keyset with device client".into(),
                    is_disabled: false,
                    is_default: false,
                },
                client: Some(CreatedClient {
                    id: client_id,
                    client_key: ViturKeyMaterial::from(vec![1, 2, 3, 4]),
                }),
            };

            let serialized = serde_json::to_value(&response).unwrap();

            // Keyset fields are flat, client is a nested object with base64-encoded key
            assert_eq!(
                serialized,
                json!({
                    "id": keyset_id,
                    "name": "device-keyset",
                    "description": "Keyset with device client",
                    "is_disabled": false,
                    "is_default": false,
                    "client": {
                        "id": client_id,
                        "client_key": "AQIDBA==",
                    },
                })
            );

            // Should round-trip
            let deserialized: CreateKeysetResponse = serde_json::from_value(serialized).unwrap();
            assert_eq!(deserialized.keyset.id, keyset_id);
            let created_client = deserialized.client.unwrap();
            assert_eq!(created_client.id, client_id);
            assert_eq!(&*created_client.client_key, &[1, 2, 3, 4]);
        }
    }

    mod create_client_request_serialization {
        use super::*;
        use crate::CreateClientRequest;

        #[test]
        fn with_keyset_id_round_trips() {
            let keyset_id = Uuid::new_v4();
            let req = CreateClientRequest {
                keyset_id: Some(IdentifiedBy::Uuid(keyset_id)),
                name: "my-client".into(),
                description: "desc".into(),
            };

            let serialized = serde_json::to_value(&req).unwrap();
            assert!(serialized.get("keyset_id").is_some());

            let deserialized: CreateClientRequest = serde_json::from_value(serialized).unwrap();
            assert_eq!(deserialized.keyset_id, Some(IdentifiedBy::Uuid(keyset_id)));
        }

        #[test]
        fn without_keyset_id_round_trips() {
            let req = CreateClientRequest {
                keyset_id: None,
                name: "my-client".into(),
                description: "desc".into(),
            };

            let serialized = serde_json::to_value(&req).unwrap();
            assert!(serialized.get("keyset_id").is_none());

            let deserialized: CreateClientRequest = serde_json::from_value(serialized).unwrap();
            assert_eq!(deserialized.keyset_id, None);
        }

        #[test]
        fn backwards_compatible_with_dataset_id() {
            let dataset_id = Uuid::new_v4();
            let json = json!({
                "dataset_id": dataset_id,
                "name": "old-client",
                "description": "old desc",
            });

            let req: CreateClientRequest = serde_json::from_value(json).unwrap();
            assert_eq!(req.keyset_id, Some(IdentifiedBy::Uuid(dataset_id)));
        }

        #[test]
        fn omitted_keyset_id_defaults_to_none() {
            let json = json!({
                "name": "no-keyset",
                "description": "no keyset",
            });

            let req: CreateClientRequest = serde_json::from_value(json).unwrap();
            assert_eq!(req.keyset_id, None);
        }
    }

    mod create_keyset_request_validation {
        use crate::CreateKeysetRequest;
        use validator::Validate;

        fn valid_request() -> CreateKeysetRequest<'static> {
            CreateKeysetRequest {
                name: "my-keyset".into(),
                description: "A test keyset".into(),
                client: None,
            }
        }

        #[test]
        fn valid_request_passes() {
            assert!(valid_request().validate().is_ok());
        }

        #[test]
        fn empty_name_fails() {
            let req = CreateKeysetRequest {
                name: "".into(),
                ..valid_request()
            };
            let errors = req.validate().unwrap_err();
            assert!(errors.field_errors().contains_key("name"));
        }

        #[test]
        fn name_over_64_chars_fails() {
            let req = CreateKeysetRequest {
                name: "a".repeat(65).into(),
                ..valid_request()
            };
            let errors = req.validate().unwrap_err();
            assert!(errors.field_errors().contains_key("name"));
        }

        #[test]
        fn reserved_default_name_fails() {
            let req = CreateKeysetRequest {
                name: "default".into(),
                ..valid_request()
            };
            let errors = req.validate().unwrap_err();
            let name_errors = &errors.field_errors()["name"];
            assert!(name_errors.iter().any(|e| e.code == "reserved_name"));
        }

        #[test]
        fn reserved_default_name_case_insensitive() {
            let req = CreateKeysetRequest {
                name: "DEFAULT".into(),
                ..valid_request()
            };
            assert!(req.validate().is_err());
        }

        #[test]
        fn name_with_invalid_characters_fails() {
            let req = CreateKeysetRequest {
                name: "has spaces".into(),
                ..valid_request()
            };
            let errors = req.validate().unwrap_err();
            let name_errors = &errors.field_errors()["name"];
            assert!(name_errors.iter().any(|e| e.code == "invalid_characters"));
        }

        #[test]
        fn name_with_special_chars_fails() {
            for name in ["test@keyset", "test!keyset", "test.keyset", "test%keyset"] {
                let req = CreateKeysetRequest {
                    name: name.into(),
                    ..valid_request()
                };
                assert!(
                    req.validate().is_err(),
                    "expected {name} to fail validation"
                );
            }
        }

        #[test]
        fn name_with_allowed_chars_passes() {
            for name in ["my-keyset", "my_keyset", "my/keyset", "MyKeyset123"] {
                let req = CreateKeysetRequest {
                    name: name.into(),
                    ..valid_request()
                };
                assert!(req.validate().is_ok(), "expected {name} to pass validation");
            }
        }

        #[test]
        fn empty_description_fails() {
            let req = CreateKeysetRequest {
                description: "".into(),
                ..valid_request()
            };
            let errors = req.validate().unwrap_err();
            assert!(errors.field_errors().contains_key("description"));
        }

        #[test]
        fn description_over_256_chars_fails() {
            let req = CreateKeysetRequest {
                description: "a".repeat(257).into(),
                ..valid_request()
            };
            let errors = req.validate().unwrap_err();
            assert!(errors.field_errors().contains_key("description"));
        }

        #[test]
        fn description_at_256_chars_passes() {
            let req = CreateKeysetRequest {
                description: "a".repeat(256).into(),
                ..valid_request()
            };
            assert!(req.validate().is_ok());
        }

        #[test]
        fn nested_client_name_validation() {
            use crate::{ClientType, CreateClientSpec};

            let req = CreateKeysetRequest {
                name: "my-keyset".into(),
                description: "desc".into(),
                client: Some(CreateClientSpec {
                    client_type: ClientType::Device,
                    name: "".into(),
                }),
            };
            let errors = req.validate().unwrap_err();
            assert!(
                errors.errors().contains_key("client"),
                "expected nested client validation error"
            );
        }
    }

    mod openapi_schema {
        use crate::{CreateClientSpec, CreateKeysetRequest};
        use utoipa::PartialSchema;

        fn schema_json<T: PartialSchema>() -> serde_json::Value {
            serde_json::to_value(T::schema()).unwrap()
        }

        #[test]
        fn create_keyset_request_name_has_constraints() {
            let schema = schema_json::<CreateKeysetRequest>();
            let name = &schema["properties"]["name"];

            assert_eq!(name["minLength"], 1);
            assert_eq!(name["maxLength"], 64);
            assert_eq!(name["pattern"], r"^[A-Za-z0-9_\-/]+$");
        }

        #[test]
        fn create_keyset_request_description_has_constraints() {
            let schema = schema_json::<CreateKeysetRequest>();
            let desc = &schema["properties"]["description"];

            assert_eq!(desc["minLength"], 1);
            assert_eq!(desc["maxLength"], 256);
        }

        #[test]
        fn create_client_spec_name_has_constraints() {
            let schema = schema_json::<CreateClientSpec>();
            let name = &schema["properties"]["name"];

            assert_eq!(name["minLength"], 1);
            assert_eq!(name["maxLength"], 64);
        }
    }

    mod backwards_compatible_deserialisation {
        use super::*;

        #[test]
        fn when_dataset_id_is_uuid() {
            let client_id = Uuid::new_v4();
            let dataset_id = Uuid::new_v4();

            let json = json!({
                "client_id": client_id,
                "dataset_id": dataset_id,
            });

            let req: LoadKeysetRequest = serde_json::from_value(json).unwrap();

            assert_eq!(
                req,
                LoadKeysetRequest {
                    client_id,
                    keyset_id: Some(IdentifiedBy::Uuid(dataset_id))
                }
            );
        }

        #[test]
        fn when_keyset_id_is_uuid() {
            let client_id = Uuid::new_v4();
            let keyset_id = Uuid::new_v4();

            let json = json!({
                "client_id": client_id,
                "keyset_id": keyset_id,
            });

            let req: LoadKeysetRequest = serde_json::from_value(json).unwrap();

            assert_eq!(
                req,
                LoadKeysetRequest {
                    client_id,
                    keyset_id: Some(IdentifiedBy::Uuid(keyset_id))
                }
            );
        }

        #[test]
        fn when_dataset_id_is_id_name() {
            let client_id = Uuid::new_v4();
            let dataset_id = IdentifiedBy::Name(Name::new_untrusted("some-dataset-name"));

            let json = json!({
                "client_id": client_id,
                "dataset_id": dataset_id,
            });

            let req: LoadKeysetRequest = serde_json::from_value(json).unwrap();

            assert_eq!(
                req,
                LoadKeysetRequest {
                    client_id,
                    keyset_id: Some(dataset_id)
                }
            );
        }
    }
}