cipherstash-client 0.42.0

use super::{DataKeyWithTag, DecryptError, RecordWithNonce, RetrieveKeyPayload};
use crate::zerokms::EncryptedRecord;
use recipher::key::Iv;
use serde::{Deserialize, Serialize};
use uuid::Uuid;
use zerokms_protocol::DecryptionPolicy;

/// The per-document key header of the SteVec envelope wire format.
///
/// A SteVec document encrypts all of its entries under ONE data key, so the
/// key-retrieval material — everything an [`EncryptedRecord`] carries besides
/// the ciphertext itself — is stored once at the document envelope (the `h`
/// field) instead of being repeated inside every entry's `c`. Entries carry
/// only their raw AEAD output; their nonces are derived from their stored
/// selectors (`nonce = selector_bytes[..12]`), so no per-entry nonce is
/// stored either.
///
/// Deliberately a dedicated struct rather than an [`EncryptedRecord`] with an
/// empty ciphertext: the header is a distinct wire concept and the blob is
/// opaque to SQL, so the field set can evolve independently under serde.
#[cfg_attr(test, derive(PartialEq, Eq))]
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct KeyHeader {
    #[serde(with = "serde_bytes")]
    pub iv: Iv,
    #[serde(with = "serde_bytes")]
    pub tag: Vec<u8>,
    pub descriptor: String,
    #[serde(default, rename = "dataset_id")]
    pub keyset_id: Option<Uuid>,
    /// Decryption policy with server-generated MAC (tag_version=1 only).
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub decryption_policy: Option<DecryptionPolicy>,
}

impl KeyHeader {
    /// Build the header for a document encrypted under `key`.
    pub fn new(
        key: &DataKeyWithTag,
        descriptor: impl Into<String>,
        keyset_id: Option<Uuid>,
    ) -> Self {
        Self {
            iv: key.key.iv,
            tag: key.tag.clone(),
            descriptor: descriptor.into(),
            keyset_id,
            decryption_policy: key.decryption_policy.clone(),
        }
    }

    /// The ZeroKMS key-retrieval payload for the document's data key — the
    /// same (iv, descriptor, tag) triple a self-contained record produces.
    pub fn retrieve_key_payload(&self) -> RetrieveKeyPayload<'_> {
        let payload = RetrieveKeyPayload::new(self.iv, self.descriptor.as_str(), &self.tag);
        if let Some(policy) = &self.decryption_policy {
            return payload.with_decryption_policy(policy.clone());
        }
        payload
    }

    /// Reassemble a decryptable record for one entry of the document: the
    /// header's key material + the entry's raw AEAD `ciphertext` + the entry's
    /// 16-byte tokenized `selector`. The selector is the source of both AEAD
    /// bindings — the nonce is its first 12 bytes and the full 16 bytes are
    /// bound into the AAD (see [`RecordWithNonce`]).
    pub fn record_with_selector(&self, ciphertext: Vec<u8>, selector: [u8; 16]) -> RecordWithNonce {
        RecordWithNonce {
            record: EncryptedRecord {
                iv: self.iv,
                ciphertext,
                tag: self.tag.clone(),
                descriptor: self.descriptor.clone(),
                keyset_id: self.keyset_id,
                decryption_policy: self.decryption_policy.clone(),
            },
            selector,
        }
    }

    /// Serialize the header to a `Vec<u8>` using MessagePack encoding.
    pub fn to_mp_bytes(&self) -> Result<Vec<u8>, DecryptError> {
        rmp_serde::to_vec(&self).map_err(DecryptError::SerializeError)
    }

    /// Serialize the header using MessagePack and convert to a base85-encoded
    /// string — the on-the-wire form of the envelope `h` field.
    pub fn to_mp_base85(&self) -> Result<String, DecryptError> {
        self.to_mp_bytes().map(|v| base85::encode(&v))
    }

    /// Deserialize a header from a slice of MessagePack-encoded bytes.
    pub fn from_mp_bytes(bytes: &[u8]) -> Result<Self, DecryptError> {
        rmp_serde::from_slice(bytes).map_err(DecryptError::DeserializeError)
    }

    /// Deserialize a header from a base85-encoded string of
    /// MessagePack-encoded bytes (the envelope `h` field).
    pub fn from_mp_base85(base85str: &str) -> Result<Self, DecryptError> {
        let bytes = base85::decode(base85str)?;
        Self::from_mp_bytes(&bytes)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn header() -> KeyHeader {
        KeyHeader {
            iv: Iv::default(),
            tag: vec![1; 16],
            descriptor: "users/profile".to_string(),
            keyset_id: Some(Uuid::new_v4()),
            decryption_policy: None,
        }
    }

    #[test]
    fn test_key_header_mp_base85_round_trip() {
        let header = header();
        let encoded = header.to_mp_base85().unwrap();
        assert_eq!(KeyHeader::from_mp_base85(&encoded).unwrap(), header);
    }

    #[test]
    fn test_record_with_selector_reassembly() {
        let header = header();
        let assembled = header.record_with_selector(vec![7; 32], [9; 16]);

        assert_eq!(assembled.selector, [9; 16]);
        assert_eq!(assembled.record.ciphertext, vec![7; 32]);
        assert_eq!(assembled.record.iv, header.iv);
        assert_eq!(assembled.record.tag, header.tag);
        assert_eq!(assembled.record.descriptor, header.descriptor);
        assert_eq!(assembled.record.keyset_id, header.keyset_id);
    }
}