net: add keepalive support to TcpSocket

Signed-off-by: Eliza Weisman <[email protected]>

Author: hawkw

src/net/mod.rs

@@ -8,7 +8,7 @@
 //! [portability guidelines]: ../struct.Poll.html#portability
 
 mod tcp;
-pub use self::tcp::{TcpListener, TcpSocket, TcpStream};
+pub use self::tcp::{TcpListener, TcpSocket, TcpStream, TcpKeepalive};
 
 mod udp;
 pub use self::udp::UdpSocket;

src/net/tcp/mod.rs

@@ -2,7 +2,7 @@ mod listener;
 pub use self::listener::TcpListener;
 
 mod socket;
-pub use self::socket::TcpSocket;
+pub use self::socket::{TcpSocket, TcpKeepalive};
 
 mod stream;
 pub use self::stream::TcpStream;

src/net/tcp/socket.rs

@@ -22,22 +22,45 @@ pub struct TcpSocket {
     sys: sys::tcp::TcpSocket,
 }
 
+/// Configures a socket's TCP keepalive parameters.
+#[derive(Debug, Default, Clone)]
+pub struct TcpKeepalive {
+    pub(crate) time: Option<Duration>,
+    #[cfg(any(
+        target_os = "linux",
+        target_os = "macos",
+        target_os = "ios",
+        target_os = "freebsd",
+        target_os = "netbsd",
+        target_os = "windows",
+    ))]
+    pub(crate) interval: Option<Duration>,
+    #[cfg(any(
+        target_os = "linux",
+        target_os = "macos",
+        target_os = "ios",
+        target_os = "freebsd",
+        target_os = "netbsd",
+    ))]
+    pub(crate) retries: Option<u32>,
+}
+
 impl TcpSocket {
     /// Create a new IPv4 TCP socket.
     ///
     /// This calls `socket(2)`.
     pub fn new_v4() -> io::Result<TcpSocket> {
-        sys::tcp::new_v4_socket().map(|sys| TcpSocket {
-            sys
+        sys::tcp::new_v4_socket().map(|sys| {
+            TcpSocket { sys }
         })
     }
 
     /// Create a new IPv6 TCP socket.
     ///
     /// This calls `socket(2)`.
     pub fn new_v6() -> io::Result<TcpSocket> {
-        sys::tcp::new_v6_socket().map(|sys| TcpSocket {
-            sys
+        sys::tcp::new_v6_socket().map(|sys| {
+            TcpSocket { sys }
         })
     }
 
@@ -168,7 +191,133 @@ impl TcpSocket {
     pub fn get_send_buffer_size(&self) -> io::Result<u32> {
         sys::tcp::get_send_buffer_size(self.sys)
     }
-    
+
+    /// Sets whether keepalive messages are enabled to be sent on this socket.
+    ///
+    /// This will set the `SO_KEEPALIVE` option on this socket.
+    pub fn set_keepalive(&self, keepalive: bool) -> io::Result<()> {
+        sys::tcp::set_keepalive(self.sys, keepalive)
+    }
+
+    /// Returns whether or not TCP keepalive probes will be sent by this socket.
+    pub fn get_keepalive(&self) -> io::Result<bool> {
+        sys::tcp::get_keepalive(self.sys)
+    }
+
+    /// Sets parameters configuring TCP keepalive probes for this socket.
+    ///
+    /// The supported parameters depend on the operating system, and are
+    /// configured using the [`TcpKeepalive`] struct. At a minimum, all systems
+    /// support configuring the [keepalive time]: the time after which the OS
+    /// will start sending keepalive messages on an idle connection.
+    ///
+    /// # Notes
+    ///
+    /// * This will enable TCP keepalive on this socket, if it is not already
+    ///   enabled.
+    /// * On some platforms, such as Windows, any keepalive parameters *not*
+    ///   configured by the `TcpKeepalive` struct passed to this function may be
+    ///   overwritten with their default values. Therefore, this function should
+    ///   either only be called once per socket, or the same parameters should
+    ///   be passed every time it is called.
+    ///
+    /// # Examples
+    /// ```
+    /// use mio::net::{TcpSocket, TcpKeepalive};
+    /// use std::time::Duration;
+    ///
+    /// # fn main() -> Result<(), std::io::Error> {
+    /// let socket = TcpSocket::new_v6()?;
+    /// let keepalive = TcpKeepalive::default()
+    ///     .with_time(Duration::from_secs(4));
+    ///     // Depending on the target operating system, we may also be able to 
+    ///     // configure the keepalive probe interval and/or the number of retries
+    ///     // here as well.
+    ///
+    /// socket.set_keepalive_params(keepalive)?;
+    /// # Ok(()) }
+    /// ```
+    ///
+    /// [`TcpKeepalive`]: ../struct.TcpKeepalive.html
+    /// [keepalive time]: ../struct.TcpKeepalive.html#method.with_time
+    pub fn set_keepalive_params(&self, keepalive: TcpKeepalive) -> io::Result<()> {
+        self.set_keepalive(true)?;
+        sys::tcp::set_keepalive_params(self.sys, keepalive)
+    }
+
+    /// Returns the amount of time after which TCP keepalive probes will be sent
+    /// on idle connections.
+    ///
+    /// If `None`, then keepalive messages are disabled.
+    ///
+    /// This returns the value of `SO_KEEPALIVE` + `IPPROTO_TCP` on OpenBSD,
+    /// NetBSD, and Haiku, `TCP_KEEPALIVE` on macOS and iOS, and `TCP_KEEPIDLE`
+    /// on all other Unix operating systems. On Windows, it is not possible to
+    /// access the value of TCP keepalive parameters after they have been set.
+    ///
+    /// Some platforms specify this value in seconds, so sub-second
+    /// specifications may be omitted.
+    #[cfg_attr(docsrs, doc(cfg(not(target_os = "windows"))))]
+    #[cfg(not(target_os = "windows"))]
+    pub fn get_keepalive_time(&self) -> io::Result<Option<Duration>> {
+        sys::tcp::get_keepalive_time(self.sys)
+    }
+
+    /// Returns the time interval between TCP keepalive probes, if TCP keepalive is
+    /// enabled on this socket.
+    ///
+    /// If `None`, then keepalive messages are disabled.
+    ///
+    /// This returns the value of `TCP_KEEPINTVL` on supported Unix operating
+    /// systems. On Windows, it is not possible to access the value of TCP
+    /// keepalive parameters after they have been set..
+    ///
+    /// Some platforms specify this value in seconds, so sub-second
+    /// specifications may be omitted.
+    #[cfg_attr(docsrs, doc(cfg(any(
+        target_os = "linux",
+        target_os = "macos",
+        target_os = "ios",
+        target_os = "freebsd",
+        target_os = "netbsd",
+    ))))]
+    #[cfg(any(
+        target_os = "linux",
+        target_os = "macos",
+        target_os = "ios",
+        target_os = "freebsd",
+        target_os = "netbsd",
+    ))]
+    pub fn get_keepalive_interval(&self) -> io::Result<Option<Duration>> {
+        sys::tcp::get_keepalive_interval(self.sys)
+    }
+
+    /// Returns the maximum number of TCP keepalive probes that will be sent before
+    /// dropping a connection, if TCP keepalive is enabled on this socket.
+    ///
+    /// If `None`, then keepalive messages are disabled.
+    ///
+    /// This returns the value of `TCP_KEEPCNT` on Unix operating systems that
+    /// support this option. On Windows, it is not possible to access the value
+    /// of TCP keepalive parameters after they have been set.
+    #[cfg_attr(docsrs, doc(cfg(any(
+        target_os = "linux",
+        target_os = "macos",
+        target_os = "ios",
+        target_os = "freebsd",
+        target_os = "netbsd",
+    ))))]
+    #[cfg(any(
+        target_os = "linux", 
+        target_os = "macos",
+        target_os = "ios",
+        target_os = "freebsd",
+        target_os = "netbsd",
+    ))]
+    pub fn get_keepalive_retries(&self) -> io::Result<Option<u32>> {
+        sys::tcp::get_keepalive_retries(self.sys)
+    }
+
     /// Returns the local address of this socket
     ///
     /// Will return `Err` result in windows if called before calling `bind`
@@ -238,3 +387,83 @@ impl FromRawSocket for TcpSocket {
         TcpSocket { sys: socket as sys::tcp::TcpSocket }
     }
 }
+
+impl TcpKeepalive {
+    // Sets the amount of time after which TCP keepalive probes will be sent
+    /// on idle connections.
+    ///
+    /// This will set the value of `SO_KEEPALIVE` + `IPPROTO_TCP` on OpenBSD,
+    /// NetBSD, and Haiku, `TCP_KEEPALIVE` on macOS and iOS, and `TCP_KEEPIDLE`
+    /// on all other Unix operating systems. On Windows, this sets the value of
+    /// the `tcp_keepalive` struct's `keepalivetime` field.
+    ///
+    /// Some platforms specify this value in seconds, so sub-second
+    /// specifications may be omitted.
+    pub fn with_time(self, time: Duration) -> Self {
+        Self {
+            time: Some(time),
+            ..self
+        }
+    }
+
+    /// Sets the time interval between TCP keepalive probes.
+    /// This sets the value of `TCP_KEEPINTVL` on supported Unix operating
+    /// systems. On Windows, this sets the value of the `tcp_keepalive` struct's
+    /// `keepaliveinterval` field.
+    ///
+    /// Some platforms specify this value in seconds, so sub-second
+    /// specifications may be omitted.
+    #[cfg_attr(docsrs, doc(cfg(any(
+        target_os = "linux",
+        target_os = "macos",
+        target_os = "ios",
+        target_os = "freebsd",
+        target_os = "netbsd",
+        target_os = "windows"
+    ))))]
+    #[cfg(any(
+        target_os = "linux",
+        target_os = "macos",
+        target_os = "ios",
+        target_os = "freebsd",
+        target_os = "netbsd",
+        target_os = "windows"
+    ))]
+    pub fn with_interval(self, interval: Duration) -> Self {
+        Self {
+            interval: Some(interval),
+            ..self
+        }
+    }
+
+    /// Sets the maximum number of TCP keepalive probes that will be sent before
+    /// dropping a connection, if TCP keepalive is enabled on this socket.
+    ///
+    /// This will set the value of `TCP_KEEPCNT` on Unix operating systems that
+    /// support this option.
+    #[cfg_attr(docsrs, doc(cfg(any(
+        target_os = "linux",
+        target_os = "macos",
+        target_os = "ios",
+        target_os = "freebsd",
+        target_os = "netbsd",
+    ))))]
+    #[cfg(any(
+        target_os = "linux",
+        target_os = "macos",
+        target_os = "ios",
+        target_os = "freebsd",
+        target_os = "netbsd",
+    ))]
+    pub fn with_retries(self, retries: u32) -> Self {
+        Self {
+            retries: Some(retries),
+            ..self
+        }
+    }
+
+    /// Returns a new, empty set of TCP keepalive parameters.
+    pub fn new() -> Self {
+        Self::default()
+    }
+}

src/sys/shell/tcp.rs

@@ -1,6 +1,7 @@
 use std::io;
 use std::net::{self, SocketAddr};
 use std::time::Duration;
+use crate::net::TcpKeepalive;
 
 pub(crate) type TcpSocket = i32;
 
@@ -70,6 +71,59 @@ pub(crate) fn get_send_buffer_size(_: TcpSocket) -> io::Result<u32> {
     os_required!();
 }
 
+pub(crate) fn set_keepalive(_: TcpSocket, _: bool) -> io::Result<()> {
+    os_required!();
+}
+
+pub(crate) fn get_keepalive(_: TcpSocket) -> io::Result<bool> {
+    os_required!();
+}
+
+#[cfg(any(
+    target_os = "linux",
+    target_os = "macos",
+    target_os = "ios",
+    target_os = "freebsd",
+    target_os = "netbsd",
+    target_os = "windows",
+))]
+pub(crate) fn set_keepalive_params(_: TcpSocket, _: TcpKeepalive) -> io::Result<()> {
+    os_required!()
+}
+
+#[cfg(any(
+    target_os = "linux",
+    target_os = "macos",
+    target_os = "ios",
+    target_os = "freebsd",
+    target_os = "netbsd",
+))]
+pub(crate) fn get_keepalive_time(_: TcpSocket) -> io::Result<Option<Duration>> {
+    os_required!()
+}
+
+#[cfg(any(
+    target_os = "linux",
+    target_os = "macos",
+    target_os = "ios",
+    target_os = "freebsd",
+    target_os = "netbsd",
+))]
+pub(crate) fn get_keepalive_interval(_: TcpSocket) -> io::Result<Option<Duration>> {
+    os_required!()
+}
+
+#[cfg(any(
+    target_os = "linux",
+    target_os = "macos",
+    target_os = "ios",
+    target_os = "freebsd",
+    target_os = "netbsd",
+))]
+pub(crate) fn get_keepalive_retries(_: TcpSocket) -> io::Result<Option<u32>> {
+    os_required!()
+}
+
 pub fn accept(_: &net::TcpListener) -> io::Result<(net::TcpStream, SocketAddr)> {
     os_required!();
 }