Handle OOM when writing to Vec

Author: a1phyr

library/std/src/io/cursor.rs

@@ -481,7 +481,7 @@ fn reserve_and_pad<A: Allocator>(
         // to have room for (pos+buf_len) bytes. Reserve allocates
         // based on additional elements from the length, so we need to
         // reserve the difference
-        vec.reserve(desired_cap - vec.len());
+        vec.try_reserve(desired_cap - vec.len())?;
     }
     // Pad if pos is above the current len.
     if pos > vec.len() {
@@ -511,7 +511,8 @@ unsafe fn vec_write_all_unchecked<A>(pos: usize, vec: &mut Vec<u8, A>, buf: &[u8
 where
     A: Allocator,
 {
-    debug_assert!(vec.capacity() >= pos + buf.len());
+    debug_assert!(pos <= vec.capacity());
+    debug_assert!(buf.len() <= vec.capacity() - pos);
     unsafe { vec.as_mut_ptr().add(pos).copy_from(buf.as_ptr(), buf.len()) };
     pos + buf.len()
 }
@@ -565,7 +566,7 @@ where
     A: Allocator,
 {
     // For safety reasons, we don't want this sum to overflow ever.
-    // If this saturates, the reserve should panic to avoid any unsound writing.
+    // If this saturates, the reserve will fail, which will avoid any unsound writing.
     let buf_len = bufs.iter().fold(0usize, |a, b| a.saturating_add(b.len()));
     let mut pos = reserve_and_pad(pos_mut, vec, buf_len)?;
 

library/std/src/io/impls.rs

@@ -480,14 +480,15 @@ impl Write for &mut [u8] {
 impl<A: Allocator> Write for Vec<u8, A> {
     #[inline]
     fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
+        self.try_reserve(buf.len())?;
         self.extend_from_slice(buf);
         Ok(buf.len())
     }
 
     #[inline]
     fn write_vectored(&mut self, bufs: &[IoSlice<'_>]) -> io::Result<usize> {
         let len = bufs.iter().map(|b| b.len()).sum();
-        self.reserve(len);
+        self.try_reserve(len)?;
         for buf in bufs {
             self.extend_from_slice(buf);
         }
@@ -501,6 +502,7 @@ impl<A: Allocator> Write for Vec<u8, A> {
 
     #[inline]
     fn write_all(&mut self, buf: &[u8]) -> io::Result<()> {
+        self.try_reserve(buf.len())?;
         self.extend_from_slice(buf);
         Ok(())
     }