io: remove unsafe from ReadToString (#2384)

1 files changed, 22 insertions, 14 deletions

diff --git a/tokio/src/io/util/read_to_string.rs b/tokio/src/io/util/read_to_string.rs
index e77d836d..cab0505a 100644
--- a/tokio/src/io/util/read_to_string.rs
+++ b/tokio/src/io/util/read_to_string.rs
@@ -4,7 +4,7 @@ use crate::io::AsyncRead;
 use std::future::Future;
 use std::pin::Pin;
 use std::task::{Context, Poll};
-use std::{io, mem, str};
+use std::{io, mem};
 
 cfg_io_util! {
     /// Future for the [`read_to_string`](super::AsyncReadExt::read_to_string) method.
@@ -25,7 +25,7 @@ where
     let start_len = buf.len();
     ReadToString {
         reader,
-        bytes: unsafe { mem::replace(buf.as_mut_vec(), Vec::new()) },
+        bytes: mem::replace(buf, String::new()).into_bytes(),
         buf,
         start_len,
     }
@@ -38,19 +38,20 @@ fn read_to_string_internal<R: AsyncRead + ?Sized>(
     bytes: &mut Vec<u8>,
     start_len: usize,
 ) -> Poll<io::Result<usize>> {
-    let ret = ready!(read_to_end_internal(reader, cx, bytes, start_len));
-    if str::from_utf8(&bytes).is_err() {
-        Poll::Ready(ret.and_then(|_| {
-            Err(io::Error::new(
+    let ret = ready!(read_to_end_internal(reader, cx, bytes, start_len))?;
+    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());
-        // Safety: `bytes` is a valid UTF-8 because `str::from_utf8` returned `Ok`.
-        mem::swap(unsafe { buf.as_mut_vec() }, bytes);
-        Poll::Ready(ret)
+            )))
+        }
     }
 }
 
@@ -67,7 +68,14 @@ where
             bytes,
             start_len,
         } = &mut *self;
-        read_to_string_internal(Pin::new(reader), cx, buf, bytes, *start_len)
+        let ret = read_to_string_internal(Pin::new(reader), cx, buf, bytes, *start_len);
+        if let Poll::Ready(Err(_)) = ret {
+            // Put back the original string.
+            bytes.truncate(*start_len);
+            **buf = String::from_utf8(mem::replace(bytes, Vec::new()))
+                .expect("original string no longer utf-8");
+        }
+        ret
     }
 }