ogg/
writing.rs

1// Ogg decoder and encoder written in Rust
2//
3// Copyright (c) 2016-2017 est31 <MTest31@outlook.com>
4// and contributors. All rights reserved.
5// Redistribution or use only under the terms
6// specified in the LICENSE file attached to this
7// source distribution.
8
9/*!
10Writing logic
11*/
12
13use std::result;
14use std::io::{self, Cursor, Write, Seek, SeekFrom};
15use byteorder::{WriteBytesExt, LittleEndian};
16use std::collections::HashMap;
17use crc::vorbis_crc32_update;
18
19
20/// Ogg version of the `std::io::Result` type.
21///
22/// We need `std::result::Result` at other points
23/// too, so we can't use `Result` as the name.
24type IoResult<T> = result::Result<T, io::Error>;
25
26/**
27Writer for packets into an Ogg stream.
28
29Note that the functionality of this struct isn't as well tested as for
30the `PacketReader` struct.
31*/
32pub struct PacketWriter<T :io::Write> {
33	wtr :T,
34
35	page_vals :HashMap<u32, CurrentPageValues>,
36}
37
38struct CurrentPageValues {
39	/// `true` if this page is the first one in the logical bitstream
40	first_page :bool,
41	/// Page counter of the current page
42	/// Increased for every page
43	sequence_num :u32,
44
45	/// Points to the first unwritten position in cur_pg_lacing.
46	segment_cnt :u8,
47	cur_pg_lacing :[u8; 255],
48	/// The data and the absgp's of the packets
49	cur_pg_data :Vec<(Box<[u8]>, u64)>,
50
51	/// Some(offs), if the last packet
52	/// couldn't make it fully into this page, and
53	/// has to be continued in the next page.
54	///
55	/// `offs` should point to the first idx in
56	/// cur_pg_data[last] that should NOT be written
57	/// in this page anymore.
58	///
59	/// None if all packets can be written nicely.
60	pck_this_overflow_idx :Option<usize>,
61
62	/// Some(offs), if the first packet
63	/// couldn't make it fully into the last page, and
64	/// has to be continued in this page.
65	///
66	/// `offs` should point to the first idx in cur_pg_data[0]
67	/// that hasn't been written.
68	///
69	/// None if all packets can be written nicely.
70	pck_last_overflow_idx :Option<usize>,
71}
72
73/// Specifies whether to end something with the write of the packet.
74///
75/// If you want to end a stream you need to inform the Ogg `PacketWriter`
76/// about this. This is the enum to do so.
77///
78/// Also, Codecs sometimes have special requirements to put
79/// the first packet of the whole stream into its own page.
80/// The `EndPage` variant can be used for this.
81#[derive(PartialEq)]
82#[derive(Clone, Copy)]
83pub enum PacketWriteEndInfo {
84	/// No ends here, just a normal packet
85	NormalPacket,
86	/// Force-end the current page
87	EndPage,
88	/// End the whole logical stream.
89	EndStream,
90}
91
92impl <T :io::Write> PacketWriter<T> {
93	pub fn new(wtr :T) -> Self {
94		return PacketWriter {
95			wtr,
96			page_vals : HashMap::new(),
97		};
98	}
99	pub fn into_inner(self) -> T {
100		self.wtr
101	}
102	/// Access the interior writer
103	///
104	/// This allows access of the writer contained inside.
105	/// No guarantees are given onto the pattern of the writes.
106	/// They may change in the future.
107	pub fn inner(&self) -> &T {
108		&self.wtr
109	}
110	/// Access the interior writer mutably
111	///
112	/// This allows access of the writer contained inside.
113	/// No guarantees are given onto the pattern of the writes.
114	/// They may change in the future.
115	pub fn inner_mut(&mut self) -> &mut T {
116		&mut self.wtr
117	}
118	/// Write a packet
119	///
120	///
121	pub fn write_packet(&mut self, pck_cont :Box<[u8]>, serial :u32,
122			inf :PacketWriteEndInfo,
123			/* TODO find a better way to design the API around
124				passing the absgp to the underlying implementation.
125				e.g. the caller passes a closure on init which gets
126				called when we encounter a new page... with the param
127				the index inside the current page, or something.
128			*/
129			absgp :u64) -> IoResult<()> {
130		let is_end_stream :bool = inf == PacketWriteEndInfo::EndStream;
131		let pg = self.page_vals.entry(serial).or_insert(
132			CurrentPageValues {
133				first_page : true,
134				sequence_num : 0,
135				segment_cnt : 0,
136				cur_pg_lacing :[0; 255],
137				cur_pg_data :Vec::with_capacity(255),
138				pck_this_overflow_idx : None,
139				pck_last_overflow_idx : None,
140			}
141		);
142
143		let cont_len = pck_cont.len();
144		pg.cur_pg_data.push((pck_cont, absgp));
145
146		let last_data_segment_size = (cont_len % 255) as u8;
147		let needed_segments :usize = (cont_len / 255) + 1;
148		let mut segment_in_page_i :u8 = pg.segment_cnt;
149		let mut at_page_end :bool = false;
150		for segment_i in 0 .. needed_segments {
151			at_page_end = false;
152			if segment_i + 1 < needed_segments {
153				// For all segments containing 255 pieces of data
154				pg.cur_pg_lacing[segment_in_page_i as usize] = 255;
155			} else {
156				// For the last segment, must contain < 255 pieces of data
157				// (including 0)
158				pg.cur_pg_lacing[segment_in_page_i as usize] = last_data_segment_size;
159			}
160			pg.segment_cnt = segment_in_page_i + 1;
161			segment_in_page_i = (segment_in_page_i + 1) % 255;
162			if segment_in_page_i == 0 {
163				if segment_i + 1 < needed_segments {
164					// We have to flush a page, but we know there are more to come...
165					pg.pck_this_overflow_idx = Some((segment_i + 1) * 255);
166					tri!(PacketWriter::write_page(&mut self.wtr, serial, pg,
167						false));
168				} else {
169					// We have to write a page end, and it's the very last
170					// we need to write
171					tri!(PacketWriter::write_page(&mut self.wtr,
172						serial, pg, is_end_stream));
173					// Not actually required
174					// (it is always None except if we set it to Some directly
175					// before we call write_page)
176					pg.pck_this_overflow_idx = None;
177					// Required (it could have been Some(offs) before)
178					pg.pck_last_overflow_idx = None;
179				}
180				at_page_end = true;
181			}
182		}
183		if (inf != PacketWriteEndInfo::NormalPacket) && !at_page_end {
184			// Write a page end
185			tri!(PacketWriter::write_page(&mut self.wtr, serial, pg,
186				is_end_stream));
187
188			pg.pck_last_overflow_idx = None;
189
190			// TODO if inf was PacketWriteEndInfo::EndStream, we have to
191			// somehow erase pg from the hashmap...
192			// any ideas? perhaps needs external scope...
193		}
194		// All went fine.
195		Ok(())
196	}
197	fn write_page(wtr :&mut T, serial :u32, pg :&mut CurrentPageValues,
198			last_page :bool)  -> IoResult<()> {
199		{
200			// The page header with everything but the lacing values:
201			let mut hdr_cur = Cursor::new(Vec::with_capacity(27));
202			tri!(hdr_cur.write_all(&[0x4f, 0x67, 0x67, 0x53, 0x00]));
203			let mut flags :u8 = 0;
204			if pg.pck_last_overflow_idx.is_some() { flags |= 0x01; }
205			if pg.first_page { flags |= 0x02; }
206			if last_page { flags |= 0x04; }
207
208			tri!(hdr_cur.write_u8(flags));
209
210			let pck_data = &pg.cur_pg_data;
211
212			let mut last_finishing_pck_absgp = (-1i64) as u64;
213			for (idx, &(_, absgp)) in pck_data.iter().enumerate() {
214				if !(idx + 1 == pck_data.len() &&
215						pg.pck_this_overflow_idx.is_some()) {
216					last_finishing_pck_absgp = absgp;
217				}
218			}
219
220			tri!(hdr_cur.write_u64::<LittleEndian>(last_finishing_pck_absgp));
221			tri!(hdr_cur.write_u32::<LittleEndian>(serial));
222			tri!(hdr_cur.write_u32::<LittleEndian>(pg.sequence_num));
223
224			// checksum, calculated later on :)
225			tri!(hdr_cur.write_u32::<LittleEndian>(0));
226
227			tri!(hdr_cur.write_u8(pg.segment_cnt));
228
229			let mut hash_calculated :u32;
230
231			let pg_lacing = &pg.cur_pg_lacing[0 .. pg.segment_cnt as usize];
232
233
234			hash_calculated = vorbis_crc32_update(0, hdr_cur.get_ref());
235			hash_calculated = vorbis_crc32_update(hash_calculated, pg_lacing);
236
237			for (idx, &(ref pck, _)) in pck_data.iter().enumerate() {
238				let mut start :usize = 0;
239				if idx == 0 { if let Some(idx) = pg.pck_last_overflow_idx {
240					start = idx;
241				}}
242				let mut end :usize = pck.len();
243				if idx + 1 == pck_data.len() {
244					if let Some(idx) = pg.pck_this_overflow_idx {
245						end = idx;
246					}
247				}
248				hash_calculated = vorbis_crc32_update(hash_calculated,
249					&pck[start .. end]);
250			}
251
252			// Go back to enter the checksum
253			// Don't do excessive checking here (that the seek
254			// succeeded & we are at the right pos now).
255			// It's hopefully not required.
256			tri!(hdr_cur.seek(SeekFrom::Start(22)));
257			tri!(hdr_cur.write_u32::<LittleEndian>(hash_calculated));
258
259			// Now all is done, write the stuff!
260			tri!(wtr.write_all(hdr_cur.get_ref()));
261			tri!(wtr.write_all(pg_lacing));
262			for (idx, &(ref pck, _)) in pck_data.iter().enumerate() {
263				let mut start :usize = 0;
264				if idx == 0 { if let Some(idx) = pg.pck_last_overflow_idx {
265					start = idx;
266				}}
267				let mut end :usize = pck.len();
268				if idx + 1 == pck_data.len() {
269					if let Some(idx) = pg.pck_this_overflow_idx {
270						end = idx;
271					}
272				}
273				tri!(wtr.write_all(&pck[start .. end]));
274			}
275		}
276
277		// Reset the page.
278		pg.first_page = false;
279		pg.sequence_num += 1;
280
281		pg.segment_cnt = 0;
282		// If we couldn't fully write the last
283		// packet, we need to keep it for the next page,
284		// otherwise just clear everything.
285		if pg.pck_this_overflow_idx.is_some() {
286			let d = pg.cur_pg_data.pop().unwrap();
287			pg.cur_pg_data.clear();
288			pg.cur_pg_data.push(d);
289		} else {
290			pg.cur_pg_data.clear();
291		}
292
293		pg.pck_last_overflow_idx = pg.pck_this_overflow_idx;
294		pg.pck_this_overflow_idx = None;
295
296		return Ok(());
297	}
298}
299
300impl<T :io::Seek + io::Write> PacketWriter<T> {
301	pub fn get_current_offs(&mut self) -> Result<u64, io::Error> {
302		self.wtr.seek(SeekFrom::Current(0))
303	}
304}
305
306// TODO once 1.18 gets released, move this
307// to the test module and make wtr pub(crate).
308#[test]
309fn test_recapture() {
310	// Test that we can deal with recapture
311	// at varying distances.
312	// This is a regression test
313	use std::io::Write;
314	use super::PacketReader;
315	let mut c = Cursor::new(Vec::new());
316	let test_arr = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10];
317	let test_arr_2 = [2, 4, 8, 16, 32, 64, 128, 127, 126, 125, 124];
318	let test_arr_3 = [3, 5, 9, 17, 33, 65, 129, 129, 127, 126, 125];
319	{
320		let np = PacketWriteEndInfo::NormalPacket;
321		let ep = PacketWriteEndInfo::EndPage;
322		{
323			let mut w = PacketWriter::new(&mut c);
324			w.write_packet(Box::new(test_arr), 0xdeadb33f, ep, 0).unwrap();
325
326			// Now, after the end of the page, put in some noise.
327			w.wtr.write_all(&[0; 38]).unwrap();
328
329			w.write_packet(Box::new(test_arr_2), 0xdeadb33f, np, 1).unwrap();
330			w.write_packet(Box::new(test_arr_3), 0xdeadb33f, ep, 2).unwrap();
331		}
332	}
333	//print_u8_slice(c.get_ref());
334	assert_eq!(c.seek(SeekFrom::Start(0)).unwrap(), 0);
335	{
336		let mut r = PacketReader::new(c);
337		let p1 = r.read_packet().unwrap().unwrap();
338		assert_eq!(test_arr, *p1.data);
339		let p2 = r.read_packet().unwrap().unwrap();
340		assert_eq!(test_arr_2, *p2.data);
341		let p3 = r.read_packet().unwrap().unwrap();
342		assert_eq!(test_arr_3, *p3.data);
343	}
344}