pal_async/
socket.rs

1// Copyright (c) Microsoft Corporation.
2// Licensed under the MIT License.
3
4//! Socket-related functionality.
5
6// UNSAFETY: Reinterpreting an initialized `&mut [u8]` as
7// `&mut [MaybeUninit<u8>]` to pass to `socket2::Socket::peek`.
8#![expect(unsafe_code)]
9
10#[cfg(unix)]
11use super::fd;
12use super::interest::InterestSlot;
13use super::interest::PollEvents;
14use crate::driver::Driver;
15use crate::driver::PollImpl;
16use futures::AsyncRead;
17use futures::AsyncWrite;
18use parking_lot::Mutex;
19use std::fmt::Debug;
20use std::future::Future;
21use std::future::poll_fn;
22use std::io;
23use std::io::Read;
24use std::io::Write;
25use std::net::Shutdown;
26#[cfg(unix)]
27use std::os::unix::prelude::*;
28#[cfg(windows)]
29use std::os::windows::prelude::*;
30use std::path::Path;
31use std::pin::Pin;
32use std::sync::Arc;
33use std::task::Context;
34use std::task::Poll;
35use unix_socket::UnixStream;
36
37/// A trait for driving socket ready polling.
38pub trait SocketReadyDriver: Unpin {
39    /// The socket ready type.
40    type SocketReady: 'static + PollSocketReady;
41
42    /// Creates a new object for polling socket readiness.
43    #[cfg(windows)]
44    fn new_socket_ready(&self, socket: RawSocket) -> io::Result<Self::SocketReady>;
45    /// Creates a new object for polling socket readiness.
46    #[cfg(unix)]
47    fn new_socket_ready(&self, socket: RawFd) -> io::Result<Self::SocketReady>;
48}
49
50#[cfg(unix)]
51impl<T: fd::FdReadyDriver> SocketReadyDriver for T {
52    type SocketReady = <Self as fd::FdReadyDriver>::FdReady;
53
54    fn new_socket_ready(&self, socket: RawFd) -> io::Result<Self::SocketReady> {
55        self.new_fd_ready(socket)
56    }
57}
58
59/// A trait for polling socket readiness.
60pub trait PollSocketReady: Unpin + Send + Sync {
61    /// Polls a socket for readiness.
62    fn poll_socket_ready(
63        &mut self,
64        cx: &mut Context<'_>,
65        slot: InterestSlot,
66        events: PollEvents,
67    ) -> Poll<PollEvents>;
68
69    /// Clears cached socket readiness so that the next call to
70    /// `poll_socket_ready` will poll the OS again.
71    fn clear_socket_ready(&mut self, slot: InterestSlot);
72}
73
74#[cfg(unix)]
75impl<T: fd::PollFdReady> PollSocketReady for T {
76    fn poll_socket_ready(
77        &mut self,
78        cx: &mut Context<'_>,
79        slot: InterestSlot,
80        events: PollEvents,
81    ) -> Poll<PollEvents> {
82        self.poll_fd_ready(cx, slot, events)
83    }
84
85    fn clear_socket_ready(&mut self, slot: InterestSlot) {
86        self.clear_fd_ready(slot)
87    }
88}
89
90/// A polled socket.
91pub struct PolledSocket<T> {
92    poll: PollImpl<dyn PollSocketReady>, // must be first--some executors require that it's dropped before socket.
93    socket: T,
94}
95
96/// Trait implemented by socket types.
97pub trait AsSockRef: Unpin {
98    /// Returns a socket reference.
99    fn as_sock_ref(&self) -> socket2::SockRef<'_>;
100}
101
102impl<T: Unpin> AsSockRef for T
103where
104    for<'a> &'a T: Into<socket2::SockRef<'a>>,
105{
106    fn as_sock_ref(&self) -> socket2::SockRef<'_> {
107        self.into()
108    }
109}
110
111impl<T: AsSockRef> PolledSocket<T> {
112    /// Creates a new polled socket.
113    pub fn new(driver: &(impl ?Sized + Driver), socket: T) -> io::Result<Self> {
114        let sock_ref = socket.as_sock_ref();
115        sock_ref.set_nonblocking(true)?;
116        #[cfg(windows)]
117        let fd = sock_ref.as_raw_socket();
118        #[cfg(unix)]
119        let fd = sock_ref.as_raw_fd();
120        Ok(Self {
121            poll: driver.new_dyn_socket_ready(fd)?,
122            socket,
123        })
124    }
125
126    /// Extracts the inner socket.
127    pub fn into_inner(self) -> T {
128        let sock_ref = self.socket.as_sock_ref();
129        sock_ref.set_nonblocking(false).unwrap();
130        self.socket
131    }
132
133    /// Polls for peeking at incoming data without consuming it.
134    ///
135    /// On success, the peeked bytes are written to the start of `buf` and the
136    /// number of bytes peeked is returned. A return value of `Ok(0)` indicates
137    /// that the peer has closed the connection (or that `buf` was empty).
138    ///
139    /// The peeked data remains in the socket's receive buffer, so a subsequent
140    /// read (or peek) will observe the same bytes.
141    pub fn poll_peek(&mut self, cx: &mut Context<'_>, buf: &mut [u8]) -> Poll<io::Result<usize>> {
142        // Short-circuit an empty buffer so the caller gets the documented
143        // `Ok(0)` immediately, rather than blocking on read readiness.
144        if buf.is_empty() {
145            return Poll::Ready(Ok(0));
146        }
147        self.poll_io(cx, InterestSlot::Read, PollEvents::IN, |this| {
148            // SAFETY: Reinterpreting an initialized `&mut [u8]` as
149            // `&mut [MaybeUninit<u8>]` is sound: every `u8` is a valid
150            // `MaybeUninit<u8>`, and `peek` only writes initialized bytes. The
151            // caller's buffer therefore remains fully initialized.
152            let uninit = unsafe {
153                std::slice::from_raw_parts_mut(
154                    buf.as_mut_ptr().cast::<std::mem::MaybeUninit<u8>>(),
155                    buf.len(),
156                )
157            };
158            this.socket.as_sock_ref().peek(uninit)
159        })
160    }
161
162    /// Peeks at incoming data without consuming it, waiting until at least one
163    /// byte is available or the peer closes the connection.
164    ///
165    /// See [`PolledSocket::poll_peek`] for details.
166    pub async fn peek(&mut self, buf: &mut [u8]) -> io::Result<usize> {
167        poll_fn(|cx| self.poll_peek(cx, buf)).await
168    }
169}
170
171impl<T> PolledSocket<T> {
172    /// Gets a reference to the inner socket.
173    pub fn get(&self) -> &T {
174        &self.socket
175    }
176
177    /// Gets a mutable reference to the inner socket.
178    pub fn get_mut(&mut self) -> &mut T {
179        &mut self.socket
180    }
181
182    /// Converts the inner socket type.
183    pub fn convert<T2: From<T>>(self) -> PolledSocket<T2> {
184        PolledSocket {
185            socket: T2::from(self.socket),
186            poll: self.poll,
187        }
188    }
189}
190
191/// Trait for objects that can be polled for readiness.
192pub trait PollReady {
193    /// Polls an object for readiness.
194    fn poll_ready(&mut self, cx: &mut Context<'_>, events: PollEvents) -> Poll<PollEvents>;
195}
196
197/// Extension methods for implementations of [`PollReady`].
198pub trait PollReadyExt {
199    /// Waits for a socket or file to hang up.
200    fn wait_ready(&mut self, events: PollEvents) -> Ready<'_, Self>
201    where
202        Self: Unpin + Sized;
203}
204
205impl<T: PollReady + Unpin> PollReadyExt for T {
206    fn wait_ready(&mut self, events: PollEvents) -> Ready<'_, Self>
207    where
208        Self: Unpin + Sized,
209    {
210        Ready(self, events)
211    }
212}
213
214/// Future for [`PollReadyExt::wait_ready`].
215pub struct Ready<'a, T>(&'a mut T, PollEvents);
216
217impl<T: Unpin + PollReady> Future for Ready<'_, T> {
218    type Output = PollEvents;
219
220    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
221        let this = self.get_mut();
222        this.0.poll_ready(cx, this.1)
223    }
224}
225
226impl<T> PolledSocket<T> {
227    /// Calls nonblocking operation `f` when the socket has least one event in
228    /// `events` ready.
229    ///
230    /// Uses interest slot `slot` to allow multiple concurrent operations.
231    ///
232    /// If `f` returns `Err(err)` with `err.kind() ==
233    /// io::ErrorKind::WouldBlock`, then this re-polls the socket for readiness
234    /// and returns `Poll::Pending`.
235    pub fn poll_io<F, R>(
236        &mut self,
237        cx: &mut Context<'_>,
238        slot: InterestSlot,
239        events: PollEvents,
240        mut f: F,
241    ) -> Poll<io::Result<R>>
242    where
243        F: FnMut(&mut Self) -> io::Result<R>,
244    {
245        loop {
246            std::task::ready!(self.poll.poll_socket_ready(cx, slot, events));
247            match f(self) {
248                Err(err) if err.kind() == io::ErrorKind::WouldBlock => {
249                    self.poll.clear_socket_ready(slot);
250                }
251                r => break Poll::Ready(r),
252            }
253        }
254    }
255}
256
257impl<T: AsSockRef> PollReady for PolledSocket<T> {
258    fn poll_ready(&mut self, cx: &mut Context<'_>, events: PollEvents) -> Poll<PollEvents> {
259        self.poll.poll_socket_ready(cx, InterestSlot::Read, events)
260    }
261}
262
263impl<T> PolledSocket<T>
264where
265    T: AsSockRef + Read + Write,
266{
267    /// Splits the socket into a read and write half that can be used
268    /// concurrently.
269    ///
270    /// This is more flexible and efficient than
271    /// [`futures::io::AsyncReadExt::split`], since it avoids holding a lock
272    /// while calling into the kernel, and it provides access to the underlying
273    /// socket for more advanced operations.
274    pub fn split(self) -> (ReadHalf<T>, WriteHalf<T>) {
275        let inner = Arc::new(SplitInner {
276            poll: Mutex::new(self.poll),
277            socket: self.socket,
278        });
279        (
280            ReadHalf {
281                inner: inner.clone(),
282            },
283            WriteHalf { inner },
284        )
285    }
286}
287
288fn is_connect_incomplete_error(err: &io::Error) -> bool {
289    // This handles the Windows and AF_UNIX case.
290    if err.kind() == io::ErrorKind::WouldBlock {
291        return true;
292    }
293    // This handles the remaining cases on Linux.
294    #[cfg(unix)]
295    if err.raw_os_error() == Some(libc::EINPROGRESS) {
296        return true;
297    }
298    false
299}
300
301impl PolledSocket<socket2::Socket> {
302    /// Connects the socket to address `addr`.
303    pub async fn connect(&mut self, addr: &socket2::SockAddr) -> io::Result<()> {
304        match self.socket.connect(addr) {
305            Ok(()) => Ok(()),
306            Err(err) if is_connect_incomplete_error(&err) => {
307                self.poll.clear_socket_ready(InterestSlot::Write);
308                poll_fn(|cx| {
309                    self.poll
310                        .poll_socket_ready(cx, InterestSlot::Write, PollEvents::OUT)
311                })
312                .await;
313                if let Some(err) = self.socket.take_error()? {
314                    return Err(err);
315                }
316                Ok(())
317            }
318            Err(err) => Err(err),
319        }
320    }
321}
322
323impl PolledSocket<UnixStream> {
324    /// Creates a new connected Unix stream socket.
325    pub async fn connect_unix(
326        driver: &(impl ?Sized + Driver),
327        addr: impl AsRef<Path>,
328    ) -> io::Result<Self> {
329        let socket = socket2::Socket::new(socket2::Domain::UNIX, socket2::Type::STREAM, None)?;
330        let mut socket = PolledSocket::new(driver, socket)?;
331        socket
332            .connect(&socket2::SockAddr::unix(addr.as_ref())?)
333            .await?;
334        Ok(socket.convert())
335    }
336}
337
338impl PolledSocket<std::net::TcpStream> {
339    /// Creates a new connected TCP stream socket.
340    pub async fn connect_tcp(
341        driver: &(impl ?Sized + Driver),
342        addr: std::net::SocketAddr,
343    ) -> io::Result<Self> {
344        let domain = match addr {
345            std::net::SocketAddr::V4(_) => socket2::Domain::IPV4,
346            std::net::SocketAddr::V6(_) => socket2::Domain::IPV6,
347        };
348        let socket =
349            socket2::Socket::new(domain, socket2::Type::STREAM, Some(socket2::Protocol::TCP))?;
350        let mut socket = PolledSocket::new(driver, socket)?;
351        socket.connect(&addr.into()).await?;
352        Ok(socket.convert())
353    }
354}
355
356impl<T: AsSockRef + Read> AsyncRead for PolledSocket<T> {
357    fn poll_read(
358        mut self: Pin<&mut Self>,
359        cx: &mut Context<'_>,
360        buf: &mut [u8],
361    ) -> Poll<io::Result<usize>> {
362        // Short-circuit an empty buffer so the caller gets `Ok(0)`
363        // immediately, rather than blocking on read readiness.
364        if buf.is_empty() {
365            return Poll::Ready(Ok(0));
366        }
367        self.poll_io(cx, InterestSlot::Read, PollEvents::IN, |this| {
368            this.socket.read(buf)
369        })
370    }
371
372    fn poll_read_vectored(
373        mut self: Pin<&mut Self>,
374        cx: &mut Context<'_>,
375        bufs: &mut [io::IoSliceMut<'_>],
376    ) -> Poll<io::Result<usize>> {
377        self.poll_io(cx, InterestSlot::Read, PollEvents::IN, |this| {
378            this.socket.read_vectored(bufs)
379        })
380    }
381}
382
383impl<T: AsSockRef + Write> AsyncWrite for PolledSocket<T> {
384    fn poll_write(
385        mut self: Pin<&mut Self>,
386        cx: &mut Context<'_>,
387        buf: &[u8],
388    ) -> Poll<io::Result<usize>> {
389        self.poll_io(cx, InterestSlot::Write, PollEvents::OUT, |this| {
390            this.socket.write(buf)
391        })
392    }
393
394    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
395        self.poll_io(cx, InterestSlot::Write, PollEvents::OUT, |this| {
396            this.socket.flush()
397        })
398    }
399
400    fn poll_close(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> {
401        Poll::Ready(self.socket.as_sock_ref().shutdown(Shutdown::Write))
402    }
403
404    fn poll_write_vectored(
405        mut self: Pin<&mut Self>,
406        cx: &mut Context<'_>,
407        bufs: &[io::IoSlice<'_>],
408    ) -> Poll<io::Result<usize>> {
409        self.poll_io(cx, InterestSlot::Write, PollEvents::OUT, |this| {
410            this.socket.write_vectored(bufs)
411        })
412    }
413}
414
415/// Trait for listening sockets.
416pub trait Listener: AsSockRef {
417    /// The socket type.
418    type Socket: AsSockRef + Read + Write + Into<socket2::Socket>;
419    /// The socket address type.
420    type Address: Debug;
421
422    /// Accepts an incoming socket.
423    fn accept(&self) -> io::Result<(Self::Socket, Self::Address)>;
424    /// Returns the local address of the listener.
425    fn local_addr(&self) -> io::Result<Self::Address>;
426}
427
428impl<'a, T> Listener for &'a T
429where
430    T: Listener,
431    &'a T: AsSockRef,
432{
433    type Socket = T::Socket;
434    type Address = T::Address;
435
436    fn accept(&self) -> io::Result<(Self::Socket, Self::Address)> {
437        (**self).accept()
438    }
439
440    fn local_addr(&self) -> io::Result<Self::Address> {
441        (**self).local_addr()
442    }
443}
444
445macro_rules! listener {
446    ($ty:ty, $socket:ty, $addr:ty) => {
447        impl Listener for $ty {
448            type Socket = $socket;
449            type Address = $addr;
450            fn accept(&self) -> io::Result<(Self::Socket, Self::Address)> {
451                <$ty>::accept(self)
452            }
453            fn local_addr(&self) -> io::Result<Self::Address> {
454                <$ty>::local_addr(self)
455            }
456        }
457    };
458}
459
460listener!(
461    std::net::TcpListener,
462    std::net::TcpStream,
463    std::net::SocketAddr
464);
465
466#[cfg(unix)]
467listener!(
468    unix_socket::UnixListener,
469    UnixStream,
470    std::os::unix::net::SocketAddr
471);
472
473#[cfg(windows)]
474impl Listener for unix_socket::UnixListener {
475    type Socket = UnixStream;
476    type Address = ();
477
478    fn accept(&self) -> io::Result<(Self::Socket, Self::Address)> {
479        self.accept()
480    }
481
482    fn local_addr(&self) -> io::Result<Self::Address> {
483        Ok(())
484    }
485}
486
487listener!(socket2::Socket, socket2::Socket, socket2::SockAddr);
488
489impl PolledSocket<socket2::Socket> {
490    /// Listens for incoming connections.
491    pub fn listen(&self, backlog: i32) -> io::Result<()> {
492        self.socket.listen(backlog)
493    }
494}
495
496impl<T: Listener> PolledSocket<T> {
497    /// Polls for a new connection.
498    pub fn poll_accept(
499        &mut self,
500        cx: &mut Context<'_>,
501    ) -> Poll<io::Result<(T::Socket, T::Address)>> {
502        self.poll_io(cx, InterestSlot::Read, PollEvents::IN, |this| {
503            this.socket.accept()
504        })
505    }
506
507    /// Accepts a new connection.
508    pub async fn accept(&mut self) -> io::Result<(T::Socket, T::Address)> {
509        poll_fn(|cx| self.poll_accept(cx)).await
510    }
511}
512
513struct SplitInner<T> {
514    poll: Mutex<PollImpl<dyn PollSocketReady>>, // must be first--some executors require that it's dropped before socket.
515    socket: T,
516}
517
518/// The read half of a socket, via [`PolledSocket::split`].
519pub struct ReadHalf<T> {
520    inner: Arc<SplitInner<T>>,
521}
522
523impl<T> ReadHalf<T> {
524    /// Gets a reference to the inner socket.
525    pub fn get(&self) -> &T {
526        &self.inner.socket
527    }
528
529    /// Calls nonblocking operation `f` when the socket is ready for read.
530    ///
531    /// If `f` returns `Err(err)` with `err.kind() ==
532    /// io::ErrorKind::WouldBlock`, then this re-polls the socket for readiness
533    /// and returns `Poll::Pending`.
534    pub fn poll_io<F, R>(&mut self, cx: &mut Context<'_>, mut f: F) -> Poll<io::Result<R>>
535    where
536        F: FnMut(&mut Self) -> io::Result<R>,
537    {
538        loop {
539            std::task::ready!(self.inner.poll.lock().poll_socket_ready(
540                cx,
541                InterestSlot::Read,
542                PollEvents::IN
543            ));
544            match f(self) {
545                Err(err) if err.kind() == io::ErrorKind::WouldBlock => {
546                    self.inner
547                        .poll
548                        .lock()
549                        .clear_socket_ready(InterestSlot::Read);
550                }
551                r => break Poll::Ready(r),
552            }
553        }
554    }
555}
556
557/// The write half of a socket, via [`PolledSocket::split`].
558pub struct WriteHalf<T> {
559    inner: Arc<SplitInner<T>>,
560}
561
562impl<T> WriteHalf<T> {
563    /// Gets a reference to the inner socket.
564    pub fn get(&self) -> &T {
565        &self.inner.socket
566    }
567
568    /// Calls nonblocking operation `f` when the socket is ready for write.
569    ///
570    /// If `f` returns `Err(err)` with `err.kind() ==
571    /// io::ErrorKind::WouldBlock`, then this re-polls the socket for readiness
572    /// and returns `Poll::Pending`.
573    pub fn poll_io<F, R>(&mut self, cx: &mut Context<'_>, mut f: F) -> Poll<io::Result<R>>
574    where
575        F: FnMut(&mut Self) -> io::Result<R>,
576    {
577        loop {
578            std::task::ready!(self.inner.poll.lock().poll_socket_ready(
579                cx,
580                InterestSlot::Write,
581                PollEvents::OUT
582            ));
583            match f(self) {
584                Err(err) if err.kind() == io::ErrorKind::WouldBlock => {
585                    self.inner
586                        .poll
587                        .lock()
588                        .clear_socket_ready(InterestSlot::Write);
589                }
590                r => break Poll::Ready(r),
591            }
592        }
593    }
594}
595
596impl<T: AsSockRef> PollReady for ReadHalf<T> {
597    fn poll_ready(&mut self, cx: &mut Context<'_>, events: PollEvents) -> Poll<PollEvents> {
598        self.inner
599            .poll
600            .lock()
601            .poll_socket_ready(cx, InterestSlot::Read, events)
602    }
603}
604
605impl<T: AsSockRef> AsyncRead for ReadHalf<T> {
606    fn poll_read(
607        mut self: Pin<&mut Self>,
608        cx: &mut Context<'_>,
609        buf: &mut [u8],
610    ) -> Poll<io::Result<usize>> {
611        // Short-circuit an empty buffer so the caller gets `Ok(0)`
612        // immediately, rather than blocking on read readiness.
613        if buf.is_empty() {
614            return Poll::Ready(Ok(0));
615        }
616        self.poll_io(cx, |this| (&*this.inner.socket.as_sock_ref()).read(buf))
617    }
618
619    fn poll_read_vectored(
620        mut self: Pin<&mut Self>,
621        cx: &mut Context<'_>,
622        bufs: &mut [io::IoSliceMut<'_>],
623    ) -> Poll<io::Result<usize>> {
624        self.poll_io(cx, |this| {
625            (&*this.inner.socket.as_sock_ref()).read_vectored(bufs)
626        })
627    }
628}
629
630impl<T: AsSockRef> PollReady for WriteHalf<T> {
631    fn poll_ready(&mut self, cx: &mut Context<'_>, events: PollEvents) -> Poll<PollEvents> {
632        self.inner
633            .poll
634            .lock()
635            .poll_socket_ready(cx, InterestSlot::Write, events)
636    }
637}
638
639impl<T: AsSockRef> AsyncWrite for WriteHalf<T> {
640    fn poll_write(
641        mut self: Pin<&mut Self>,
642        cx: &mut Context<'_>,
643        buf: &[u8],
644    ) -> Poll<io::Result<usize>> {
645        self.poll_io(cx, |this| (&*this.inner.socket.as_sock_ref()).write(buf))
646    }
647
648    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
649        self.poll_io(cx, |this| (&*this.inner.socket.as_sock_ref()).flush())
650    }
651
652    fn poll_close(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> {
653        Poll::Ready(self.inner.socket.as_sock_ref().shutdown(Shutdown::Write))
654    }
655
656    fn poll_write_vectored(
657        mut self: Pin<&mut Self>,
658        cx: &mut Context<'_>,
659        bufs: &[io::IoSlice<'_>],
660    ) -> Poll<io::Result<usize>> {
661        self.poll_io(cx, |this| {
662            (&*this.inner.socket.as_sock_ref()).write_vectored(bufs)
663        })
664    }
665}
666
667#[cfg(test)]
668mod tests {
669    use super::PolledSocket;
670    use crate::DefaultDriver;
671    use futures::AsyncReadExt;
672    use futures::AsyncWriteExt;
673    use pal_async_test::async_test;
674    use unix_socket::UnixStream;
675
676    #[async_test]
677    async fn split(driver: DefaultDriver) {
678        let (a, b) = UnixStream::pair().unwrap();
679        let a = PolledSocket::new(&driver, a).unwrap();
680        let b = PolledSocket::new(&driver, b).unwrap();
681        let (mut ar, mut aw) = a.split();
682        let (br, mut bw) = b.split();
683        let copy = async {
684            futures::io::copy(br, &mut bw).await.unwrap();
685            bw.close().await.unwrap();
686        };
687        let rest = async {
688            aw.write_all(b"abc").await.unwrap();
689            let mut v = vec![0; 3];
690            ar.read_exact(&mut v).await.unwrap();
691            aw.write_all(b"def").await.unwrap();
692            aw.close().await.unwrap();
693            ar.read_to_end(&mut v).await.unwrap();
694            assert_eq!(&v, b"abcdef");
695        };
696        futures::future::join(copy, rest).await;
697    }
698
699    #[async_test]
700    async fn peek(driver: DefaultDriver) {
701        let (a, b) = UnixStream::pair().unwrap();
702        let mut a = PolledSocket::new(&driver, a).unwrap();
703        let mut b = PolledSocket::new(&driver, b).unwrap();
704
705        b.write_all(b"abc").await.unwrap();
706
707        // Peeking does not consume the data.
708        let mut buf = [0; 2];
709        let n = a.peek(&mut buf).await.unwrap();
710        assert_eq!(n, 2);
711        assert_eq!(&buf, b"ab");
712
713        // A second peek observes the same bytes.
714        let mut buf = [0; 3];
715        let n = a.peek(&mut buf).await.unwrap();
716        assert_eq!(n, 3);
717        assert_eq!(&buf, b"abc");
718
719        // A subsequent read still observes the peeked bytes.
720        let mut buf = [0; 3];
721        a.read_exact(&mut buf).await.unwrap();
722        assert_eq!(&buf, b"abc");
723
724        // Peeking after the peer closes returns 0.
725        drop(b);
726        let mut buf = [0; 1];
727        let n = a.peek(&mut buf).await.unwrap();
728        assert_eq!(n, 0);
729    }
730}
pal_async/socket.rs

pal_async/
socket.rs
