wolfcose 0.1.0

use crate::{
    ByteBuf, CborDeserialize, CborDeserializer, CborMajorType, CborValue, Error, HeaderMap,
    ProtectedHeader, Result, UnprotectedHeader,
};

/// Known COSE message kind.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum CoseMessageKind {
    /// COSE_Sign1.
    Sign1,
    /// COSE_Sign.
    Sign,
    /// COSE_Encrypt0.
    Encrypt0,
    /// COSE_Encrypt.
    Encrypt,
    /// COSE_Mac0.
    Mac0,
    /// COSE_Mac.
    Mac,
}

/// Parsed COSE_Sign1 envelope.
#[derive(Clone, Debug, PartialEq)]
pub struct CoseSign1Message<'a> {
    raw: &'a [u8],
    protected: ProtectedHeader,
    unprotected: UnprotectedHeader,
    payload_attached: bool,
}

/// Parsed COSE_Sign envelope.
#[derive(Clone, Debug, PartialEq)]
pub struct CoseSignMessage<'a> {
    raw: &'a [u8],
    protected: ProtectedHeader,
    unprotected: UnprotectedHeader,
    payload_attached: bool,
    signer_count: usize,
}

/// Parsed COSE_Encrypt0 envelope.
#[derive(Clone, Debug, PartialEq)]
pub struct CoseEncrypt0Message<'a> {
    raw: &'a [u8],
    protected: ProtectedHeader,
    unprotected: UnprotectedHeader,
    ciphertext_attached: bool,
}

/// Parsed COSE_Encrypt envelope.
#[derive(Clone, Debug, PartialEq)]
pub struct CoseEncryptMessage<'a> {
    raw: &'a [u8],
    protected: ProtectedHeader,
    unprotected: UnprotectedHeader,
    ciphertext_attached: bool,
    recipient_count: usize,
}

/// Parsed COSE_Mac0 envelope.
#[derive(Clone, Debug, PartialEq)]
pub struct CoseMac0Message<'a> {
    raw: &'a [u8],
    protected: ProtectedHeader,
    unprotected: UnprotectedHeader,
    payload_attached: bool,
}

/// Parsed COSE_Mac envelope.
#[derive(Clone, Debug, PartialEq)]
pub struct CoseMacMessage<'a> {
    raw: &'a [u8],
    protected: ProtectedHeader,
    unprotected: UnprotectedHeader,
    payload_attached: bool,
    recipient_count: usize,
}

macro_rules! common_methods {
    ($ty:ty, $attached:ident, $attached_name:ident) => {
        impl<'a> $ty {
            /// Original encoded message bytes.
            pub fn raw(&self) -> &'a [u8] {
                self.raw
            }

            /// Protected headers.
            pub fn protected(&self) -> &HeaderMap {
                &self.protected.0
            }

            /// Unprotected headers.
            pub fn unprotected(&self) -> &HeaderMap {
                &self.unprotected.0
            }

            /// Whether the payload or ciphertext is attached in the message.
            pub fn $attached_name(&self) -> bool {
                self.$attached
            }
        }
    };
}

common_methods!(CoseSign1Message<'a>, payload_attached, payload_attached);
common_methods!(CoseSignMessage<'a>, payload_attached, payload_attached);
common_methods!(
    CoseEncrypt0Message<'a>,
    ciphertext_attached,
    ciphertext_attached
);
common_methods!(
    CoseEncryptMessage<'a>,
    ciphertext_attached,
    ciphertext_attached
);
common_methods!(CoseMac0Message<'a>, payload_attached, payload_attached);
common_methods!(CoseMacMessage<'a>, payload_attached, payload_attached);

impl<'a> CoseSign1Message<'a> {
    /// Parse a COSE_Sign1 envelope.
    pub fn parse(input: &'a [u8]) -> Result<Self> {
        let parsed = parse_envelope(input, CoseMessageKind::Sign1, 4)?;
        Ok(Self {
            raw: input,
            protected: parsed.protected,
            unprotected: parsed.unprotected,
            payload_attached: parsed.body_attached,
        })
    }
}

impl<'a> CoseSignMessage<'a> {
    /// Parse a COSE_Sign envelope.
    pub fn parse(input: &'a [u8]) -> Result<Self> {
        let parsed = parse_envelope(input, CoseMessageKind::Sign, 4)?;
        Ok(Self {
            raw: input,
            protected: parsed.protected,
            unprotected: parsed.unprotected,
            payload_attached: parsed.body_attached,
            signer_count: parsed.trailing_count,
        })
    }

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

impl<'a> CoseEncrypt0Message<'a> {
    /// Parse a COSE_Encrypt0 envelope.
    pub fn parse(input: &'a [u8]) -> Result<Self> {
        let parsed = parse_envelope(input, CoseMessageKind::Encrypt0, 3)?;
        Ok(Self {
            raw: input,
            protected: parsed.protected,
            unprotected: parsed.unprotected,
            ciphertext_attached: parsed.body_attached,
        })
    }
}

impl<'a> CoseEncryptMessage<'a> {
    /// Parse a COSE_Encrypt envelope.
    pub fn parse(input: &'a [u8]) -> Result<Self> {
        let parsed = parse_envelope(input, CoseMessageKind::Encrypt, 4)?;
        Ok(Self {
            raw: input,
            protected: parsed.protected,
            unprotected: parsed.unprotected,
            ciphertext_attached: parsed.body_attached,
            recipient_count: parsed.trailing_count,
        })
    }

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

impl<'a> CoseMac0Message<'a> {
    /// Parse a COSE_Mac0 envelope.
    pub fn parse(input: &'a [u8]) -> Result<Self> {
        let parsed = parse_envelope(input, CoseMessageKind::Mac0, 4)?;
        Ok(Self {
            raw: input,
            protected: parsed.protected,
            unprotected: parsed.unprotected,
            payload_attached: parsed.body_attached,
        })
    }
}

impl<'a> CoseMacMessage<'a> {
    /// Parse a COSE_Mac envelope.
    pub fn parse(input: &'a [u8]) -> Result<Self> {
        let parsed = parse_envelope(input, CoseMessageKind::Mac, 5)?;
        Ok(Self {
            raw: input,
            protected: parsed.protected,
            unprotected: parsed.unprotected,
            payload_attached: parsed.body_attached,
            recipient_count: parsed.trailing_count,
        })
    }

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

struct ParsedEnvelope {
    protected: ProtectedHeader,
    unprotected: UnprotectedHeader,
    body_attached: bool,
    trailing_count: usize,
}

fn parse_envelope(input: &[u8], kind: CoseMessageKind, array_len: usize) -> Result<ParsedEnvelope> {
    let mut deserializer = CborDeserializer::new(input);
    consume_optional_expected_tag(&mut deserializer, kind)?;
    let len = deserializer.array()?;
    if len != array_len {
        return Err(Error::CborMalformed);
    }

    let protected =
        ProtectedHeader::from_bstr(ByteBuf::deserialize(&mut deserializer)?.0.as_slice())?;
    let unprotected = UnprotectedHeader(HeaderMap::deserialize(&mut deserializer)?);
    let body_attached = !consume_null_or_item(&mut deserializer)?;

    let trailing_count = match kind {
        CoseMessageKind::Sign1 | CoseMessageKind::Encrypt0 | CoseMessageKind::Mac0 => {
            deserializer.decoder_mut().skip()?;
            0
        }
        CoseMessageKind::Sign | CoseMessageKind::Encrypt => {
            let count = deserializer.array()?;
            for _ in 0..count {
                deserializer.decoder_mut().skip()?;
            }
            count
        }
        CoseMessageKind::Mac => {
            deserializer.decoder_mut().skip()?;
            let count = deserializer.array()?;
            for _ in 0..count {
                deserializer.decoder_mut().skip()?;
            }
            count
        }
    };

    if !deserializer.is_finished() {
        return Err(Error::CborMalformed);
    }

    Ok(ParsedEnvelope {
        protected,
        unprotected,
        body_attached,
        trailing_count,
    })
}

fn consume_optional_expected_tag(
    deserializer: &mut CborDeserializer<'_>,
    kind: CoseMessageKind,
) -> Result<()> {
    if deserializer.decoder_mut().peek_type() != Some(CborMajorType::TAG) {
        return Ok(());
    }
    let tag = deserializer.tag()?;
    let expected = match kind {
        CoseMessageKind::Sign1 => crate::raw::WOLFCOSE_TAG_SIGN1 as u64,
        CoseMessageKind::Sign => crate::raw::WOLFCOSE_TAG_SIGN as u64,
        CoseMessageKind::Encrypt0 => crate::raw::WOLFCOSE_TAG_ENCRYPT0 as u64,
        CoseMessageKind::Encrypt => crate::raw::WOLFCOSE_TAG_ENCRYPT as u64,
        CoseMessageKind::Mac0 => crate::raw::WOLFCOSE_TAG_MAC0 as u64,
        CoseMessageKind::Mac => crate::raw::WOLFCOSE_TAG_MAC as u64,
    };
    if tag == expected {
        Ok(())
    } else {
        Err(Error::CoseBadTag)
    }
}

fn consume_null_or_item(deserializer: &mut CborDeserializer<'_>) -> Result<bool> {
    if deserializer.decoder_mut().peek_initial_byte() == Some(crate::raw::WOLFCOSE_CBOR_NULL as u8)
    {
        deserializer.null()?;
        Ok(true)
    } else {
        CborValue::deserialize(deserializer)?;
        Ok(false)
    }
}