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

use std::fmt;
use std::marker::PhantomData;
use std::mem::ManuallyDrop;
use std::ops::Deref;
#[cfg(unix)]
use std::os::unix::io::{AsRawFd, FromRawFd};
#[cfg(windows)]
use std::os::windows::io::{AsRawSocket, FromRawSocket};

use crate::Socket;

/// A reference to a [`Socket`] that can be used to configure socket types other
/// than the `Socket` type itself.
///
/// This allows for example a [`TcpStream`], found in the standard library, to
/// be configured using all the additional methods found in the [`Socket`] API.
///
/// `SockRef` can be created from any socket type that implements [`AsRawFd`]
/// (Unix) or [`AsRawSocket`] (Windows) using the [`From`] implementation, but
/// the caller must ensure the file descriptor/socket is a valid.
///
/// [`TcpStream`]: std::net::TcpStream
/// [`AsRawFd`]: std::os::unix::io::AsRawFd
/// [`AsRawSocket`]: std::os::windows::io::AsRawSocket
///
/// # Examples
///
/// ```
/// use std::net::{TcpStream, SocketAddr};
///
/// use socket2::SockRef;
///
/// # fn main() -> Result<(), Box<dyn std::error::Error>> {
/// // Create `TcpStream` from the standard library.
/// let address: SocketAddr = "127.0.0.1:1234".parse()?;
/// # let b1 = std::sync::Arc::new(std::sync::Barrier::new(2));
/// # let b2 = b1.clone();
/// # let handle = std::thread::spawn(move || {
/// #    let listener = std::net::TcpListener::bind(address).unwrap();
/// #    b2.wait();
/// #    let (stream, _) = listener.accept().unwrap();
/// #    std::thread::sleep(std::time::Duration::from_millis(10));
/// #    drop(stream);
/// # });
/// # b1.wait();
/// let stream = TcpStream::connect(address)?;
///
/// // Create a `SockRef`erence to the stream.
/// let socket_ref = SockRef::from(&stream);
/// // Use `Socket::set_nodelay` on the stream.
/// socket_ref.set_nodelay(true)?;
/// drop(socket_ref);
///
/// assert_eq!(stream.nodelay()?, true);
/// # handle.join().unwrap();
/// # Ok(())
/// # }
/// ```
pub struct SockRef<'s> {
    /// Because this is a reference we don't own the `Socket`, however `Socket`
    /// closes itself when dropped, so we use `ManuallyDrop` to prevent it from
    /// closing itself.
    socket: ManuallyDrop<Socket>,
    /// Because we don't own the socket we need to ensure the socket remains
    /// open while we have a "reference" to it, the lifetime `'s` ensures this.
    _lifetime: PhantomData<&'s Socket>,
}

impl<'s> Deref for SockRef<'s> {
    type Target = Socket;

    fn deref(&self) -> &Self::Target {
        &self.socket
    }
}

#[cfg(unix)]
impl<'s, S> From<&'s S> for SockRef<'s>
where
    S: AsRawFd,
{
    /// The caller must ensure `S` is actually a socket.
    fn from(socket: &'s S) -> Self {
        SockRef {
            socket: ManuallyDrop::new(unsafe { Socket::from_raw_fd(socket.as_raw_fd()) }),
            _lifetime: PhantomData,
        }
    }
}

#[cfg(windows)]
impl<'s, S> From<&'s S> for SockRef<'s>
where
    S: AsRawSocket,
{
    /// See the `From<AsRawFd>` implementation.
    fn from(socket: &'s S) -> Self {
        SockRef {
            socket: ManuallyDrop::new(unsafe { Socket::from_raw_socket(socket.as_raw_socket()) }),
            _lifetime: PhantomData,
        }
    }
}

impl fmt::Debug for SockRef<'_> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("SockRef")
            .field("raw", &self.socket.inner)
            .field("local_addr", &self.socket.local_addr().ok())
            .field("peer_addr", &self.socket.peer_addr().ok())
            .finish()
    }
}