mptcp/

std.rs

use std::{
    io,
    net::{SocketAddr, TcpListener, TcpStream, ToSocketAddrs},
};

use crate::{sys::MptcpSocketBuilder, MptcpExt, MptcpOpt, MptcpSocket};

/// Extension trait for std::net::TcpStream to support MPTCP.
pub trait MptcpStreamExt {
    type Output;

    /// Establishes an MPTCP connection with the given address and MptcpOpt.
    ///
    /// # Arguments
    ///
    /// * `addr` - The address to connect to.
    /// * `opt` - The MptcpOpt options for the connection.
    ///
    /// # Returns
    ///
    /// Returns an `io::Result` containing the MptcpSocket if the connection is successful,
    /// or an `io::Error` if an error occurs during the connection.
    fn connect_mptcp_opt<A: ToSocketAddrs>(
        addr: A,
        opt: MptcpOpt,
    ) -> io::Result<MptcpSocket<Self::Output>>;

    /// Establishes an MPTCP connection with the given address. If MPTCP cannot be used
    /// the connection will fallback to a regular TCP connection.
    ///
    /// # Arguments
    ///
    /// * `addr` - The address to connect to.
    ///
    /// # Returns
    ///
    /// Returns an `io::Result` containing the MptcpSocket if the connection is successful,
    /// or an `io::Error` if an error occurs during the connection.
    fn connect_mptcp<A: ToSocketAddrs>(addr: A) -> io::Result<MptcpSocket<Self::Output>> {
        Self::connect_mptcp_opt(addr, MptcpOpt::Fallback)
    }

    /// Establishes an MPTCP connection with the given address. Returns an error even if
    /// MPTCP cannot be used. See `connect_mptcp` for a version that falls back to TCP.
    ///
    /// # Arguments
    ///
    /// * `addr` - The address to connect to.
    ///
    /// # Returns
    ///
    /// Returns an `io::Result` containing the MptcpSocket if the connection is successful,
    /// or an `io::Error` if an error occurs during the connection.
    fn connect_mptcp_force<A: ToSocketAddrs>(addr: A) -> io::Result<Self::Output> {
        Ok(Self::connect_mptcp_opt(addr, MptcpOpt::NoFallback)?.into_socket())
    }
}

/// Extension trait for std::net::TcpListener.
pub trait MptcpListenerExt {
    type Output;

    /// Binds an MPTCP socket to the specified address with the given MptcpOpt.
    ///
    /// # Arguments
    ///
    /// * `addr` - The address to bind the socket to.
    /// * `opt` - The MptcpOpt to use for the socket.
    ///
    /// # Returns
    ///
    /// Returns an `io::Result` containing the MptcpSocket with the specified MptcpOpt.
    fn bind_mptcp_opt<A: ToSocketAddrs>(
        addr: A,
        opt: MptcpOpt,
    ) -> io::Result<MptcpSocket<Self::Output>>;

    /// Binds an MPTCP socket to the specified address. If MPTCP cannot be used
    /// the connection will fallback to a regular TCP connection.
    ///
    /// # Arguments
    ///
    /// * `addr` - The address to bind the socket to.
    ///
    /// # Returns
    ///
    /// Returns an `io::Result` containing the MptcpSocket with the default MptcpOpt (Fallback).
    fn bind_mptcp<A: ToSocketAddrs>(addr: A) -> io::Result<MptcpSocket<Self::Output>> {
        Self::bind_mptcp_opt(addr, MptcpOpt::Fallback)
    }

    /// Binds an MPTCP socket to the specified address. Returns an error even if
    /// MPTCP cannot be used. See `bind_mptcp` for a version that falls back to TCP.
    ///
    /// # Arguments
    ///
    /// * `addr` - The address to bind the socket to.
    ///
    /// # Returns
    ///
    /// Returns an `io::Result` containing the MptcpSocket with the MptcpOpt set to NoFallback.
    fn bind_mptcp_force<A: ToSocketAddrs>(addr: A) -> io::Result<Self::Output> {
        Ok(Self::bind_mptcp_opt(addr, MptcpOpt::NoFallback)?.into_socket())
    }
}

fn resolve_each_addr<A: ToSocketAddrs, F, T>(addr: &A, mut f: F) -> io::Result<T>
where
    F: FnMut(SocketAddr) -> io::Result<T>,
{
    let addrs = addr.to_socket_addrs()?;
    let mut last_err = None;
    for addr in addrs {
        match f(addr) {
            Ok(l) => return Ok(l),
            Err(e) => last_err = Some(e),
        }
    }
    Err(last_err.unwrap_or_else(|| {
        io::Error::new(
            io::ErrorKind::InvalidInput,
            "could not resolve to any address",
        )
    }))
}

impl MptcpStreamExt for TcpStream {
    type Output = Self;

    fn connect_mptcp_opt<A: ToSocketAddrs>(
        addr: A,
        opt: MptcpOpt,
    ) -> io::Result<MptcpSocket<Self::Output>> {
        match resolve_each_addr(&addr, |addr| {
            MptcpSocketBuilder::new_for_addr(addr)?.connect(addr)
        }) {
            Ok(sock) => Ok(MptcpSocket::Mptcp(sock.into())),
            Err(_) if matches!(opt, MptcpOpt::Fallback) => {
                Ok(MptcpSocket::Tcp(Self::connect(addr)?))
            }
            Err(err) => Err(err),
        }
    }
}

impl MptcpExt for TcpStream {}

impl From<MptcpSocket<TcpStream>> for TcpStream {
    fn from(socket: MptcpSocket<TcpStream>) -> Self {
        socket.into_socket()
    }
}

impl MptcpListenerExt for TcpListener {
    type Output = Self;

    fn bind_mptcp_opt<A: ToSocketAddrs>(
        addr: A,
        opt: MptcpOpt,
    ) -> io::Result<MptcpSocket<Self::Output>> {
        match resolve_each_addr(&addr, |addr| {
            MptcpSocketBuilder::new_for_addr(addr)?.bind(addr)
        }) {
            Ok(sock) => Ok(MptcpSocket::Mptcp(sock.into())),
            Err(_) if matches!(opt, MptcpOpt::Fallback) => Ok(MptcpSocket::Tcp(Self::bind(addr)?)),
            Err(err) => Err(err),
        }
    }
}

impl From<MptcpSocket<TcpListener>> for TcpListener {
    fn from(socket: MptcpSocket<TcpListener>) -> Self {
        socket.into_socket()
    }
}

#[cfg(all(test, target_os = "linux"))]
mod test {
    use std::net::{IpAddr, Ipv4Addr};

    use crate::sys::is_mptcp_enabled;

    use super::*;

    #[test]
    fn test_resolve_each_addr() {
        let addr = "127.0.0.1:80";
        let result = resolve_each_addr(&addr, |addr| {
            assert_eq!(addr.port(), 80);
            assert_eq!(addr.ip(), IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)));
            Ok(())
        });
        assert!(result.is_ok());
    }

    #[test]
    fn test_resolve_each_addr_error() {
        let addr = "thisisanerror";
        let result = resolve_each_addr(&addr, |_| Ok(()));
        assert!(result.is_err());
    }

    #[test]
    fn test_mptcp_socket() {
        let mptcp_enabled = is_mptcp_enabled();

        let listener = TcpListener::bind_mptcp("127.0.0.1:0");
        if mptcp_enabled {
            assert!(matches!(listener, Ok(MptcpSocket::Mptcp(..))));
        } else {
            assert!(matches!(listener, Ok(MptcpSocket::Tcp(..))));
        }

        let listener = listener.unwrap().into_socket();
        let local_addr = listener.local_addr().unwrap();

        let stream = TcpStream::connect_mptcp(local_addr);
        if mptcp_enabled {
            assert!(matches!(stream, Ok(MptcpSocket::Mptcp(..))));
        } else {
            assert!(matches!(stream, Ok(MptcpSocket::Tcp(..))));
        }
    }

    #[test]
    fn test_mptcp_no_fallback() {
        let mptcp_enabled = is_mptcp_enabled();

        if mptcp_enabled {
            // If the system supports MPTCP, we cannot test the no fallback option
            return;
        }

        let listener = TcpListener::bind_mptcp_force("127.0.0.1:0");
        assert!(listener.is_err());

        let stream = TcpStream::connect_mptcp_force("127.0.0.1:0");
        assert!(stream.is_err());
    }
}