Skip to main content

mp4_edit/
writer.rs

1use derive_more::Display;
2use futures_io::AsyncWrite;
3use futures_util::AsyncWriteExt;
4use thiserror::Error;
5
6use crate::{atom::FourCC, Atom, AtomData};
7
8#[derive(Debug, Error)]
9#[error("{kind}{}", self.source.as_ref().map(|e| format!(" ({e})")).unwrap_or_default())]
10pub struct WriteError {
11    /// The kind of error that occurred during writing.
12    kind: WriteErrorKind,
13    /// The source error that caused this error.
14    #[source]
15    source: Option<Box<dyn std::error::Error + Send + Sync>>,
16}
17
18#[derive(Debug, Display)]
19pub enum WriteErrorKind {
20    #[display("I/O error")]
21    Io,
22}
23
24pub trait SerializeAtom: Sized {
25    /// [`FourCC`] representing atom type
26    fn atom_type(&self) -> FourCC;
27
28    /// Serialize an atom's body
29    fn into_body_bytes(self) -> Vec<u8>;
30
31    /// Serialize an atom into bytes
32    fn into_bytes(self) -> Vec<u8> {
33        let atom_type = self.atom_type();
34        let mut body = self.into_body_bytes();
35        let mut header = serialize_atom_header(atom_type, body.len() as u64);
36        header.append(&mut body);
37        header
38    }
39}
40
41pub struct Mp4Writer<W> {
42    writer: W,
43    offset: usize,
44}
45
46impl<W: AsyncWrite + Unpin> Mp4Writer<W> {
47    pub fn new(writer: W) -> Self {
48        Self { writer, offset: 0 }
49    }
50
51    pub fn current_offset(&self) -> usize {
52        self.offset
53    }
54
55    pub async fn flush(&mut self) -> Result<(), WriteError> {
56        self.writer.flush().await.map_err(|e| WriteError {
57            kind: WriteErrorKind::Io,
58            source: Some(Box::new(e)),
59        })
60    }
61
62    pub async fn write_atom_header(
63        &mut self,
64        atom_type: FourCC,
65        data_size: usize,
66    ) -> Result<(), WriteError> {
67        // Write atom header
68        let header_bytes = serialize_atom_header(atom_type, data_size as u64);
69        self.writer
70            .write_all(&header_bytes)
71            .await
72            .map_err(|e| WriteError {
73                kind: WriteErrorKind::Io,
74                source: Some(Box::new(e)),
75            })?;
76        self.offset += header_bytes.len();
77        Ok(())
78    }
79
80    pub async fn write_leaf_atom(
81        &mut self,
82        atom_type: FourCC,
83        data: AtomData,
84    ) -> Result<(), WriteError> {
85        let data_bytes: Vec<u8> = data.into_body_bytes();
86        self.write_atom_header(atom_type, data_bytes.len()).await?;
87        self.writer
88            .write_all(&data_bytes)
89            .await
90            .map_err(|e| WriteError {
91                kind: WriteErrorKind::Io,
92                source: Some(Box::new(e)),
93            })?;
94        self.offset += data_bytes.len();
95        Ok(())
96    }
97
98    pub async fn write_atom(&mut self, atom: Atom) -> Result<(), WriteError> {
99        // Serialize the entire atom tree into bytes
100        let bytes = atom.into_bytes();
101
102        // Write all bytes at once
103        self.writer
104            .write_all(&bytes)
105            .await
106            .map_err(|e| WriteError {
107                kind: WriteErrorKind::Io,
108                source: Some(Box::new(e)),
109            })?;
110
111        self.offset += bytes.len();
112
113        Ok(())
114    }
115
116    pub async fn write_raw(&mut self, data: &[u8]) -> Result<(), WriteError> {
117        self.writer.write_all(data).await.map_err(|e| WriteError {
118            kind: WriteErrorKind::Io,
119            source: Some(Box::new(e)),
120        })?;
121
122        self.offset += data.len();
123        Ok(())
124    }
125}
126
127fn serialize_atom_header(atom_type: FourCC, data_size: u64) -> Vec<u8> {
128    let mut result = Vec::new();
129
130    // Determine if we need 64-bit size
131    let total_size_with_32bit_header = 8u64 + data_size;
132    let use_64bit = total_size_with_32bit_header > u64::from(u32::MAX);
133
134    if use_64bit {
135        // Extended 64-bit size format: size=1 (4 bytes) + type (4 bytes) + extended_size (8 bytes) + content
136        let total_size = 16u64 + data_size;
137
138        // Write size=1 to indicate extended format
139        result.extend_from_slice(&1u32.to_be_bytes());
140
141        // Write atom type
142        result.extend_from_slice(&atom_type.0);
143
144        // Write extended size
145        result.extend_from_slice(&total_size.to_be_bytes());
146    } else {
147        // Standard 32-bit size format: size (4 bytes) + type (4 bytes) + content
148        let total_size = total_size_with_32bit_header as u32;
149
150        // Write size
151        result.extend_from_slice(&total_size.to_be_bytes());
152
153        // Write atom type
154        result.extend_from_slice(&atom_type.0);
155    }
156
157    result
158}