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
149
150

//! Non-blocking TCP or Unix connect.
//!
//! This crate allows you to create a [`TcpStream`] or a [`UnixStream`] in a non-blocking way,
//! without waiting for the connection to become fully established.
//!
//! [`TcpStream`]: https://doc.rust-lang.org/stable/std/net/struct.TcpStream.html
//! [`UnixStream`]: https://doc.rust-lang.org/stable/std/os/unix/net/struct.UnixStream.html
//!
//! # Examples
//!
//! ```
//! use polling::{Event, Poller};
//! use std::time::Duration;
//!
//! // Create a pending TCP connection.
//! let stream = nb_connect::tcp(([127, 0, 0, 1], 80))?;
//!
//! // Create a poller that waits for the stream to become writable.
//! let poller = Poller::new()?;
//! poller.add(&stream, Event::writable(0))?;
//!
//! // Wait for at most 1 second.
//! if poller.wait(&mut Vec::new(), Some(Duration::from_secs(1)))? == 0 {
//!     println!("timeout");
//! } else if let Some(err) = stream.take_error()? {
//!     println!("error: {}", err);
//! } else {
//!     println!("connected");
//! }
//! # std::io::Result::Ok(())
//! ```

#![warn(missing_docs, missing_debug_implementations, rust_2018_idioms)]

use std::io;
use std::net::{SocketAddr, TcpStream};

use socket2::{Domain, Protocol, SockAddr, Socket, Type};

#[cfg(unix)]
use std::{os::unix::net::UnixStream, path::Path};

fn connect(addr: SockAddr, domain: Domain, protocol: Option<Protocol>) -> io::Result<Socket> {
    let sock_type = Type::STREAM;
    #[cfg(any(
        target_os = "android",
        target_os = "dragonfly",
        target_os = "freebsd",
        target_os = "fuchsia",
        target_os = "illumos",
        target_os = "linux",
        target_os = "netbsd",
        target_os = "openbsd"
    ))]
    // If we can, set nonblocking at socket creation for unix
    let sock_type = sock_type.nonblocking();
    // This automatically handles cloexec on unix, no_inherit on windows and nosigpipe on macos
    let socket = Socket::new(domain, sock_type, protocol)?;
    #[cfg(not(any(
        target_os = "android",
        target_os = "dragonfly",
        target_os = "freebsd",
        target_os = "fuchsia",
        target_os = "illumos",
        target_os = "linux",
        target_os = "netbsd",
        target_os = "openbsd"
    )))]
    // If the current platform doesn't support nonblocking at creation, enable it after creation
    socket.set_nonblocking(true)?;
    match socket.connect(&addr) {
        Ok(_) => {}
        #[cfg(unix)]
        Err(err) if err.raw_os_error() == Some(libc::EINPROGRESS) => {}
        Err(err) if err.kind() == io::ErrorKind::WouldBlock => {}
        Err(err) => return Err(err),
    }
    Ok(socket)
}

/// Creates a pending Unix connection to the specified path.
///
/// The returned Unix stream will be in non-blocking mode and in the process of connecting to the
/// specified path.
///
/// The stream becomes writable when connected.
///
/// # Examples
///
/// ```no_run
/// use polling::{Event, Poller};
/// use std::time::Duration;
///
/// // Create a pending Unix connection.
/// let stream = nb_connect::unix("/tmp/socket")?;
///
/// // Create a poller that waits for the stream to become writable.
/// let poller = Poller::new()?;
/// poller.add(&stream, Event::writable(0))?;
///
/// // Wait for at most 1 second.
/// if poller.wait(&mut Vec::new(), Some(Duration::from_secs(1)))? == 0 {
///     println!("timeout");
/// } else {
///     println!("connected");
/// }
/// # std::io::Result::Ok(())
/// ```
#[cfg(unix)]
pub fn unix<P: AsRef<Path>>(path: P) -> io::Result<UnixStream> {
    let socket = connect(SockAddr::unix(path)?, Domain::UNIX, None)?;
    Ok(socket.into())
}

/// Creates a pending TCP connection to the specified address.
///
/// The returned TCP stream will be in non-blocking mode and in the process of connecting to the
/// specified address.
///
/// The stream becomes writable when connected.
///
/// # Examples
///
/// ```
/// use polling::{Event, Poller};
/// use std::time::Duration;
///
/// // Create a pending TCP connection.
/// let stream = nb_connect::tcp(([127, 0, 0, 1], 80))?;
///
/// // Create a poller that waits for the stream to become writable.
/// let poller = Poller::new()?;
/// poller.add(&stream, Event::writable(0))?;
///
/// // Wait for at most 1 second.
/// if poller.wait(&mut Vec::new(), Some(Duration::from_secs(1)))? == 0 {
///     println!("timeout");
/// } else if let Some(err) = stream.take_error()? {
///     println!("error: {}", err);
/// } else {
///     println!("connected");
/// }
/// # std::io::Result::Ok(())
/// ```
pub fn tcp<A: Into<SocketAddr>>(addr: A) -> io::Result<TcpStream> {
    let addr = addr.into();
    let domain = Domain::for_address(addr);
    let socket = connect(addr.into(), domain, Some(Protocol::TCP))?;
    Ok(socket.into())
}