From 621790e1656d15cd46c971e046cbcf3a7d0d8351 Mon Sep 17 00:00:00 2001
From: Alice Ryhl <alice@ryhl.io>
Date: Thu, 27 Jan 2022 11:45:04 +0100
Subject: [PATCH] io: fix `take` when using evil reader (#4428)

---
 tokio/src/io/util/take.rs |  5 +++++
 tokio/tests/io_take.rs    | 30 +++++++++++++++++++++++++++++-
 2 files changed, 34 insertions(+), 1 deletion(-)

diff --git a/tokio/src/io/util/take.rs b/tokio/src/io/util/take.rs
index b5e90c936fa..d73512bdfaf 100644
--- a/tokio/src/io/util/take.rs
+++ b/tokio/src/io/util/take.rs
@@ -84,9 +84,14 @@ impl<R: AsyncRead> AsyncRead for Take<R> {
             return Poll::Ready(Ok(()));
         }
 
+        let buf_ptr = buf.filled().as_ptr();
+
         let me = self.project();
         let mut b = buf.take(*me.limit_ as usize);
+
         ready!(me.inner.poll_read(cx, &mut b))?;
+        assert_eq!(b.filled().as_ptr(), buf_ptr);
+
         let n = b.filled().len();
 
         // We need to update the original ReadBuf
diff --git a/tokio/tests/io_take.rs b/tokio/tests/io_take.rs
index 683606f7272..45c61f276b1 100644
--- a/tokio/tests/io_take.rs
+++ b/tokio/tests/io_take.rs
@@ -1,7 +1,9 @@
 #![warn(rust_2018_idioms)]
 #![cfg(feature = "full")]
 
-use tokio::io::AsyncReadExt;
+use std::pin::Pin;
+use std::task::{Context, Poll};
+use tokio::io::{self, AsyncRead, AsyncReadExt, ReadBuf};
 use tokio_test::assert_ok;
 
 #[tokio::test]
@@ -14,3 +16,29 @@ async fn take() {
     assert_eq!(n, 4);
     assert_eq!(&buf, &b"hell\0\0"[..]);
 }
+
+struct BadReader;
+
+impl AsyncRead for BadReader {
+    fn poll_read(
+        self: Pin<&mut Self>,
+        _cx: &mut Context<'_>,
+        read_buf: &mut ReadBuf<'_>,
+    ) -> Poll<io::Result<()>> {
+        let vec = vec![0; 10];
+
+        let mut buf = ReadBuf::new(vec.leak());
+        buf.put_slice(&[123; 10]);
+        *read_buf = buf;
+
+        Poll::Ready(Ok(()))
+    }
+}
+
+#[tokio::test]
+#[should_panic]
+async fn bad_reader_fails() {
+    let mut buf = Vec::with_capacity(10);
+
+    BadReader.take(10).read_buf(&mut buf).await.unwrap();
+}