Skip to content

Commit

Permalink
Handle out of memory errors in io:Read::read_to_end()
Browse files Browse the repository at this point in the history
  • Loading branch information
kornelski committed Nov 29, 2023
1 parent abe34e9 commit 561e5b8
Show file tree
Hide file tree
Showing 3 changed files with 30 additions and 3 deletions.
4 changes: 2 additions & 2 deletions library/std/src/fs.rs
Original file line number Diff line number Diff line change
Expand Up @@ -770,14 +770,14 @@ impl Read for &File {
// Reserves space in the buffer based on the file size when available.
fn read_to_end(&mut self, buf: &mut Vec<u8>) -> io::Result<usize> {
let size = buffer_capacity_required(self);
buf.reserve(size.unwrap_or(0));
buf.try_reserve(size.unwrap_or(0)).map_err(|_| io::ErrorKind::OutOfMemory)?;
io::default_read_to_end(self, buf, size)
}

// Reserves space in the buffer based on the file size when available.
fn read_to_string(&mut self, buf: &mut String) -> io::Result<usize> {
let size = buffer_capacity_required(self);
buf.reserve(size.unwrap_or(0));
buf.try_reserve(size.unwrap_or(0)).map_err(|_| io::ErrorKind::OutOfMemory)?;
io::default_read_to_string(self, buf, size)
}
}
Expand Down
1 change: 1 addition & 0 deletions library/std/src/io/buffered/bufreader.rs
Original file line number Diff line number Diff line change
Expand Up @@ -345,6 +345,7 @@ impl<R: ?Sized + Read> Read for BufReader<R> {
// delegate to the inner implementation.
fn read_to_end(&mut self, buf: &mut Vec<u8>) -> io::Result<usize> {
let inner_buf = self.buffer();
buf.try_reserve(inner_buf.len()).map_err(|_| io::ErrorKind::OutOfMemory)?;
buf.extend_from_slice(inner_buf);
let nread = inner_buf.len();
self.discard_buffer();
Expand Down
28 changes: 27 additions & 1 deletion library/std/src/io/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -430,6 +430,7 @@ pub(crate) fn default_read_to_end<R: Read + ?Sized>(
loop {
match r.read(&mut probe) {
Ok(n) => {
buf.try_reserve(n).map_err(|_| io::ErrorKind::OutOfMemory)?;
buf.extend_from_slice(&probe[..n]);
return Ok(n);
}
Expand Down Expand Up @@ -462,7 +463,8 @@ pub(crate) fn default_read_to_end<R: Read + ?Sized>(
}

if buf.len() == buf.capacity() {
buf.reserve(PROBE_SIZE); // buf is full, need more space
// buf is full, need more space
buf.try_reserve(PROBE_SIZE).map_err(|_| ErrorKind::OutOfMemory)?;
}

let mut spare = buf.spare_capacity_mut();
Expand Down Expand Up @@ -815,6 +817,30 @@ pub trait Read {
/// file.)
///
/// [`std::fs::read`]: crate::fs::read
///
/// ## Implementing `read_to_end`
///
/// When implementing the `io::Read` trait, it is recommended to allocate
/// memory using [`Vec::try_reserve`]. However, this behavior is not guaranteed
/// by all implementations, and `read_to_end` may not handle out-of-memory
/// situations gracefully.
///
/// ```no_run
/// # use std::io;
/// # struct Example; impl Example {
/// # fn read_some_data_for_the_example(&self) -> &'static [u8] { &[] }
/// fn read_to_end(&mut self, buf: &mut Vec<u8>) -> io::Result<usize> {
/// let data_read = self.read_some_data_for_the_example();
///
/// buf.try_reserve(data_read.len()).map_err(|_| io::ErrorKind::OutOfMemory)?;
/// buf.extend_from_slice(data_read);
///
/// Ok(data_read.len())
/// }
/// # }
/// ```
///
/// [`Vec::try_reserve`]: crate::vec::Vec::try_reserve
#[stable(feature = "rust1", since = "1.0.0")]
fn read_to_end(&mut self, buf: &mut Vec<u8>) -> Result<usize> {
default_read_to_end(self, buf, None)
Expand Down

0 comments on commit 561e5b8

Please sign in to comment.