mp4_atom/
header.rs

1use std::io::{Cursor, Read};
2
3use crate::*;
4
5/// A atom header, which contains the atom's kind and size.
6#[derive(Debug, Clone, Copy)]
7pub struct Header {
8    /// The name of the atom, always 4 bytes.
9    pub kind: FourCC,
10
11    /// The size of the atom, **excluding** the header.
12    /// This is optional when the atom extends to the end of the file.
13    pub size: Option<usize>,
14}
15
16impl Encode for Header {
17    fn encode<B: BufMut>(&self, buf: &mut B) -> Result<()> {
18        match self.size.map(|size| size + 8) {
19            Some(size) if size > u32::MAX as usize => {
20                1u32.encode(buf)?;
21                self.kind.encode(buf)?;
22
23                // Have to include the size of this extra field
24                ((size + 8) as u64).encode(buf)
25            }
26            Some(size) => {
27                (size as u32).encode(buf)?;
28                self.kind.encode(buf)
29            }
30            None => {
31                0u32.encode(buf)?;
32                self.kind.encode(buf)
33            }
34        }
35    }
36}
37
38impl Decode for Header {
39    fn decode<B: Buf>(buf: &mut B) -> Result<Self> {
40        let size = u32::decode(buf)?;
41        let kind = FourCC::decode(buf)?;
42
43        let size = match size {
44            0 => None,
45            1 => {
46                // Read another 8 bytes
47                let size = u64::decode(buf)?;
48                Some(size.checked_sub(16).ok_or(Error::InvalidSize)? as usize)
49            }
50            _ => Some(size.checked_sub(8).ok_or(Error::InvalidSize)? as usize),
51        };
52
53        Ok(Self { kind, size })
54    }
55}
56
57impl DecodeMaybe for Header {
58    fn decode_maybe<B: Buf>(buf: &mut B) -> Result<Option<Self>> {
59        if buf.remaining() < 8 {
60            return Ok(None);
61        }
62
63        let size = u32::from_be_bytes(buf.slice(4).try_into().unwrap());
64        if size == 1 && buf.remaining() < 16 {
65            return Ok(None);
66        }
67
68        Ok(Some(Self::decode(buf)?))
69    }
70}
71
72impl ReadFrom for Header {
73    fn read_from<R: Read>(r: &mut R) -> Result<Self> {
74        <Option<Header> as ReadFrom>::read_from(r)?.ok_or(Error::UnexpectedEof)
75    }
76}
77
78impl ReadFrom for Option<Header> {
79    fn read_from<R: Read>(r: &mut R) -> Result<Self> {
80        let mut buf = [0u8; 8];
81        let n = r.read(&mut buf)?;
82        if n == 0 {
83            return Ok(None);
84        }
85
86        r.read_exact(&mut buf[n..])?;
87
88        let size = u32::from_be_bytes(buf[0..4].try_into().unwrap());
89        let kind = u32::from_be_bytes(buf[4..8].try_into().unwrap()).into();
90
91        let size = match size {
92            0 => None,
93            1 => {
94                // Read another 8 bytes
95                r.read_exact(&mut buf)?;
96                let size = u64::from_be_bytes(buf);
97                let size = size.checked_sub(16).ok_or(Error::InvalidSize)?;
98
99                Some(size as usize)
100            }
101            _ => Some(size.checked_sub(8).ok_or(Error::InvalidSize)? as usize),
102        };
103
104        Ok(Some(Header { kind, size }))
105    }
106}
107
108// Utility methods
109impl Header {
110    pub(crate) fn read_body<R: Read>(&self, r: &mut R) -> Result<Cursor<Vec<u8>>> {
111        // TODO This allocates on the heap.
112        // Ideally, we should use ReadFrom instead of Decode to avoid this.
113
114        // Don't use `with_capacity` on an untrusted size
115        // We allocate at most 4096 bytes upfront and grow as needed
116        let cap = self.size.unwrap_or(0).min(4096);
117        let mut buf = Vec::with_capacity(cap);
118
119        match self.size {
120            Some(size) => {
121                let n = std::io::copy(&mut r.take(size as _), &mut buf)? as _;
122                if size != n {
123                    return Err(Error::OutOfBounds);
124                }
125            }
126            None => {
127                std::io::copy(r, &mut buf)?;
128            }
129        };
130
131        Ok(Cursor::new(buf))
132    }
133
134    #[cfg(feature = "tokio")]
135    pub(crate) async fn read_body_tokio<R: ::tokio::io::AsyncRead + Unpin>(
136        &self,
137        r: &mut R,
138    ) -> Result<Cursor<Vec<u8>>> {
139        use ::tokio::io::AsyncReadExt;
140
141        // TODO This allocates on the heap.
142        // Ideally, we should use ReadFrom instead of Decode to avoid this.
143
144        // Don't use `with_capacity` on an untrusted size
145        // We allocate at most 4096 bytes upfront and grow as needed
146        let cap = self.size.unwrap_or(0).min(4096);
147        let mut buf = Vec::with_capacity(cap);
148
149        match self.size {
150            Some(size) => {
151                let n = ::tokio::io::copy(&mut r.take(size as _), &mut buf).await? as _;
152                if size != n {
153                    return Err(Error::OutOfBounds);
154                }
155            }
156            None => {
157                ::tokio::io::copy(r, &mut buf).await?;
158            }
159        };
160
161        Ok(Cursor::new(buf))
162    }
163}
mp4_atom/header.rs

mp4_atom/
header.rs
