Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Handle out of memory errors in io:Read::read_to_end() #117925

Merged
merged 2 commits into from
Jan 30, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 6 additions & 4 deletions library/std/src/fs.rs
Original file line number Diff line number Diff line change
Expand Up @@ -260,7 +260,8 @@ pub fn read<P: AsRef<Path>>(path: P) -> io::Result<Vec<u8>> {
fn inner(path: &Path) -> io::Result<Vec<u8>> {
let mut file = File::open(path)?;
let size = file.metadata().map(|m| m.len() as usize).ok();
let mut bytes = Vec::with_capacity(size.unwrap_or(0));
let mut bytes = Vec::new();
bytes.try_reserve_exact(size.unwrap_or(0)).map_err(|_| io::ErrorKind::OutOfMemory)?;
io::default_read_to_end(&mut file, &mut bytes, size)?;
Ok(bytes)
}
Expand Down Expand Up @@ -302,7 +303,8 @@ pub fn read_to_string<P: AsRef<Path>>(path: P) -> io::Result<String> {
fn inner(path: &Path) -> io::Result<String> {
let mut file = File::open(path)?;
let size = file.metadata().map(|m| m.len() as usize).ok();
let mut string = String::with_capacity(size.unwrap_or(0));
let mut string = String::new();
string.try_reserve_exact(size.unwrap_or(0)).map_err(|_| io::ErrorKind::OutOfMemory)?;
io::default_read_to_string(&mut file, &mut string, size)?;
Ok(string)
}
Expand Down Expand Up @@ -774,14 +776,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_exact(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_exact(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
5 changes: 3 additions & 2 deletions library/std/src/io/impls.rs
Original file line number Diff line number Diff line change
Expand Up @@ -303,8 +303,9 @@ impl Read for &[u8] {

#[inline]
fn read_to_end(&mut self, buf: &mut Vec<u8>) -> io::Result<usize> {
buf.extend_from_slice(*self);
let len = self.len();
buf.try_reserve(len).map_err(|_| ErrorKind::OutOfMemory)?;
buf.extend_from_slice(*self);
*self = &self[len..];
Ok(len)
}
Expand Down Expand Up @@ -451,7 +452,7 @@ impl<A: Allocator> Read for VecDeque<u8, A> {
fn read_to_end(&mut self, buf: &mut Vec<u8>) -> io::Result<usize> {
// The total len is known upfront so we can reserve it in a single call.
let len = self.len();
buf.reserve(len);
buf.try_reserve(len).map_err(|_| ErrorKind::OutOfMemory)?;

let (front, back) = self.as_slices();
buf.extend_from_slice(front);
Expand Down
38 changes: 37 additions & 1 deletion library/std/src/io/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -430,6 +430,8 @@ pub(crate) fn default_read_to_end<R: Read + ?Sized>(
loop {
match r.read(&mut probe) {
Ok(n) => {
// there is no way to recover from allocation failure here
// because the data has already been read.
buf.extend_from_slice(&probe[..n]);
return Ok(n);
}
Expand Down Expand Up @@ -462,7 +464,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 +818,39 @@ 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.
the8472 marked this conversation as resolved.
Show resolved Hide resolved
///
/// ```no_run
/// # use std::io::{self, BufRead};
/// # struct Example { example_datasource: io::Empty } impl Example {
/// # fn get_some_data_for_the_example(&self) -> &'static [u8] { &[] }
/// fn read_to_end(&mut self, dest_vec: &mut Vec<u8>) -> io::Result<usize> {
/// let initial_vec_len = dest_vec.len();
/// loop {
/// let src_buf = self.example_datasource.fill_buf()?;
/// if src_buf.is_empty() {
/// break;
/// }
/// dest_vec.try_reserve(src_buf.len()).map_err(|_| io::ErrorKind::OutOfMemory)?;
/// dest_vec.extend_from_slice(src_buf);
///
/// // Any irreversible side effects should happen after `try_reserve` succeeds,
/// // to avoid losing data on allocation error.
/// let read = src_buf.len();
/// self.example_datasource.consume(read);
/// }
/// Ok(dest_vec.len() - initial_vec_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
Loading