feature(media/ogg_writer): add packet segments to page (#499)

* feature(media/ogg_writer): add packet segments to page

* fix(media/ogg_writer): handle payload sizes that are multiples of 255

Author: yinshanyang

media/src/io/ogg_writer/mod.rs

@@ -127,8 +127,10 @@ impl<W: Write + Seek> OggWriter<W> {
     ) -> Result<()> {
         self.last_payload_size = payload.len();
         self.last_payload = payload.clone();
+        let n_segments = (self.last_payload_size + 255 - 1) / 255;
 
-        let mut page = Vec::with_capacity(PAGE_HEADER_SIZE + 1 + self.last_payload_size);
+        let mut page =
+            Vec::with_capacity(PAGE_HEADER_SIZE + 1 + self.last_payload_size + n_segments);
         {
             let mut header_writer = BufWriter::new(&mut page);
             header_writer.write_all(PAGE_HEADER_SIGNATURE)?; // page headers starts with 'OggS'//0-3
@@ -138,8 +140,16 @@ impl<W: Write + Seek> OggWriter<W> {
             header_writer.write_u32::<LittleEndian>(self.serial)?; // Bitstream serial number//14-17
             header_writer.write_u32::<LittleEndian>(page_index)?; // Page sequence number//18-21
             header_writer.write_u32::<LittleEndian>(0)?; //Checksum reserve //22-25
-            header_writer.write_u8(1)?; // Number of segments in page, giving always 1 segment //26
-            header_writer.write_u8(self.last_payload_size as u8)?; // Segment Table inserting at 27th position since page header length is 27
+            header_writer.write_u8(n_segments as u8)?; // Number of segments in page //26
+
+            // Filling the segment table with the lacing values.
+            // First (n_segments - 1) values will always be 255.
+            for _ in 0..n_segments - 1 {
+                header_writer.write_u8(255)?;
+            }
+            // The last value will be the remainder.
+            header_writer.write_u8((self.last_payload_size - (n_segments * 255 - 255)) as u8)?;
+
             header_writer.write_all(payload)?; // inserting at 28th since Segment Table(1) + header length(27)
         }
 

media/src/io/ogg_writer/ogg_writer_test.rs

@@ -63,3 +63,167 @@ fn test_ogg_writer_add_packet_and_close() -> Result<()> {
 
     Ok(())
 }
+
+#[test]
+fn test_ogg_writer_add_packet() -> Result<()> {
+    let raw_pkt = Bytes::from_iter(std::iter::repeat(0x45).take(235));
+
+    let mut valid_packet = rtp::packet::Packet {
+        header: rtp::header::Header {
+            marker: true,
+            extension: true,
+            extension_profile: 1,
+            version: 2,
+            payload_type: 111,
+            sequence_number: 27023,
+            timestamp: 3653407706,
+            ssrc: 476325762,
+            csrc: vec![],
+            padding: false,
+            extensions: vec![],
+            extensions_padding: 0,
+        },
+        payload: raw_pkt,
+    };
+    valid_packet
+        .header
+        .set_extension(0, Bytes::from_static(&[0xFF, 0xFF, 0xFF, 0xFF]))?;
+
+    let buffer = Cursor::new(Vec::<u8>::new());
+    let mut writer = OggWriter::new(buffer, 48000, 2)?;
+    let result = writer.write_rtp(&valid_packet);
+
+    assert!(
+        result.is_ok(),
+        "OggWriter should be able to write an Opus packet smaller than 255 bytes"
+    );
+    assert!(
+        writer.writer.into_inner()[126..128] == [1, 235],
+        "OggWriter should be able to write an Opus packet smaller than 255 bytes"
+    );
+
+    Ok(())
+}
+
+#[test]
+fn test_ogg_writer_add_packet_of_255() -> Result<()> {
+    let raw_pkt = Bytes::from_iter(std::iter::repeat(0x45).take(255));
+
+    let mut valid_packet = rtp::packet::Packet {
+        header: rtp::header::Header {
+            marker: true,
+            extension: true,
+            extension_profile: 1,
+            version: 2,
+            payload_type: 111,
+            sequence_number: 27023,
+            timestamp: 3653407706,
+            ssrc: 476325762,
+            csrc: vec![],
+            padding: false,
+            extensions: vec![],
+            extensions_padding: 0,
+        },
+        payload: raw_pkt,
+    };
+    valid_packet
+        .header
+        .set_extension(0, Bytes::from_static(&[0xFF, 0xFF, 0xFF, 0xFF]))?;
+
+    let buffer = Cursor::new(Vec::<u8>::new());
+    let mut writer = OggWriter::new(buffer, 48000, 2)?;
+    let result = writer.write_rtp(&valid_packet);
+
+    assert!(
+        result.is_ok(),
+        "OggWriter should be able to write an Opus packet of exactly 255"
+    );
+    assert!(
+        writer.writer.into_inner()[126..128] == [1, 255],
+        "OggWriter should be able to write an Opus packet of exactly 255"
+    );
+
+    Ok(())
+}
+
+#[test]
+fn test_ogg_writer_add_large_packet() -> Result<()> {
+    let raw_pkt = Bytes::from_iter(std::iter::repeat(0x45).take(1000));
+
+    let mut valid_packet = rtp::packet::Packet {
+        header: rtp::header::Header {
+            marker: true,
+            extension: true,
+            extension_profile: 1,
+            version: 2,
+            payload_type: 111,
+            sequence_number: 27023,
+            timestamp: 3653407706,
+            ssrc: 476325762,
+            csrc: vec![],
+            padding: false,
+            extensions: vec![],
+            extensions_padding: 0,
+        },
+        payload: raw_pkt,
+    };
+    valid_packet
+        .header
+        .set_extension(0, Bytes::from_static(&[0xFF, 0xFF, 0xFF, 0xFF]))?;
+
+    let buffer = Cursor::new(Vec::<u8>::new());
+    let mut writer = OggWriter::new(buffer, 48000, 2)?;
+    let result = writer.write_rtp(&valid_packet);
+
+    assert!(
+        result.is_ok(),
+        "OggWriter should be able to write a large (> 255 bytes) Opus packet"
+    );
+    assert!(
+        writer.writer.into_inner()[126..131] == [4, 255, 255, 255, 235],
+        "OggWriter should be able to write multiple segments per page, for 1000 bytes, 4 segments of 255, 255, 255 and 235 long"
+    );
+
+    Ok(())
+}
+
+#[test]
+fn test_ogg_writer_add_large_packet_with_multiple_of_255() -> Result<()> {
+    let raw_pkt = Bytes::from_iter(std::iter::repeat(0x45).take(255 * 4));
+
+    let mut valid_packet = rtp::packet::Packet {
+        header: rtp::header::Header {
+            marker: true,
+            extension: true,
+            extension_profile: 1,
+            version: 2,
+            payload_type: 111,
+            sequence_number: 27023,
+            timestamp: 3653407706,
+            ssrc: 476325762,
+            csrc: vec![],
+            padding: false,
+            extensions: vec![],
+            extensions_padding: 0,
+        },
+        payload: raw_pkt,
+    };
+    valid_packet
+        .header
+        .set_extension(0, Bytes::from_static(&[0xFF, 0xFF, 0xFF, 0xFF]))?;
+
+    let buffer = Cursor::new(Vec::<u8>::new());
+    let mut writer = OggWriter::new(buffer, 48000, 2)?;
+    let result = writer.write_rtp(&valid_packet);
+
+    assert!(
+        result.is_ok(),
+        "OggWriter should be able to write a large (> 255 bytes) Opus packet"
+    );
+    assert!(
+        writer.writer.into_inner()[126..131] == [4, 255, 255, 255, 255],
+        "OggWriter should be able to write multiple segments per page, for 1020 bytes, 4 segments of 255 each"
+    );
+
+    Ok(())
+}