use std::io::{self, Read, Write};
pub const FILE_MAGIC: [u8; 4] = *b"DKIF";
pub const FILE_HEADER_BYTES: usize = 32;
pub const FRAME_HEADER_BYTES: usize = 12;
pub const AV1_FOURCC: [u8; 4] = *b"AV01";
pub fn write_file_header<W: Write>(
writer: &mut W,
width: u16,
height: u16,
fps_num: u32,
fps_den: u32,
frame_count: u32,
) -> io::Result<()> {
let mut buf = [0u8; FILE_HEADER_BYTES];
buf[0..4].copy_from_slice(&FILE_MAGIC);
buf[6..8].copy_from_slice(&(FILE_HEADER_BYTES as u16).to_le_bytes());
buf[8..12].copy_from_slice(&AV1_FOURCC);
buf[12..14].copy_from_slice(&width.to_le_bytes());
buf[14..16].copy_from_slice(&height.to_le_bytes());
buf[16..20].copy_from_slice(&fps_num.to_le_bytes());
buf[20..24].copy_from_slice(&fps_den.to_le_bytes());
buf[24..28].copy_from_slice(&frame_count.to_le_bytes());
writer.write_all(&buf)
}
pub fn write_frame<W: Write>(writer: &mut W, frame_bytes: &[u8], pts: u64) -> io::Result<()> {
let mut header = [0u8; FRAME_HEADER_BYTES];
let size = u32::try_from(frame_bytes.len())
.ok()
.ok_or_else(|| io::Error::new(io::ErrorKind::InvalidInput, "IVF frame > 4 GiB"))?;
header[0..4].copy_from_slice(&size.to_le_bytes());
header[4..12].copy_from_slice(&pts.to_le_bytes());
writer.write_all(&header)?;
writer.write_all(frame_bytes)
}
pub struct IvfReader<R> {
reader: R,
}
impl<R: Read> IvfReader<R> {
pub fn open(mut reader: R) -> io::Result<Self> {
let mut buf = [0u8; FILE_HEADER_BYTES];
reader.read_exact(&mut buf)?;
if buf[0..4] != FILE_MAGIC {
return Err(io::Error::new(
io::ErrorKind::InvalidData,
format!("IVF: missing DKIF magic (got {:02x?})", &buf[0..4],),
));
}
Ok(Self { reader })
}
pub fn read_frame(&mut self) -> io::Result<Option<Vec<u8>>> {
let mut header = [0u8; FRAME_HEADER_BYTES];
if self.reader.read(&mut header[..1])? == 0 {
return Ok(None);
}
self.reader.read_exact(&mut header[1..])?;
let size = u32::from_le_bytes([header[0], header[1], header[2], header[3]]) as usize;
let mut frame = vec![0u8; size];
self.reader.read_exact(&mut frame)?;
Ok(Some(frame))
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::io::Cursor;
#[test]
fn round_trip_two_frames() {
let mut buf = Vec::new();
buf.extend_from_slice(&FILE_MAGIC); buf.extend_from_slice(&0u16.to_le_bytes()); buf.extend_from_slice(&32u16.to_le_bytes()); buf.extend_from_slice(b"AV01");
buf.extend_from_slice(&64u16.to_le_bytes()); buf.extend_from_slice(&64u16.to_le_bytes()); buf.extend_from_slice(&1u32.to_le_bytes()); buf.extend_from_slice(&1u32.to_le_bytes()); buf.extend_from_slice(&2u32.to_le_bytes()); buf.extend_from_slice(&0u32.to_le_bytes()); for (i, body) in [b"FRAME-ONE".as_slice(), b"BODY-2"].iter().enumerate() {
buf.extend_from_slice(&(body.len() as u32).to_le_bytes());
buf.extend_from_slice(&(i as u64).to_le_bytes());
buf.extend_from_slice(body);
}
let mut r = IvfReader::open(Cursor::new(buf)).expect("open");
assert_eq!(
r.read_frame().unwrap().as_deref(),
Some(b"FRAME-ONE".as_slice())
);
assert_eq!(
r.read_frame().unwrap().as_deref(),
Some(b"BODY-2".as_slice())
);
assert!(
r.read_frame().unwrap().is_none(),
"clean EOF after last frame"
);
}
#[test]
fn rejects_bad_magic() {
let mut buf = vec![0u8; FILE_HEADER_BYTES];
buf[0..4].copy_from_slice(b"NOPE");
let err = IvfReader::open(Cursor::new(buf)).err().expect("bad magic");
assert_eq!(err.kind(), io::ErrorKind::InvalidData);
}
}