From 621790e1656d15cd46c971e046cbcf3a7d0d8351 Mon Sep 17 00:00:00 2001 From: Alice Ryhl 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 AsyncRead for Take { 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> { + 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(); +}