flv: Impl `Seek` for `FlvReader`

flv/src/reader.rs

@@ -1,3 +1,5 @@
+use std::io::{Error, ErrorKind, Result, Seek, SeekFrom};
+
 /// A reader that allows demuxing an FLV container.
 pub struct FlvReader<'a> {
     source: &'a [u8],
@@ -95,8 +97,57 @@ impl<'a> FlvReader<'a> {
             self.read(8)?.try_into().expect("eight bytes"),
         ))
     }
+}
 
-    pub fn position(&self) -> usize {
-        self.position
+impl<'a> Seek for FlvReader<'a> {
+    fn seek(&mut self, pos: SeekFrom) -> Result<u64> {
+        let newpos = match pos {
+            SeekFrom::Start(pos) => pos,
+            SeekFrom::Current(pos) => (self.position as i64 + pos)
+                .try_into()
+                .map_err(|e| Error::new(ErrorKind::InvalidInput, e))?,
+            SeekFrom::End(pos) => (self.source.len() as i64 - pos)
+                .try_into()
+                .map_err(|e| Error::new(ErrorKind::InvalidInput, e))?,
+        };
+
+        self.position = newpos as usize;
+
+        Ok(self.position as u64)
+    }
+}
+
+#[cfg(test)]
+#[allow(clippy::seek_from_current, clippy::seek_to_start_instead_of_rewind)]
+mod tests {
+    use crate::reader::FlvReader;
+    use std::io::{Seek, SeekFrom};
+
+    #[test]
+    fn valid_position_seek() {
+        let data = vec![0; 4000];
+        let mut reader = FlvReader::from_source(&data);
+
+        assert_eq!(reader.seek(SeekFrom::Current(0)).unwrap(), 0);
+        assert_eq!(reader.seek(SeekFrom::Current(4000)).unwrap(), 4000);
+        assert_eq!(reader.seek(SeekFrom::Current(-2000)).unwrap(), 2000);
+        assert_eq!(reader.seek(SeekFrom::Start(0)).unwrap(), 0);
+        assert_eq!(reader.seek(SeekFrom::Start(4000)).unwrap(), 4000);
+        assert_eq!(reader.seek(SeekFrom::End(0)).unwrap(), 4000);
+        assert_eq!(reader.seek(SeekFrom::End(4000)).unwrap(), 0);
+    }
+
+    #[test]
+    fn invalid_position_seek() {
+        let data = vec![];
+        let mut reader = FlvReader::from_parts(&data, 12000);
+
+        assert_eq!(reader.seek(SeekFrom::Current(0)).unwrap(), 12000);
+        assert_eq!(reader.seek(SeekFrom::Current(4000)).unwrap(), 16000);
+        assert_eq!(reader.seek(SeekFrom::Current(-2000)).unwrap(), 14000);
+        assert_eq!(reader.seek(SeekFrom::Start(0)).unwrap(), 0);
+        assert_eq!(reader.seek(SeekFrom::Start(4000)).unwrap(), 4000);
+        assert_eq!(reader.seek(SeekFrom::End(0)).unwrap(), 0);
+        assert!(reader.seek(SeekFrom::End(4000)).is_err());
     }
 }

flv/src/sound.rs

@@ -1,4 +1,5 @@
 use crate::FlvReader;
+use std::io::Seek;
 
 #[repr(u8)]
 #[derive(PartialEq, Eq, Debug)]
@@ -129,7 +130,7 @@ impl<'a> AudioData<'a> {
     ///
     /// If `None` is yielded, the data stream is not a valid audio header.
     pub fn parse(reader: &mut FlvReader<'a>, data_size: u32) -> Option<Self> {
-        let start = reader.position();
+        let start = reader.stream_position().expect("current position") as usize;
         let format_spec = reader.read_u8()?;
 
         let format = SoundFormat::try_from(format_spec & 0x0F).ok()?;
@@ -137,7 +138,7 @@ impl<'a> AudioData<'a> {
         let size = SoundSize::try_from((format_spec >> 6) & 0x01).ok()?;
         let sound_type = SoundType::try_from((format_spec >> 7) & 0x01).ok()?;
 
-        let header_size = reader.position() - start;
+        let header_size = reader.stream_position().expect("current position") as usize - start;
         if (data_size as usize) < header_size {
             return None;
         }

flv/src/video.rs

@@ -1,4 +1,5 @@
 use crate::reader::FlvReader;
+use std::io::Seek;
 
 #[repr(u8)]
 #[derive(PartialEq, Eq, Debug, Copy, Clone)]
@@ -102,13 +103,13 @@ impl<'a> VideoData<'a> {
     ///
     /// If `None` is yielded, the data stream is not a valid video header.
     pub fn parse(reader: &mut FlvReader<'a>, data_size: u32) -> Option<Self> {
-        let start = reader.position();
+        let start = reader.stream_position().expect("current position") as usize;
         let format_spec = reader.read_u8()?;
 
         let frame_type = FrameType::try_from(format_spec & 0x0F).ok()?;
         let codec_id = CodecId::try_from(format_spec >> 4).ok()?;
 
-        let header_size = reader.position() - start;
+        let header_size = reader.stream_position().expect("current position") as usize - start;
         if (data_size as usize) < header_size {
             return None;
         }