From d5f5b0ab21848a7928dcb0823066e09d874790ed Mon Sep 17 00:00:00 2001
From: Roman Kashitsyn <roman@dfinity.org>
Date: Wed, 10 Mar 2021 16:21:23 +0100
Subject: [PATCH] feat(tcp): add support for custom socket options

This change extends to AddrIncoming API to support setting custom
options before binding the server socket.  This change introduces one
such option: `reuse_port`, which enables the `SO_REUSEPORT` option on
the socket on UNIX platforms.

Example:

    use hyper::server::conn::{AddrIncoming, SocketOptions};
    let incoming = AddrIncoming::bind_with_option(
        "0.0.0.0:80".parse().unwrap(),
        SocketOptions::new().reuse_port(true)
    );
---
 Cargo.toml        |  2 +-
 src/server/tcp.rs | 55 +++++++++++++++++++++++++++++++++++++++++++++--
 2 files changed, 54 insertions(+), 3 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index 31ac0356f5..d044454d3d 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -42,7 +42,7 @@ want = "0.3"
 # Optional
 
 libc = { version = "0.2", optional = true }
-socket2 = { version = "0.3.16", optional = true }
+socket2 = { version = "0.3.16", optional = true, features = ["reuseport"] }
 
 [dev-dependencies]
 futures-util = { version = "0.3", default-features = false, features = ["alloc"] }
diff --git a/src/server/tcp.rs b/src/server/tcp.rs
index 16a663719c..429565a946 100644
--- a/src/server/tcp.rs
+++ b/src/server/tcp.rs
@@ -12,6 +12,34 @@ use crate::common::{task, Future, Pin, Poll};
 pub use self::addr_stream::AddrStream;
 use super::accept::Accept;
 
+/// Options to set on the server socket before binding it.
+#[derive(Clone, Debug)]
+pub struct SocketOptions {
+    reuse_port: bool,
+}
+
+impl Default for SocketOptions {
+    fn default() -> Self {
+        Self { reuse_port: false }
+    }
+}
+
+impl SocketOptions {
+    /// Constructs default socket options.
+    pub fn new() -> Self {
+        Self::default()
+    }
+
+    /// Requests setting the `SO_REUSEPORT` option on the server
+    /// socket before binding it.
+    ///
+    /// Only has effect on UNIX platforms.
+    pub fn reuse_port(mut self, enable: bool) -> Self {
+        self.reuse_port = enable;
+        self
+    }
+}
+
 /// A stream of connections from binding to an address.
 #[must_use = "streams do nothing unless polled"]
 pub struct AddrIncoming {
@@ -25,9 +53,27 @@ pub struct AddrIncoming {
 
 impl AddrIncoming {
     pub(super) fn new(addr: &SocketAddr) -> crate::Result<Self> {
-        let std_listener = StdTcpListener::bind(addr).map_err(crate::Error::new_listen)?;
+        AddrIncoming::new_with_options(addr, SocketOptions::default())
+    }
 
-        AddrIncoming::from_std(std_listener)
+    fn new_with_options(addr: &SocketAddr, _opts: SocketOptions) -> crate::Result<Self> {
+        use socket2::{Domain, Protocol, SockAddr, Type};
+        let domain = match &addr {
+            SocketAddr::V4(_) => Domain::ipv4(),
+            SocketAddr::V6(_) => Domain::ipv6(),
+        };
+        let socket = socket2::Socket::new(domain, Type::stream(), Some(Protocol::tcp()))
+            .map_err(crate::Error::new_listen)?;
+        #[cfg(all(unix, not(any(target_os = "solaris", target_os = "illumos"))))]
+        if _opts.reuse_port {
+            socket
+                .set_reuse_port(true)
+                .map_err(crate::Error::new_listen)?;
+        }
+        socket
+            .bind(&SockAddr::from(*addr))
+            .map_err(crate::Error::new_listen)?;
+        AddrIncoming::from_std(socket.into_tcp_listener())
     }
 
     pub(super) fn from_std(std_listener: StdTcpListener) -> crate::Result<Self> {
@@ -44,6 +90,11 @@ impl AddrIncoming {
         AddrIncoming::new(addr)
     }
 
+    /// Creates a new `AddrIncoming` binding the provided socket address with specified options.
+    pub fn bind_with_options(addr: &SocketAddr, opts: SocketOptions) -> crate::Result<Self> {
+        AddrIncoming::new_with_options(addr, opts)
+    }
+
     /// Creates a new `AddrIncoming` from an existing `tokio::net::TcpListener`.
     pub fn from_listener(listener: TcpListener) -> crate::Result<Self> {
         let addr = listener.local_addr().map_err(crate::Error::new_listen)?;