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

use std::io::{self, Read, Write};
use std::cmp::min;
use seek_forward::{SeekForward, Tell};

/// Wraps around a stream to limit the length of the underlying stream.
///
/// This implementation differs from `std::io::Take` in that it also allows writes,
/// and seeking forward is allowed if the underlying stream supports it.
pub struct Take<T> {
    inner: T,
    limit: u64,
}

impl<T> Take<T> {
    /// Creates a new `Take` with `limit` bytes
    pub fn new(inner: T, limit: u64) -> Self {
        Take {
            inner: inner,
            limit: limit,
        }
    }
}

impl<T: Write> Write for Take<T> {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        let limit = min(self.limit, buf.len() as u64);

        if limit == 0 {
            return Ok(0)
        }

        let buf = &buf[..limit as usize];
        let inner = try!(self.inner.write(buf));
        self.limit -= inner as u64;
        Ok(inner)
    }

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

impl<T: Read> Read for Take<T> {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        let limit = min(self.limit, buf.len() as u64);

        if limit == 0 {
            return Ok(0)
        }

        let buf = &mut buf[..limit as usize];
        let inner = try!(self.inner.read(buf));
        self.limit -= inner as u64;
        Ok(inner)
    }
}

impl<T: SeekForward> SeekForward for Take<T> {
    fn seek_forward(&mut self, offset: u64) -> io::Result<u64> {
        let res = try!(self.inner.seek_forward(min(offset, self.limit)));
        self.limit -= res;
        Ok(res)
    }
}

impl<T: Tell> Tell for Take<T> {
    fn tell(&mut self) -> io::Result<u64> {
        self.inner.tell()
    }
}