ogg_pager: Restrict page sizes to 8KB; Fix segment table generation when paginating

Author: Serial-ATA

ogg_pager/src/paginate.rs

@@ -1,7 +1,8 @@
 use crate::error::Result;
 use crate::{
 	segment_table, Page, PageHeader, CONTAINS_FIRST_PAGE_OF_BITSTREAM,
-	CONTAINS_LAST_PAGE_OF_BITSTREAM, CONTINUED_PACKET, MAX_CONTENT_SIZE,
+	CONTAINS_LAST_PAGE_OF_BITSTREAM, CONTINUED_PACKET, MAX_WRITTEN_CONTENT_SIZE,
+	MAX_WRITTEN_SEGMENT_COUNT,
 };
 
 use std::io::Read;
@@ -33,13 +34,13 @@ impl PaginateContext {
 			},
 			pos: 0,
 			idx: 0,
-			remaining_page_size: MAX_CONTENT_SIZE,
+			remaining_page_size: MAX_WRITTEN_CONTENT_SIZE,
 			current_packet_len: 0,
 		}
 	}
 
-	fn flush(&mut self, content: &mut Vec<u8>, segments: &mut Vec<u8>) {
-		let header = PageHeader {
+	fn flush(&mut self, content: &mut Vec<u8>, segment_table: &mut Vec<u8>) {
+		let mut header = PageHeader {
 			start: self.pos,
 			header_type_flag: {
 				match self.flags.first_page {
@@ -60,31 +61,44 @@ impl PaginateContext {
 			},
 			stream_serial: self.stream_serial,
 			sequence_number: self.idx as u32,
+			segments: Vec::new(),
 			// No need to calculate this yet
 			checksum: 0,
 		};
 
 		let content = core::mem::take(content);
-		let segments = core::mem::take(segments);
-
 		let content_len = content.len();
 		self.pos += content_len as u64;
 
+		// Moving on to a new packet
+		if self.pos > self.current_packet_len as u64 {
+			self.flags.packet_spans_multiple_pages = false;
+		}
+
+		// We need to determine how many segments our page content takes up.
+		// If it takes up the remainder of the segment table for the entire packet,
+		// we'll just consume it as is.
+		let segments_occupied = if content_len >= 255 {
+			content_len / 255
+		} else {
+			1
+		};
+
+		if self.flags.packet_spans_multiple_pages {
+			header.segments = segment_table.drain(..segments_occupied).collect();
+		} else {
+			header.segments = core::mem::take(segment_table);
+		}
+
 		self.pages.push(Page {
 			content,
-			segments,
 			header,
 			end: self.pos,
 		});
 
 		self.idx += 1;
 		self.flags.packet_finished_on_page = false;
-		self.remaining_page_size = MAX_CONTENT_SIZE;
-
-		// Moving on to a new packet
-		if self.pos > self.current_packet_len as u64 {
-			self.flags.packet_spans_multiple_pages = false;
-		}
+		self.remaining_page_size = MAX_WRITTEN_CONTENT_SIZE;
 	}
 }
 
@@ -117,8 +131,6 @@ pub fn paginate<'a, I: 'a>(
 where
 	I: IntoIterator<Item = &'a [u8]>,
 {
-	const MAX_SEGMENT_COUNT: usize = 255;
-
 	let mut ctx = PaginateContext::new(abgp, stream_serial, flags);
 
 	let mut packets_iter = packets.into_iter();
@@ -129,14 +141,15 @@ where
 	};
 	ctx.current_packet_len = packet.len();
 
-	let mut segments = Vec::with_capacity(255);
+	let mut segments = segment_table(packet.len());
 	let mut page_content = Vec::new();
 
 	loop {
 		if !ctx.flags.packet_spans_multiple_pages && !ctx.flags.first_page {
 			match packets_iter.next() {
 				Some(packet_) => {
 					packet = packet_;
+					segments.append(&mut segment_table(packet.len()));
 					ctx.current_packet_len = packet.len();
 					ctx.flags.fresh_packet = true;
 				},
@@ -150,30 +163,30 @@ where
 			.take(ctx.remaining_page_size as u64)
 			.read_to_end(&mut page_content)?;
 		ctx.remaining_page_size -= bytes_read;
-		packet = &packet[bytes_read..];
-
-		segments.append(&mut segment_table(bytes_read)?);
-		let remaining_segments = MAX_SEGMENT_COUNT - segments.len();
-
-		// We have a maximum of 255 segments available per page, if we require more than
-		// is left in the segment table, we'll have to split the packet into multiple pages.
-		let segments_required = (packet.len() / 255) + 1;
-		ctx.flags.packet_spans_multiple_pages = segments_required > remaining_segments;
 
+		packet = &packet[bytes_read..];
 		// We need to indicate whether or not any packet was finished on this page.
 		// This is used for the absolute granule position.
 		if packet.is_empty() {
 			ctx.flags.packet_finished_on_page = true;
 		}
 
 		// The first packet of the bitstream must have its own page, unlike any other packet.
-		let first_packet_finished_on_page = ctx.flags.first_page
-			&& ctx.header_flags & CONTAINS_FIRST_PAGE_OF_BITSTREAM != 0
-			&& ctx.flags.packet_finished_on_page;
+		let first_page_of_bitstream = ctx.header_flags & CONTAINS_FIRST_PAGE_OF_BITSTREAM != 0;
+		let first_packet_finished_on_page =
+			ctx.flags.first_page && first_page_of_bitstream && ctx.flags.packet_finished_on_page;
+
+		// We have a maximum of `MAX_WRITTEN_SEGMENT_COUNT` segments available per page, if we require more than
+		// is left in the segment table, we'll have to split the packet into multiple pages.
+		let segments_required = (packet.len() / MAX_WRITTEN_SEGMENT_COUNT) + 1;
+		let remaining_segments = MAX_WRITTEN_SEGMENT_COUNT.saturating_sub(segments.len());
+		ctx.flags.packet_spans_multiple_pages = segments_required > remaining_segments;
 
 		if first_packet_finished_on_page
 			// We've completely filled this page, we need to flush before moving on
 			|| (ctx.remaining_page_size == 0 || remaining_segments == 0)
+			// We've read all this packet has to offer
+			|| packet.is_empty()
 		{
 			ctx.flush(&mut page_content, &mut segments);
 		}