async_peek/
lib.rs

1/**************************************************************************************************
2 *                                                                                                *
3 * This Source Code Form is subject to the terms of the Mozilla Public                            *
4 * License, v. 2.0. If a copy of the MPL was not distributed with this                            *
5 * file, You can obtain one at http://mozilla.org/MPL/2.0/.                                       *
6 *                                                                                                *
7 **************************************************************************************************/
8
9// ======================================== Documentation ======================================= \\
10
11//! This crate provides [`AsyncPeek`], a trait to read data asynchronously without removing it
12//! from the queue (like when using the blocking methods [`std::net::TcpStream::peek()`] and
13//! [`std::net::UdpSocket::peek()`]).
14
15// =========================================== Imports ========================================== \\
16
17use cfg_if::cfg_if;
18use core::cmp;
19use core::future::Future;
20use core::ops::{Deref, DerefMut};
21use core::pin::Pin;
22use core::task::{Context, Poll};
23use std::io::Error;
24
25#[cfg(feature = "async-io")]
26use std::net::{TcpStream, UdpSocket};
27
28// ========================================= Interfaces ========================================= \\
29
30/// Read data asynchronously without removing it from the queue.
31pub trait AsyncPeek {
32    /// Attempts to read data into `buf` without removing it from the queue.
33    ///
34    /// Returns the number of bytes read on success, or [`io::Error`] if an error is encountered.
35    ///
36    /// If no data is available, the current task is registered to be notified when data becomes
37    /// available or the stream is closed, and `Poll::Pending` is returned.
38    ///
39    /// [`io::Error`]: std::io::Error
40    fn poll_peek(
41        self: Pin<&mut Self>,
42        ctx: &mut Context,
43        buf: &mut [u8],
44    ) -> Poll<Result<usize, Error>>;
45}
46
47/// An extension trait which adds utility methods to [`AsyncPeek`] types.
48pub trait AsyncPeekExt: AsyncPeek {
49    /// Tries to read data into `buf` without removing it from the queue.
50    ///
51    /// Returns the number of bytes read, or [`io::Error`] if an error is encountered.
52    ///
53    /// ## Examples
54    ///
55    /// ```rust
56    /// # smol::block_on(async {
57    /// #
58    /// use async_peek::{AsyncPeek, AsyncPeekExt};
59    /// use smol::Async;
60    /// use std::net::{TcpStream, ToSocketAddrs};
61    ///
62    /// let addr = "127.0.0.1:0".to_socket_addrs()?.next().unwrap();
63    /// let mut stream = Async::<TcpStream>::connect(addr).await?;
64    /// let mut buf = [0; 64];
65    ///
66    /// let n = stream.peek(&mut buf).await?;
67    /// println!("Peeked {} bytes", n);
68    /// #
69    /// # std::io::Result::Ok(()) });
70    /// ```
71    ///
72    /// [`io::Error`]: std::io::Error
73    fn peek<'peek>(&'peek mut self, buf: &'peek mut [u8]) -> Peek<'peek, Self>
74    where
75        Self: Unpin,
76    {
77        Peek { peek: self, buf }
78    }
79}
80
81// ============================================ Types =========================================== \\
82
83/// Future for the [`peek()`] function.
84///
85/// [`peek()`]: AsyncPeekExt::peek()
86pub struct Peek<'peek, P: ?Sized> {
87    peek: &'peek mut P,
88    buf: &'peek mut [u8],
89}
90
91// ======================================== macro_rules! ======================================== \\
92
93#[allow(unused)]
94macro_rules! impl_for_net {
95    ($net:ty) => {
96        impl AsyncPeek for $net {
97            fn poll_peek(
98                self: Pin<&mut Self>,
99                ctx: &mut Context,
100                buf: &mut [u8],
101            ) -> Poll<Result<usize, Error>> {
102                let fut = (&*self).peek(buf);
103                ufut::pin!(fut);
104
105                fut.poll(ctx)
106            }
107        }
108    };
109}
110
111// ======================================= impl AsyncPeek ======================================= \\
112
113impl AsyncPeek for &[u8] {
114    fn poll_peek(
115        self: Pin<&mut Self>,
116        _: &mut Context,
117        buf: &mut [u8],
118    ) -> Poll<Result<usize, Error>> {
119        let len = cmp::min(buf.len(), self.len());
120        buf[0..len].copy_from_slice(&self[0..len]);
121
122        Poll::Ready(Ok(len))
123    }
124}
125
126impl<T> AsyncPeek for &mut T
127where
128    T: AsyncPeek + Unpin + ?Sized,
129{
130    fn poll_peek(
131        mut self: Pin<&mut Self>,
132        ctx: &mut Context,
133        buf: &mut [u8],
134    ) -> Poll<Result<usize, Error>> {
135        Pin::new(&mut **self).poll_peek(ctx, buf)
136    }
137}
138
139impl<T> AsyncPeek for Box<T>
140where
141    T: AsyncPeek + Unpin + ?Sized,
142{
143    fn poll_peek(
144        mut self: Pin<&mut Self>,
145        ctx: &mut Context,
146        buf: &mut [u8],
147    ) -> Poll<Result<usize, Error>> {
148        Pin::new(&mut **self).poll_peek(ctx, buf)
149    }
150}
151
152impl<T> AsyncPeek for Pin<T>
153where
154    T: DerefMut + Unpin,
155    <T as Deref>::Target: AsyncPeek,
156{
157    fn poll_peek(
158        self: Pin<&mut Self>,
159        ctx: &mut Context,
160        buf: &mut [u8],
161    ) -> Poll<Result<usize, Error>> {
162        self.get_mut().as_mut().poll_peek(ctx, buf)
163    }
164}
165
166
167
168cfg_if! {
169    if #[cfg(feature = "async-io")] {
170        impl_for_net!(async_io::Async<TcpStream>);
171        impl_for_net!(async_io::Async<UdpSocket>);
172        impl_for_net!(&async_io::Async<TcpStream>);
173        impl_for_net!(&async_io::Async<UdpSocket>);
174    }
175}
176
177cfg_if! {
178    if #[cfg(feature = "async-net")] {
179        impl_for_net!(async_net::TcpStream);
180        impl_for_net!(async_net::UdpSocket);
181        impl_for_net!(&async_net::TcpStream);
182        impl_for_net!(&async_net::UdpSocket);
183    }
184}
185
186// ====================================== impl AsyncPeekExt ===================================== \\
187
188impl<P: AsyncPeek + ?Sized> AsyncPeekExt for P {}
189
190// ========================================= impl Future ======================================== \\
191
192impl<P: AsyncPeek + Unpin> Future for Peek<'_, P> {
193    type Output = Result<usize, Error>;
194
195    fn poll(mut self: Pin<&mut Self>, ctx: &mut Context) -> Poll<Self::Output> {
196        let this = &mut *self;
197        Pin::new(&mut this.peek).poll_peek(ctx, this.buf)
198    }
199}
async_peek/lib.rs

async_peek/
lib.rs
