wolfcose 0.1.0

use crate::{
    Algorithm, ByteBuf, CborDeserialize, CborDeserializer, CborSerialize, CborValue, Error, Result,
};
use alloc::{string::String, vec::Vec};

/// COSE header label.
#[derive(Clone, Debug, Eq, PartialEq, Ord, PartialOrd, Hash)]
pub enum HeaderLabel {
    /// `alg` header.
    Algorithm,
    /// `crit` header.
    Critical,
    /// `content_type` header.
    ContentType,
    /// `kid` header.
    Kid,
    /// `iv` header.
    Iv,
    /// `partial_iv` header.
    PartialIv,
    /// Ephemeral key header.
    EphemeralKey,
    /// Raw integer label.
    Int(i64),
    /// Raw text label.
    Text(String),
}

impl HeaderLabel {
    /// Convert a raw integer label.
    pub fn from_i64(label: i64) -> Self {
        match label {
            1 => Self::Algorithm,
            2 => Self::Critical,
            3 => Self::ContentType,
            4 => Self::Kid,
            5 => Self::Iv,
            6 => Self::PartialIv,
            -1 => Self::EphemeralKey,
            other => Self::Int(other),
        }
    }

    /// Return a raw integer label for standard integer labels.
    pub fn as_i64(&self) -> Option<i64> {
        match self {
            Self::Algorithm => Some(1),
            Self::Critical => Some(2),
            Self::ContentType => Some(3),
            Self::Kid => Some(4),
            Self::Iv => Some(5),
            Self::PartialIv => Some(6),
            Self::EphemeralKey => Some(-1),
            Self::Int(value) => Some(*value),
            Self::Text(_) => None,
        }
    }
}

impl CborSerialize for HeaderLabel {
    fn serialize(&self, serializer: &mut crate::CborSerializer<'_>) -> Result<()> {
        match self {
            Self::Text(value) => value.serialize(serializer),
            _ => self
                .as_i64()
                .ok_or(Error::InvalidArgument)?
                .serialize(serializer),
        }
    }
}

impl<'de> CborDeserialize<'de> for HeaderLabel {
    fn deserialize(deserializer: &mut CborDeserializer<'de>) -> Result<Self> {
        match deserializer.decoder_mut().peek_type() {
            Some(crate::CborMajorType::TEXT) => Ok(Self::Text(String::deserialize(deserializer)?)),
            Some(crate::CborMajorType::UINT) | Some(crate::CborMajorType::NEGATIVE_INT) => {
                Ok(Self::from_i64(i64::deserialize(deserializer)?))
            }
            _ => Err(Error::CborType),
        }
    }
}

/// COSE header value.
#[derive(Clone, Debug, PartialEq)]
pub enum HeaderValue {
    /// Algorithm value.
    Algorithm(Algorithm),
    /// Integer value.
    Int(i64),
    /// Text value.
    Text(String),
    /// Byte-string value.
    Bytes(Vec<u8>),
    /// Array value.
    Array(Vec<CborValue>),
    /// Nested map or extension value.
    Raw(CborValue),
}

impl CborSerialize for HeaderValue {
    fn serialize(&self, serializer: &mut crate::CborSerializer<'_>) -> Result<()> {
        match self {
            Self::Algorithm(value) => value.id().serialize(serializer),
            Self::Int(value) => value.serialize(serializer),
            Self::Text(value) => value.serialize(serializer),
            Self::Bytes(value) => ByteBuf(value.clone()).serialize(serializer),
            Self::Array(value) => value.serialize(serializer),
            Self::Raw(value) => value.serialize(serializer),
        }
    }
}

impl<'de> CborDeserialize<'de> for HeaderValue {
    fn deserialize(deserializer: &mut CborDeserializer<'de>) -> Result<Self> {
        let value = CborValue::deserialize(deserializer)?;
        Ok(match value {
            CborValue::Integer(value) => Self::Int(value),
            CborValue::Unsigned(value) => {
                let value = i64::try_from(value).map_err(|_| Error::CborOverflow)?;
                Self::Int(value)
            }
            CborValue::Text(value) => Self::Text(value),
            CborValue::Bytes(value) => Self::Bytes(value),
            CborValue::Array(value) => Self::Array(value),
            value => Self::Raw(value),
        })
    }
}

/// Typed COSE header map preserving unknown labels.
#[derive(Clone, Debug, Default, PartialEq)]
pub struct HeaderMap {
    entries: Vec<(HeaderLabel, HeaderValue)>,
}

impl HeaderMap {
    /// Create an empty header map.
    pub fn new() -> Self {
        Self::default()
    }

    /// Number of entries.
    pub fn len(&self) -> usize {
        self.entries.len()
    }

    /// Whether there are no entries.
    pub fn is_empty(&self) -> bool {
        self.entries.is_empty()
    }

    /// Insert a header value, replacing an existing matching label.
    pub fn insert(&mut self, label: HeaderLabel, value: HeaderValue) {
        if let Some((_, existing)) = self.entries.iter_mut().find(|(entry, _)| *entry == label) {
            *existing = value;
        } else {
            self.entries.push((label, value));
        }
    }

    /// Borrow a value.
    pub fn get(&self, label: &HeaderLabel) -> Option<&HeaderValue> {
        self.entries
            .iter()
            .find_map(|(entry, value)| (entry == label).then_some(value))
    }

    /// Iterate over entries.
    pub fn iter(&self) -> impl Iterator<Item = &(HeaderLabel, HeaderValue)> {
        self.entries.iter()
    }

    /// Algorithm header, if present and integer-like.
    pub fn algorithm(&self) -> Option<Algorithm> {
        match self.get(&HeaderLabel::Algorithm) {
            Some(HeaderValue::Algorithm(value)) => Some(*value),
            Some(HeaderValue::Int(value)) => i32::try_from(*value).ok().map(Algorithm::from_id),
            _ => None,
        }
    }

    /// Key identifier, if present.
    pub fn kid(&self) -> Option<&[u8]> {
        match self.get(&HeaderLabel::Kid) {
            Some(HeaderValue::Bytes(value)) => Some(value),
            _ => None,
        }
    }
}

impl CborSerialize for HeaderMap {
    fn serialize(&self, serializer: &mut crate::CborSerializer<'_>) -> Result<()> {
        serializer.map(self.entries.len())?;
        for (label, value) in &self.entries {
            label.serialize(serializer)?;
            value.serialize(serializer)?;
        }
        Ok(())
    }
}

impl<'de> CborDeserialize<'de> for HeaderMap {
    fn deserialize(deserializer: &mut CborDeserializer<'de>) -> Result<Self> {
        let len = deserializer.map()?;
        let mut map = Self::new();
        for _ in 0..len {
            let label = HeaderLabel::deserialize(deserializer)?;
            let mut value = HeaderValue::deserialize(deserializer)?;
            if label == HeaderLabel::Algorithm {
                if let HeaderValue::Int(id) = value {
                    value = HeaderValue::Algorithm(Algorithm::from_id(
                        i32::try_from(id).map_err(|_| Error::CborOverflow)?,
                    ));
                }
            }
            map.insert(label, value);
        }
        Ok(map)
    }
}

/// Protected COSE header map.
#[derive(Clone, Debug, Default, PartialEq)]
pub struct ProtectedHeader(pub HeaderMap);

/// Unprotected COSE header map.
#[derive(Clone, Debug, Default, PartialEq)]
pub struct UnprotectedHeader(pub HeaderMap);

impl ProtectedHeader {
    /// Decode from a protected header byte string.
    pub fn from_bstr(input: &[u8]) -> Result<Self> {
        if input.is_empty() {
            return Ok(Self(HeaderMap::new()));
        }
        crate::from_slice(input).map(Self)
    }
}