libp2p_mdns/behaviour/
socket.rs

1// Copyright 2018 Parity Technologies (UK) Ltd.
2//
3// Permission is hereby granted, free of charge, to any person obtaining a
4// copy of this software and associated documentation files (the "Software"),
5// to deal in the Software without restriction, including without limitation
6// the rights to use, copy, modify, merge, publish, distribute, sublicense,
7// and/or sell copies of the Software, and to permit persons to whom the
8// Software is furnished to do so, subject to the following conditions:
9//
10// The above copyright notice and this permission notice shall be included in
11// all copies or substantial portions of the Software.
12//
13// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
14// OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
15// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
16// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
17// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
18// FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
19// DEALINGS IN THE SOFTWARE.
20
21use std::{
22    io::Error,
23    net::{SocketAddr, UdpSocket},
24    task::{Context, Poll},
25};
26
27/// Interface that must be implemented by the different runtimes to use the [`UdpSocket`] in async
28/// mode
29#[allow(unreachable_pub)] // Users should not depend on this.
30pub trait AsyncSocket: Unpin + Send + 'static {
31    /// Create the async socket from the [`std::net::UdpSocket`]
32    fn from_std(socket: UdpSocket) -> std::io::Result<Self>
33    where
34        Self: Sized;
35
36    /// Attempts to receive a single packet on the socket
37    /// from the remote address to which it is connected.
38    fn poll_read(
39        &mut self,
40        _cx: &mut Context,
41        _buf: &mut [u8],
42    ) -> Poll<Result<(usize, SocketAddr), Error>>;
43
44    /// Attempts to send data on the socket to a given address.
45    fn poll_write(
46        &mut self,
47        _cx: &mut Context,
48        _packet: &[u8],
49        _to: SocketAddr,
50    ) -> Poll<Result<(), Error>>;
51}
52
53#[cfg(feature = "tokio")]
54pub(crate) mod tokio {
55    use ::tokio::{io::ReadBuf, net::UdpSocket as TkUdpSocket};
56
57    use super::*;
58
59    /// Tokio ASync Socket`
60    pub(crate) type TokioUdpSocket = TkUdpSocket;
61    impl AsyncSocket for TokioUdpSocket {
62        fn from_std(socket: UdpSocket) -> std::io::Result<Self> {
63            socket.set_nonblocking(true)?;
64            TokioUdpSocket::from_std(socket)
65        }
66
67        fn poll_read(
68            &mut self,
69            cx: &mut Context,
70            buf: &mut [u8],
71        ) -> Poll<Result<(usize, SocketAddr), Error>> {
72            let mut rbuf = ReadBuf::new(buf);
73            match self.poll_recv_from(cx, &mut rbuf) {
74                Poll::Pending => Poll::Pending,
75                Poll::Ready(Err(err)) => Poll::Ready(Err(err)),
76                Poll::Ready(Ok(addr)) => Poll::Ready(Ok((rbuf.filled().len(), addr))),
77            }
78        }
79
80        fn poll_write(
81            &mut self,
82            cx: &mut Context,
83            packet: &[u8],
84            to: SocketAddr,
85        ) -> Poll<Result<(), Error>> {
86            match self.poll_send_to(cx, packet, to) {
87                Poll::Pending => Poll::Pending,
88                Poll::Ready(Err(err)) => Poll::Ready(Err(err)),
89                Poll::Ready(Ok(_len)) => Poll::Ready(Ok(())),
90            }
91        }
92    }
93}