Remove unsafety from setting TcpStream keepalive (#1187)

We now have a way to convert `tokio` -> `std` `TcpStream`s and
use `socket2`'s `Socket::From<TcpStream>` impl to safely set a
socket's keepalive option. This eliminates the need for us to use
`unsafe` in the keepalive setter.

Signed-off-by: Kevin Leimkuhler <[email protected]>

Author: kleimkuhler

linkerd/proxy/transport/src/connect.rs

@@ -37,7 +37,7 @@ impl<T: Param<Remote<ServerAddr>>> tower::Service<T> for ConnectTcp {
         Box::pin(async move {
             let io = TcpStream::connect(&addr).await?;
             super::set_nodelay_or_warn(&io);
-            super::set_keepalive_or_warn(&io, keepalive);
+            let io = super::set_keepalive_or_warn(io, keepalive)?;
             debug!(
                 local.addr = %io.local_addr().expect("cannot load local addr"),
                 ?keepalive,

linkerd/proxy/transport/src/lib.rs

@@ -1,9 +1,9 @@
 //! Utilities for use TCP servers & clients.
 //!
-//! Uses unsafe code to interact with socket options for keepalive and SO_ORIGINAL_DST.
+//! Uses unsafe code to interact with socket options for SO_ORIGINAL_DST.
 
 #![deny(warnings, rust_2018_idioms)]
-//#![forbid(unsafe_code)]
+// #![forbid(unsafe_code)]
 
 pub mod addrs;
 mod connect;
@@ -17,6 +17,7 @@ pub use self::{
     listen::{Bind, BindTcp},
     orig_dst::BindWithOrigDst,
 };
+use linkerd_io as io;
 use socket2::TcpKeepalive;
 use std::time::Duration;
 use tokio::net::TcpStream;
@@ -38,47 +39,20 @@ fn set_nodelay_or_warn(socket: &TcpStream) {
     }
 }
 
-fn set_keepalive_or_warn(tcp: &TcpStream, keepalive_duration: Option<Duration>) {
-    // TODO(eliza): when https://github.com/tokio-rs/tokio/pull/3189 merges
-    // upstream, we will be able to convert the Tokio `TcpStream` into a
-    // `socket2::Socket` without unsafe, by converting it to a
-    // `std::net::TcpStream` (as `socket2::Socket` has a
-    // `From<std::net::TcpStream>`). What we're doing now is more or less
-    // equivalent, but this would use a safe interface...
-    #[cfg(unix)]
-    let sock = unsafe {
-        // Safety: `from_raw_fd` takes ownership of the underlying file
-        // descriptor, and will close it when dropped. However, we obtain the
-        // file descriptor via `as_raw_fd` rather than `into_raw_fd`, so the
-        // Tokio `TcpStream` *also* retains ownership of the socket --- which is
-        // what we want. Instead of letting the `socket2` socket returned by
-        // `from_raw_fd` close the fd, we `mem::forget` the `Socket`, so that
-        // its `Drop` impl will not run. This ensures the fd is not closed
-        // prematurely.
-        use std::os::unix::io::{AsRawFd, FromRawFd};
-        socket2::Socket::from_raw_fd(tcp.as_raw_fd())
+fn set_keepalive_or_warn(
+    tcp: TcpStream,
+    keepalive_duration: Option<Duration>,
+) -> io::Result<TcpStream> {
+    let sock = {
+        let stream = tokio::net::TcpStream::into_std(tcp)?;
+        socket2::Socket::from(stream)
     };
-    #[cfg(windows)]
-    let sock = unsafe {
-        // Safety: `from_raw_socket` takes ownership of the underlying Windows
-        // SOCKET, and will close it when dropped. However, we obtain the
-        // SOCKET via `as_raw_socket` rather than `into_raw_socket`, so the
-        // Tokio `TcpStream` *also* retains ownership of the socket --- which is
-        // what we want. Instead of letting the `socket2` socket returned by
-        // `from_raw_socket` close the SOCKET, we `mem::forget` the `Socket`, so
-        // that its `Drop` impl will not run. This ensures the socket is not
-        // closed prematurely.
-        use std::os::windows::io::{AsRawSocket, FromRawSocket};
-        socket2::Socket::from_raw_socket(tcp.as_raw_socket())
-    };
-
     let ka = keepalive_duration
         .into_iter()
         .fold(TcpKeepalive::new(), |k, t| k.with_time(t));
     if let Err(e) = sock.set_tcp_keepalive(&ka) {
         tracing::warn!("failed to set keepalive: {}", e);
     }
-
-    // Don't let the socket2 socket close the fd on drop!
-    std::mem::forget(sock);
+    let stream: std::net::TcpStream = socket2::Socket::into(sock);
+    tokio::net::TcpStream::from_std(stream)
 }

linkerd/proxy/transport/src/listen.rs

@@ -66,7 +66,7 @@ where
         let accept = TcpListenerStream::new(listen).map(move |res| {
             let tcp = res?;
             super::set_nodelay_or_warn(&tcp);
-            super::set_keepalive_or_warn(&tcp, keepalive);
+            let tcp = super::set_keepalive_or_warn(tcp, keepalive)?;
             let client = Remote(ClientAddr(tcp.peer_addr()?));
             Ok((Addrs { server, client }, tcp))
         });