immeta/common/
riff.rs

1use std::io::{self, Read, Take};
2use std::str;
3use std::result;
4use std::fmt;
5
6use byteorder::{ReadBytesExt, LittleEndian};
7
8use types::Result;
9use utils::ReadExt;
10
11#[derive(Copy, Clone, Eq, PartialEq, Debug)]
12pub struct ChunkId(pub [u8; 4]);
13
14impl ChunkId {
15    // TODO: add a new() method and hide public field after const fns become stable
16
17    #[inline]
18    pub fn as_str(&self) -> Option<&str> {
19        str::from_utf8(&self.0).ok()
20    }
21
22    #[inline]
23    pub fn as_bytes(&self) -> &[u8] {
24        &self.0
25    }
26}
27
28impl fmt::Display for ChunkId {
29    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
30        match self.as_str() {
31            Some(s) => f.write_str(s),
32            None => write!(f, "{:?}", self.as_bytes())
33        }
34    }
35}
36
37pub struct RiffReader<R: Read> {
38    source: R
39}
40
41impl<R: Read> RiffReader<R> {
42    pub fn new(source: R) -> RiffReader<R> {
43        RiffReader {
44            source: source
45        }
46    }
47
48    pub fn root(&mut self) -> Result<RiffListChunk> {
49        let (id, len) = match try!(read_id_and_len(&mut self.source)) {
50            Some(t) => t,
51            None => return Err(unexpected_eof!())
52        };
53
54        if id.as_bytes() != b"RIFF" {
55            return Err(invalid_format!("RIFF file header is invalid"));
56        }
57
58        Ok(try!(RiffChunk {
59            data: Counter {
60                delegate: (&mut self.source as &mut Read).take(len as u64),
61                counter: None
62            },
63            tainted: false,
64            chunk_id: id,
65            len: len
66        }.into_list_unchecked()))
67    }
68}
69
70struct Counter<'a, R> {
71    delegate: R,
72    counter: Option<&'a mut u32>
73}
74
75impl<'a, R: Read> Read for Counter<'a, R> {
76    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
77        self.delegate.read(buf).map(|n| {
78            if let Some(ref mut counter) = self.counter {
79                **counter += n as u32;
80            }
81            n
82        })
83    }
84}
85
86pub struct RiffChunk<'a> {
87    chunk_id: ChunkId,
88    len: u32,
89    tainted: bool,
90    data: Counter<'a, Take<&'a mut Read>>
91}
92
93impl<'a> RiffChunk<'a> {
94    #[inline]
95    pub fn chunk_id(&self) -> ChunkId { self.chunk_id }
96
97    #[inline]
98    pub fn len(&self) -> u32 { self.len }
99
100    #[inline]
101    pub fn contents(&mut self) -> &mut Read { 
102        self.tainted = true;
103        &mut self.data
104    }
105
106    #[inline]
107    pub fn can_have_subchunks(&self) -> bool {
108        !self.tainted && match &self.chunk_id.0 {
109            b"RIFF" | b"LIST" => true,
110            _ => false
111        }
112    }
113
114    #[inline]
115    pub fn into_list(self) -> result::Result<Result<RiffListChunk<'a>>, RiffChunk<'a>> {
116        if self.can_have_subchunks() {
117            Ok(self.into_list_unchecked())
118        } else {
119            Err(self)
120        }
121    }
122
123    fn into_list_unchecked(mut self) -> Result<RiffListChunk<'a>> {
124        let mut chunk_type = [0u8; 4];
125
126        if try!(self.data.read_exact_0(&mut chunk_type)) != 4 {
127            return Err(unexpected_eof!("when reading chunk type of chunk {}", self.chunk_id));
128        }
129
130        Ok(RiffListChunk {
131            chunk_id: self.chunk_id,
132            len: self.len,
133            chunk_type: ChunkId(chunk_type),
134            data: self.data,
135            cur_chunk_len: 0,
136            cur_chunk_read: 0
137        })
138    }
139}
140
141pub struct RiffListChunk<'a> {
142    chunk_id: ChunkId,
143    len: u32,
144    chunk_type: ChunkId,
145    data: Counter<'a, Take<&'a mut Read>>,
146    cur_chunk_len: u32,
147    cur_chunk_read: u32
148}
149
150impl<'a> RiffListChunk<'a> {
151    #[inline]
152    pub fn chunk_id(&self) -> ChunkId { self.chunk_id }
153
154    #[inline]
155    pub fn len(&self) -> u32 { self.len }
156
157    #[inline]
158    pub fn chunk_type(&self) -> ChunkId { self.chunk_type }
159
160    #[inline]
161    pub fn next(&mut self) -> Option<Result<RiffChunk>> {
162        if self.cur_chunk_read < self.cur_chunk_len {
163            let to_skip = (self.cur_chunk_len - self.cur_chunk_read) as u64;
164            match self.data.skip_exact_0(to_skip) {
165                Ok(n) if n == to_skip => {}
166                Ok(_) => return Some(Err(unexpected_eof!())),
167                Err(e) => return Some(Err(e.into()))
168            }
169        }
170
171        let (id, len) = match read_id_and_len(&mut self.data) {
172            Ok(Some(t)) => t,
173            Ok(None) => return None,
174            Err(e) => return Some(Err(e.into()))
175        };
176
177        self.cur_chunk_read = 0;
178        self.cur_chunk_len = len;
179
180        Some(Ok(RiffChunk {
181            chunk_id: id,
182            len: len,
183            tainted: false,
184            data: Counter {
185                delegate: (&mut self.data as &mut Read).take(len as u64),
186                counter: Some(&mut self.cur_chunk_read)
187            }
188        }))
189    }
190}
191
192fn read_id_and_len<R: Read>(source: &mut R) -> Result<Option<(ChunkId, u32)>> {
193    let mut id = [0u8; 4];
194
195    match try!(source.read_exact_0(&mut id)) {
196        0 => return Ok(None),
197        4 => {}
198        _ => return Err(unexpected_eof!())
199    }
200
201    let len = try!(source.read_u32::<LittleEndian>());
202
203    Ok(Some((ChunkId(id), len)))
204}
205
206#[cfg(test)]
207mod tests {
208    use std::io::{Read, Write};
209
210    use byteorder::{WriteBytesExt, LittleEndian};
211
212    use utils::ReadExt;
213
214    use super::{RiffReader, RiffListChunk, ChunkId};
215
216    macro_rules! build {
217        ($($arg:expr),+) => {{
218            let mut data = Vec::new();
219            $(data.write($arg).unwrap();)+
220            data
221        }}
222    }
223
224    fn n(n: u32) -> [u8; 4] {
225        let mut r = [0u8; 4];
226        (&mut r as &mut [u8]).write_u32::<LittleEndian>(n).unwrap();
227        r
228    }
229
230    fn check_next_chunk<'a>(cl: &mut RiffListChunk<'a>, id: ChunkId, len: u32, data: &[u8]) {
231        let chunk = cl.next();
232        let mut chunk = chunk.unwrap().unwrap();
233
234        assert_eq!(chunk.chunk_id(), id);
235        assert_eq!(chunk.len(), len);
236        let contents = chunk.contents().read_to_vec().unwrap();
237        assert_eq!(&*contents, data);
238        assert!(!chunk.can_have_subchunks());
239    }
240
241    #[test]
242    fn test_invalid_header() {
243        let mut data: &[u8] = b"XXXX\x04abcd";
244
245        let mut r = RiffReader::new(&mut data);
246
247        let root = r.root();
248        assert!(root.is_err());
249    }
250
251    #[test]
252    fn test_flat_chunks() {
253        let data = build! {
254            b"RIFF", &n(37), b"abcd",
255            b"A   ", &n(4), b"1234",
256            b"B   ", &n(5), b"56789"
257        };
258        let mut data: &[u8] = &data;
259
260        let mut r = RiffReader::new(&mut data);
261
262        let mut root = r.root().unwrap();
263
264        assert_eq!(root.chunk_id(), ChunkId(*b"RIFF"));
265        assert_eq!(root.len(), 37);
266        assert_eq!(root.chunk_type(), ChunkId(*b"abcd"));
267
268        check_next_chunk(&mut root, ChunkId(*b"A   "), 4, b"1234");
269        check_next_chunk(&mut root, ChunkId(*b"B   "), 5, b"56789");
270
271        assert!(root.next().is_none());
272    }
273
274    #[test]
275    fn test_nested_chunks() {
276        let data = build! {
277            b"RIFF", &n(77), b"abcd",
278            b"A   ", &n(1), b"z",
279            b"LIST", &n(56), b"wxyz",
280                b" B  ", &n(3), b"123",
281                b"LIST", &n(22), b"hi  ",
282                    b"  C ", &n(0),
283                    b"   D", &n(2), b"op",
284                b"E   ", &n(3), b"fuz"
285        };
286        let mut data: &[u8] = &data;
287
288        let mut r = RiffReader::new(&mut data);
289
290        let mut root = r.root().unwrap();
291
292        assert_eq!(root.chunk_id(), ChunkId(*b"RIFF"));
293        assert_eq!(root.len(), 77);
294        assert_eq!(root.chunk_type(), ChunkId(*b"abcd"));
295
296        check_next_chunk(&mut root, ChunkId(*b"A   "), 1, b"z");
297
298        {
299            let chunk = root.next().unwrap().unwrap();
300            assert_eq!(chunk.chunk_id(), ChunkId(*b"LIST"));
301            assert_eq!(chunk.len(), 56);
302            assert!(chunk.can_have_subchunks());
303
304            let chunk = chunk.into_list();
305            assert!(chunk.is_ok());
306            let mut chunk = chunk.ok().unwrap().unwrap();
307
308            check_next_chunk(&mut chunk, ChunkId(*b" B  "), 3, b"123");
309
310            {
311                let sublist = chunk.next().unwrap().unwrap();
312                assert_eq!(sublist.chunk_id(), ChunkId(*b"LIST"));
313                assert_eq!(sublist.len(), 22);
314                assert!(sublist.can_have_subchunks());
315
316                let sublist = sublist.into_list();
317                assert!(sublist.is_ok());
318                let mut sublist = sublist.ok().unwrap().unwrap();
319
320                check_next_chunk(&mut sublist, ChunkId(*b"  C "), 0, b"");
321                check_next_chunk(&mut sublist, ChunkId(*b"   D"), 2, b"op");
322            }
323
324            check_next_chunk(&mut chunk, ChunkId(*b"E   "), 3, b"fuz");
325
326            assert!(chunk.next().is_none());
327        }
328
329        assert!(root.next().is_none());
330    }
331
332    #[test]
333    fn test_skip_chunk_data() {
334        let data = build! {
335            b"RIFF", &n(77), b"abcd",
336            b"A   ", &n(10), b"abcdefghij",
337            b" B  ", &n(12), b"123456789012",
338            b"  C ", &n(8),  b"ABCDEFGH"
339        };
340        let mut data: &[u8] = &data;
341
342        let mut r = RiffReader::new(&mut data);
343
344        let mut root = r.root().unwrap();
345
346        {
347            let mut chunk = root.next().unwrap().unwrap();
348            assert_eq!(chunk.chunk_id(), ChunkId(*b"A   "));
349            assert_eq!(chunk.len(), 10);
350            assert_eq!(
351                (&mut chunk.contents() as &mut Read).take(5).read_to_vec().unwrap(),
352                b"abcde".to_owned()
353            );
354        }
355
356        {
357            let chunk = root.next().unwrap().unwrap();
358            assert_eq!(chunk.chunk_id(), ChunkId(*b" B  "));
359            assert_eq!(chunk.len(), 12);
360        }
361
362        {
363            let mut chunk = root.next().unwrap().unwrap();
364            assert_eq!(chunk.chunk_id(), ChunkId(*b"  C "));
365            assert_eq!(chunk.len(), 8);
366            assert_eq!(chunk.contents().read_to_vec().unwrap(), b"ABCDEFGH".to_owned());
367        }
368
369        assert!(root.next().is_none());
370    }
371}
immeta/common/riff.rs

immeta/common/
riff.rs
