image4 0.8.2

//! Contains the [`RestoreInfo`] type that represents Image4 restore info (an IM4R file).

use crate::{
    property::{DictRef, ValueRef},
    util, Tag,
};
use core::ops::Add;
use der::{
    asn1::Ia5StringRef, Decode, Encode, EncodeValue, FixedTag, Length, Reader, Tagged, Writer,
};
#[cfg(any(feature = "alloc", test))]
use {
    crate::property::{Dict, Value},
    alloc::collections::BTreeMap,
    der::referenced::{OwnedToRef, RefToOwned},
};

/// Restore info that might be present in an Image4 file.
#[derive(Clone, Debug)]
#[cfg_attr(test, derive(Eq, PartialEq))]
pub struct RestoreInfoRef<'a>(DictRef<'a>);

impl<'a> RestoreInfoRef<'a> {
    /// Decodes the part of encoded restore info after the magic string.
    ///
    /// May be used when it is required to first identify what kind of image you're looking at by
    /// checking the magic string.
    pub fn decode_after_magic<R: Reader<'a>>(decoder: &mut R) -> der::Result<Self> {
        Ok(Self(DictRef::decode(decoder)?))
    }

    /// Decodes body until `"BNCN"` tag is found, then returns the integer value.
    ///
    /// # Errors
    ///
    /// An error may be returned if a decoding error occurs while searching for the tag or if the tag value is
    /// not an integer.
    ///
    /// # Security
    ///
    /// This function may not be used to validate that Image4 encoding inside restore info is valid. As soon
    /// as the boot nonce tag is found decoding stops and whatever follows the tag may be invalid. In case it
    /// is required to validate the encoding, [`RestoreInfoRef::decode_owned`] method should be used.
    pub fn decode_boot_nonce(&self) -> der::Result<Option<u64>> {
        for result in self.0.iter()? {
            let (tag, value) = result?;

            if tag != Tag::from_bytes(*b"BNCN") {
                continue;
            }

            let ValueRef::Integer(nonce) = value else {
                return Err(der::Error::new(
                    der::ErrorKind::TagUnexpected {
                        expected: Some(der::Tag::Integer),
                        actual: value.tag(),
                    },
                    Length::ZERO,
                ));
            };

            return Ok(Some(nonce));
        }

        Ok(None)
    }

    /// Decodes the contents of the body into an owned type.
    ///
    /// This function completely validates the contents and currently is the only way to do
    /// this.
    #[cfg(any(feature = "alloc", test))]
    pub fn decode_owned(&self) -> der::Result<BTreeMap<Tag, Value>> {
        self.0.decode_owned()
    }
}

impl<'a> Decode<'a> for RestoreInfoRef<'a> {
    fn decode<R: Reader<'a>>(decoder: &mut R) -> der::Result<Self> {
        decoder.sequence(|nested_decoder| {
            util::decode_and_check_magic(nested_decoder, b"IM4R")?;
            Self::decode_after_magic(nested_decoder)
        })
    }
}

impl EncodeValue for RestoreInfoRef<'_> {
    fn value_len(&self) -> der::Result<Length> {
        Ia5StringRef::new(&b"IM4R")?
            .encoded_len()?
            .add(self.0.encoded_len()?)
    }

    fn encode_value(&self, encoder: &mut impl Writer) -> der::Result<()> {
        Ia5StringRef::new(&b"IM4R")?.encode(encoder)?;
        self.0.encode(encoder)
    }
}

impl FixedTag for RestoreInfoRef<'_> {
    const TAG: der::Tag = der::Tag::Sequence;
}

#[cfg(any(feature = "alloc", test))]
impl<'a> RefToOwned<'a> for RestoreInfoRef<'a> {
    type Owned = RestoreInfo;

    fn ref_to_owned(&self) -> Self::Owned {
        RestoreInfo(self.0.ref_to_owned())
    }
}

/// Restore info that might be present in an Image4 file.
#[derive(Clone, Debug)]
#[cfg_attr(test, derive(Eq, PartialEq))]
#[cfg(any(feature = "alloc", test))]
pub struct RestoreInfo(Dict);

#[cfg(any(feature = "alloc", test))]
impl RestoreInfo {
    /// Encodes a [`RestoreInfo`] containing only a boot nonce.
    pub fn with_boot_nonce(nonce: u64) -> der::Result<Self> {
        let mut map = BTreeMap::new();
        map.insert(Tag::from_bytes(*b"BNCN"), Value::Integer(nonce));

        Self::encode_from(&map)
    }

    /// Encodes a [`BTreeMap`] describing a dictionary of Image4 properties into a [`Dict`] and wraps
    /// it into a [`RestoreInfo`].
    pub fn encode_from(value: &BTreeMap<Tag, Value>) -> der::Result<Self> {
        Dict::encode_from(value).map(Self)
    }

    /// Returns a byte slice for the encoded wrapped dictionary.
    pub fn as_bytes(&self) -> &[u8] {
        self.0.as_bytes()
    }

    /// Converts restore info into the dictionary it wraps.
    pub fn into_dict(self) -> Dict {
        self.0
    }

    /// Returns a borrowed version of the underlying dictionary.
    pub fn as_dict_ref(&self) -> DictRef<'_> {
        self.0.owned_to_ref()
    }
}

#[cfg(any(feature = "alloc", test))]
impl From<Dict> for RestoreInfo {
    fn from(value: Dict) -> Self {
        Self(value)
    }
}

#[cfg(any(feature = "alloc", test))]
impl From<&'_ DictRef<'_>> for RestoreInfo {
    fn from(value: &'_ DictRef<'_>) -> Self {
        Self(value.into())
    }
}

#[cfg(any(feature = "alloc", test))]
impl From<DictRef<'_>> for RestoreInfo {
    fn from(value: DictRef<'_>) -> Self {
        Self(value.into())
    }
}

#[cfg(any(feature = "alloc", test))]
impl From<RestoreInfo> for Dict {
    fn from(value: RestoreInfo) -> Self {
        value.0
    }
}

#[cfg(any(feature = "alloc", test))]
impl<'a> Decode<'a> for RestoreInfo {
    fn decode<R: Reader<'a>>(decoder: &mut R) -> der::Result<Self> {
        Dict::decode(decoder).map(Self)
    }
}

#[cfg(any(feature = "alloc", test))]
impl Encode for RestoreInfo {
    fn encoded_len(&self) -> der::Result<Length> {
        self.owned_to_ref().encoded_len()
    }

    fn encode(&self, encoder: &mut impl Writer) -> der::Result<()> {
        self.owned_to_ref().encode(encoder)
    }
}

#[cfg(any(feature = "alloc", test))]
impl OwnedToRef for RestoreInfo {
    type Borrowed<'a> = RestoreInfoRef<'a>;

    fn owned_to_ref(&self) -> Self::Borrowed<'_> {
        RestoreInfoRef(self.0.owned_to_ref())
    }
}