Skip to content

Commit

Permalink
io: fix panic in read_line
Browse files Browse the repository at this point in the history
The changes in this PR are inspired by tokio-rs#2384.

Fixes: tokio-rs#2532
  • Loading branch information
Darksonn committed May 16, 2020
1 parent f480659 commit f893e2e
Show file tree
Hide file tree
Showing 3 changed files with 121 additions and 19 deletions.
36 changes: 21 additions & 15 deletions tokio/src/io/util/read_line.rs
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,6 @@ use std::future::Future;
use std::io;
use std::mem;
use std::pin::Pin;
use std::str;
use std::task::{Context, Poll};

cfg_io_util! {
Expand All @@ -26,7 +25,7 @@ where
{
ReadLine {
reader,
bytes: unsafe { mem::replace(buf.as_mut_vec(), Vec::new()) },
bytes: mem::replace(buf, String::new()).into_bytes(),
buf,
read: 0,
}
Expand All @@ -39,20 +38,20 @@ pub(super) fn read_line_internal<R: AsyncBufRead + ?Sized>(
bytes: &mut Vec<u8>,
read: &mut usize,
) -> Poll<io::Result<usize>> {
let ret = ready!(read_until_internal(reader, cx, b'\n', bytes, read));
if str::from_utf8(&bytes).is_err() {
Poll::Ready(ret.and_then(|_| {
Err(io::Error::new(
let ret = ready!(read_until_internal(reader, cx, b'\n', bytes, read))?;
match String::from_utf8(mem::replace(bytes, Vec::new())) {
Ok(string) => {
debug_assert!(buf.is_empty());
*buf = string;
Poll::Ready(Ok(ret))
}
Err(e) => {
*bytes = e.into_bytes();
Poll::Ready(Err(io::Error::new(
io::ErrorKind::InvalidData,
"stream did not contain valid UTF-8",
))
}))
} else {
debug_assert!(buf.is_empty());
debug_assert_eq!(*read, 0);
// Safety: `bytes` is a valid UTF-8 because `str::from_utf8` returned `Ok`.
mem::swap(unsafe { buf.as_mut_vec() }, bytes);
Poll::Ready(ret)
)))
}
}
}

Expand All @@ -66,7 +65,14 @@ impl<R: AsyncBufRead + ?Sized + Unpin> Future for ReadLine<'_, R> {
bytes,
read,
} = &mut *self;
read_line_internal(Pin::new(reader), cx, buf, bytes, read)
let ret = ready!(read_line_internal(Pin::new(reader), cx, buf, bytes, read));
if ret.is_err() {
// If an error occurred, put the data back, but only if it is valid utf-8.
if let Ok(string) = String::from_utf8(mem::replace(bytes, Vec::new())) {
**buf = string;
}
}
Poll::Ready(ret)
}
}

Expand Down
52 changes: 50 additions & 2 deletions tokio/tests/io_read_line.rs
Original file line number Diff line number Diff line change
@@ -1,8 +1,9 @@
#![warn(rust_2018_idioms)]
#![cfg(feature = "full")]

use tokio::io::AsyncBufReadExt;
use tokio_test::assert_ok;
use std::io::ErrorKind;
use tokio::io::{AsyncBufReadExt, BufReader, Error};
use tokio_test::{assert_ok, io::Builder};

use std::io::Cursor;

Expand All @@ -27,3 +28,50 @@ async fn read_line() {
assert_eq!(n, 0);
assert_eq!(buf, "");
}

#[tokio::test]
async fn read_line_not_all_ready() {
let mock = Builder::new()
.read(b"Hello Wor")
.read(b"ld\nFizzBuz")
.read(b"z\n1\n2")
.build();

let mut read = BufReader::new(mock);

let mut line = "We say ".to_string();
let bytes = read.read_line(&mut line).await.unwrap();
assert_eq!(bytes, "Hello World\n".len());
assert_eq!(line.as_str(), "We say Hello World\n");

line = "I solve ".to_string();
let bytes = read.read_line(&mut line).await.unwrap();
assert_eq!(bytes, "FizzBuzz\n".len());
assert_eq!(line.as_str(), "I solve FizzBuzz\n");

line.clear();
let bytes = read.read_line(&mut line).await.unwrap();
assert_eq!(bytes, 2);
assert_eq!(line.as_str(), "1\n");

line.clear();
let bytes = read.read_line(&mut line).await.unwrap();
assert_eq!(bytes, 1);
assert_eq!(line.as_str(), "2");
}

#[tokio::test]
async fn read_line_fail() {
let mock = Builder::new()
.read(b"Hello Wor")
.read_error(Error::new(ErrorKind::Other, "The world has no end"))
.build();

let mut read = BufReader::new(mock);

let mut line = "Foo".to_string();
let err = read.read_line(&mut line).await.expect_err("Should fail");
assert_eq!(err.kind(), ErrorKind::Other);
assert_eq!(err.to_string(), "The world has no end");
assert_eq!(line.as_str(), "FooHello Wor");
}
52 changes: 50 additions & 2 deletions tokio/tests/io_read_until.rs
Original file line number Diff line number Diff line change
@@ -1,8 +1,9 @@
#![warn(rust_2018_idioms)]
#![cfg(feature = "full")]

use tokio::io::AsyncBufReadExt;
use tokio_test::assert_ok;
use std::io::ErrorKind;
use tokio::io::{AsyncBufReadExt, BufReader, Error};
use tokio_test::{assert_ok, io::Builder};

#[tokio::test]
async fn read_until() {
Expand All @@ -21,3 +22,50 @@ async fn read_until() {
assert_eq!(n, 0);
assert_eq!(buf, []);
}

#[tokio::test]
async fn read_until_not_all_ready() {
let mock = Builder::new()
.read(b"Hello Wor")
.read(b"ld#FizzBuz")
.read(b"z#1#2")
.build();

let mut read = BufReader::new(mock);

let mut chunk = b"We say ".to_vec();
let bytes = read.read_until(b'#', &mut chunk).await.unwrap();
assert_eq!(bytes, b"Hello World#".len());
assert_eq!(chunk, b"We say Hello World#");

chunk = b"I solve ".to_vec();
let bytes = read.read_until(b'#', &mut chunk).await.unwrap();
assert_eq!(bytes, "FizzBuzz\n".len());
assert_eq!(chunk, b"I solve FizzBuzz#");

chunk.clear();
let bytes = read.read_until(b'#', &mut chunk).await.unwrap();
assert_eq!(bytes, 2);
assert_eq!(chunk, b"1#");

chunk.clear();
let bytes = read.read_until(b'#', &mut chunk).await.unwrap();
assert_eq!(bytes, 1);
assert_eq!(chunk, b"2");
}

#[tokio::test]
async fn read_until_fail() {
let mock = Builder::new()
.read(b"Hello Wor")
.read_error(Error::new(ErrorKind::Other, "The world has no end"))
.build();

let mut read = BufReader::new(mock);

let mut chunk = b"Foo".to_vec();
let err = read.read_until(b'#', &mut chunk).await.expect_err("Should fail");
assert_eq!(err.kind(), ErrorKind::Other);
assert_eq!(err.to_string(), "The world has no end");
assert_eq!(chunk, b"FooHello Wor");
}

0 comments on commit f893e2e

Please sign in to comment.