sdlan-lib-rs/src/utils/socks.rs

use default_net::Interface;
use sdlan_sn_rs::{
    config::{AF_INET, AF_INET6},
    utils::{Result, SDLanError},
};
use socket2::{Domain, Protocol, SockAddr};
use std::{net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr, SocketAddrV4}, num::NonZeroU32};
use tokio::net::ToSocketAddrs;
use tracing::{debug, error};

use sdlan_sn_rs::peer::SdlanSock;
use tokio::net::UdpSocket;

use crate::network::Node;

#[allow(unused)]
pub struct SocketV6 {
    ipv6: Option<Ipv4Addr>,
    port: u16,
    has_v6: bool,
}

pub struct Socket {
    udp: UdpSocket,
}

impl Socket {
    pub async fn send_to<A: ToSocketAddrs>(&self, buf: &[u8], target: A) -> Result<usize> {
        let m = self.udp.send_to(buf, target).await?;
        Ok(m)
    }

    pub fn get_local_port(&self) -> u16 {
        match self.udp.local_addr() {
            Ok(addr) => addr.port(),
            Err(_e) => 0,
        }
    }

    pub async fn recv_from(&self, buf: &mut [u8]) -> Result<(usize, SocketAddr)> {
        let m = self.udp.recv_from(buf).await?;
        Ok(m)
    }

    pub fn ttl(&self) -> Result<u32> {
        if let Ok(v) = self.udp.ttl() {
            Ok(v)
        } else {
            Err(SDLanError::NormalError("no ttl found"))
        }
    }

    pub fn set_ttl(&self, ttl: u32) -> Result<()> {
        if let Ok(_) = self.udp.set_ttl(ttl) {
            Ok(())
        } else {
            Err(SDLanError::NormalError("failed to set ttl"))
        }
    }

    pub async fn build_v6(v6: Ipv6Addr, port: u16) -> Result<Self> {
        let udp = UdpSocket::bind(format!("[{}]:{}", v6, port)).await?;
        Ok(Self { udp })
    }

    pub async fn build(
        port: u16,
        bind_any: bool,
        join_multicast: bool,
        tos: u32,
        default_interface: &Option<Interface>,
    ) -> Result<Self> {
        let addr = match bind_any {
            true => "0.0.0.0",
            false => "127.0.0.1",
        };
        let addr = format!("{}:{}", addr, port);

        let socket2 =
            socket2::Socket::new(Domain::IPV4, socket2::Type::DGRAM, Some(Protocol::UDP))?;

        socket2.bind(&addr.parse::<SocketAddr>()?.into());

        if let Some(iface) = default_interface {
            set_unicast_if_v4(&socket2, iface);
        }

        let std_socket: std::net::UdpSocket = socket2.into();
        std_socket.set_nonblocking(true);

        let udp = UdpSocket::from_std(std_socket)?;

        // let udp = UdpSocket::bind(&addr).await?;
        if join_multicast {
            if let Err(e) =
                udp.join_multicast_v4(Ipv4Addr::new(224, 0, 0, 69), Ipv4Addr::new(0, 0, 0, 0))
            {
                error!("failed to join multicast: {}", e.to_string());
            } else {
                debug!("{} joined multicast ok", addr);
            }
        }
        if tos != 0 {
            if let Err(e) = udp.set_tos_v4(tos) {
                error!("failed to set tos: {}", e.to_string());
            }
        }
        Ok(Self { udp })
    }
}

#[cfg(not(target_os = "windows"))]
pub fn set_unicast_if_v4(sock: &socket2::Socket, interface: &Interface) -> Result<()> {
    sock.bind_device_by_index_v4(NonZeroU32::new(interface.index))?;
    Ok(())
}

#[cfg(target_os = "windows")]
pub fn set_unicast_if_v4(sock: &socket2::Socket, interface: &Interface) -> Result<()> {
    use std::os::windows::io::AsRawSocket;
    use windows_sys::Win32::Networking::WinSock::{setsockopt, IPPROTO_IP, IP_UNICAST_IF, SOCKET};

    let idx = interface.index.to_be();
    let ret = unsafe {
        setsockopt(
            sock.as_raw_socket() as SOCKET,
            IPPROTO_IP as i32,
            IP_UNICAST_IF as i32,
            &idx as *const _ as *const _,
            std::mem::size_of_val(&idx) as i32,
        )
    };
    if ret == 0 {
        Ok(())
    } else {
        Err(SDLanError::IOError(
            std::io::Error::last_os_error().to_string(),
        ))
    }
}

/*
pub async fn send_to_sock_v4_and_v6(
    // sk: &Socket,
    eee: &Node,
    content: &[u8],
    sock: &SdlanSock,
    v6: &Option<V6Info>,
) -> Result<()> {
    let _ = send_to_sock(&eee, content, sock).await;
    if let Some(v6) = v6 {
        // let sk6 = eee.udp_sock_v6.read().unwrap().clone();
        let sock = SdlanSock {
            family: AF_INET6,
            port: v6.port,
            v4: [0; 4],
            v6: v6.v6,
        };
        let _ = send_to_sock(eee, content, &sock).await;
    }
    Ok(())
}
*/

#[allow(unused)]
pub fn get_socketaddr_from_sock(s: &SdlanSock) -> SocketAddr {
    let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::from(s.v4)), s.port);
    addr
}

pub async fn send_to_sock(
    // sk: &Socket,
    eee: &Node,
    content: &[u8],
    sock: &SdlanSock,
    // v6: &Option<V6Info>,
) -> Result<()> {
    match sock.family {
        AF_INET => {
            // sockv4
            let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::from(sock.v4)), sock.port);
            // println!("udp sock v4 send to {}", addr);
            eee.udp_sock_v4.send_to(content, addr).await?;
            return Ok(());
        }
        AF_INET6 => {
            // sock v6
            let sk = eee.udp_sock_v6.read().unwrap().clone();
            match sk.as_ref() {
                None => {
                    error!("ipv6 not opened, not responding");
                    return Ok(());
                }
                Some(sk) => {
                    let addr = SocketAddr::new(IpAddr::V6(Ipv6Addr::from(sock.v6)), sock.port);
                    debug!("send with ipv6");
                    sk.send_to(content, addr).await?;
                    return Ok(());
                }
            }
        }
        other => {
            error!("unknown family {}, aborting", other);
            return Err(SDLanError::NormalError("unknown family"));
        }
    }
}

/*
pub async fn send_to_sock(sk: &Socket, content: &[u8], sock: &SdlanSock) -> Result<usize> {
    match sock.family {
        AF_INET => {
            let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::from(sock.v4)), sock.port);
            let n = sk.send_to(content, addr).await?;
            return Ok(n);
        }
        AF_INET6 => {
            let addr = SocketAddr::new(IpAddr::V6(Ipv6Addr::from(sock.v6)), sock.port);
            let n = sk.send_to(content, addr).await?;
            return Ok(n);
        }
        other => {
            error!("AF family {} not implemented", other);
            return Ok(0);
        }
    }
}
*/