image4 0.8.2

use crate::{
    property::{ConstraintRef, ValueRef},
    util, Tag,
};
use der::{asn1::Ia5StringRef, Decode, Encode, ErrorKind, Length, Reader, SliceReader};
#[cfg(any(feature = "alloc", test))]
use {
    crate::property::{Constraint, Value},
    alloc::collections::BTreeMap,
    der::{Header, Writer},
};

const DER_TAG_NUMBER_MASK: u8 = 0x1f;
const DER_TAG_CLASS_AND_CP_MASK: u8 = 0xe0;

fn read_multibyte_tag<'a, const CLASS_AND_CP: u8, R: Reader<'a>>(
    decoder: &mut R,
) -> der::Result<u32> {
    let first_byte = decoder.read_byte()?;
    if first_byte & DER_TAG_CLASS_AND_CP_MASK != CLASS_AND_CP {
        return Err(decoder.error(ErrorKind::TagNumberInvalid));
    }

    let tag_number = if (first_byte & DER_TAG_NUMBER_MASK) != DER_TAG_NUMBER_MASK {
        (first_byte & DER_TAG_NUMBER_MASK) as u32
    } else {
        let mut value = 0u32;

        loop {
            let byte = decoder.read_byte()?;
            let lower_bits = (byte & 0x7f) as u32;
            if value == 0 && lower_bits == 0 {
                return Err(decoder.error(ErrorKind::TagNumberInvalid));
            }

            value = lower_bits
                | value
                    .checked_shl(7)
                    .ok_or_else(|| decoder.error(ErrorKind::TagNumberInvalid))?;

            if byte & 0x80 == 0 {
                break;
            }
        }

        value
    };

    Ok(tag_number)
}

#[cfg(any(feature = "alloc", test))]
fn write_multibyte_tag<const CLASS_AND_CP: u8, W: Writer>(
    encoder: &mut W,
    tag: u32,
) -> der::Result<()> {
    if tag & (DER_TAG_NUMBER_MASK as u32) == tag && tag != (DER_TAG_NUMBER_MASK as u32) {
        encoder.write_byte(tag as u8 | CLASS_AND_CP)
    } else {
        encoder.write_byte(DER_TAG_NUMBER_MASK | CLASS_AND_CP)?;

        let mut shift = (32 - tag.leading_zeros()) / 7 * 7;
        loop {
            let byte = (tag.wrapping_shr(shift) & 0x7f) as u8 | if shift > 0 { 0x80 } else { 0x00 };

            encoder.write_byte(byte)?;

            if shift == 0 {
                break Ok(());
            } else {
                shift -= 7;
            }
        }
    }
}

#[cfg(any(feature = "alloc", test))]
fn multibyte_tag_encoded_len(tag: u32) -> Length {
    let bits = 32 - tag.leading_zeros();

    Length::new(if bits <= 5 && tag != 0x1f {
        1u32
    } else {
        1 + (bits + 6) / 7
    } as u16)
}

pub fn decode_property<'a, R: Reader<'a>, V: Decode<'a>>(decoder: &mut R) -> der::Result<(Tag, V)> {
    let tag = read_multibyte_tag::<0xe0, _>(decoder)?;
    let value = decode_len_and_body(tag, decoder)?;

    Ok((Tag::from(tag), value))
}

#[cfg(any(feature = "alloc", test))]
pub fn decode_dict_body<'a, R: Reader<'a>, V: Decode<'a>>(
    decoder: &mut R,
    header: Header,
) -> der::Result<BTreeMap<Tag, V>> {
    let mut dict = BTreeMap::new();

    let slice = decoder.read_slice(header.length)?;
    let mut decoder = SliceReader::new(slice)?;

    while !decoder.is_finished() {
        let position = decoder.position();

        let (tag, value) = decode_property(&mut decoder)?;
        if dict.insert(tag, value).is_some() {
            return Err(ErrorKind::TagNumberInvalid.at(position));
        }
    }

    decoder.finish(dict)
}

fn decode_len_and_body<'a, R: Reader<'a>, V: Decode<'a>>(
    tag: u32,
    decoder: &mut R,
) -> der::Result<V> {
    let len = Length::decode(decoder)?;
    let slice = decoder.read_slice(len)?;
    let seq_position = decoder.position();
    let mut decoder = SliceReader::new(slice)?;

    let value = decoder
        .sequence(|decoder| {
            let tag_str = Ia5StringRef::decode(decoder)?;
            if tag_str.as_bytes() != tag.to_be_bytes() {
                return Err(decoder.error(ErrorKind::Value {
                    tag: der::Tag::Ia5String,
                }));
            }

            V::decode(decoder)
        })
        .map_err(|e| util::shift_error_position(e, seq_position))?;

    decoder.finish(value)
}

#[cfg(any(feature = "alloc", test))]
fn encoded_property_len<V: Encode>(tag: Tag, value: &V) -> der::Result<Length> {
    let tag_str_bytes = tag.to_bytes();
    let tag_str = Ia5StringRef::new(&tag_str_bytes)?;

    let seq_body_len = (tag_str.encoded_len() + value.encoded_len()?)?;
    let seq_header = Header {
        tag: der::Tag::Sequence,
        length: seq_body_len,
    };
    let seq_len = (seq_header.encoded_len() + seq_body_len)?;

    multibyte_tag_encoded_len(tag.into()) + seq_len.encoded_len()? + seq_len
}

#[cfg(any(feature = "alloc", test))]
pub fn encoded_dict_body_len<V: Encode>(dict: &BTreeMap<Tag, V>) -> der::Result<Length> {
    dict.iter().try_fold(Length::ZERO, |r, (&tag, value)| {
        r + encoded_property_len(tag, value)?
    })
}

#[cfg(any(feature = "alloc", test))]
pub fn encode_dict_body<W: Writer, V: Encode>(
    dict: &BTreeMap<Tag, V>,
    encoder: &mut W,
) -> der::Result<()> {
    for (tag, value) in dict {
        let tag_str_bytes = tag.to_bytes();
        let tag_str = Ia5StringRef::new(&tag_str_bytes)?;

        let seq_body_len = (tag_str.encoded_len() + value.encoded_len()?)?;
        let seq_header = Header {
            tag: der::Tag::Sequence,
            length: seq_body_len,
        };
        let seq_len = (seq_header.encoded_len() + seq_body_len)?;

        write_multibyte_tag::<0xe0, _>(encoder, (*tag).into())?;
        seq_len.encode(encoder)?;
        seq_header.encode(encoder)?;
        tag_str.encode(encoder)?;
        value.encode(encoder)?;
    }

    Ok(())
}

/// A marker trait for types that can be decoded as either an Image4 property constraint or a
/// regular property.
pub trait DecodeProperty<'a>: Decode<'a> + sealed::Sealed {}

/// A marker trait for types that represent either an Image4 property constraint or a regular
/// property and may contain undecoded dictionaries and arrays.
pub trait PropertyRef<'a>: DecodeProperty<'a> + Encode {
    /// An associated type that represents the same flavor of properties that are fully decoded and
    /// own their contents.
    #[cfg(any(feature = "alloc", test))]
    type Owned: PropertyOwned<Ref<'a> = Self>;
}

/// A marker trait for types that represent either an Image4 property constraint or a regular
/// property, are fully decoded and own their contents.
#[cfg(any(feature = "alloc", test))]
pub trait PropertyOwned: for<'a> DecodeProperty<'a> + Encode {
    /// An associated type that represents the same flavor of properties that may be not fully
    /// decoded.
    type Ref<'a>: PropertyRef<'a, Owned = Self>;
}

mod sealed {
    pub trait Sealed {}
}

impl sealed::Sealed for ValueRef<'_> {}

impl<'a> DecodeProperty<'a> for ValueRef<'a> {}

impl<'a> PropertyRef<'a> for ValueRef<'a> {
    #[cfg(any(feature = "alloc", test))]
    type Owned = Value;
}

impl sealed::Sealed for ConstraintRef<'_> {}

impl<'a> DecodeProperty<'a> for ConstraintRef<'a> {}

impl<'a> PropertyRef<'a> for ConstraintRef<'a> {
    #[cfg(any(feature = "alloc", test))]
    type Owned = Constraint;
}

#[cfg(any(feature = "alloc", test))]
impl sealed::Sealed for Value {}

#[cfg(any(feature = "alloc", test))]
impl DecodeProperty<'_> for Value {}

#[cfg(any(feature = "alloc", test))]
impl PropertyOwned for Value {
    type Ref<'a> = ValueRef<'a>;
}

#[cfg(any(feature = "alloc", test))]
impl sealed::Sealed for Constraint {}

#[cfg(any(feature = "alloc", test))]
impl DecodeProperty<'_> for Constraint {}

#[cfg(any(feature = "alloc", test))]
impl PropertyOwned for Constraint {
    type Ref<'a> = ConstraintRef<'a>;
}