From c95f8a6966d6a7b0946c5d5e64e0cd25a8839c0a Mon Sep 17 00:00:00 2001
From: Paolo Barbolini <paolo.barbolini@m4ss.net>
Date: Thu, 2 Jan 2025 10:54:38 +0100
Subject: [PATCH] io: extend `Buf` length only after having read into it

---
 tokio/src/fs/file.rs     |  5 +++--
 tokio/src/io/blocking.rs | 32 +++++++++++++++-----------------
 2 files changed, 18 insertions(+), 19 deletions(-)

diff --git a/tokio/src/fs/file.rs b/tokio/src/fs/file.rs
index 63dd8af3e98..a02faafb407 100644
--- a/tokio/src/fs/file.rs
+++ b/tokio/src/fs/file.rs
@@ -7,6 +7,7 @@ use crate::io::blocking::{Buf, DEFAULT_MAX_BUF_SIZE};
 use crate::io::{AsyncRead, AsyncSeek, AsyncWrite, ReadBuf};
 use crate::sync::Mutex;
 
+use std::cmp;
 use std::fmt;
 use std::fs::{Metadata, Permissions};
 use std::future::Future;
@@ -600,11 +601,11 @@ impl AsyncRead for File {
                         return Poll::Ready(Ok(()));
                     }
 
-                    buf.ensure_capacity_for(dst, me.max_buf_size);
                     let std = me.std.clone();
 
+                    let max_buf_size = cmp::min(dst.remaining(), me.max_buf_size);
                     inner.state = State::Busy(spawn_blocking(move || {
-                        let res = buf.read_from(&mut &*std);
+                        let res = unsafe { buf.read_from(&mut &*std, max_buf_size) };
                         (Operation::Read(res), buf)
                     }));
                 }
diff --git a/tokio/src/io/blocking.rs b/tokio/src/io/blocking.rs
index f189136b52e..aa6046496c9 100644
--- a/tokio/src/io/blocking.rs
+++ b/tokio/src/io/blocking.rs
@@ -5,6 +5,7 @@ use std::cmp;
 use std::future::Future;
 use std::io;
 use std::io::prelude::*;
+use std::mem::MaybeUninit;
 use std::pin::Pin;
 use std::task::{ready, Context, Poll};
 
@@ -64,11 +65,11 @@ where
                         return Poll::Ready(Ok(()));
                     }
 
-                    buf.ensure_capacity_for(dst, DEFAULT_MAX_BUF_SIZE);
                     let mut inner = self.inner.take().unwrap();
 
+                    let max_buf_size = cmp::min(dst.remaining(), DEFAULT_MAX_BUF_SIZE);
                     self.state = State::Busy(sys::run(move || {
-                        let res = buf.read_from(&mut inner);
+                        let res = unsafe { buf.read_from(&mut inner, max_buf_size) };
                         (res, buf, inner)
                     }));
                 }
@@ -227,25 +228,22 @@ impl Buf {
         &self.buf[self.pos..]
     }
 
-    pub(crate) fn ensure_capacity_for(&mut self, bytes: &ReadBuf<'_>, max_buf_size: usize) {
+    /// Safety: `rd` must not read from the buffer passed to `read` and
+    /// must correctly report the length of the written portion of the buffer.
+    pub(crate) unsafe fn read_from<T: Read>(
+        &mut self,
+        rd: &mut T,
+        max_buf_size: usize,
+    ) -> io::Result<usize> {
         assert!(self.is_empty());
+        self.buf.reserve(max_buf_size);
 
-        let len = cmp::min(bytes.remaining(), max_buf_size);
-
-        if self.buf.len() < len {
-            self.buf.reserve(len - self.buf.len());
-        }
-
-        unsafe {
-            self.buf.set_len(len);
-        }
-    }
-
-    pub(crate) fn read_from<T: Read>(&mut self, rd: &mut T) -> io::Result<usize> {
-        let res = uninterruptibly!(rd.read(&mut self.buf));
+        let buf = &mut self.buf.spare_capacity_mut()[..max_buf_size];
+        let buf = unsafe { &mut *(buf as *mut [MaybeUninit<u8>] as *mut [u8]) };
+        let res = uninterruptibly!(rd.read(buf));
 
         if let Ok(n) = res {
-            self.buf.truncate(n);
+            unsafe { self.buf.set_len(n) }
         } else {
             self.buf.clear();
         }