tokio/io/util/
read_line.rs

1use crate::io::util::read_until::read_until_internal;
2use crate::io::AsyncBufRead;
3
4use pin_project_lite::pin_project;
5use std::future::Future;
6use std::io;
7use std::marker::PhantomPinned;
8use std::mem;
9use std::pin::Pin;
10use std::string::FromUtf8Error;
11use std::task::{ready, Context, Poll};
12
13pin_project! {
14    /// Future for the [`read_line`](crate::io::AsyncBufReadExt::read_line) method.
15    #[derive(Debug)]
16    #[must_use = "futures do nothing unless you `.await` or poll them"]
17    pub struct ReadLine<'a, R: ?Sized> {
18        reader: &'a mut R,
19        // This is the buffer we were provided. It will be replaced with an empty string
20        // while reading to postpone utf-8 handling until after reading.
21        output: &'a mut String,
22        // The actual allocation of the string is moved into this vector instead.
23        buf: Vec<u8>,
24        // The number of bytes appended to buf. This can be less than buf.len() if
25        // the buffer was not empty when the operation was started.
26        read: usize,
27        // Make this future `!Unpin` for compatibility with async trait methods.
28        #[pin]
29        _pin: PhantomPinned,
30    }
31}
32
33pub(crate) fn read_line<'a, R>(reader: &'a mut R, string: &'a mut String) -> ReadLine<'a, R>
34where
35    R: AsyncBufRead + ?Sized + Unpin,
36{
37    ReadLine {
38        reader,
39        buf: mem::take(string).into_bytes(),
40        output: string,
41        read: 0,
42        _pin: PhantomPinned,
43    }
44}
45
46fn put_back_original_data(output: &mut String, mut vector: Vec<u8>, num_bytes_read: usize) {
47    let original_len = vector.len() - num_bytes_read;
48    vector.truncate(original_len);
49    *output = String::from_utf8(vector).expect("The original data must be valid utf-8.");
50}
51
52/// This handles the various failure cases and puts the string back into `output`.
53///
54/// The `truncate_on_io_error` `bool` is necessary because `read_to_string` and `read_line`
55/// disagree on what should happen when an IO error occurs.
56pub(super) fn finish_string_read(
57    io_res: io::Result<usize>,
58    utf8_res: Result<String, FromUtf8Error>,
59    read: usize,
60    output: &mut String,
61    truncate_on_io_error: bool,
62) -> Poll<io::Result<usize>> {
63    match (io_res, utf8_res) {
64        (Ok(num_bytes), Ok(string)) => {
65            debug_assert_eq!(read, 0);
66            *output = string;
67            Poll::Ready(Ok(num_bytes))
68        }
69        (Err(io_err), Ok(string)) => {
70            *output = string;
71            if truncate_on_io_error {
72                let original_len = output.len() - read;
73                output.truncate(original_len);
74            }
75            Poll::Ready(Err(io_err))
76        }
77        (Ok(num_bytes), Err(utf8_err)) => {
78            debug_assert_eq!(read, 0);
79            put_back_original_data(output, utf8_err.into_bytes(), num_bytes);
80
81            Poll::Ready(Err(io::Error::new(
82                io::ErrorKind::InvalidData,
83                "stream did not contain valid UTF-8",
84            )))
85        }
86        (Err(io_err), Err(utf8_err)) => {
87            put_back_original_data(output, utf8_err.into_bytes(), read);
88
89            Poll::Ready(Err(io_err))
90        }
91    }
92}
93
94pub(super) fn read_line_internal<R: AsyncBufRead + ?Sized>(
95    reader: Pin<&mut R>,
96    cx: &mut Context<'_>,
97    output: &mut String,
98    buf: &mut Vec<u8>,
99    read: &mut usize,
100) -> Poll<io::Result<usize>> {
101    let io_res = ready!(read_until_internal(reader, cx, b'\n', buf, read));
102    let utf8_res = String::from_utf8(mem::take(buf));
103
104    // At this point both buf and output are empty. The allocation is in utf8_res.
105
106    debug_assert!(buf.is_empty());
107    debug_assert!(output.is_empty());
108    finish_string_read(io_res, utf8_res, *read, output, false)
109}
110
111impl<R: AsyncBufRead + ?Sized + Unpin> Future for ReadLine<'_, R> {
112    type Output = io::Result<usize>;
113
114    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
115        let me = self.project();
116
117        read_line_internal(Pin::new(*me.reader), cx, me.output, me.buf, me.read)
118    }
119}
tokio/io/util/read_line.rs

tokio/io/util/
read_line.rs
