1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122

//! The `dribble` library helps you test implementations of the traits
//! `std::io::Read` and `std::io::Write` by passing data to them in small,
//! random-sized chunks.  This allows you to stress-test the code you run
//! near buffer boundaries.

extern crate rand;

use rand::distributions::{IndependentSample, Range};
use std::cmp::min;
use std::io::{self, Read, Write};

/// Wrap an implementation of `Read`, and return bytes in small,
/// random-sized chunks when `read` is called.
///
/// ```
/// use std::io::{Cursor, Read};
/// use dribble::DribbleReader;
///
/// let input = b"This is my test data";
/// let mut cursor = Cursor::new(input as &[u8]);
/// let mut dribble = DribbleReader::new(&mut cursor);
/// let mut output = vec!();
/// dribble.read_to_end(&mut output).unwrap();
/// assert_eq!(input as &[u8], &output as &[u8]);
/// ```
pub struct DribbleReader<R: Read> {
    source: R,
    buffer: Vec<u8>,
    used: usize,
    available: usize
}

impl<R: Read> DribbleReader<R> {
    /// Create a new `DribbleReader`.  The `read` function will only return
    /// 0 if `source.read` returns 0.
    pub fn new(source: R) -> Self {
        DribbleReader{source: source, buffer: vec![0; 64],
                      used: 0, available: 0}
    }
}

impl<R: Read> Read for DribbleReader<R> {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        assert!(self.used <= self.available);

        // If we're out of buffered data, try to read some.
        if self.used == self.available {
            self.used = 0;
            self.available = 0;
            self.available = try!(self.source.read(&mut self.buffer));
        }

        if self.available == 0 {
            // We tried to read bytes and didn't get any, so pass it
            // through.
            Ok(0)
        } else {
            // Decide how many bytes to copy, limiting it to the smallest
            // of (1) the destination space available, (2) a random number
            // between 1 and 4 inclusive, and (3) the number of bytes
            // currently buffered.
            let mut rng = rand::thread_rng();
            let bytes = min(buf.len(),
                            min(Range::new(1, 5).ind_sample(&mut rng),
                                self.available -  self.used));
            assert!(1 <= bytes && bytes <= 4);

            // Copy bytes, mark them as used, and return.
            for i in 0..bytes { buf[i] = self.buffer[self.used+i]; }
            self.used += bytes;
            Ok(bytes)
        }
    }
}

/// Wrap an implementation of `Write`, and pass through bytes in small,
/// random-sized chunks when `write` is called.
///
/// ```
/// use std::io::Write;
/// use dribble::DribbleWriter;
///
/// let input = b"This is my test data";
/// let mut output = vec!();
/// {
///     let mut dribble = DribbleWriter::new(&mut output);
///     dribble.write(input).unwrap();
/// }
/// assert_eq!(input as &[u8], &output as &[u8]);        
/// ```
pub struct DribbleWriter<W: Write> {
    dest: W
}

impl<W: Write> DribbleWriter<W> {
    /// Create a new `DribbleWriter`.
    pub fn new(dest: W) -> Self {
        DribbleWriter{dest: dest}
    }
}

impl<W: Write> Write for DribbleWriter<W> {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        let mut rng = rand::thread_rng();
        let mut written = 0;
        while written < buf.len() {
            // Decide how many bytes to write, not exceeding the number
            // available.
            let bytes = min(buf.len() - written,
                            Range::new(0, 5).ind_sample(&mut rng));

            // Write the bytes.
            try!(self.dest.write(&buf[written..written+bytes]));
            written += bytes;
        }
        Ok(written)
    }

    fn flush(&mut self) -> io::Result<()> {
        self.dest.flush()
    }
}