seqknock-common 0.3.1

/*
 * Copyright 2023 Jonas Eriksson
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

use libc::{
    __errno_location,
    __u32,
    AF_INET,
    AF_INET6,
    EINPROGRESS,
    SOCK_STREAM,
    SOL_TCP,
    //TCP_SEND_QUEUE,
    TCP_QUEUE_SEQ,
    TCP_REPAIR,
    TCP_REPAIR_QUEUE,
    c_int,
    c_void,
    close as libc_close,
    connect as libc_connect,
    in_addr,
    in6_addr,
    setsockopt,
    sockaddr,
    sockaddr_in,
    sockaddr_in6,
    socket,
};

// Define the TCP_SEND_QUEUE constant here since it's missing from the libc-crate. It also seems to
// be missing from musl, so not sure how that would work. In the end, it's defined by kernel.
static TCP_SEND_QUEUE: __u32 = 2;

use std::io::{Error, ErrorKind, Result};
use std::mem::size_of;
use std::net::SocketAddr;
use std::net::ToSocketAddrs;
use std::os::unix::io::FromRawFd;

use log::debug;

#[cfg(feature = "async")]
use async_io;
#[cfg(feature = "async")]
use async_net;

#[derive(Clone, Copy, PartialEq)]
pub enum Family {
    V4,
    V6,
}

fn family_matches(socket_addr: &SocketAddr, family: Option<Family>) -> bool {
    if let Some(f) = family {
        if f == Family::V4 && !socket_addr.is_ipv4() {
            return false;
        }
        if f == Family::V6 && !socket_addr.is_ipv6() {
            return false;
        }
    }
    true
}

pub fn connect<A: ToSocketAddrs>(
    sequence_no: u32,
    addr: A,
    force_family: Option<Family>,
) -> Result<std::net::TcpStream> {
    unsafe {
        // Lookup addresses
        let socket_addrs = addr.to_socket_addrs()?;

        // Try to connect to all addresses
        let mut maybe_err = None;
        for socket_addr in socket_addrs {
            if !family_matches(&socket_addr, force_family) {
                debug!("skipping {}, not of requested family", socket_addr);
                continue;
            }
            debug!("Trying to connect to {}", socket_addr);

            let sock = create_socket(family_of(&socket_addr), sequence_no)?;

            match connect_socket(sock, socket_addr, true) {
                Ok(()) => {
                    debug!("Connected to {}", socket_addr);
                    // Create stream from socket fd
                    return Ok(std::net::TcpStream::from_raw_fd(sock));
                }
                Err(e) => maybe_err = Some(e),
            }

            libc_close(sock);
        }

        if let Some(e) = maybe_err {
            return Err(e);
        }
    }
    Err(Error::new(
        ErrorKind::AddrNotAvailable,
        "No address entries for hostname",
    ))
}

#[cfg(feature = "async")]
pub async fn connect_async<A>(
    sequence_no: u32,
    socket_addr: A,
    force_family: Option<Family>,
) -> Result<async_net::TcpStream>
where
    A: async_net::AsyncToSocketAddrs,
{
    let socket_addrs = async_net::resolve(socket_addr).await?;

    unsafe {
        // Try to connect to all addresses
        let mut maybe_err = None;
        for socket_addr in socket_addrs {
            if !family_matches(&socket_addr, force_family) {
                debug!("skipping {}, not of requested family", socket_addr);
                continue;
            }
            debug!("Trying to connect to {}", socket_addr);

            let sock = create_socket(family_of(&socket_addr), sequence_no)?;

            match connect_socket(sock, socket_addr, true) {
                Ok(()) => {
                    let stream = match async_io::Async::new(std::net::TcpStream::from_raw_fd(sock))
                    {
                        Ok(s) => s,
                        Err(e) => {
                            maybe_err = Some(e);
                            continue;
                        }
                    };
                    match stream.writable().await {
                        Ok(_) => match stream.get_ref().take_error()? {
                            None => {
                                debug!("Connected to {}", socket_addr);
                                return Ok(stream.into());
                            }
                            Some(e) => {
                                maybe_err = Some(e);
                                continue;
                            }
                        },
                        Err(e) => maybe_err = Some(e),
                    }
                }
                Err(e) => maybe_err = Some(e),
            }

            libc_close(sock);
        }

        if let Some(e) = maybe_err {
            return Err(e);
        }
    }

    Err(Error::new(
        ErrorKind::AddrNotAvailable,
        "No address entries for hostname",
    ))
}

unsafe fn connect_socket(
    sock: c_int,
    socket_addr: std::net::SocketAddr,
    blocking: bool,
) -> Result<()> {
    unsafe {
        match socket_addr {
            SocketAddr::V4(v4addr) => {
                let octets = v4addr.ip().octets();
                let u32_addr: u32 = (octets[0] as u32)
                    | (octets[1] as u32) << 8
                    | (octets[2] as u32) << 16
                    | (octets[3] as u32) << 24;
                let saddr = sockaddr_in {
                    sin_family: AF_INET as u16,
                    sin_port: v4addr.port().to_be(),
                    sin_addr: in_addr { s_addr: u32_addr },
                    sin_zero: [0; 8],
                };
                let result = libc_connect(
                    sock,
                    &saddr as *const sockaddr_in as *const sockaddr,
                    size_of::<sockaddr_in>() as u32,
                );
                if result < 0 && (blocking || (*__errno_location()) != EINPROGRESS) {
                    return Err(Error::last_os_error());
                }
            }
            SocketAddr::V6(v6addr) => {
                let saddr = sockaddr_in6 {
                    sin6_family: AF_INET6 as u16,
                    sin6_port: v6addr.port().to_be(),
                    sin6_flowinfo: 0,
                    sin6_addr: in6_addr {
                        s6_addr: v6addr.ip().octets(),
                    },
                    sin6_scope_id: 0,
                };
                let result = libc_connect(
                    sock,
                    &saddr as *const sockaddr_in6 as *const sockaddr,
                    size_of::<sockaddr_in6>() as u32,
                );
                if result < 0 && (blocking || (*__errno_location()) != EINPROGRESS) {
                    return Err(Error::last_os_error());
                }
            }
        }

        Ok(())
    }
}

unsafe fn sso_tcp_wrapper(sock: c_int, cmd: c_int, data: u32) -> Result<()> {
    unsafe {
        let dataptr = &data as *const __u32 as *const c_void;
        if setsockopt(sock, SOL_TCP, cmd, dataptr, 4) < 0 {
            return Err(Error::last_os_error());
        }
        Ok(())
    }
}

unsafe fn create_socket(family: c_int, sequence_no: u32) -> Result<c_int> {
    unsafe {
        // Create socket
        let sock: c_int = socket(family, SOCK_STREAM, 0);
        if sock < 0 {
            return Err(Error::last_os_error());
        }

        // Enter repair mode
        sso_tcp_wrapper(sock, TCP_REPAIR, 1)?;
        // Enter repair queue mode for the send queue
        sso_tcp_wrapper(sock, TCP_REPAIR_QUEUE, TCP_SEND_QUEUE)?;
        // Set sequence number
        sso_tcp_wrapper(sock, TCP_QUEUE_SEQ, sequence_no)?;
        // Exit repair mode
        sso_tcp_wrapper(sock, TCP_REPAIR, 0)?;

        Ok(sock)
    }
}

fn family_of(socket_addr: &std::net::SocketAddr) -> c_int {
    match socket_addr {
        SocketAddr::V4(_) => AF_INET,
        SocketAddr::V6(_) => AF_INET6,
    }
}