mp4-track 0.1.0

use crate::{Error, FourCc, Limits};
use mp4_atom::ReadFrom as _;

pub(crate) const FTYP: FourCc = FourCc(*b"ftyp");
pub(crate) const MOOV: FourCc = FourCc(*b"moov");
pub(crate) const MOOF: FourCc = FourCc(*b"moof");
pub(crate) const MDAT: FourCc = FourCc(*b"mdat");

pub(crate) fn validate_typed_atom(kind: FourCc, body: &[u8]) -> Result<(), Error> {
    let size = body.len().checked_add(8).ok_or(Error::LimitExceeded {
        limit: "atom_size",
        value: u64::MAX,
        max: u64::from(u32::MAX),
    })?;
    let size = u32::try_from(size).map_err(|_| Error::LimitExceeded {
        limit: "atom_size",
        value: u64::try_from(body.len()).unwrap_or(u64::MAX),
        max: u64::from(u32::MAX).saturating_sub(8),
    })?;
    let mut encoded = Vec::with_capacity(usize::try_from(size).unwrap_or(0));
    encoded.extend_from_slice(&size.to_be_bytes());
    encoded.extend_from_slice(&kind.0);
    encoded.extend_from_slice(body);
    let mut cursor = std::io::Cursor::new(encoded);
    match mp4_atom::Any::read_from(&mut cursor) {
        Ok(_) => Ok(()),
        // This crate tolerates trailing bytes inside a box it has already
        // parsed, so a partial decode by the box layer is not a failure.
        Err(mp4_atom::Error::UnderDecode(_)) => Ok(()),
        Err(mp4_atom::Error::OverDecode(inner)) => Err(Error::Malformed {
            kind: FourCc(<[u8; 4]>::from(inner)),
            reason: "declared content exceeds the box body",
        }),
        Err(error) => Err(Error::Atom(error)),
    }
}

#[derive(Clone, Copy, Debug)]
pub(crate) struct BoxHeader {
    pub(crate) kind: FourCc,
    pub(crate) start: usize,
    pub(crate) header_len: usize,
    pub(crate) size: usize,
}

impl BoxHeader {
    pub(crate) fn body_start(self) -> Result<usize, Error> {
        self.start
            .checked_add(self.header_len)
            .ok_or(Error::Malformed {
                kind: self.kind,
                reason: "box body offset overflow",
            })
    }

    pub(crate) fn end(self) -> Result<usize, Error> {
        self.start.checked_add(self.size).ok_or(Error::Malformed {
            kind: self.kind,
            reason: "box end overflow",
        })
    }

    pub(crate) fn body(self, data: &[u8]) -> Result<&[u8], Error> {
        let start = self.body_start()?;
        let end = self.end()?;
        data.get(start..end).ok_or(Error::Malformed {
            kind: self.kind,
            reason: "box body is outside its parent",
        })
    }
}

fn u32_at(data: &[u8], at: usize) -> Option<u32> {
    let bytes: [u8; 4] = data.get(at..at.checked_add(4)?)?.try_into().ok()?;
    Some(u32::from_be_bytes(bytes))
}

fn u64_at(data: &[u8], at: usize) -> Option<u64> {
    let bytes: [u8; 8] = data.get(at..at.checked_add(8)?)?.try_into().ok()?;
    Some(u64::from_be_bytes(bytes))
}

pub(crate) fn parse_header(
    data: &[u8],
    at: usize,
    parent: Option<FourCc>,
    absolute_base: u64,
    allow_zero: bool,
) -> Result<BoxHeader, Error> {
    let absolute = absolute_base
        .checked_add(u64::try_from(at).map_err(|_| Error::Malformed {
            kind: parent.unwrap_or(FourCc(*b"root")),
            reason: "offset does not fit u64",
        })?)
        .ok_or(Error::Malformed {
            kind: parent.unwrap_or(FourCc(*b"root")),
            reason: "absolute offset overflow",
        })?;
    let remaining = data.len().saturating_sub(at);
    if remaining < 8 {
        return Err(Error::Truncated {
            at: absolute,
            needed: 8,
        });
    }
    let size32 = u32_at(data, at).ok_or(Error::Truncated {
        at: absolute,
        needed: 4,
    })?;
    let kind_at = at.checked_add(4).ok_or(Error::Truncated {
        at: absolute,
        needed: 8,
    })?;
    let kind_bytes: [u8; 4] = data
        .get(
            kind_at..kind_at.checked_add(4).ok_or(Error::Truncated {
                at: absolute,
                needed: 8,
            })?,
        )
        .ok_or(Error::Truncated {
            at: absolute,
            needed: 8,
        })?
        .try_into()
        .map_err(|_| Error::Truncated {
            at: absolute,
            needed: 8,
        })?;
    let kind = FourCc(kind_bytes);
    let (mut size, mut header_len) = match size32 {
        0 => {
            if !allow_zero {
                return Err(Error::ZeroSizeNested { kind, at: absolute });
            }
            (remaining, 8usize)
        }
        1 => {
            if remaining < 16 {
                return Err(Error::Truncated {
                    at: absolute,
                    needed: 16,
                });
            }
            let large_at = at.checked_add(8).ok_or(Error::Truncated {
                at: absolute,
                needed: 16,
            })?;
            let large = u64_at(data, large_at).ok_or(Error::Truncated {
                at: absolute,
                needed: 16,
            })?;
            let converted = usize::try_from(large).map_err(|_| Error::BoxExceedsFile {
                kind,
                at: absolute,
                declared_end: u64::MAX,
                file_len: u64::try_from(data.len()).unwrap_or(u64::MAX),
            })?;
            (converted, 16usize)
        }
        value => (
            usize::try_from(value).map_err(|_| Error::BoxSizeTooSmall {
                kind,
                size: u64::from(value),
                at: absolute,
            })?,
            8usize,
        ),
    };
    if kind == FourCc(*b"uuid") {
        header_len = header_len.checked_add(16).ok_or(Error::BoxSizeTooSmall {
            kind,
            size: u64::try_from(size).unwrap_or(u64::MAX),
            at: absolute,
        })?;
    }
    if size < header_len {
        return Err(Error::BoxSizeTooSmall {
            kind,
            size: u64::try_from(size).unwrap_or(u64::MAX),
            at: absolute,
        });
    }
    if size32 == 0 {
        size = remaining;
    }
    Ok(BoxHeader {
        kind,
        start: at,
        header_len,
        size,
    })
}

pub(crate) fn children(
    body: &[u8],
    parent: FourCc,
    absolute_base: u64,
) -> Result<Vec<(BoxHeader, &[u8])>, Error> {
    let mut result = Vec::new();
    let mut at = 0usize;
    while at < body.len() {
        let header = parse_header(body, at, Some(parent), absolute_base, false)?;
        let end = header.end()?;
        if end > body.len() {
            let child_at = absolute_base.saturating_add(u64::try_from(at).unwrap_or(u64::MAX));
            return Err(Error::ChildExceedsParent {
                parent,
                child: header.kind,
                at: child_at,
            });
        }
        result.push((header, header.body(body)?));
        at = end;
    }
    Ok(result)
}

fn is_container(kind: FourCc) -> bool {
    [
        FourCc(*b"moov"),
        FourCc(*b"trak"),
        FourCc(*b"mdia"),
        FourCc(*b"minf"),
        FourCc(*b"stbl"),
        FourCc(*b"edts"),
        FourCc(*b"dinf"),
        FourCc(*b"mvex"),
        FourCc(*b"moof"),
        FourCc(*b"traf"),
        FourCc(*b"udta"),
        FourCc(*b"ilst"),
    ]
    .contains(&kind)
}

pub(crate) fn validate_tree(
    data: &[u8],
    root: FourCc,
    absolute_base: u64,
    limits: Limits,
) -> Result<(), Error> {
    validate_children(data, root, absolute_base, 1, limits)
}

fn validate_children(
    data: &[u8],
    parent: FourCc,
    absolute_base: u64,
    depth: u8,
    limits: Limits,
) -> Result<(), Error> {
    if depth > limits.max_depth {
        return Err(Error::NestingTooDeep {
            depth,
            at: absolute_base,
        });
    }
    let mut at = 0usize;
    while at < data.len() {
        let header = parse_header(data, at, Some(parent), absolute_base, false)?;
        let end = header.end()?;
        if end > data.len() {
            return Err(Error::ChildExceedsParent {
                parent,
                child: header.kind,
                at: absolute_base.saturating_add(u64::try_from(at).unwrap_or(u64::MAX)),
            });
        }
        if is_container(header.kind) {
            let body = header.body(data)?;
            let child_base = absolute_base
                .checked_add(u64::try_from(header.body_start()?).unwrap_or(u64::MAX))
                .ok_or(Error::Malformed {
                    kind: header.kind,
                    reason: "absolute child offset overflow",
                })?;
            validate_children(
                body,
                header.kind,
                child_base,
                depth.checked_add(1).ok_or(Error::NestingTooDeep {
                    depth: u8::MAX,
                    at: child_base,
                })?,
                limits,
            )?;
        } else if header.kind == FourCc(*b"meta") {
            let body = header.body(data)?;
            let nested = body.get(4..).ok_or(Error::Malformed {
                kind: header.kind,
                reason: "missing full-box fields",
            })?;
            let child_base = absolute_base
                .checked_add(
                    u64::try_from(header.body_start()?.saturating_add(4)).unwrap_or(u64::MAX),
                )
                .ok_or(Error::Malformed {
                    kind: header.kind,
                    reason: "absolute child offset overflow",
                })?;
            validate_children(
                nested,
                header.kind,
                child_base,
                depth.checked_add(1).ok_or(Error::NestingTooDeep {
                    depth: u8::MAX,
                    at: child_base,
                })?,
                limits,
            )?;
        }
        at = end;
    }
    Ok(())
}

pub(crate) struct Cursor<'a> {
    data: &'a [u8],
    pos: usize,
    kind: FourCc,
}

impl<'a> Cursor<'a> {
    pub(crate) fn new(data: &'a [u8], kind: FourCc) -> Self {
        Self { data, pos: 0, kind }
    }

    pub(crate) fn position(&self) -> usize {
        self.pos
    }

    pub(crate) fn remaining(&self) -> usize {
        self.data.len().saturating_sub(self.pos)
    }

    pub(crate) fn skip(&mut self, count: usize) -> Result<(), Error> {
        self.take(count).map(|_| ())
    }

    pub(crate) fn take(&mut self, count: usize) -> Result<&'a [u8], Error> {
        let end = self.pos.checked_add(count).ok_or(Error::Malformed {
            kind: self.kind,
            reason: "field offset overflow",
        })?;
        let value = self.data.get(self.pos..end).ok_or(Error::Malformed {
            kind: self.kind,
            reason: "box field is truncated",
        })?;
        self.pos = end;
        Ok(value)
    }

    pub(crate) fn rest(&mut self) -> &'a [u8] {
        let value = self.data.get(self.pos..).unwrap_or_default();
        self.pos = self.data.len();
        value
    }

    pub(crate) fn u8(&mut self) -> Result<u8, Error> {
        self.take(1)?.first().copied().ok_or(Error::Malformed {
            kind: self.kind,
            reason: "missing byte",
        })
    }

    pub(crate) fn u16(&mut self) -> Result<u16, Error> {
        let bytes: [u8; 2] = self.take(2)?.try_into().map_err(|_| Error::Malformed {
            kind: self.kind,
            reason: "missing u16",
        })?;
        Ok(u16::from_be_bytes(bytes))
    }

    pub(crate) fn i16(&mut self) -> Result<i16, Error> {
        let bytes: [u8; 2] = self.take(2)?.try_into().map_err(|_| Error::Malformed {
            kind: self.kind,
            reason: "missing i16",
        })?;
        Ok(i16::from_be_bytes(bytes))
    }

    pub(crate) fn u24(&mut self) -> Result<u32, Error> {
        let bytes = self.take(3)?;
        let a = u32::from(bytes.first().copied().unwrap_or_default());
        let b = u32::from(bytes.get(1).copied().unwrap_or_default());
        let c = u32::from(bytes.get(2).copied().unwrap_or_default());
        Ok((a << 16) | (b << 8) | c)
    }

    pub(crate) fn u32(&mut self) -> Result<u32, Error> {
        let bytes: [u8; 4] = self.take(4)?.try_into().map_err(|_| Error::Malformed {
            kind: self.kind,
            reason: "missing u32",
        })?;
        Ok(u32::from_be_bytes(bytes))
    }

    pub(crate) fn i32(&mut self) -> Result<i32, Error> {
        let bytes: [u8; 4] = self.take(4)?.try_into().map_err(|_| Error::Malformed {
            kind: self.kind,
            reason: "missing i32",
        })?;
        Ok(i32::from_be_bytes(bytes))
    }

    pub(crate) fn u64(&mut self) -> Result<u64, Error> {
        let bytes: [u8; 8] = self.take(8)?.try_into().map_err(|_| Error::Malformed {
            kind: self.kind,
            reason: "missing u64",
        })?;
        Ok(u64::from_be_bytes(bytes))
    }

    pub(crate) fn i64(&mut self) -> Result<i64, Error> {
        let bytes: [u8; 8] = self.take(8)?.try_into().map_err(|_| Error::Malformed {
            kind: self.kind,
            reason: "missing i64",
        })?;
        Ok(i64::from_be_bytes(bytes))
    }

    pub(crate) fn fourcc(&mut self) -> Result<FourCc, Error> {
        let bytes: [u8; 4] = self.take(4)?.try_into().map_err(|_| Error::Malformed {
            kind: self.kind,
            reason: "missing fourcc",
        })?;
        Ok(FourCc(bytes))
    }

    pub(crate) fn full_box(&mut self) -> Result<(u8, u32), Error> {
        Ok((self.u8()?, self.u24()?))
    }
}

pub(crate) fn one<'a>(
    items: &[(BoxHeader, &'a [u8])],
    kind: FourCc,
    parent: FourCc,
    required: bool,
) -> Result<Option<&'a [u8]>, Error> {
    let mut found = items.iter().filter(|(header, _)| header.kind == kind);
    let first = found.next().map(|(_, body)| *body);
    if found.next().is_some() {
        return Err(Error::DuplicateBox { kind, parent });
    }
    if required && first.is_none() {
        return Err(Error::MissingBox {
            kind,
            parent: Some(parent),
        });
    }
    Ok(first)
}