diff --git a/library/std/src/io/impls.rs b/library/std/src/io/impls.rs index e5048dcc8acd9..927136df3c6ad 100644 --- a/library/std/src/io/impls.rs +++ b/library/std/src/io/impls.rs @@ -377,18 +377,28 @@ impl Write for &mut [u8] { /// Write is implemented for `Vec` by appending to the vector. /// The vector will grow as needed. +/// +/// # Panics +/// +/// In case of allocation error or capacity overflow, write operations will panic. +/// The panicking behavior is not guaranteed. In the future, it may become +/// a regular `io::Error`. #[stable(feature = "rust1", since = "1.0.0")] impl Write for Vec { #[inline] fn write(&mut self, buf: &[u8]) -> io::Result { + fallible_reserve(self, buf.len()); self.extend_from_slice(buf); Ok(buf.len()) } #[inline] fn write_vectored(&mut self, bufs: &[IoSlice<'_>]) -> io::Result { - let len = bufs.iter().map(|b| b.len()).sum(); - self.reserve(len); + let len = bufs + .iter() + .try_fold(0usize, |len, b| len.checked_add(b.len())) + .expect("capacity overflow"); + fallible_reserve(self, len); for buf in bufs { self.extend_from_slice(buf); } @@ -402,6 +412,7 @@ impl Write for Vec { #[inline] fn write_all(&mut self, buf: &[u8]) -> io::Result<()> { + fallible_reserve(self, buf.len()); self.extend_from_slice(buf); Ok(()) } @@ -412,6 +423,20 @@ impl Write for Vec { } } +#[inline] +fn fallible_reserve(vec: &mut Vec, len: usize) { + if len > vec.capacity().wrapping_sub(vec.len()) { + do_reserve_and_handle(vec, len); + } +} + +#[cold] +fn do_reserve_and_handle(vec: &mut Vec, len: usize) { + if let Err(_) = vec.try_reserve(len) { + panic!("out of memory"); + } +} + /// Read is implemented for `VecDeque` by consuming bytes from the front of the `VecDeque`. #[stable(feature = "vecdeque_read_write", since = "1.63.0")] impl Read for VecDeque {