From 03545161e68c1b8c3fe052243e372ee75e4ec513 Mon Sep 17 00:00:00 2001 From: Kornel Date: Mon, 9 Oct 2023 14:25:31 +0100 Subject: [PATCH 1/2] Handle out of memory errors in io:Read::read_to_end() --- library/std/src/io/buffered/bufreader.rs | 1 + library/std/src/io/impls.rs | 5 ++-- library/std/src/io/mod.rs | 38 +++++++++++++++++++++++- 3 files changed, 41 insertions(+), 3 deletions(-) diff --git a/library/std/src/io/buffered/bufreader.rs b/library/std/src/io/buffered/bufreader.rs index 6c7494a6a6ff9..e920500d7d07f 100644 --- a/library/std/src/io/buffered/bufreader.rs +++ b/library/std/src/io/buffered/bufreader.rs @@ -345,6 +345,7 @@ impl Read for BufReader { // delegate to the inner implementation. fn read_to_end(&mut self, buf: &mut Vec) -> io::Result { 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(); diff --git a/library/std/src/io/impls.rs b/library/std/src/io/impls.rs index d8c8d933eb403..557e64dc8674e 100644 --- a/library/std/src/io/impls.rs +++ b/library/std/src/io/impls.rs @@ -303,8 +303,9 @@ impl Read for &[u8] { #[inline] fn read_to_end(&mut self, buf: &mut Vec) -> io::Result { - 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) } @@ -451,7 +452,7 @@ impl Read for VecDeque { fn read_to_end(&mut self, buf: &mut Vec) -> io::Result { // 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); diff --git a/library/std/src/io/mod.rs b/library/std/src/io/mod.rs index e3aa973741f1f..8fca66fa17c5e 100644 --- a/library/std/src/io/mod.rs +++ b/library/std/src/io/mod.rs @@ -430,6 +430,8 @@ pub(crate) fn default_read_to_end( 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); } @@ -462,7 +464,8 @@ pub(crate) fn default_read_to_end( } 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(); @@ -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. + /// + /// ```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) -> io::Result { + /// 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) -> Result { default_read_to_end(self, buf, None) From 60f46289cf55759dc474d796edbf9d38cbab86d4 Mon Sep 17 00:00:00 2001 From: Kornel Date: Mon, 9 Oct 2023 14:32:15 +0100 Subject: [PATCH 2/2] Handle out of memory errors in fs::read/read_to_string --- library/std/src/fs.rs | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/library/std/src/fs.rs b/library/std/src/fs.rs index c4a92927937a1..80d369eb067d4 100644 --- a/library/std/src/fs.rs +++ b/library/std/src/fs.rs @@ -260,7 +260,8 @@ pub fn read>(path: P) -> io::Result> { fn inner(path: &Path) -> io::Result> { 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) } @@ -302,7 +303,8 @@ pub fn read_to_string>(path: P) -> io::Result { fn inner(path: &Path) -> io::Result { 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) } @@ -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) -> io::Result { 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 { 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) } }