cose2 0.5.0

//! COSE_KDF_Context and its sub-structures (RFC 9053 §5.2).

use std::borrow::Cow;

use cbor2::Cbor;
use serde::{
    de::{Error as DeError, IgnoredAny, SeqAccess, Visitor},
    ser::{Error as _, SerializeSeq},
    Deserialize, Deserializer, Serialize, Serializer,
};

use crate::{
    header::{decode_protected, encode_protected},
    Error, Header, Label,
};

/// A `PartyInfo` nonce: `bstr / int` (RFC 9053 §5.2, where the field is
/// `bstr / int / nil`; absence is modeled with `Option`).
///
/// Integer nonces retain the complete CBOR integer range in an `i128`.
#[derive(Clone, Debug, PartialEq)]
pub enum PartyNonce {
    /// A byte-string nonce.
    Bytes(Vec<u8>),
    /// An integer nonce.
    Int(i128),
}

impl From<Vec<u8>> for PartyNonce {
    fn from(value: Vec<u8>) -> Self {
        PartyNonce::Bytes(value)
    }
}

impl From<&[u8]> for PartyNonce {
    fn from(value: &[u8]) -> Self {
        PartyNonce::Bytes(value.to_vec())
    }
}

impl From<i64> for PartyNonce {
    fn from(value: i64) -> Self {
        PartyNonce::Int(i128::from(value))
    }
}

impl From<u64> for PartyNonce {
    fn from(value: u64) -> Self {
        PartyNonce::Int(i128::from(value))
    }
}

impl Serialize for PartyNonce {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: Serializer,
    {
        match self {
            PartyNonce::Bytes(b) => serializer.serialize_bytes(b),
            PartyNonce::Int(i) => {
                validate_party_nonce_integer(*i).map_err(serde::ser::Error::custom)?;
                serializer.serialize_i128(*i)
            }
        }
    }
}

impl<'de> Deserialize<'de> for PartyNonce {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: Deserializer<'de>,
    {
        struct NonceVisitor;

        impl Visitor<'_> for NonceVisitor {
            type Value = PartyNonce;

            fn expecting(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
                f.write_str("a byte string or integer PartyInfo nonce")
            }

            fn visit_bytes<E: DeError>(self, v: &[u8]) -> Result<PartyNonce, E> {
                Ok(PartyNonce::Bytes(v.to_vec()))
            }

            fn visit_byte_buf<E: DeError>(self, v: Vec<u8>) -> Result<PartyNonce, E> {
                Ok(PartyNonce::Bytes(v))
            }

            fn visit_i64<E: DeError>(self, v: i64) -> Result<PartyNonce, E> {
                Ok(PartyNonce::Int(i128::from(v)))
            }

            fn visit_u64<E: DeError>(self, v: u64) -> Result<PartyNonce, E> {
                Ok(PartyNonce::Int(i128::from(v)))
            }

            fn visit_i128<E: DeError>(self, v: i128) -> Result<PartyNonce, E> {
                validate_party_nonce_integer(v).map_err(E::custom)?;
                Ok(PartyNonce::Int(v))
            }

            fn visit_u128<E: DeError>(self, v: u128) -> Result<PartyNonce, E> {
                let value =
                    i128::try_from(v).map_err(|_| E::custom("integer nonce out of CBOR range"))?;
                validate_party_nonce_integer(value).map_err(E::custom)?;
                Ok(PartyNonce::Int(value))
            }
        }

        deserializer.deserialize_any(NonceVisitor)
    }
}

fn validate_party_nonce_integer(value: i128) -> Result<(), Error> {
    if value < -1 - i128::from(u64::MAX) || value > i128::from(u64::MAX) {
        Err(Error::UnexpectedType(
            "integer nonce out of CBOR integer range".into(),
        ))
    } else {
        Ok(())
    }
}

/// A `PartyInfo` array `[identity, nonce, other]` (RFC 9053 §5.2).
///
/// `identity` and `other` are byte strings or `null`; `nonce` is a byte
/// string, an integer, or `null`.
#[derive(Clone, Debug, Default, PartialEq, Cbor)]
#[cbor(array)]
pub struct PartyInfo {
    /// Party identity information.
    #[serde(with = "crate::strict::optional_bytes")]
    pub identity: Option<Vec<u8>>,
    /// Party-provided nonce.
    pub nonce: Option<PartyNonce>,
    /// Other party-provided information.
    #[serde(with = "crate::strict::optional_bytes")]
    pub other: Option<Vec<u8>>,
}

/// A `SuppPubInfo` structure: `[keyDataLength, protected, ?other]`.
#[derive(Clone, Debug, Default)]
pub struct SuppPubInfo {
    /// Length of the derived key material, in bits.
    pub key_data_length: u64,
    /// The protected header of the enclosing structure.
    pub protected: Header,
    /// Optional other supplemental public information.
    pub other: Option<Vec<u8>>,
    /// Exact decoded protected bytes. Leave as `None` when constructing.
    #[doc(hidden)]
    pub protected_raw: Option<Vec<u8>>,
}

impl PartialEq for SuppPubInfo {
    fn eq(&self, other: &Self) -> bool {
        self.key_data_length == other.key_data_length
            && self.protected == other.protected
            && self.other == other.other
    }
}

impl SuppPubInfo {
    /// Returns the decoded protected-header bytes, when available.
    pub fn protected_raw(&self) -> Option<&[u8]> {
        self.protected_raw.as_deref()
    }
}

impl Serialize for SuppPubInfo {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: Serializer,
    {
        let protected_raw: Cow<'_, [u8]> = match &self.protected_raw {
            Some(raw) => {
                crate::header::validate_protected_state(&self.protected, raw)
                    .map_err(S::Error::custom)?;
                Cow::Borrowed(raw)
            }
            None => Cow::Owned(encode_protected(&self.protected).map_err(S::Error::custom)?),
        };
        let len = if self.other.is_some() { 3 } else { 2 };
        let mut seq = serializer.serialize_seq(Some(len))?;
        seq.serialize_element(&self.key_data_length)?;
        seq.serialize_element(serde_bytes::Bytes::new(protected_raw.as_ref()))?;
        if let Some(other) = &self.other {
            seq.serialize_element(serde_bytes::Bytes::new(other))?;
        }
        seq.end()
    }
}

impl<'de> Deserialize<'de> for SuppPubInfo {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: Deserializer<'de>,
    {
        struct SuppPubInfoVisitor;

        impl<'de> Visitor<'de> for SuppPubInfoVisitor {
            type Value = SuppPubInfo;

            fn expecting(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
                f.write_str("a SuppPubInfo array of 2 or 3 elements")
            }

            fn visit_seq<A>(self, mut seq: A) -> Result<SuppPubInfo, A::Error>
            where
                A: SeqAccess<'de>,
            {
                let key_data_length: u64 = seq
                    .next_element()?
                    .ok_or_else(|| A::Error::custom("missing keyDataLength"))?;
                let protected_raw: crate::strict::StrictBytes = seq
                    .next_element()?
                    .ok_or_else(|| A::Error::custom("missing protected header"))?;
                let other = seq.next_element::<crate::strict::StrictBytes>()?;
                if seq.next_element::<IgnoredAny>()?.is_some() {
                    return Err(A::Error::invalid_length(4, &self));
                }
                let protected = decode_protected(&protected_raw.0).map_err(A::Error::custom)?;
                Ok(SuppPubInfo {
                    key_data_length,
                    protected,
                    other: other.map(|o| o.0),
                    protected_raw: Some(protected_raw.0),
                })
            }
        }

        deserializer.deserialize_seq(SuppPubInfoVisitor)
    }
}

/// A COSE_KDF_Context structure (RFC 9053 §5.2):
/// `[AlgorithmID, PartyUInfo, PartyVInfo, SuppPubInfo, ?SuppPrivInfo]`.
#[derive(Clone, Debug, PartialEq)]
pub struct KdfContext {
    /// Identifier of the algorithm the derived key is used with
    /// (`int / tstr`, RFC 9053 §5.2).
    pub algorithm_id: Label,
    /// Information about party U.
    pub party_u_info: PartyInfo,
    /// Information about party V.
    pub party_v_info: PartyInfo,
    /// Supplemental public information.
    pub supp_pub_info: SuppPubInfo,
    /// Optional supplemental private information.
    pub supp_priv_info: Option<Vec<u8>>,
}

impl KdfContext {
    /// Decodes a context from CBOR bytes.
    pub fn from_slice(data: &[u8]) -> Result<Self, Error> {
        crate::strict::validate_array(data)?;
        Ok(cbor2::from_slice(data)?)
    }

    /// Encodes the context to canonical CBOR bytes.
    pub fn to_vec(&self) -> Result<Vec<u8>, Error> {
        Ok(cbor2::to_canonical_vec(self)?)
    }
}

impl Serialize for KdfContext {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: Serializer,
    {
        let len = if self.supp_priv_info.is_some() { 5 } else { 4 };
        let mut seq = serializer.serialize_seq(Some(len))?;
        seq.serialize_element(&self.algorithm_id)?;
        seq.serialize_element(&self.party_u_info)?;
        seq.serialize_element(&self.party_v_info)?;
        seq.serialize_element(&self.supp_pub_info)?;
        if let Some(priv_info) = &self.supp_priv_info {
            seq.serialize_element(serde_bytes::Bytes::new(priv_info))?;
        }
        seq.end()
    }
}

impl<'de> Deserialize<'de> for KdfContext {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: Deserializer<'de>,
    {
        struct KdfContextVisitor;

        impl<'de> Visitor<'de> for KdfContextVisitor {
            type Value = KdfContext;

            fn expecting(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
                f.write_str("a COSE_KDF_Context array of 4 or 5 elements")
            }

            fn visit_seq<A>(self, mut seq: A) -> Result<KdfContext, A::Error>
            where
                A: SeqAccess<'de>,
            {
                let algorithm_id: Label = seq
                    .next_element()?
                    .ok_or_else(|| A::Error::custom("missing AlgorithmID"))?;
                let party_u_info: PartyInfo = seq
                    .next_element()?
                    .ok_or_else(|| A::Error::custom("missing PartyUInfo"))?;
                let party_v_info: PartyInfo = seq
                    .next_element()?
                    .ok_or_else(|| A::Error::custom("missing PartyVInfo"))?;
                let supp_pub_info: SuppPubInfo = seq
                    .next_element()?
                    .ok_or_else(|| A::Error::custom("missing SuppPubInfo"))?;
                let supp_priv_info = seq.next_element::<crate::strict::StrictBytes>()?;
                if seq.next_element::<IgnoredAny>()?.is_some() {
                    return Err(A::Error::invalid_length(6, &self));
                }
                Ok(KdfContext {
                    algorithm_id,
                    party_u_info,
                    party_v_info,
                    supp_pub_info,
                    supp_priv_info: supp_priv_info.map(|p| p.0),
                })
            }
        }

        deserializer.deserialize_seq(KdfContextVisitor)
    }
}