anni_flac/
header.rs

1use crate::blocks::*;
2use crate::error::FlacError;
3use crate::prelude::*;
4use crate::utils::*;
5use byteorder::{BigEndian, ReadBytesExt, WriteBytesExt};
6use std::fmt;
7use std::fs::{File, OpenOptions};
8use std::io::{Read, Seek, SeekFrom, Write};
9use std::path::{Path, PathBuf};
10
11pub struct FlacHeader {
12    pub blocks: Vec<MetadataBlock>,
13    pub path: PathBuf,
14    frame_offset: usize,
15}
16
17impl FlacHeader {
18    pub fn parse<R: Read>(reader: &mut R, path: PathBuf) -> Result<FlacHeader> {
19        if reader.read_u8()? != b'f'
20            || reader.read_u8()? != b'L'
21            || reader.read_u8()? != b'a'
22            || reader.read_u8()? != b'C'
23        {
24            return Err(FlacError::InvalidMagicNumber);
25        }
26
27        let stream_info = MetadataBlock::from_reader(reader)?;
28        match stream_info.data {
29            MetadataBlockData::StreamInfo(_) => {}
30            _ => return Err(FlacError::InvalidFirstBlock),
31        }
32
33        let mut is_last = stream_info.is_last;
34        let mut blocks = vec![stream_info];
35        let mut frame_offset = 4 + 4 + 34;
36        while !is_last {
37            let block = MetadataBlock::from_reader(reader)?;
38            frame_offset += 4 + block.length;
39            is_last = block.is_last;
40            blocks.push(block);
41        }
42        Ok(FlacHeader {
43            blocks,
44            path: path,
45            frame_offset,
46        })
47    }
48
49    #[cfg(feature = "async")]
50    pub async fn parse_async<R>(reader: &mut R, path: PathBuf) -> Result<FlacHeader>
51    where
52        R: AsyncRead + Unpin + Send,
53    {
54        if reader.read_u8().await? != b'f'
55            || reader.read_u8().await? != b'L'
56            || reader.read_u8().await? != b'a'
57            || reader.read_u8().await? != b'C'
58        {
59            return Err(FlacError::InvalidMagicNumber);
60        }
61
62        let stream_info = MetadataBlock::from_async_reader(reader).await?;
63        match stream_info.data {
64            MetadataBlockData::StreamInfo(_) => {}
65            _ => return Err(FlacError::InvalidFirstBlock),
66        }
67
68        let mut is_last = stream_info.is_last;
69        let mut blocks = vec![stream_info];
70        let mut frame_offset = 4 + 4 + 34;
71        while !is_last {
72            let block = MetadataBlock::from_async_reader(reader).await?;
73            frame_offset += 4 + block.length;
74            is_last = block.is_last;
75            blocks.push(block);
76        }
77        Ok(FlacHeader {
78            blocks,
79            path: path,
80            frame_offset,
81        })
82    }
83
84    pub fn from_file<P: AsRef<Path>>(filename: P) -> Result<FlacHeader> {
85        let mut file = File::open(filename.as_ref())?;
86        let header = Self::parse(&mut file, filename.as_ref().to_path_buf())?;
87        Ok(header)
88    }
89
90    pub fn stream_info(&self) -> &BlockStreamInfo {
91        let block = self.blocks.get(0).unwrap();
92        match &block.data {
93            MetadataBlockData::StreamInfo(i) => i,
94            _ => panic!("First block is not stream info!"),
95        }
96    }
97
98    fn block_of(&self, id: u8) -> Option<&MetadataBlock> {
99        self.blocks
100            .iter()
101            .find(|&block| u8::from(&block.data) == id)
102    }
103
104    fn block_of_mut(&mut self, id: u8) -> Option<&mut MetadataBlock> {
105        self.blocks
106            .iter_mut()
107            .find(|block| u8::from(&block.data) == id)
108    }
109
110    pub fn comments(&self) -> Option<&BlockVorbisComment> {
111        self.block_of(4).map(|b| match &b.data {
112            MetadataBlockData::Comment(c) => c,
113            _ => unreachable!(),
114        })
115    }
116
117    /// Get a mutable comments blocks for edit
118    ///
119    /// If VorbisComment block does not exist, a new block would be appended to header.
120    pub fn comments_mut(&mut self) -> &mut BlockVorbisComment {
121        let is_none = self.block_of_mut(4).is_none();
122        if is_none {
123            let comment = BlockVorbisComment {
124                vendor_string: format!("anni-flac v{}", env!("CARGO_PKG_VERSION")),
125                comments: vec![],
126            };
127            self.blocks
128                .push(MetadataBlock::new(MetadataBlockData::Comment(comment)));
129        }
130        self.block_of_mut(4)
131            .map(|b| match &mut b.data {
132                MetadataBlockData::Comment(c) => c,
133                _ => unreachable!(),
134            })
135            .unwrap()
136    }
137
138    fn frame_offset_now(&self) -> usize {
139        let mut frame_offset_now = 4;
140        for block in self.blocks.iter() {
141            frame_offset_now += 4 + block.data.len(); // block header + data
142        }
143        frame_offset_now
144    }
145
146    pub fn save<P: AsRef<Path>>(&mut self, output: Option<P>) -> Result<()> {
147        let input_path = self.path.to_path_buf();
148        let output_path = match output {
149            Some(p) => p.as_ref().to_path_buf(),
150            None => input_path.clone(),
151        };
152
153        self.format();
154        if input_path != output_path {
155            // save to another file
156            let mut file = File::create(output_path)?;
157
158            // write magic number
159            file.write_all(b"fLaC")?;
160            // write header blocks
161            for block in self.blocks.iter() {
162                block.write_to(&mut file)?;
163            }
164            // write frames
165            let mut file_input = File::open(input_path)?;
166            file_input.seek(SeekFrom::Start(self.frame_offset as u64))?;
167            std::io::copy(&mut file_input, &mut file)?;
168        } else {
169            // recalculate frame offset after header modify
170            let frame_offset_now = self.frame_offset_now();
171            log::debug!(
172                "frame_offset_now = {}, flac.frame_offset = {}",
173                frame_offset_now,
174                self.frame_offset
175            );
176
177            let need_new_file = frame_offset_now > self.frame_offset || {
178                // if header is smaller than / the same size as previous header
179                // means we do not need more space
180                // just need to write all data to the header
181                let space_to_add = self.frame_offset - frame_offset_now;
182
183                // try to get last block for padding
184                let last = self.blocks.last_mut().unwrap();
185                if let MetadataBlockData::Padding(size) = &mut last.data {
186                    // padding block exists, modify padding size directly
187                    last.length += space_to_add;
188                    log::debug!(
189                        "padding.size_original = {}, adjusted to {}",
190                        size,
191                        last.length
192                    );
193                    *size = last.length;
194                    false
195                } else if space_to_add >= 4 {
196                    // padding block does not exist, add a new padding block
197                    let space_to_add = space_to_add - 4;
198                    // make the last block not last
199                    self.blocks.last_mut().unwrap().is_last = false;
200                    // add padding block
201                    self.blocks.push(MetadataBlock {
202                        is_last: true,
203                        length: space_to_add,
204                        data: MetadataBlockData::Padding(space_to_add),
205                    });
206                    false
207                } else {
208                    // a new padding block needs at least 4 bytes
209                    // so if the space left is less than 4 bytes
210                    // padding block can not be created
211                    // we handle this situation as frame_offset_now > frame_offset_old
212                    true
213                }
214            };
215            if need_new_file {
216                // write to filename.anni
217                let output_new_path = output_path.with_extension("anni");
218                self.save(Some(output_new_path.as_path()))?;
219
220                let original_backup_path = output_path.with_extension("anni.bak");
221                // move original file to filename.anni.bak
222                std::fs::rename(&output_path, &original_backup_path)?;
223                // move new file to original file
224                std::fs::rename(output_new_path, output_path)?;
225                // remove backup of original file
226                std::fs::remove_file(original_backup_path)?;
227            } else {
228                // write back to input directly
229                // so we only need to write header blocks to override the original header
230                let mut file = OpenOptions::new().write(true).open(input_path)?;
231                // skip magic number b"fLaC"
232                file.seek(SeekFrom::Start(4))?;
233                // write header blocks
234                for block in self.blocks.iter() {
235                    block.write_to(&mut file)?;
236                }
237            }
238        }
239        Ok(())
240    }
241
242    // TODO: make this method private
243    pub fn format(&mut self) {
244        // recalculate frame offset after header modify
245        let frame_offset_now = self.frame_offset_now();
246
247        // merge padding blocks
248        let mut padding_size: Option<usize> = None;
249        self.blocks.retain(|block| match &block.data {
250            MetadataBlockData::Padding(size) => {
251                // update padding block size
252                padding_size = Some(padding_size.unwrap_or_default() + size);
253                // remove all padding blocks
254                false
255            }
256            // keep all other blocks
257            _ => true,
258        });
259
260        // insert padding block if necessary
261        if let Some(mut padding_block_size) = padding_size {
262            let need_padding = frame_offset_now != self.frame_offset
263                && if frame_offset_now > self.frame_offset {
264                    // need more space
265                    let needed = frame_offset_now - self.frame_offset;
266                    if needed <= padding_block_size {
267                        // have enough space
268                        padding_block_size -= frame_offset_now - self.frame_offset;
269                        true
270                    } else if needed == padding_block_size + 4 {
271                        // space needed == padding size + padding header size
272                        padding_block_size = 0;
273                        false
274                    } else {
275                        // padding space not enough
276                        padding_block_size = 8192;
277                        true
278                    }
279                } else {
280                    // expand padding space
281                    let expanded = self.frame_offset - frame_offset_now;
282                    padding_block_size += expanded;
283                    true
284                };
285
286            if need_padding {
287                self.blocks.push(MetadataBlock {
288                    is_last: true,
289                    length: padding_block_size,
290                    data: MetadataBlockData::Padding(padding_block_size),
291                })
292            }
293        }
294
295        self.fix_is_last()
296    }
297
298    fn fix_is_last(&mut self) {
299        let last = self.blocks.len() - 1;
300        for (index, block) in self.blocks.iter_mut().enumerate() {
301            block.is_last = index == last
302        }
303    }
304}
305
306pub struct MetadataBlock {
307    /// Whether the block is the last block in header.
308    ///
309    /// Must be fixed using `fix_is_last` before writing.
310    is_last: bool,
311    /// length of the block at **read time**
312    ///
313    /// Not trustable if any changes has been made
314    length: usize,
315    pub data: MetadataBlockData,
316}
317
318impl Decode for MetadataBlock {
319    fn from_reader<R: Read>(reader: &mut R) -> Result<Self> {
320        let first_byte = reader.read_u8()?;
321        let block_type = first_byte & 0b01111111;
322        let length = reader.read_u24::<BigEndian>()? as usize;
323        Ok(MetadataBlock {
324            is_last: first_byte & 0b10000000 > 0,
325            length,
326            data: match block_type {
327                0 => MetadataBlockData::StreamInfo(BlockStreamInfo::from_reader(
328                    &mut reader.take(length as u64),
329                )?),
330                1 => MetadataBlockData::Padding(skip(reader, length)? as usize),
331                2 => MetadataBlockData::Application(BlockApplication::from_reader(
332                    &mut reader.take(length as u64),
333                )?),
334                3 => MetadataBlockData::SeekTable(BlockSeekTable::from_reader(
335                    &mut reader.take(length as u64),
336                )?),
337                4 => MetadataBlockData::Comment(BlockVorbisComment::from_reader(
338                    &mut reader.take(length as u64),
339                )?),
340                5 => MetadataBlockData::CueSheet(BlockCueSheet::from_reader(
341                    &mut reader.take(length as u64),
342                )?),
343                6 => MetadataBlockData::Picture(BlockPicture::from_reader(
344                    &mut reader.take(length as u64),
345                )?),
346                _ => MetadataBlockData::Reserved((block_type, take(reader, length)?)),
347            },
348        })
349    }
350}
351
352#[cfg(feature = "async")]
353#[async_trait::async_trait]
354impl AsyncDecode for MetadataBlock {
355    async fn from_async_reader<R>(reader: &mut R) -> Result<Self>
356    where
357        R: AsyncRead + Unpin + Send,
358    {
359        let first_byte = reader.read_u8().await?;
360        let block_type = first_byte & 0b01111111;
361        let length = read_u24_async(reader).await? as usize;
362        Ok(MetadataBlock {
363            is_last: first_byte & 0b10000000 > 0,
364            length,
365            data: match block_type {
366                0 => MetadataBlockData::StreamInfo(
367                    BlockStreamInfo::from_async_reader(&mut reader.take(length as u64)).await?,
368                ),
369                1 => MetadataBlockData::Padding(skip_async(reader, length).await? as usize),
370                2 => MetadataBlockData::Application(
371                    BlockApplication::from_async_reader(&mut reader.take(length as u64)).await?,
372                ),
373                3 => MetadataBlockData::SeekTable(
374                    BlockSeekTable::from_async_reader(&mut reader.take(length as u64)).await?,
375                ),
376                4 => MetadataBlockData::Comment(
377                    BlockVorbisComment::from_async_reader(&mut reader.take(length as u64)).await?,
378                ),
379                5 => MetadataBlockData::CueSheet(
380                    BlockCueSheet::from_async_reader(&mut reader.take(length as u64)).await?,
381                ),
382                6 => MetadataBlockData::Picture(
383                    BlockPicture::from_async_reader(&mut reader.take(length as u64)).await?,
384                ),
385                _ => MetadataBlockData::Reserved((block_type, take_async(reader, length).await?)),
386            },
387        })
388    }
389}
390
391impl Encode for MetadataBlock {
392    fn write_to<W: Write>(&self, writer: &mut W) -> Result<()> {
393        writer.write_u8((if self.is_last { 0b10000000 } else { 0 }) + u8::from(&self.data))?;
394        writer.write_u24::<BigEndian>(self.data.len() as u32)?;
395        match &self.data {
396            MetadataBlockData::StreamInfo(s) => s.write_to(writer)?,
397            MetadataBlockData::Padding(p) => writer.write_all(&vec![0; *p])?, // FIXME: Why does writing zero needs to allocate memory?!
398            MetadataBlockData::Application(a) => a.write_to(writer)?,
399            MetadataBlockData::SeekTable(s) => s.write_to(writer)?,
400            MetadataBlockData::Comment(c) => c.write_to(writer)?,
401            MetadataBlockData::CueSheet(c) => c.write_to(writer)?,
402            MetadataBlockData::Picture(p) => p.write_to(writer)?,
403            MetadataBlockData::Reserved((_, data)) => writer.write_all(data)?,
404        }
405        Ok(())
406    }
407}
408
409impl MetadataBlock {
410    pub fn new(data: MetadataBlockData) -> Self {
411        MetadataBlock {
412            is_last: false,
413            length: 0,
414            data,
415        }
416    }
417
418    pub fn write(&self, dst: &mut dyn Write, i: usize) -> std::result::Result<(), std::io::Error> {
419        let data = &self.data;
420        writeln!(dst, "METADATA block #{}", i)?;
421        writeln!(dst, "  type: {} ({})", u8::from(data), data.as_str())?;
422        writeln!(dst, "  is last: {}", self.is_last)?;
423        writeln!(dst, "  length: {}", self.length)?;
424        writeln!(dst, "{:2?}", data)
425    }
426}
427
428pub enum MetadataBlockData {
429    StreamInfo(BlockStreamInfo),
430    Padding(usize),
431    Application(BlockApplication),
432    SeekTable(BlockSeekTable),
433    Comment(BlockVorbisComment),
434    CueSheet(BlockCueSheet),
435    Picture(BlockPicture),
436    Reserved((u8, Vec<u8>)),
437}
438
439impl From<&MetadataBlockData> for u8 {
440    fn from(data: &MetadataBlockData) -> Self {
441        match data {
442            MetadataBlockData::StreamInfo(_) => 0,
443            MetadataBlockData::Padding(_) => 1,
444            MetadataBlockData::Application(_) => 2,
445            MetadataBlockData::SeekTable(_) => 3,
446            MetadataBlockData::Comment(_) => 4,
447            MetadataBlockData::CueSheet(_) => 5,
448            MetadataBlockData::Picture(_) => 6,
449            MetadataBlockData::Reserved((t, _)) => *t,
450        }
451    }
452}
453
454impl MetadataBlockData {
455    pub fn as_str(&self) -> &'static str {
456        match self {
457            MetadataBlockData::StreamInfo(_) => "STREAMINFO",
458            MetadataBlockData::Padding(_) => "PADDING",
459            MetadataBlockData::Application(_) => "APPLICATION",
460            MetadataBlockData::SeekTable(_) => "SEEKTABLE",
461            MetadataBlockData::Comment(_) => "VORBIS_COMMENT",
462            MetadataBlockData::CueSheet(_) => "CUESHEET",
463            MetadataBlockData::Picture(_) => "PICTURE",
464            _ => "RESERVED",
465        }
466    }
467
468    #[allow(clippy::len_without_is_empty)]
469    pub fn len(&self) -> usize {
470        match self {
471            MetadataBlockData::StreamInfo(_) => 34,
472            MetadataBlockData::Padding(p) => *p,
473            MetadataBlockData::Application(a) => a.data.len() + 4,
474            MetadataBlockData::SeekTable(t) => t.seek_points.len() * 18,
475            MetadataBlockData::Comment(c) => {
476                8 + c.vendor_string.len() + c.comments.iter().map(|c| 4 + c.len()).sum::<usize>()
477            }
478            MetadataBlockData::CueSheet(c) => {
479                396 + c
480                    .tracks
481                    .iter()
482                    .map(|t| 36 + t.track_index.len() * 12)
483                    .sum::<usize>()
484            }
485            MetadataBlockData::Picture(p) => {
486                32 + p.mime_type.len() + p.description.len() + p.data.len()
487            }
488            MetadataBlockData::Reserved((_, arr)) => arr.len(),
489        }
490    }
491}
492
493impl fmt::Debug for MetadataBlockData {
494    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
495        let prefix = if let Some(prefix) = f.width() {
496            prefix
497        } else {
498            0
499        };
500        match self {
501            MetadataBlockData::Padding(_) => Ok(()),
502            MetadataBlockData::Reserved(_) => Ok(()),
503            MetadataBlockData::StreamInfo(s) => write!(f, "{:prefix$?}", s, prefix = prefix),
504            MetadataBlockData::Application(s) => write!(f, "{:prefix$?}", s, prefix = prefix),
505            MetadataBlockData::SeekTable(s) => write!(f, "{:prefix$?}", s, prefix = prefix),
506            MetadataBlockData::Comment(s) => write!(f, "{:prefix$?}", s, prefix = prefix),
507            MetadataBlockData::CueSheet(s) => write!(f, "{:prefix$?}", s, prefix = prefix),
508            MetadataBlockData::Picture(s) => write!(f, "{:prefix$?}", s, prefix = prefix),
509        }
510    }
511}