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 kind: WriteErrorKind,
13 #[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 fn atom_type(&self) -> FourCC;
27
28 fn into_body_bytes(self) -> Vec<u8>;
30
31 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 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 let bytes = atom.into_bytes();
101
102 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 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 let total_size = 16u64 + data_size;
137
138 result.extend_from_slice(&1u32.to_be_bytes());
140
141 result.extend_from_slice(&atom_type.0);
143
144 result.extend_from_slice(&total_size.to_be_bytes());
146 } else {
147 let total_size = total_size_with_32bit_header as u32;
149
150 result.extend_from_slice(&total_size.to_be_bytes());
152
153 result.extend_from_slice(&atom_type.0);
155 }
156
157 result
158}