tokio_dual_stack 0.1.5

//! [![git]](https://git.philomathiclife.com/tokio_dual_stack/log.html) [![crates-io]](https://crates.io/crates/tokio_dual_stack) [![docs-rs]](crate)
//!
//! [git]: https://git.philomathiclife.com/git_badge.svg
//! [crates-io]: https://img.shields.io/badge/crates.io-fc8d62?style=for-the-badge&labelColor=555555&logo=rust
//! [docs-rs]: https://img.shields.io/badge/docs.rs-66c2a5?style=for-the-badge&labelColor=555555&logo=docs.rs
//!
//! `tokio_dual_stack` is a library that adds a "dual-stack" [`TcpListener`].
//!
//! ## Why is this useful?
//!
//! Only certain platforms offer the ability for one socket to handle both IPv6 and IPv4 requests
//! (e.g., OpenBSD does not). For the platforms that do, it is often dependent on runtime configuration
//! (e.g., [`IPV6_V6ONLY`](https://www.man7.org/linux/man-pages/man7/ipv6.7.html)). Additionally those platforms
//! that support it often require the "wildcard" IPv6 address to be used (i.e., `::`) which has the unfortunate
//! consequence of preventing other services from using the same protocol port.
//!
//! There are a few ways to work around this issue. One is to deploy the same service twice: one that uses
//! an IPv6 socket and the other that uses an IPv4 socket. This can complicate deployments (e.g., the application
//! may not have been written with the expectation that multiple deployments could be running at the same time) in
//! addition to using more resources. Another is for the application to manually handle each socket (e.g.,
//! [`select`](https://docs.rs/tokio/latest/tokio/macro.select.html)/[`join`](https://docs.rs/tokio/latest/tokio/macro.join.html)
//! each [`TcpListener::accept`]).
//!
//! [`DualStackTcpListener`] chooses an implementation similar to what the equivalent `select` would do while
//! also ensuring that one socket does not "starve" another by ensuring each socket is fairly given an opportunity
//! to `TcpListener::accept` a connection. This has the nice benefit of having a similar API to what a single
//! `TcpListener` would have as well as having similar performance to a socket that does handle both IPv6 and
//! IPv4 requests.
use core::{
    net::{SocketAddr, SocketAddrV4, SocketAddrV6},
    pin::Pin,
    sync::atomic::{AtomicBool, Ordering},
    task::{Context, Poll},
};
use pin_project_lite::pin_project;
use std::io::{Error, ErrorKind, Result};
use tokio::net::{self, TcpListener, TcpSocket, TcpStream, ToSocketAddrs};
/// Prevents [`Sealed`] from being publicly implementable.
mod private {
    /// Marker trait to prevent [`super::Tcp`] from being publicly implementable.
    #[expect(unnameable_types, reason = "want Tcp to be 'sealed'")]
    pub trait Sealed {}
}
use private::Sealed;
/// TCP "listener".
///
/// This `trait` is sealed and cannot be implemented for types outside of `tokio_dual_stack`.
///
/// This exists primarily as a way to define type constructors or polymorphic functions
/// that can user either a [`TcpListener`] or [`DualStackTcpListener`].
///
/// # Examples
///
/// ```no_run
/// # use core::convert::Infallible;
/// # use tokio_dual_stack::Tcp;
/// async fn main_loop<T: Tcp>(listener: T) -> Infallible {
///     loop {
///         match listener.accept().await {
///             Ok((_, socket)) => println!("Client socket: {socket}"),
///             Err(e) => println!("TCP connection failure: {e}"),
///         }
///     }
/// }
/// ```
pub trait Tcp: Sealed + Sized {
    /// Creates a new TCP listener, which will be bound to the specified address(es).
    ///
    /// The returned listener is ready for accepting connections.
    ///
    /// Binding with a port number of 0 will request that the OS assigns a port to this listener.
    /// The port allocated can be queried via the `local_addr` method.
    ///
    /// The address type can be any implementor of the [`ToSocketAddrs`] trait. If `addr` yields
    /// multiple addresses, bind will be attempted with each of the addresses until one succeeds
    /// and returns the listener. If none of the addresses succeed in creating a listener, the
    /// error returned from the last attempt (the last address) is returned.
    ///
    /// This function sets the `SO_REUSEADDR` option on the socket.
    ///
    /// # Examples
    ///
    /// ```no_run
    /// # use core::net::{Ipv4Addr, Ipv6Addr, SocketAddr, SocketAddrV4, SocketAddrV6};
    /// # use std::io::Result;
    /// # use tokio_dual_stack::{DualStackTcpListener, Tcp as _};
    /// #[tokio::main(flavor = "current_thread")]
    /// async fn main() -> Result<()> {
    ///     let listener = DualStackTcpListener::bind(
    ///         [
    ///             SocketAddr::V6(SocketAddrV6::new(Ipv6Addr::LOCALHOST, 8080, 0, 0)),
    ///             SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::LOCALHOST, 8080)),
    ///         ]
    ///         .as_slice(),
    ///     )
    ///     .await?;
    ///     Ok(())
    /// }
    /// ```
    fn bind<A: ToSocketAddrs>(addr: A) -> impl Future<Output = Result<Self>>;
    /// Accepts a new incoming connection from this listener.
    ///
    /// This function will yield once a new TCP connection is established. When established,
    /// the corresponding `TcpStream` and the remote peer’s address will be returned.
    ///
    /// # Cancel safety
    ///
    /// This method is cancel safe. If the method is used as the event in a
    /// [`tokio::select!`](https://docs.rs/tokio/latest/tokio/macro.select.html)
    /// statement and some other branch completes first, then it is guaranteed that no new
    /// connections were accepted by this method.
    ///
    /// # Examples
    ///
    /// ```no_run
    /// # use core::net::{Ipv4Addr, Ipv6Addr, SocketAddr, SocketAddrV4, SocketAddrV6};
    /// # use std::io::Result;
    /// # use tokio_dual_stack::{DualStackTcpListener, Tcp as _};
    /// #[tokio::main(flavor = "current_thread")]
    /// async fn main() -> Result<()> {
    ///     match DualStackTcpListener::bind(
    ///         [
    ///             SocketAddr::V6(SocketAddrV6::new(Ipv6Addr::LOCALHOST, 8080, 0, 0)),
    ///             SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::LOCALHOST, 8080)),
    ///         ]
    ///         .as_slice(),
    ///     )
    ///     .await?.accept().await {
    ///         Ok((_, addr)) => println!("new client: {addr}"),
    ///         Err(e) => println!("couldn't get client: {e}"),
    ///     }
    ///     Ok(())
    /// }
    /// ```
    fn accept(&self) -> impl Future<Output = Result<(TcpStream, SocketAddr)>> + Send + Sync;
    /// Polls to accept a new incoming connection to this listener.
    ///
    /// If there is no connection to accept, `Poll::Pending` is returned and the current task will be notified by
    /// a waker. Note that on multiple calls to `poll_accept`, only the `Waker` from the `Context` passed to the
    /// most recent call is scheduled to receive a wakeup.
    fn poll_accept(&self, cx: &mut Context<'_>) -> Poll<Result<(TcpStream, SocketAddr)>>;
}
impl Sealed for TcpListener {}
impl Tcp for TcpListener {
    #[inline]
    fn bind<A: ToSocketAddrs>(addr: A) -> impl Future<Output = Result<Self>> {
        Self::bind(addr)
    }
    #[inline]
    fn accept(&self) -> impl Future<Output = Result<(TcpStream, SocketAddr)>> + Send + Sync {
        self.accept()
    }
    #[inline]
    fn poll_accept(&self, cx: &mut Context<'_>) -> Poll<Result<(TcpStream, SocketAddr)>> {
        self.poll_accept(cx)
    }
}
/// "Dual-stack" TCP listener.
///
/// IPv6 and IPv4 TCP listener.
#[derive(Debug)]
pub struct DualStackTcpListener {
    /// IPv6 TCP listener.
    ip6: TcpListener,
    /// IPv4 TCP listener.
    ip4: TcpListener,
    /// `true` iff [`Self::ip6::accept`] should be `poll`ed first; otherwise [`Self::ip4::accept`] is `poll`ed
    /// first.
    ///
    /// This exists to prevent one IP version from "starving" another. Each time [`Self::accept`] or
    /// [`Self::poll_accept`] is called, it's overwritten with the opposite `bool`.
    ///
    /// Note we could make this a `core::cell::Cell`; but for maximal flexibility and consistency with `TcpListener`,
    /// we use an `AtomicBool`. This among other things means `DualStackTcpListener` will implement `Sync`.
    ip6_first: AtomicBool,
}
impl DualStackTcpListener {
    /// Creates `Self` using the [`TcpListener`]s returned from [`TcpSocket::listen`].
    ///
    /// [`Self::bind`] is useful when the behavior of [`TcpListener::bind`] is sufficient; however if the underlying
    /// `TcpSocket`s need to be configured differently, then one must call this function instead.
    ///
    /// # Errors
    ///
    /// Errors iff [`TcpSocket::local_addr`] does for either socket, the underlying sockets use the same IP version,
    /// or [`TcpSocket::listen`] errors for either socket.
    ///
    /// Note on Windows-based platforms `TcpSocket::local_addr` will error if [`TcpSocket::bind`] was not called.
    ///
    /// # Examples
    ///
    /// ```no_run
    /// # use core::net::{Ipv4Addr, Ipv6Addr, SocketAddr, SocketAddrV4, SocketAddrV6};
    /// # use std::io::Result;
    /// # use tokio_dual_stack::DualStackTcpListener;
    /// # use tokio::net::TcpSocket;
    /// #[tokio::main(flavor = "current_thread")]
    /// async fn main() -> Result<()> {
    ///     let ip6 = TcpSocket::new_v6()?;
    ///     ip6.bind(SocketAddr::V6(SocketAddrV6::new(Ipv6Addr::LOCALHOST, 8080, 0, 0)))?;
    ///     let ip4 = TcpSocket::new_v4()?;
    ///     ip4.bind(SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::LOCALHOST, 8080)))?;
    ///     let listener = DualStackTcpListener::from_sockets((ip6, 1024), (ip4, 1024))?;
    ///     Ok(())
    /// }
    /// ```
    #[inline]
    pub fn from_sockets(
        (socket_1, backlog_1): (TcpSocket, u32),
        (socket_2, backlog_2): (TcpSocket, u32),
    ) -> Result<Self> {
        socket_1.local_addr().and_then(|sock| {
            socket_2.local_addr().and_then(|sock_2| {
                if sock.is_ipv6() {
                    if sock_2.is_ipv4() {
                        socket_1.listen(backlog_1).and_then(|ip6| {
                            socket_2.listen(backlog_2).map(|ip4| Self {
                                ip6,
                                ip4,
                                ip6_first: AtomicBool::new(true),
                            })
                        })
                    } else {
                        Err(Error::new(
                            ErrorKind::InvalidData,
                            "TcpSockets are the same IP version",
                        ))
                    }
                } else if sock_2.is_ipv6() {
                    socket_1.listen(backlog_1).and_then(|ip4| {
                        socket_2.listen(backlog_2).map(|ip6| Self {
                            ip6,
                            ip4,
                            ip6_first: AtomicBool::new(true),
                        })
                    })
                } else {
                    Err(Error::new(
                        ErrorKind::InvalidData,
                        "TcpSockets are the same IP version",
                    ))
                }
            })
        })
    }
    /// Returns the local address of each socket that the listeners are bound to.
    ///
    /// This can be useful, for example, when binding to port 0 to figure out which port was actually bound.
    ///
    /// # Errors
    ///
    /// Errors iff [`TcpListener::local_addr`] does for either listener.
    ///
    /// # Examples
    ///
    /// ```no_run
    /// # use core::net::{Ipv4Addr, Ipv6Addr, SocketAddr, SocketAddrV4, SocketAddrV6};
    /// # use std::io::Result;
    /// # use tokio_dual_stack::{DualStackTcpListener, Tcp as _};
    /// #[tokio::main(flavor = "current_thread")]
    /// async fn main() -> Result<()> {
    ///     let ip6 = SocketAddrV6::new(Ipv6Addr::LOCALHOST, 8080, 0, 0);
    ///     let ip4 = SocketAddrV4::new(Ipv4Addr::LOCALHOST, 8080);
    ///     assert_eq!(
    ///         DualStackTcpListener::bind([SocketAddr::V6(ip6), SocketAddr::V4(ip4)].as_slice())
    ///             .await?
    ///             .local_addr()?,
    ///         (ip6, ip4)
    ///     );
    ///     Ok(())
    /// }
    /// ```
    #[expect(clippy::unreachable, reason = "we want to crash when there is a bug")]
    #[inline]
    pub fn local_addr(&self) -> Result<(SocketAddrV6, SocketAddrV4)> {
        self.ip6.local_addr().and_then(|ip6| {
            self.ip4.local_addr().map(|ip4| {
                (
                    if let SocketAddr::V6(sock6) = ip6 {
                        sock6
                    } else {
                        unreachable!("there is a bug in DualStackTcpListener::bind")
                    },
                    if let SocketAddr::V4(sock4) = ip4 {
                        sock4
                    } else {
                        unreachable!("there is a bug in DualStackTcpListener::bind")
                    },
                )
            })
        })
    }
    /// Sets the value for the `IP_TTL` option on both sockets.
    ///
    /// This value sets the time-to-live field that is used in every packet sent from each socket.
    /// `ttl_ip6` is the `IP_TTL` value for the IPv6 socket and `ttl_ip4` is the `IP_TTL` value for the
    /// IPv4 socket.
    ///
    /// # Errors
    ///
    /// Errors iff [`TcpListener::set_ttl`] does for either listener.
    ///
    /// # Examples
    ///
    /// ```no_run
    /// # use core::net::{Ipv4Addr, Ipv6Addr, SocketAddr, SocketAddrV4, SocketAddrV6};
    /// # use std::io::Result;
    /// # use tokio_dual_stack::{DualStackTcpListener, Tcp as _};
    /// #[tokio::main(flavor = "current_thread")]
    /// async fn main() -> Result<()> {
    ///     DualStackTcpListener::bind(
    ///         [
    ///             SocketAddr::V6(SocketAddrV6::new(Ipv6Addr::LOCALHOST, 8080, 0, 0)),
    ///             SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::LOCALHOST, 8080)),
    ///         ]
    ///         .as_slice(),
    ///     )
    ///     .await?.set_ttl(100, 100).expect("could not set TTL");
    ///     Ok(())
    /// }
    /// ```
    #[inline]
    pub fn set_ttl(&self, ttl_ip6: u32, ttl_ip4: u32) -> Result<()> {
        self.ip6
            .set_ttl(ttl_ip6)
            .and_then(|()| self.ip4.set_ttl(ttl_ip4))
    }
    /// Gets the values of the `IP_TTL` option for both sockets.
    ///
    /// The first `u32` represents the `IP_TTL` value for the IPv6 socket and the second `u32` is the
    /// `IP_TTL` value for the IPv4 socket. For more information about this option, see [`Self::set_ttl`].
    ///
    /// # Errors
    ///
    /// Errors iff [`TcpListener::ttl`] does for either listener.
    ///
    /// # Examples
    ///
    /// ```no_run
    /// # use core::net::{Ipv4Addr, Ipv6Addr, SocketAddr, SocketAddrV4, SocketAddrV6};
    /// # use std::io::Result;
    /// # use tokio_dual_stack::{DualStackTcpListener, Tcp as _};
    /// #[tokio::main(flavor = "current_thread")]
    /// async fn main() -> Result<()> {
    ///     let listener = DualStackTcpListener::bind(
    ///         [
    ///             SocketAddr::V6(SocketAddrV6::new(Ipv6Addr::LOCALHOST, 8080, 0, 0)),
    ///             SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::LOCALHOST, 8080)),
    ///         ]
    ///         .as_slice(),
    ///     )
    ///     .await?;
    ///     listener.set_ttl(100, 100).expect("could not set TTL");
    ///     assert_eq!(listener.ttl()?, (100, 100));
    ///     Ok(())
    /// }
    /// ```
    #[inline]
    pub fn ttl(&self) -> Result<(u32, u32)> {
        self.ip6
            .ttl()
            .and_then(|ip6| self.ip4.ttl().map(|ip4| (ip6, ip4)))
    }
}
pin_project! {
    /// `Future` returned by [`DualStackTcpListener::accept]`.
    struct AcceptFut<
        F: Future<Output = Result<(TcpStream, SocketAddr)>>,
        F2: Future<Output = Result<(TcpStream, SocketAddr)>>,
    > {
        // Accept future for one `TcpListener`.
        #[pin]
        fut_1: F,
        // Accept future for the other `TcpListener`.
        #[pin]
        fut_2: F2,
    }
}
impl<
    F: Future<Output = Result<(TcpStream, SocketAddr)>>,
    F2: Future<Output = Result<(TcpStream, SocketAddr)>>,
> Future for AcceptFut<F, F2>
{
    type Output = Result<(TcpStream, SocketAddr)>;
    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let this = self.project();
        match this.fut_1.poll(cx) {
            Poll::Ready(res) => Poll::Ready(res),
            Poll::Pending => this.fut_2.poll(cx),
        }
    }
}
impl Sealed for DualStackTcpListener {}
impl Tcp for DualStackTcpListener {
    #[inline]
    async fn bind<A: ToSocketAddrs>(addr: A) -> Result<Self> {
        match net::lookup_host(addr).await {
            Ok(socks) => {
                let mut last_err = None;
                let mut ip6_opt = None;
                let mut ip4_opt = None;
                for sock in socks {
                    match ip6_opt {
                        None => match ip4_opt {
                            None => {
                                let is_ip6 = sock.is_ipv6();
                                match TcpListener::bind(sock).await {
                                    Ok(ip) => {
                                        if is_ip6 {
                                            ip6_opt = Some(ip);
                                        } else {
                                            ip4_opt = Some(ip);
                                        }
                                    }
                                    Err(err) => last_err = Some(err),
                                }
                            }
                            Some(ip4) => {
                                if sock.is_ipv6() {
                                    match TcpListener::bind(sock).await {
                                        Ok(ip6) => {
                                            return Ok(Self {
                                                ip6,
                                                ip4,
                                                ip6_first: AtomicBool::new(true),
                                            });
                                        }
                                        Err(err) => last_err = Some(err),
                                    }
                                }
                                ip4_opt = Some(ip4);
                            }
                        },
                        Some(ip6) => {
                            if sock.is_ipv4() {
                                match TcpListener::bind(sock).await {
                                    Ok(ip4) => {
                                        return Ok(Self {
                                            ip6,
                                            ip4,
                                            ip6_first: AtomicBool::new(true),
                                        });
                                    }
                                    Err(err) => last_err = Some(err),
                                }
                            }
                            ip6_opt = Some(ip6);
                        }
                    }
                }
                Err(last_err.unwrap_or_else(|| {
                    Error::new(
                        ErrorKind::InvalidInput,
                        "could not resolve to an IPv6 and IPv4 address",
                    )
                }))
            }
            Err(err) => Err(err),
        }
    }
    #[inline]
    fn accept(&self) -> impl Future<Output = Result<(TcpStream, SocketAddr)>> + Send + Sync {
        // The correctness of code does not depend on `self.ip6_first`; therefore
        // we elect for the most performant `Ordering`.
        if self.ip6_first.swap(false, Ordering::Relaxed) {
            AcceptFut {
                fut_1: self.ip6.accept(),
                fut_2: self.ip4.accept(),
            }
        } else {
            // The correctness of code does not depend on `self.ip6_first`; therefore
            // we elect for the most performant `Ordering`.
            self.ip6_first.store(true, Ordering::Relaxed);
            AcceptFut {
                fut_1: self.ip4.accept(),
                fut_2: self.ip6.accept(),
            }
        }
    }
    #[inline]
    fn poll_accept(&self, cx: &mut Context<'_>) -> Poll<Result<(TcpStream, SocketAddr)>> {
        // The correctness of code does not depend on `self.ip6_first`; therefore
        // we elect for the most performant `Ordering`.
        if self.ip6_first.swap(false, Ordering::Relaxed) {
            self.ip6.poll_accept(cx)
        } else {
            // The correctness of code does not depend on `self.ip6_first`; therefore
            // we elect for the most performant `Ordering`.
            self.ip6_first.store(true, Ordering::Relaxed);
            self.ip4.poll_accept(cx)
        }
    }
}