mz_ore/netio/
read_exact.rs

Help
1// Copyright 2019 Tokio Contributors.
2// Copyright Materialize, Inc. and contributors. All rights reserved.
3//
4// Licensed under the Apache License, Version 2.0 (the "License");
5// you may not use this file except in compliance with the License.
6// You may obtain a copy of the License in the LICENSE file at the
7// root of this repository, or online at
8//
9//     http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing, software
12// distributed under the License is distributed on an "AS IS" BASIS,
13// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14// See the License for the specific language governing permissions and
15// limitations under the License.
16//
17// Portions of this file are derived from the ReadExact combinator in the Tokio
18// project. The original source code was retrieved on March 1, 2019 from:
19//
20//     https://github.com/tokio-rs/tokio/blob/195c4b04963742ecfff202ee9d0b72cc923aee81/tokio-io/src/io/read_exact.rs
21//
22// The original source code is subject to the terms of the MIT license, a copy
23// of which can be found in the LICENSE file at the root of this repository.
24
25use std::future::Future;
26use std::pin::Pin;
27use std::task::{Context, Poll};
28
29use futures::ready;
30use tokio::io::{self, AsyncRead, ReadBuf};
31
32/// A future which reads exactly enough bytes to fill a buffer, unless EOF is
33/// reached first.
34///
35/// Create a `ReadExactOrEof` struct by calling the [`read_exact_or_eof`]
36/// function.
37#[derive(Debug)]
38pub struct ReadExactOrEof<'a, A> {
39    reader: &'a mut A,
40    buf: &'a mut [u8],
41    pos: usize,
42}
43
44/// Creates a future which will read exactly enough bytes to fill `buf`, unless
45/// EOF is reached first. If a short read should be considered an error, use
46/// [`tokio::io::AsyncReadExt::read_exact`] instead.
47///
48/// The returned future will resolve to the number of bytes read.
49///
50/// In the case of an error the contents of the buffer are unspecified.
51pub fn read_exact_or_eof<'a, A>(reader: &'a mut A, buf: &'a mut [u8]) -> ReadExactOrEof<'a, A>
52where
53    A: AsyncRead,
54{
55    ReadExactOrEof {
56        reader,
57        buf,
58        pos: 0,
59    }
60}
61
62impl<A> Future for ReadExactOrEof<'_, A>
63where
64    A: AsyncRead + Unpin,
65{
66    type Output = io::Result<usize>;
67
68    fn poll(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output> {
69        while self.pos < self.buf.len() {
70            let me = &mut *self;
71            let mut buf = ReadBuf::new(&mut me.buf[me.pos..]);
72            ready!(Pin::new(&mut me.reader).poll_read(cx, &mut buf))?;
73            me.pos += buf.filled().len();
74            if buf.filled().len() == 0 {
75                break;
76            }
77        }
78        Poll::Ready(Ok(self.pos))
79    }
80}
mz_ore/netio/read_exact.rs

mz_ore/netio/
read_exact.rs
