1 extern crate rand;
2
3 use super::EncoderWriter;
4 use tests::random_config;
5 use {encode_config, encode_config_buf, STANDARD_NO_PAD, URL_SAFE};
6
7 use std::io::{Cursor, Write};
8 use std::{cmp, io, str};
9
10 use self::rand::Rng;
11
12 #[test]
encode_three_bytes()13 fn encode_three_bytes() {
14 let mut c = Cursor::new(Vec::new());
15 {
16 let mut enc = EncoderWriter::new(&mut c, URL_SAFE);
17
18 let sz = enc.write(b"abc").unwrap();
19 assert_eq!(sz, 3);
20 }
21 assert_eq!(&c.get_ref()[..], encode_config("abc", URL_SAFE).as_bytes());
22 }
23
24 #[test]
encode_nine_bytes_two_writes()25 fn encode_nine_bytes_two_writes() {
26 let mut c = Cursor::new(Vec::new());
27 {
28 let mut enc = EncoderWriter::new(&mut c, URL_SAFE);
29
30 let sz = enc.write(b"abcdef").unwrap();
31 assert_eq!(sz, 6);
32 let sz = enc.write(b"ghi").unwrap();
33 assert_eq!(sz, 3);
34 }
35 assert_eq!(
36 &c.get_ref()[..],
37 encode_config("abcdefghi", URL_SAFE).as_bytes()
38 );
39 }
40
41 #[test]
encode_one_then_two_bytes()42 fn encode_one_then_two_bytes() {
43 let mut c = Cursor::new(Vec::new());
44 {
45 let mut enc = EncoderWriter::new(&mut c, URL_SAFE);
46
47 let sz = enc.write(b"a").unwrap();
48 assert_eq!(sz, 1);
49 let sz = enc.write(b"bc").unwrap();
50 assert_eq!(sz, 2);
51 }
52 assert_eq!(&c.get_ref()[..], encode_config("abc", URL_SAFE).as_bytes());
53 }
54
55 #[test]
encode_one_then_five_bytes()56 fn encode_one_then_five_bytes() {
57 let mut c = Cursor::new(Vec::new());
58 {
59 let mut enc = EncoderWriter::new(&mut c, URL_SAFE);
60
61 let sz = enc.write(b"a").unwrap();
62 assert_eq!(sz, 1);
63 let sz = enc.write(b"bcdef").unwrap();
64 assert_eq!(sz, 5);
65 }
66 assert_eq!(
67 &c.get_ref()[..],
68 encode_config("abcdef", URL_SAFE).as_bytes()
69 );
70 }
71
72 #[test]
encode_1_2_3_bytes()73 fn encode_1_2_3_bytes() {
74 let mut c = Cursor::new(Vec::new());
75 {
76 let mut enc = EncoderWriter::new(&mut c, URL_SAFE);
77
78 let sz = enc.write(b"a").unwrap();
79 assert_eq!(sz, 1);
80 let sz = enc.write(b"bc").unwrap();
81 assert_eq!(sz, 2);
82 let sz = enc.write(b"def").unwrap();
83 assert_eq!(sz, 3);
84 }
85 assert_eq!(
86 &c.get_ref()[..],
87 encode_config("abcdef", URL_SAFE).as_bytes()
88 );
89 }
90
91 #[test]
encode_with_padding()92 fn encode_with_padding() {
93 let mut c = Cursor::new(Vec::new());
94 {
95 let mut enc = EncoderWriter::new(&mut c, URL_SAFE);
96
97 enc.write_all(b"abcd").unwrap();
98
99 enc.flush().unwrap();
100 }
101 assert_eq!(&c.get_ref()[..], encode_config("abcd", URL_SAFE).as_bytes());
102 }
103
104 #[test]
encode_with_padding_multiple_writes()105 fn encode_with_padding_multiple_writes() {
106 let mut c = Cursor::new(Vec::new());
107 {
108 let mut enc = EncoderWriter::new(&mut c, URL_SAFE);
109
110 assert_eq!(1, enc.write(b"a").unwrap());
111 assert_eq!(2, enc.write(b"bc").unwrap());
112 assert_eq!(3, enc.write(b"def").unwrap());
113 assert_eq!(1, enc.write(b"g").unwrap());
114
115 enc.flush().unwrap();
116 }
117 assert_eq!(
118 &c.get_ref()[..],
119 encode_config("abcdefg", URL_SAFE).as_bytes()
120 );
121 }
122
123 #[test]
finish_writes_extra_byte()124 fn finish_writes_extra_byte() {
125 let mut c = Cursor::new(Vec::new());
126 {
127 let mut enc = EncoderWriter::new(&mut c, URL_SAFE);
128
129 assert_eq!(6, enc.write(b"abcdef").unwrap());
130
131 // will be in extra
132 assert_eq!(1, enc.write(b"g").unwrap());
133
134 // 1 trailing byte = 2 encoded chars
135 let _ = enc.finish().unwrap();
136 }
137 assert_eq!(
138 &c.get_ref()[..],
139 encode_config("abcdefg", URL_SAFE).as_bytes()
140 );
141 }
142
143 #[test]
write_partial_chunk_encodes_partial_chunk()144 fn write_partial_chunk_encodes_partial_chunk() {
145 let mut c = Cursor::new(Vec::new());
146 {
147 let mut enc = EncoderWriter::new(&mut c, STANDARD_NO_PAD);
148
149 // nothing encoded yet
150 assert_eq!(2, enc.write(b"ab").unwrap());
151 // encoded here
152 let _ = enc.finish().unwrap();
153 }
154 assert_eq!(
155 &c.get_ref()[..],
156 encode_config("ab", STANDARD_NO_PAD).as_bytes()
157 );
158 assert_eq!(3, c.get_ref().len());
159 }
160
161 #[test]
write_1_chunk_encodes_complete_chunk()162 fn write_1_chunk_encodes_complete_chunk() {
163 let mut c = Cursor::new(Vec::new());
164 {
165 let mut enc = EncoderWriter::new(&mut c, STANDARD_NO_PAD);
166
167 assert_eq!(3, enc.write(b"abc").unwrap());
168 let _ = enc.finish().unwrap();
169 }
170 assert_eq!(
171 &c.get_ref()[..],
172 encode_config("abc", STANDARD_NO_PAD).as_bytes()
173 );
174 assert_eq!(4, c.get_ref().len());
175 }
176
177 #[test]
write_1_chunk_and_partial_encodes_only_complete_chunk()178 fn write_1_chunk_and_partial_encodes_only_complete_chunk() {
179 let mut c = Cursor::new(Vec::new());
180 {
181 let mut enc = EncoderWriter::new(&mut c, STANDARD_NO_PAD);
182
183 // "d" not written
184 assert_eq!(3, enc.write(b"abcd").unwrap());
185 let _ = enc.finish().unwrap();
186 }
187 assert_eq!(
188 &c.get_ref()[..],
189 encode_config("abc", STANDARD_NO_PAD).as_bytes()
190 );
191 assert_eq!(4, c.get_ref().len());
192 }
193
194 #[test]
write_2_partials_to_exactly_complete_chunk_encodes_complete_chunk()195 fn write_2_partials_to_exactly_complete_chunk_encodes_complete_chunk() {
196 let mut c = Cursor::new(Vec::new());
197 {
198 let mut enc = EncoderWriter::new(&mut c, STANDARD_NO_PAD);
199
200 assert_eq!(1, enc.write(b"a").unwrap());
201 assert_eq!(2, enc.write(b"bc").unwrap());
202 let _ = enc.finish().unwrap();
203 }
204 assert_eq!(
205 &c.get_ref()[..],
206 encode_config("abc", STANDARD_NO_PAD).as_bytes()
207 );
208 assert_eq!(4, c.get_ref().len());
209 }
210
211 #[test]
write_partial_then_enough_to_complete_chunk_but_not_complete_another_chunk_encodes_complete_chunk_without_consuming_remaining()212 fn write_partial_then_enough_to_complete_chunk_but_not_complete_another_chunk_encodes_complete_chunk_without_consuming_remaining() {
213 let mut c = Cursor::new(Vec::new());
214 {
215 let mut enc = EncoderWriter::new(&mut c, STANDARD_NO_PAD);
216
217 assert_eq!(1, enc.write(b"a").unwrap());
218 // doesn't consume "d"
219 assert_eq!(2, enc.write(b"bcd").unwrap());
220 let _ = enc.finish().unwrap();
221 }
222 assert_eq!(
223 &c.get_ref()[..],
224 encode_config("abc", STANDARD_NO_PAD).as_bytes()
225 );
226 assert_eq!(4, c.get_ref().len());
227 }
228
229 #[test]
write_partial_then_enough_to_complete_chunk_and_another_chunk_encodes_complete_chunks()230 fn write_partial_then_enough_to_complete_chunk_and_another_chunk_encodes_complete_chunks() {
231 let mut c = Cursor::new(Vec::new());
232 {
233 let mut enc = EncoderWriter::new(&mut c, STANDARD_NO_PAD);
234
235 assert_eq!(1, enc.write(b"a").unwrap());
236 // completes partial chunk, and another chunk
237 assert_eq!(5, enc.write(b"bcdef").unwrap());
238 let _ = enc.finish().unwrap();
239 }
240 assert_eq!(
241 &c.get_ref()[..],
242 encode_config("abcdef", STANDARD_NO_PAD).as_bytes()
243 );
244 assert_eq!(8, c.get_ref().len());
245 }
246
247 #[test]
write_partial_then_enough_to_complete_chunk_and_another_chunk_and_another_partial_chunk_encodes_only_complete_chunks()248 fn write_partial_then_enough_to_complete_chunk_and_another_chunk_and_another_partial_chunk_encodes_only_complete_chunks() {
249 let mut c = Cursor::new(Vec::new());
250 {
251 let mut enc = EncoderWriter::new(&mut c, STANDARD_NO_PAD);
252
253 assert_eq!(1, enc.write(b"a").unwrap());
254 // completes partial chunk, and another chunk, with one more partial chunk that's not
255 // consumed
256 assert_eq!(5, enc.write(b"bcdefe").unwrap());
257 let _ = enc.finish().unwrap();
258 }
259 assert_eq!(
260 &c.get_ref()[..],
261 encode_config("abcdef", STANDARD_NO_PAD).as_bytes()
262 );
263 assert_eq!(8, c.get_ref().len());
264 }
265
266 #[test]
drop_calls_finish_for_you()267 fn drop_calls_finish_for_you() {
268 let mut c = Cursor::new(Vec::new());
269 {
270 let mut enc = EncoderWriter::new(&mut c, STANDARD_NO_PAD);
271 assert_eq!(1, enc.write(b"a").unwrap());
272 }
273 assert_eq!(
274 &c.get_ref()[..],
275 encode_config("a", STANDARD_NO_PAD).as_bytes()
276 );
277 assert_eq!(2, c.get_ref().len());
278 }
279
280 #[test]
every_possible_split_of_input()281 fn every_possible_split_of_input() {
282 let mut rng = rand::thread_rng();
283 let mut orig_data = Vec::<u8>::new();
284 let mut stream_encoded = Vec::<u8>::new();
285 let mut normal_encoded = String::new();
286
287 let size = 5_000;
288
289 for i in 0..size {
290 orig_data.clear();
291 stream_encoded.clear();
292 normal_encoded.clear();
293
294 for _ in 0..size {
295 orig_data.push(rng.gen());
296 }
297
298 let config = random_config(&mut rng);
299 encode_config_buf(&orig_data, config, &mut normal_encoded);
300
301 {
302 let mut stream_encoder = EncoderWriter::new(&mut stream_encoded, config);
303 // Write the first i bytes, then the rest
304 stream_encoder.write_all(&orig_data[0..i]).unwrap();
305 stream_encoder.write_all(&orig_data[i..]).unwrap();
306 }
307
308 assert_eq!(normal_encoded, str::from_utf8(&stream_encoded).unwrap());
309 }
310 }
311
312 #[test]
encode_random_config_matches_normal_encode_reasonable_input_len()313 fn encode_random_config_matches_normal_encode_reasonable_input_len() {
314 // choose up to 2 * buf size, so ~half the time it'll use a full buffer
315 do_encode_random_config_matches_normal_encode(super::encoder::BUF_SIZE * 2)
316 }
317
318 #[test]
encode_random_config_matches_normal_encode_tiny_input_len()319 fn encode_random_config_matches_normal_encode_tiny_input_len() {
320 do_encode_random_config_matches_normal_encode(10)
321 }
322
323 #[test]
retrying_writes_that_error_with_interrupted_works()324 fn retrying_writes_that_error_with_interrupted_works() {
325 let mut rng = rand::thread_rng();
326 let mut orig_data = Vec::<u8>::new();
327 let mut stream_encoded = Vec::<u8>::new();
328 let mut normal_encoded = String::new();
329
330 for _ in 0..1_000 {
331 orig_data.clear();
332 stream_encoded.clear();
333 normal_encoded.clear();
334
335 let orig_len: usize = rng.gen_range(100, 20_000);
336 for _ in 0..orig_len {
337 orig_data.push(rng.gen());
338 }
339
340 // encode the normal way
341 let config = random_config(&mut rng);
342 encode_config_buf(&orig_data, config, &mut normal_encoded);
343
344 // encode via the stream encoder
345 {
346 let mut interrupt_rng = rand::thread_rng();
347 let mut interrupting_writer = InterruptingWriter {
348 w: &mut stream_encoded,
349 rng: &mut interrupt_rng,
350 fraction: 0.8,
351 };
352
353 let mut stream_encoder = EncoderWriter::new(&mut interrupting_writer, config);
354 let mut bytes_consumed = 0;
355 while bytes_consumed < orig_len {
356 // use short inputs since we want to use `extra` a lot as that's what needs rollback
357 // when errors occur
358 let input_len: usize = cmp::min(rng.gen_range(0, 10), orig_len - bytes_consumed);
359
360 retry_interrupted_write_all(
361 &mut stream_encoder,
362 &orig_data[bytes_consumed..bytes_consumed + input_len],
363 ).unwrap();
364
365 bytes_consumed += input_len;
366 }
367
368 loop {
369 let res = stream_encoder.finish();
370 match res {
371 Ok(_) => break,
372 Err(e) => match e.kind() {
373 io::ErrorKind::Interrupted => continue,
374 _ => Err(e).unwrap(), // bail
375 },
376 }
377 }
378
379 assert_eq!(orig_len, bytes_consumed);
380 }
381
382 assert_eq!(normal_encoded, str::from_utf8(&stream_encoded).unwrap());
383 }
384 }
385
386 #[test]
writes_that_only_write_part_of_input_and_sometimes_interrupt_produce_correct_encoded_data()387 fn writes_that_only_write_part_of_input_and_sometimes_interrupt_produce_correct_encoded_data() {
388 let mut rng = rand::thread_rng();
389 let mut orig_data = Vec::<u8>::new();
390 let mut stream_encoded = Vec::<u8>::new();
391 let mut normal_encoded = String::new();
392
393 for _ in 0..1_000 {
394 orig_data.clear();
395 stream_encoded.clear();
396 normal_encoded.clear();
397
398 let orig_len: usize = rng.gen_range(100, 20_000);
399 for _ in 0..orig_len {
400 orig_data.push(rng.gen());
401 }
402
403 // encode the normal way
404 let config = random_config(&mut rng);
405 encode_config_buf(&orig_data, config, &mut normal_encoded);
406
407 // encode via the stream encoder
408 {
409 let mut partial_rng = rand::thread_rng();
410 let mut partial_writer = PartialInterruptingWriter {
411 w: &mut stream_encoded,
412 rng: &mut partial_rng,
413 full_input_fraction: 0.1,
414 no_interrupt_fraction: 0.1
415 };
416
417 let mut stream_encoder = EncoderWriter::new(&mut partial_writer, config);
418 let mut bytes_consumed = 0;
419 while bytes_consumed < orig_len {
420 // use at most medium-length inputs to exercise retry logic more aggressively
421 let input_len: usize = cmp::min(rng.gen_range(0, 100), orig_len - bytes_consumed);
422
423 let res = stream_encoder.write(&orig_data[bytes_consumed..bytes_consumed + input_len]);
424
425 // retry on interrupt
426 match res {
427 Ok(len) => bytes_consumed += len,
428 Err(e) => match e.kind() {
429 io::ErrorKind::Interrupted => continue,
430 _ => { panic!("should not see other errors"); }
431 },
432 }
433 };
434
435 stream_encoder.finish().unwrap();
436
437 assert_eq!(orig_len, bytes_consumed);
438 }
439
440 assert_eq!(normal_encoded, str::from_utf8(&stream_encoded).unwrap());
441 }
442 }
443
444
445 /// Retry writes until all the data is written or an error that isn't Interrupted is returned.
retry_interrupted_write_all<W: Write>(w: &mut W, buf: &[u8]) -> io::Result<()>446 fn retry_interrupted_write_all<W: Write>(w: &mut W, buf: &[u8]) -> io::Result<()> {
447 let mut bytes_consumed = 0;
448
449 while bytes_consumed < buf.len() {
450 let res = w.write(&buf[bytes_consumed..]);
451
452 match res {
453 Ok(len) => bytes_consumed += len,
454 Err(e) => match e.kind() {
455 io::ErrorKind::Interrupted => continue,
456 _ => { return Err(e) }
457 },
458 }
459 }
460
461 Ok(())
462 }
463
do_encode_random_config_matches_normal_encode(max_input_len: usize)464 fn do_encode_random_config_matches_normal_encode(max_input_len: usize) {
465 let mut rng = rand::thread_rng();
466 let mut orig_data = Vec::<u8>::new();
467 let mut stream_encoded = Vec::<u8>::new();
468 let mut normal_encoded = String::new();
469
470 for _ in 0..1_000 {
471 orig_data.clear();
472 stream_encoded.clear();
473 normal_encoded.clear();
474
475 let orig_len: usize = rng.gen_range(100, 20_000);
476 for _ in 0..orig_len {
477 orig_data.push(rng.gen());
478 }
479
480 // encode the normal way
481 let config = random_config(&mut rng);
482 encode_config_buf(&orig_data, config, &mut normal_encoded);
483
484 // encode via the stream encoder
485 {
486 let mut stream_encoder = EncoderWriter::new(&mut stream_encoded, config);
487 let mut bytes_consumed = 0;
488 while bytes_consumed < orig_len {
489 let input_len: usize =
490 cmp::min(rng.gen_range(0, max_input_len), orig_len - bytes_consumed);
491
492 // write a little bit of the data
493 stream_encoder
494 .write_all(&orig_data[bytes_consumed..bytes_consumed + input_len])
495 .unwrap();
496
497 bytes_consumed += input_len;
498 }
499
500 stream_encoder.finish().unwrap();
501
502 assert_eq!(orig_len, bytes_consumed);
503 }
504
505 assert_eq!(normal_encoded, str::from_utf8(&stream_encoded).unwrap());
506 }
507 }
508
509 /// A `Write` implementation that returns Interrupted some fraction of the time, randomly.
510 struct InterruptingWriter<'a, W: 'a + Write, R: 'a + Rng> {
511 w: &'a mut W,
512 rng: &'a mut R,
513 /// In [0, 1]. If a random number in [0, 1] is `<= threshold`, `Write` methods will return
514 /// an `Interrupted` error
515 fraction: f64,
516 }
517
518 impl<'a, W: Write, R: Rng> Write for InterruptingWriter<'a, W, R> {
write(&mut self, buf: &[u8]) -> io::Result<usize>519 fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
520 if self.rng.gen_range(0.0, 1.0) <= self.fraction {
521 return Err(io::Error::new(io::ErrorKind::Interrupted, "interrupted"));
522 }
523
524 self.w.write(buf)
525 }
526
flush(&mut self) -> io::Result<()>527 fn flush(&mut self) -> io::Result<()> {
528 if self.rng.gen_range(0.0, 1.0) <= self.fraction {
529 return Err(io::Error::new(io::ErrorKind::Interrupted, "interrupted"));
530 }
531
532 self.w.flush()
533 }
534 }
535
536 /// A `Write` implementation that sometimes will only write part of its input.
537 struct PartialInterruptingWriter<'a, W: 'a + Write, R: 'a + Rng> {
538 w: &'a mut W,
539 rng: &'a mut R,
540 /// In [0, 1]. If a random number in [0, 1] is `<= threshold`, `write()` will write all its
541 /// input. Otherwise, it will write a random substring
542 full_input_fraction: f64,
543 no_interrupt_fraction: f64
544 }
545
546 impl<'a, W: Write, R: Rng> Write for PartialInterruptingWriter<'a, W, R> {
write(&mut self, buf: &[u8]) -> io::Result<usize>547 fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
548 if self.rng.gen_range(0.0, 1.0) > self.no_interrupt_fraction{
549 return Err(io::Error::new(io::ErrorKind::Interrupted, "interrupted"));
550 }
551
552 if self.rng.gen_range(0.0, 1.0) <= self.full_input_fraction || buf.len() == 0 {
553 // pass through the buf untouched
554 self.w.write(buf)
555 } else {
556 // only use a prefix of it
557 self.w.write(&buf[0..(self.rng.gen_range(0, buf.len() - 1))])
558 }
559 }
560
flush(&mut self) -> io::Result<()>561 fn flush(&mut self) -> io::Result<()> {
562 self.w.flush()
563 }
564 }