diff options
| -rw-r--r-- | src/libstd/io/mod.rs | 35 |
1 files changed, 35 insertions, 0 deletions
diff --git a/src/libstd/io/mod.rs b/src/libstd/io/mod.rs index 5b319f4c687..3b4e15953c4 100644 --- a/src/libstd/io/mod.rs +++ b/src/libstd/io/mod.rs @@ -669,6 +669,11 @@ impl<T> Take<T> { impl<T: Read> Read for Take<T> { fn read(&mut self, buf: &mut [u8]) -> Result<usize> { + // Don't call into inner reader at all at EOF because it may still block + if self.limit == 0 { + return Ok(0); + } + let max = cmp::min(buf.len() as u64, self.limit) as usize; let n = try!(self.inner.read(&mut buf[..max])); self.limit -= n as u64; @@ -676,6 +681,21 @@ impl<T: Read> Read for Take<T> { } } +impl<T: BufRead> BufRead for Take<T> { + fn fill_buf(&mut self) -> Result<&[u8]> { + let buf = try!(self.inner.fill_buf()); + let cap = cmp::min(buf.len() as u64, self.limit) as usize; + Ok(&buf[..cap]) + } + + fn consume(&mut self, amt: usize) { + // Don't let callers reset the limit by passing an overlarge value + let amt = cmp::min(amt as u64, self.limit) as usize; + self.limit -= amt as u64; + self.inner.consume(amt); + } +} + /// An adaptor which will emit all read data to a specified writer as well. /// /// For more information see `ReadExt::tee` @@ -846,6 +866,7 @@ impl<B: BufRead> Iterator for Lines<B> { mod tests { use prelude::v1::*; use io::prelude::*; + use io; use super::Cursor; #[test] @@ -943,4 +964,18 @@ mod tests { let mut v = String::new(); assert!(c.read_to_string(&mut v).is_err()); } + + #[test] + fn take_eof() { + struct R; + + impl Read for R { + fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> { + Err(io::Error::new(io::ErrorKind::Other, "", None)) + } + } + + let mut buf = [0; 1]; + assert_eq!(Ok(0), R.take(0).read(&mut buf)); + } } |