Add TCP functionality from net2

4 files changed, 209 insertions, 0 deletions

diff --git a/src/libstd/sys/common/net.rs b/src/libstd/sys/common/net.rs
index 1cb9303a9fc..0ac4056de8e 100644
--- a/src/libstd/sys/common/net.rs
+++ b/src/libstd/sys/common/net.rs
@@ -228,6 +228,54 @@ impl TcpStream {
     pub fn duplicate(&self) -> io::Result<TcpStream> {
         self.inner.duplicate().map(|s| TcpStream { inner: s })
     }
+
+    pub fn set_nodelay(&self, nodelay: bool) -> io::Result<()> {
+        setsockopt(&self.inner, c::IPPROTO_TCP, c::TCP_NODELAY, nodelay as c_int)
+    }
+
+    pub fn nodelay(&self) -> io::Result<bool> {
+        let raw: c_int = try!(getsockopt(&self.inner, c::IPPROTO_TCP, c::TCP_NODELAY));
+        Ok(raw != 0)
+    }
+
+    pub fn set_keepalive(&self, keepalive: Option<Duration>) -> io::Result<()> {
+        self.inner.set_keepalive(keepalive)
+    }
+
+    pub fn keepalive(&self) -> io::Result<Option<Duration>> {
+        self.inner.keepalive()
+    }
+
+    pub fn set_ttl(&self, ttl: u32) -> io::Result<()> {
+        setsockopt(&self.inner, c::IPPROTO_IP, c::IP_TTL, ttl as c_int)
+    }
+
+    pub fn ttl(&self) -> io::Result<u32> {
+        let raw: c_int = try!(getsockopt(&self.inner, c::IPPROTO_IP, c::IP_TTL));
+        Ok(raw as u32)
+    }
+
+    pub fn set_only_v6(&self, only_v6: bool) -> io::Result<()> {
+        setsockopt(&self.inner, c::IPPROTO_IPV6, c::IPV6_V6ONLY, only_v6 as c_int)
+    }
+
+    pub fn only_v6(&self) -> io::Result<bool> {
+        let raw: c_int = try!(getsockopt(&self.inner, c::IPPROTO_IPV6, c::IPV6_V6ONLY));
+        Ok(raw != 0)
+    }
+
+    pub fn take_error(&self) -> io::Result<Option<io::Error>> {
+        let raw: c_int = try!(getsockopt(&self.inner, c::SOL_SOCKET, c::SO_ERROR));
+        if raw == 0 {
+            Ok(None)
+        } else {
+            Ok(Some(io::Error::from_raw_os_error(raw as i32)))
+        }
+    }
+
+    pub fn set_nonblocking(&self, nonblocking: bool) -> io::Result<()> {
+        self.inner.set_nonblocking(nonblocking)
+    }
 }
 
 impl FromInner<Socket> for TcpStream {
@@ -307,6 +355,37 @@ impl TcpListener {
     pub fn duplicate(&self) -> io::Result<TcpListener> {
         self.inner.duplicate().map(|s| TcpListener { inner: s })
     }
+
+    pub fn set_ttl(&self, ttl: u32) -> io::Result<()> {
+        setsockopt(&self.inner, c::IPPROTO_IP, c::IP_TTL, ttl as c_int)
+    }
+
+    pub fn ttl(&self) -> io::Result<u32> {
+        let raw: c_int = try!(getsockopt(&self.inner, c::IPPROTO_IP, c::IP_TTL));
+        Ok(raw as u32)
+    }
+
+    pub fn set_only_v6(&self, only_v6: bool) -> io::Result<()> {
+        setsockopt(&self.inner, c::IPPROTO_IPV6, c::IPV6_V6ONLY, only_v6 as c_int)
+    }
+
+    pub fn only_v6(&self) -> io::Result<bool> {
+        let raw: c_int = try!(getsockopt(&self.inner, c::IPPROTO_IPV6, c::IPV6_V6ONLY));
+        Ok(raw != 0)
+    }
+
+    pub fn take_error(&self) -> io::Result<Option<io::Error>> {
+        let raw: c_int = try!(getsockopt(&self.inner, c::SOL_SOCKET, c::SO_ERROR));
+        if raw == 0 {
+            Ok(None)
+        } else {
+            Ok(Some(io::Error::from_raw_os_error(raw as i32)))
+        }
+    }
+
+    pub fn set_nonblocking(&self, nonblocking: bool) -> io::Result<()> {
+        self.inner.set_nonblocking(nonblocking)
+    }
 }
 
 impl FromInner<Socket> for TcpListener {
diff --git a/src/libstd/sys/unix/net.rs b/src/libstd/sys/unix/net.rs
index 16c369674f0..7a2ac7257af 100644
--- a/src/libstd/sys/unix/net.rs
+++ b/src/libstd/sys/unix/net.rs
@@ -35,6 +35,16 @@ use libc::SOCK_CLOEXEC;
 #[cfg(not(target_os = "linux"))]
 const SOCK_CLOEXEC: c_int = 0;
 
+#[cfg(any(target_os = "openbsd", taret_os = "freebsd"))]
+use libc::SO_KEEPALIVE as TCP_KEEPALIVE;
+#[cfg(any(target_os = "macos", taret_os = "ios"))]
+use libc::TCP_KEEPALIVE;
+#[cfg(not(any(target_os = "openbsd",
+              target_os = "freebsd",
+              target_os = "macos",
+              target_os = "ios")))]
+use libc::TCP_KEEPIDLE as TCP_KEEPALIVE;
+
 pub struct Socket(FileDesc);
 
 pub fn init() {}
@@ -168,6 +178,44 @@ impl Socket {
         try!(cvt(unsafe { libc::shutdown(self.0.raw(), how) }));
         Ok(())
     }
+
+    pub fn set_keepalive(&self, keepalive: Option<Duration>) -> io::Result<()> {
+        try!(setsockopt(self,
+                        libc::SOL_SOCKET,
+                        libc::SO_KEEPALIVE,
+                        keepalive.is_some() as libc::c_int));
+        if let Some(dur) = keepalive {
+            let mut raw = dur.as_secs();
+            if dur.subsec_nanos() > 0 {
+                raw = raw.saturating_add(1);
+            }
+
+            let raw = if raw > libc::c_int::max_value() as u64 {
+                libc::c_int::max_value()
+            } else {
+                raw as libc::c_int
+            };
+
+            try!(setsockopt(self, libc::IPPROTO_TCP, TCP_KEEPALIVE, raw));
+        }
+
+        Ok(())
+    }
+
+    pub fn keepalive(&self) -> io::Result<Option<Duration>> {
+        let raw: c_int = try!(getsockopt(self, libc::SOL_SOCKET, libc::SO_KEEPALIVE));
+        if raw == 0 {
+            return Ok(None);
+        }
+
+        let raw: c_int = try!(getsockopt(self, libc::IPPROTO_TCP, TCP_KEEPALIVE));
+        Ok(Some(Duration::from_secs(raw as u64)))
+    }
+
+    pub fn set_nonblocking(&self, nonblocking: bool) -> io::Result<()> {
+        let mut nonblocking = nonblocking as libc::c_ulong;
+        cvt(unsafe { libc::ioctl(*self.as_inner(), libc::FIONBIO, &mut nonblocking) }).map(|_| ())
+    }
 }
 
 impl AsInner<c_int> for Socket {
diff --git a/src/libstd/sys/windows/c.rs b/src/libstd/sys/windows/c.rs
index 5cbfec01bed..cc420763fd7 100644
--- a/src/libstd/sys/windows/c.rs
+++ b/src/libstd/sys/windows/c.rs
@@ -78,6 +78,13 @@ pub type SOCKET = ::os::windows::raw::SOCKET;
 pub type socklen_t = c_int;
 pub type ADDRESS_FAMILY = USHORT;
 
+pub type LPWSAOVERLAPPED_COMPLETION_ROUTINE =
+    Option<unsafe extern "system" fn(dwError: DWORD,
+                                     cbTransferred: DWORD,
+                                     lpOverlapped: LPWSAOVERLAPPED,
+                                     dwFlags: DWORD)>;
+pub type LPWSAOVERLAPPED = *mut OVERLAPPED;
+
 pub const TRUE: BOOL = 1;
 pub const FALSE: BOOL = 0;
 
@@ -114,6 +121,9 @@ pub const FILE_FLAG_OPEN_REPARSE_POINT: DWORD = 0x00200000;
 pub const FILE_FLAG_BACKUP_SEMANTICS: DWORD = 0x02000000;
 pub const SECURITY_SQOS_PRESENT: DWORD = 0x00100000;
 
+pub const SIO_KEEPALIVE_VALS: DWORD = 0x98000004;
+pub const FIONBIO: c_ulong = 0x8004667e;
+
 #[repr(C)]
 #[derive(Copy)]
 pub struct WIN32_FIND_DATAW {
@@ -775,6 +785,13 @@ pub struct in6_addr {
     pub s6_addr: [u8; 16],
 }
 
+#[repr(C)]
+pub struct tcp_keepalive {
+    pub onoff: c_ulong,
+    pub keepalivetime: c_ulong,
+    pub keepaliveinterval: c_ulong,
+}
+
 #[cfg(all(target_arch = "x86_64", target_env = "gnu"))]
 pub enum UNWIND_HISTORY_TABLE {}
 
@@ -833,6 +850,17 @@ extern "system" {
                       lpProtocolInfo: LPWSAPROTOCOL_INFO,
                       g: GROUP,
                       dwFlags: DWORD) -> SOCKET;
+    pub fn WSAIoctl(s: SOCKET,
+                    dwIoControlCode: DWORD,
+                    lpvInBuffer: LPVOID,
+                    cbInBuffer: DWORD,
+                    lpvOutBuffer: LPVOID,
+                    cbOutBuffer: DWORD,
+                    lpcbBytesReturned: LPDWORD,
+                    lpOverlapped: LPWSAOVERLAPPED,
+                    lpCompletionRoutine: LPWSAOVERLAPPED_COMPLETION_ROUTINE)
+                    -> c_int;
+    pub fn ioctlsocket(s: SOCKET, cmd: c_long, argp: *mut u_long) -> c_int;
     pub fn InitializeCriticalSection(CriticalSection: *mut CRITICAL_SECTION);
     pub fn EnterCriticalSection(CriticalSection: *mut CRITICAL_SECTION);
     pub fn TryEnterCriticalSection(CriticalSection: *mut CRITICAL_SECTION) -> BOOLEAN;
diff --git a/src/libstd/sys/windows/net.rs b/src/libstd/sys/windows/net.rs
index 49ba8e9c659..be13657aaf4 100644
--- a/src/libstd/sys/windows/net.rs
+++ b/src/libstd/sys/windows/net.rs
@@ -185,6 +185,60 @@ impl Socket {
         try!(cvt(unsafe { c::shutdown(self.0, how) }));
         Ok(())
     }
+
+    pub fn set_keepalive(&self, keepalive: Option<Duration>) -> io::Result<()> {
+        let ms = keepalive.map(sys::dur2timeout).unwrap_or(c::INFINITE);
+        let ka = c::tcp_keepalive {
+            onoff: keepalive.is_some() as c::c_ulong,
+            keepalivetime: ms as c::c_ulong,
+            keepaliveinterval: ms as c::c_ulong,
+        };
+        sys::cvt(unsafe {
+            c::WSAIoctl(self.0,
+                        c::SIO_KEEPALIVE_VALS,
+                        &ka as *const _ as *mut _,
+                        mem::size_of_val(&ka) as c::DWORD,
+                        0 as *mut _,
+                        0,
+                        0 as *mut _,
+                        0 as *mut _,
+                        None)
+        }).map(|_| ())
+    }
+
+    pub fn keepalive(&self) -> io::Result<Option<Duration>> {
+        let mut ka = c::tcp_keepalive {
+            onoff: 0,
+            keepalivetime: 0,
+            keepaliveinterval: 0,
+        };
+        try!(sys::cvt(unsafe {
+            WSAIoctl(self.0,
+                     c::SIO_KEEPALIVE_VALS,
+                     0 as *mut _,
+                     0,
+                     &mut ka as *mut _ as *mut _,
+                     mem::size_of_val(&ka) as c::DWORD,
+                     0 as *mut _,
+                     0 as *mut _,
+                     None)
+        }));
+
+        if ka.onoff == 0 {
+            Ok(None)
+        } else {
+            let secs = ka.keepaliveinterval / 1000;
+            let nsec = (ka.keepaliveinterval % 1000) * 1000000;
+            Ok(Some(Duration::new(secs as u64, nsec as u32)))
+        }
+    }
+
+    pub fn set_nonblocking(&self, nonblocking: bool) -> io::Result<()> {
+        let mut nonblocking = nonblocking as c::c_ulong;
+        sys::cvt(unsafe {
+            c::ioctlsocket(self.0, c::FIONBIO as c::c_int, &mut nonblocking)
+        }).map(|_| ())
+    }
 }
 
 impl Drop for Socket {