diff --git a/futures-util/src/io/lines.rs b/futures-util/src/io/lines.rs index 8c4d17c58..0a1abf4bf 100644 --- a/futures-util/src/io/lines.rs +++ b/futures-util/src/io/lines.rs @@ -35,6 +35,7 @@ impl Stream for Lines { fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { let this = self.project(); let n = ready!(read_line_internal(this.reader, cx, this.buf, this.bytes, this.read))?; + *this.read = 0; if n == 0 && this.buf.is_empty() { return Poll::Ready(None); } diff --git a/futures-util/src/io/read_line.rs b/futures-util/src/io/read_line.rs index b4483223d..43942add5 100644 --- a/futures-util/src/io/read_line.rs +++ b/futures-util/src/io/read_line.rs @@ -18,13 +18,14 @@ pub struct ReadLine<'a, R: ?Sized> { buf: &'a mut String, bytes: Vec, read: usize, + finished: bool, } impl Unpin for ReadLine<'_, R> {} impl<'a, R: AsyncBufRead + ?Sized + Unpin> ReadLine<'a, R> { pub(super) fn new(reader: &'a mut R, buf: &'a mut String) -> Self { - Self { reader, bytes: mem::take(buf).into_bytes(), buf, read: 0 } + Self { reader, bytes: mem::take(buf).into_bytes(), buf, read: 0, finished: false } } } @@ -35,26 +36,42 @@ pub(super) fn read_line_internal( bytes: &mut Vec, read: &mut usize, ) -> Poll> { - let ret = ready!(read_until_internal(reader, cx, b'\n', bytes, read)); - if str::from_utf8(bytes).is_err() { - bytes.clear(); - Poll::Ready(ret.and_then(|_| { - 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) + let mut ret = ready!(read_until_internal(reader, cx, b'\n', bytes, read)); + if str::from_utf8(&bytes[bytes.len() - *read..bytes.len()]).is_err() { + bytes.truncate(bytes.len() - *read); + if ret.is_ok() { + ret = Err(io::Error::new( + io::ErrorKind::InvalidData, + "stream did not contain valid UTF-8", + )); + } } + *read = 0; + // Safety: `bytes` is valid UTF-8 because it was taken from a String + // and the newly read bytes are either valid UTF-8 or have been removed. + mem::swap(unsafe { buf.as_mut_vec() }, bytes); + Poll::Ready(ret) } impl Future for ReadLine<'_, R> { type Output = io::Result; fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { - let Self { reader, buf, bytes, read } = &mut *self; - read_line_internal(Pin::new(reader), cx, buf, bytes, read) + let Self { reader, buf, bytes, read, finished: _ } = &mut *self; + let ret = ready!(read_line_internal(Pin::new(reader), cx, buf, bytes, read)); + self.finished = true; + Poll::Ready(ret) + } +} + +impl Drop for ReadLine<'_, R> { + fn drop(&mut self) { + // restore old string contents + if !self.finished { + self.bytes.truncate(self.bytes.len() - self.read); + // Safety: `bytes` is valid UTF-8 because it was taken from a String + // and the newly read bytes have been removed. + mem::swap(unsafe { self.buf.as_mut_vec() }, &mut self.bytes); + } } } diff --git a/futures-util/src/io/read_until.rs b/futures-util/src/io/read_until.rs index d6121d6f0..adc359db5 100644 --- a/futures-util/src/io/read_until.rs +++ b/futures-util/src/io/read_until.rs @@ -3,7 +3,6 @@ use futures_core::ready; use futures_core::task::{Context, Poll}; use futures_io::AsyncBufRead; use std::io; -use std::mem; use std::pin::Pin; use std::vec::Vec; @@ -46,7 +45,7 @@ pub(super) fn read_until_internal( reader.as_mut().consume(used); *read += used; if done || used == 0 { - return Poll::Ready(Ok(mem::replace(read, 0))); + return Poll::Ready(Ok(*read)); } } }