use std::io::{self, Read, Take};
use std::str;
use std::result;
use std::fmt;
use byteorder::{ReadBytesExt, LittleEndian};
use types::Result;
use utils::ReadExt;
#[derive(Copy, Clone, Eq, PartialEq, Debug)]
pub struct ChunkId(pub [u8; 4]);
impl ChunkId {
#[inline]
pub fn as_str(&self) -> Option<&str> {
str::from_utf8(&self.0).ok()
}
#[inline]
pub fn as_bytes(&self) -> &[u8] {
&self.0
}
}
impl fmt::Display for ChunkId {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
match self.as_str() {
Some(s) => f.write_str(s),
None => write!(f, "{:?}", self.as_bytes())
}
}
}
pub struct RiffReader<R: Read> {
source: R
}
impl<R: Read> RiffReader<R> {
pub fn new(source: R) -> RiffReader<R> {
RiffReader {
source: source
}
}
pub fn root(&mut self) -> Result<RiffListChunk> {
let (id, len) = match try!(read_id_and_len(&mut self.source)) {
Some(t) => t,
None => return Err(unexpected_eof!())
};
if id.as_bytes() != b"RIFF" {
return Err(invalid_format!("RIFF file header is invalid"));
}
Ok(try!(RiffChunk {
data: Counter {
delegate: (&mut self.source as &mut Read).take(len as u64),
counter: None
},
tainted: false,
chunk_id: id,
len: len
}.into_list_unchecked()))
}
}
struct Counter<'a, R> {
delegate: R,
counter: Option<&'a mut u32>
}
impl<'a, R: Read> Read for Counter<'a, R> {
fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
self.delegate.read(buf).map(|n| {
if let Some(ref mut counter) = self.counter {
**counter += n as u32;
}
n
})
}
}
pub struct RiffChunk<'a> {
chunk_id: ChunkId,
len: u32,
tainted: bool,
data: Counter<'a, Take<&'a mut Read>>
}
impl<'a> RiffChunk<'a> {
#[inline]
pub fn chunk_id(&self) -> ChunkId { self.chunk_id }
#[inline]
pub fn len(&self) -> u32 { self.len }
#[inline]
pub fn contents(&mut self) -> &mut Read {
self.tainted = true;
&mut self.data
}
#[inline]
pub fn can_have_subchunks(&self) -> bool {
!self.tainted && match &self.chunk_id.0 {
b"RIFF" | b"LIST" => true,
_ => false
}
}
#[inline]
pub fn into_list(self) -> result::Result<Result<RiffListChunk<'a>>, RiffChunk<'a>> {
if self.can_have_subchunks() {
Ok(self.into_list_unchecked())
} else {
Err(self)
}
}
fn into_list_unchecked(mut self) -> Result<RiffListChunk<'a>> {
let mut chunk_type = [0u8; 4];
if try!(self.data.read_exact_0(&mut chunk_type)) != 4 {
return Err(unexpected_eof!("when reading chunk type of chunk {}", self.chunk_id));
}
Ok(RiffListChunk {
chunk_id: self.chunk_id,
len: self.len,
chunk_type: ChunkId(chunk_type),
data: self.data,
cur_chunk_len: 0,
cur_chunk_read: 0
})
}
}
pub struct RiffListChunk<'a> {
chunk_id: ChunkId,
len: u32,
chunk_type: ChunkId,
data: Counter<'a, Take<&'a mut Read>>,
cur_chunk_len: u32,
cur_chunk_read: u32
}
impl<'a> RiffListChunk<'a> {
#[inline]
pub fn chunk_id(&self) -> ChunkId { self.chunk_id }
#[inline]
pub fn len(&self) -> u32 { self.len }
#[inline]
pub fn chunk_type(&self) -> ChunkId { self.chunk_type }
#[inline]
pub fn next(&mut self) -> Option<Result<RiffChunk>> {
if self.cur_chunk_read < self.cur_chunk_len {
let to_skip = (self.cur_chunk_len - self.cur_chunk_read) as u64;
match self.data.skip_exact_0(to_skip) {
Ok(n) if n == to_skip => {}
Ok(_) => return Some(Err(unexpected_eof!())),
Err(e) => return Some(Err(e.into()))
}
}
let (id, len) = match read_id_and_len(&mut self.data) {
Ok(Some(t)) => t,
Ok(None) => return None,
Err(e) => return Some(Err(e.into()))
};
self.cur_chunk_read = 0;
self.cur_chunk_len = len;
Some(Ok(RiffChunk {
chunk_id: id,
len: len,
tainted: false,
data: Counter {
delegate: (&mut self.data as &mut Read).take(len as u64),
counter: Some(&mut self.cur_chunk_read)
}
}))
}
}
fn read_id_and_len<R: Read>(source: &mut R) -> Result<Option<(ChunkId, u32)>> {
let mut id = [0u8; 4];
match try!(source.read_exact_0(&mut id)) {
0 => return Ok(None),
4 => {}
_ => return Err(unexpected_eof!())
}
let len = try!(source.read_u32::<LittleEndian>());
Ok(Some((ChunkId(id), len)))
}
#[cfg(test)]
mod tests {
use std::io::{Read, Write};
use byteorder::{WriteBytesExt, LittleEndian};
use utils::ReadExt;
use super::{RiffReader, RiffListChunk, ChunkId};
macro_rules! build {
($($arg:expr),+) => {{
let mut data = Vec::new();
$(data.write($arg).unwrap();)+
data
}}
}
fn n(n: u32) -> [u8; 4] {
let mut r = [0u8; 4];
(&mut r as &mut [u8]).write_u32::<LittleEndian>(n).unwrap();
r
}
fn check_next_chunk<'a>(cl: &mut RiffListChunk<'a>, id: ChunkId, len: u32, data: &[u8]) {
let chunk = cl.next();
let mut chunk = chunk.unwrap().unwrap();
assert_eq!(chunk.chunk_id(), id);
assert_eq!(chunk.len(), len);
let contents = chunk.contents().read_to_vec().unwrap();
assert_eq!(&*contents, data);
assert!(!chunk.can_have_subchunks());
}
#[test]
fn test_invalid_header() {
let mut data: &[u8] = b"XXXX\x04abcd";
let mut r = RiffReader::new(&mut data);
let root = r.root();
assert!(root.is_err());
}
#[test]
fn test_flat_chunks() {
let data = build! {
b"RIFF", &n(37), b"abcd",
b"A ", &n(4), b"1234",
b"B ", &n(5), b"56789"
};
let mut data: &[u8] = &data;
let mut r = RiffReader::new(&mut data);
let mut root = r.root().unwrap();
assert_eq!(root.chunk_id(), ChunkId(*b"RIFF"));
assert_eq!(root.len(), 37);
assert_eq!(root.chunk_type(), ChunkId(*b"abcd"));
check_next_chunk(&mut root, ChunkId(*b"A "), 4, b"1234");
check_next_chunk(&mut root, ChunkId(*b"B "), 5, b"56789");
assert!(root.next().is_none());
}
#[test]
fn test_nested_chunks() {
let data = build! {
b"RIFF", &n(77), b"abcd",
b"A ", &n(1), b"z",
b"LIST", &n(56), b"wxyz",
b" B ", &n(3), b"123",
b"LIST", &n(22), b"hi ",
b" C ", &n(0),
b" D", &n(2), b"op",
b"E ", &n(3), b"fuz"
};
let mut data: &[u8] = &data;
let mut r = RiffReader::new(&mut data);
let mut root = r.root().unwrap();
assert_eq!(root.chunk_id(), ChunkId(*b"RIFF"));
assert_eq!(root.len(), 77);
assert_eq!(root.chunk_type(), ChunkId(*b"abcd"));
check_next_chunk(&mut root, ChunkId(*b"A "), 1, b"z");
{
let chunk = root.next().unwrap().unwrap();
assert_eq!(chunk.chunk_id(), ChunkId(*b"LIST"));
assert_eq!(chunk.len(), 56);
assert!(chunk.can_have_subchunks());
let chunk = chunk.into_list();
assert!(chunk.is_ok());
let mut chunk = chunk.ok().unwrap().unwrap();
check_next_chunk(&mut chunk, ChunkId(*b" B "), 3, b"123");
{
let sublist = chunk.next().unwrap().unwrap();
assert_eq!(sublist.chunk_id(), ChunkId(*b"LIST"));
assert_eq!(sublist.len(), 22);
assert!(sublist.can_have_subchunks());
let sublist = sublist.into_list();
assert!(sublist.is_ok());
let mut sublist = sublist.ok().unwrap().unwrap();
check_next_chunk(&mut sublist, ChunkId(*b" C "), 0, b"");
check_next_chunk(&mut sublist, ChunkId(*b" D"), 2, b"op");
}
check_next_chunk(&mut chunk, ChunkId(*b"E "), 3, b"fuz");
assert!(chunk.next().is_none());
}
assert!(root.next().is_none());
}
#[test]
fn test_skip_chunk_data() {
let data = build! {
b"RIFF", &n(77), b"abcd",
b"A ", &n(10), b"abcdefghij",
b" B ", &n(12), b"123456789012",
b" C ", &n(8), b"ABCDEFGH"
};
let mut data: &[u8] = &data;
let mut r = RiffReader::new(&mut data);
let mut root = r.root().unwrap();
{
let mut chunk = root.next().unwrap().unwrap();
assert_eq!(chunk.chunk_id(), ChunkId(*b"A "));
assert_eq!(chunk.len(), 10);
assert_eq!(
(&mut chunk.contents() as &mut Read).take(5).read_to_vec().unwrap(),
b"abcde".to_owned()
);
}
{
let chunk = root.next().unwrap().unwrap();
assert_eq!(chunk.chunk_id(), ChunkId(*b" B "));
assert_eq!(chunk.len(), 12);
}
{
let mut chunk = root.next().unwrap().unwrap();
assert_eq!(chunk.chunk_id(), ChunkId(*b" C "));
assert_eq!(chunk.len(), 8);
assert_eq!(chunk.contents().read_to_vec().unwrap(), b"ABCDEFGH".to_owned());
}
assert!(root.next().is_none());
}
}