1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148

//! [`UdpSocket`](crate::net::UdpSocket) split support.
//!
//! The [`split`](method@crate::net::UdpSocket::split) method splits a
//! `UdpSocket` into a receive half and a send half, which can be used to
//! receive and send datagrams concurrently, even from two different tasks.
//!
//! The halves provide access to the underlying socket, implementing
//! `AsRef<UdpSocket>`. This allows you to call `UdpSocket` methods that takes
//! `&self`, e.g., to get local address, to get and set socket options, to join
//! or leave multicast groups, etc.
//!
//! The halves can be reunited to the original socket with their `reunite`
//! methods.

use crate::future::poll_fn;
use crate::net::udp::UdpSocket;

use std::error::Error;
use std::fmt;
use std::io;
use std::net::SocketAddr;
use std::sync::Arc;

/// The send half after [`split`](super::UdpSocket::split).
///
/// Use [`send_to`](method@Self::send_to) or [`send`](method@Self::send) to send
/// datagrams.
#[derive(Debug)]
pub struct SendHalf(Arc<UdpSocket>);

/// The recv half after [`split`](super::UdpSocket::split).
///
/// Use [`recv_from`](method@Self::recv_from) or [`recv`](method@Self::recv) to receive
/// datagrams.
#[derive(Debug)]
pub struct RecvHalf(Arc<UdpSocket>);

pub(crate) fn split(socket: UdpSocket) -> (RecvHalf, SendHalf) {
    let shared = Arc::new(socket);
    let send = shared.clone();
    let recv = shared;
    (RecvHalf(recv), SendHalf(send))
}

/// Error indicating that two halves were not from the same socket, and thus could
/// not be `reunite`d.
#[derive(Debug)]
pub struct ReuniteError(pub SendHalf, pub RecvHalf);

impl fmt::Display for ReuniteError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(
            f,
            "tried to reunite halves that are not from the same socket"
        )
    }
}

impl Error for ReuniteError {}

fn reunite(s: SendHalf, r: RecvHalf) -> Result<UdpSocket, ReuniteError> {
    if Arc::ptr_eq(&s.0, &r.0) {
        drop(r);
        // Only two instances of the `Arc` are ever created, one for the
        // receiver and one for the sender, and those `Arc`s are never exposed
        // externally. And so when we drop one here, the other one must be the
        // only remaining one.
        Ok(Arc::try_unwrap(s.0).expect("udp: try_unwrap failed in reunite"))
    } else {
        Err(ReuniteError(s, r))
    }
}

impl RecvHalf {
    /// Attempts to put the two "halves" of a `UdpSocket` back together and
    /// recover the original socket. Succeeds only if the two "halves"
    /// originated from the same call to `UdpSocket::split`.
    pub fn reunite(self, other: SendHalf) -> Result<UdpSocket, ReuniteError> {
        reunite(other, self)
    }

    /// Returns a future that receives a single datagram on the socket. On success,
    /// the future resolves to the number of bytes read and the origin.
    ///
    /// The function must be called with valid byte array `buf` of sufficient size
    /// to hold the message bytes. If a message is too long to fit in the supplied
    /// buffer, excess bytes may be discarded.
    pub async fn recv_from(&mut self, buf: &mut [u8]) -> io::Result<(usize, SocketAddr)> {
        poll_fn(|cx| self.0.poll_recv_from(cx, buf)).await
    }

    /// Returns a future that receives a single datagram message on the socket from
    /// the remote address to which it is connected. On success, the future will resolve
    /// to the number of bytes read.
    ///
    /// The function must be called with valid byte array `buf` of sufficient size to
    /// hold the message bytes. If a message is too long to fit in the supplied buffer,
    /// excess bytes may be discarded.
    ///
    /// The [`connect`] method will connect this socket to a remote address. The future
    /// will fail if the socket is not connected.
    ///
    /// [`connect`]: super::UdpSocket::connect
    pub async fn recv(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        poll_fn(|cx| self.0.poll_recv(cx, buf)).await
    }
}

impl SendHalf {
    /// Attempts to put the two "halves" of a `UdpSocket` back together and
    /// recover the original socket. Succeeds only if the two "halves"
    /// originated from the same call to `UdpSocket::split`.
    pub fn reunite(self, other: RecvHalf) -> Result<UdpSocket, ReuniteError> {
        reunite(self, other)
    }

    /// Returns a future that sends data on the socket to the given address.
    /// On success, the future will resolve to the number of bytes written.
    ///
    /// The future will resolve to an error if the IP version of the socket does
    /// not match that of `target`.
    pub async fn send_to(&mut self, buf: &[u8], target: &SocketAddr) -> io::Result<usize> {
        poll_fn(|cx| self.0.poll_send_to(cx, buf, target)).await
    }

    /// Returns a future that sends data on the socket to the remote address to which it is connected.
    /// On success, the future will resolve to the number of bytes written.
    ///
    /// The [`connect`] method will connect this socket to a remote address. The future
    /// will resolve to an error if the socket is not connected.
    ///
    /// [`connect`]: super::UdpSocket::connect
    pub async fn send(&mut self, buf: &[u8]) -> io::Result<usize> {
        poll_fn(|cx| self.0.poll_send(cx, buf)).await
    }
}

impl AsRef<UdpSocket> for SendHalf {
    fn as_ref(&self) -> &UdpSocket {
        &self.0
    }
}

impl AsRef<UdpSocket> for RecvHalf {
    fn as_ref(&self) -> &UdpSocket {
        &self.0
    }
}