1use 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
20type IoResult<T> = result::Result<T, io::Error>;
25
26pub struct PacketWriter<T :io::Write> {
33 wtr :T,
34
35 page_vals :HashMap<u32, CurrentPageValues>,
36}
37
38struct CurrentPageValues {
39 first_page :bool,
41 sequence_num :u32,
44
45 segment_cnt :u8,
47 cur_pg_lacing :[u8; 255],
48 cur_pg_data :Vec<(Box<[u8]>, u64)>,
50
51 pck_this_overflow_idx :Option<usize>,
61
62 pck_last_overflow_idx :Option<usize>,
71}
72
73#[derive(PartialEq)]
82#[derive(Clone, Copy)]
83pub enum PacketWriteEndInfo {
84 NormalPacket,
86 EndPage,
88 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 pub fn inner(&self) -> &T {
108 &self.wtr
109 }
110 pub fn inner_mut(&mut self) -> &mut T {
116 &mut self.wtr
117 }
118 pub fn write_packet(&mut self, pck_cont :Box<[u8]>, serial :u32,
122 inf :PacketWriteEndInfo,
123 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 pg.cur_pg_lacing[segment_in_page_i as usize] = 255;
155 } else {
156 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 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 tri!(PacketWriter::write_page(&mut self.wtr,
172 serial, pg, is_end_stream));
173 pg.pck_this_overflow_idx = None;
177 pg.pck_last_overflow_idx = None;
179 }
180 at_page_end = true;
181 }
182 }
183 if (inf != PacketWriteEndInfo::NormalPacket) && !at_page_end {
184 tri!(PacketWriter::write_page(&mut self.wtr, serial, pg,
186 is_end_stream));
187
188 pg.pck_last_overflow_idx = None;
189
190 }
194 Ok(())
196 }
197 fn write_page(wtr :&mut T, serial :u32, pg :&mut CurrentPageValues,
198 last_page :bool) -> IoResult<()> {
199 {
200 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 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 tri!(hdr_cur.seek(SeekFrom::Start(22)));
257 tri!(hdr_cur.write_u32::<LittleEndian>(hash_calculated));
258
259 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 pg.first_page = false;
279 pg.sequence_num += 1;
280
281 pg.segment_cnt = 0;
282 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#[test]
309fn test_recapture() {
310 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 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 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}