pistol 4.0.18

/// Code obtained by analyzing the traffic generated by the Linux `traceroute` command.
#[cfg(any(feature = "trace", feature = "os"))]
use rand::Rng;
#[cfg(any(feature = "trace", feature = "os"))]
use std::net::IpAddr;
#[cfg(any(feature = "trace", feature = "os"))]
use std::time::Duration;
#[cfg(any(feature = "trace", feature = "os"))]
use tracing::debug;

#[cfg(any(feature = "trace", feature = "os"))]
use crate::error::PistolError;
#[cfg(feature = "trace")]
use crate::utils::random_port_range;

#[cfg(any(feature = "trace", feature = "os"))]
pub mod icmp;
#[cfg(any(feature = "trace", feature = "os"))]
pub mod icmpv6;
#[cfg(any(feature = "trace", feature = "os"))]
pub mod tcp;
#[cfg(any(feature = "trace", feature = "os"))]
pub mod tcp6;
#[cfg(any(feature = "trace", feature = "os"))]
pub mod udp;
#[cfg(any(feature = "trace", feature = "os"))]
pub mod udp6;

#[cfg(any(feature = "trace", feature = "os"))]
const UDP_DATA_SIZE: usize = 32;
#[cfg(feature = "trace")]
const START_PORT: u16 = 33434;

#[cfg(any(feature = "trace", feature = "os"))]
#[derive(Debug, Clone, Copy, PartialEq)]
pub enum HopStatus {
    TimeExceeded(IpAddr),
    NoResponse,
    Unreachable(IpAddr),
    RecvReply(IpAddr),
}

/// The recommended timeout value is 5 seconds, which is consistent with the default value of traceroute.
#[cfg(feature = "trace")]
pub fn syn_trace(
    dst_addr: IpAddr,
    dst_port: Option<u16>, // default is 80
    src_addr: IpAddr,
    timeout: Option<Duration>,
) -> Result<u8, PistolError> {
    let dst_port = match dst_port {
        Some(p) => p,
        None => 80,
    };
    match dst_addr {
        IpAddr::V4(dst_ipv4) => {
            if let IpAddr::V4(src_ipv4) = src_addr {
                let mut rng = rand::rng();
                let mut ip_id = rng.random();
                // ensure that this u16 does not overflow
                if ip_id > u16::MAX - 30 {
                    ip_id -= 30
                }
                let mut last_response_ttl = 0;
                for ttl in 1..=30 {
                    let random_src_port = random_port_range(1000, 65535);
                    // let random_src_port = 61234; // debug use
                    let (hop_status, rtt) = tcp::send_syn_trace_packet(
                        dst_ipv4,
                        dst_port,
                        src_ipv4,
                        random_src_port,
                        ip_id,
                        ttl,
                        timeout,
                    )?;
                    ip_id += 1;
                    match hop_status {
                        HopStatus::TimeExceeded(_) => {
                            last_response_ttl = ttl;
                        }
                        HopStatus::RecvReply(_) => {
                            // tcp rst packet, means packet arrive the target machine
                            debug!("ttl: {}, {:?} - {:.2}s", ttl, hop_status, rtt.as_secs_f64());
                            return Ok(ttl);
                        }
                        _ => (),
                    }
                }
                debug!("last ttl: {}", last_response_ttl);
                Ok(last_response_ttl)
            } else {
                return Err(PistolError::AddressProtocolError { dst_addr, src_addr });
            }
        }
        IpAddr::V6(dst_ipv6) => {
            if let IpAddr::V6(src_ipv6) = src_addr {
                let mut last_response_hop_limit = 0;
                for hop_limit in 1..=30 {
                    let random_src_port = random_port_range(1000, 65535);
                    let (hop_status, rtt) = tcp6::send_syn_trace_packet(
                        dst_ipv6,
                        dst_port,
                        src_ipv6,
                        random_src_port,
                        hop_limit,
                        timeout,
                    )?;

                    match hop_status {
                        HopStatus::TimeExceeded(_) => {
                            last_response_hop_limit = hop_limit;
                        }
                        HopStatus::RecvReply(_) => {
                            debug!(
                                "hop_limit: {}, {:?} - {:.2}s",
                                hop_limit,
                                hop_status,
                                rtt.as_secs_f64()
                            );
                            // tcp rst packet, means packet arrive the target machine
                            return Ok(hop_limit);
                        }
                        _ => (),
                    }
                }
                debug!("last hop limit: {}", last_response_hop_limit);
                Ok(last_response_hop_limit)
            } else {
                return Err(PistolError::AddressProtocolError { dst_addr, src_addr });
            }
        }
    }
}

/// The recommended timeout value is 5 seconds, which is consistent with the default value of traceroute.
#[cfg(any(feature = "trace", feature = "os"))]
pub fn icmp_trace(
    dst_addr: IpAddr,
    src_addr: IpAddr,
    timeout: Option<Duration>,
) -> Result<u8, PistolError> {
    let mut rng = rand::rng();
    let icmp_id = rng.random();
    match dst_addr {
        IpAddr::V4(dst_ipv4) => {
            if let IpAddr::V4(src_ipv4) = src_addr {
                let mut ip_id = rng.random();
                if ip_id > u16::MAX - 30 {
                    ip_id -= 30;
                }
                let mut last_response_ttl = 0;
                for ttl in 1..=30 {
                    let (hop_status, rtt) = icmp::send_icmp_trace_packet(
                        dst_ipv4, src_ipv4, ip_id, ttl, icmp_id, ttl as u16, timeout,
                    )?;
                    ip_id += 1;
                    match hop_status {
                        HopStatus::TimeExceeded(_) => last_response_ttl = ttl,
                        // recv echo reply, means packet arrive the target machine
                        HopStatus::RecvReply(_) => {
                            debug!("ttl: {}, {:?} - {:.2}s", ttl, hop_status, rtt.as_secs_f64());
                            return Ok(ttl);
                        }
                        _ => (),
                    }
                }
                debug!("last ttl: {}", last_response_ttl);
                Ok(last_response_ttl)
            } else {
                return Err(PistolError::AddressProtocolError { dst_addr, src_addr });
            }
        }
        IpAddr::V6(dst_ipv6) => {
            if let IpAddr::V6(src_ipv6) = src_addr {
                let mut last_response_hop_limit = 0;
                for hop_limit in 1..=30 {
                    let (hop_status, rtt) = icmpv6::send_icmpv6_trace_packet(
                        dst_ipv6,
                        src_ipv6,
                        hop_limit,
                        icmp_id,
                        hop_limit as u16,
                        timeout,
                    )?;

                    match hop_status {
                        HopStatus::TimeExceeded(_) => last_response_hop_limit = hop_limit,
                        // recv echo reply, means packet arrive the target machine
                        HopStatus::RecvReply(_) => {
                            debug!(
                                "hop_limit: {}, {:?} - {:.2}s",
                                hop_limit,
                                hop_status,
                                rtt.as_secs_f64()
                            );
                            return Ok(hop_limit);
                        }
                        _ => (),
                    }
                }
                debug!("last hop limit: {}", last_response_hop_limit);
                Ok(last_response_hop_limit)
            } else {
                return Err(PistolError::AddressProtocolError { dst_addr, src_addr });
            }
        }
    }
}

/// The recommended timeout value is 5 seconds, which is consistent with the default value of traceroute.
#[cfg(feature = "trace")]
pub fn udp_trace(
    dst_addr: IpAddr,
    src_addr: IpAddr,
    timeout: Option<Duration>,
) -> Result<u8, PistolError> {
    match dst_addr {
        IpAddr::V4(dst_ipv4) => {
            if let IpAddr::V4(src_ipv4) = src_addr {
                let mut rng = rand::rng();
                let mut ip_id = rng.random();
                if ip_id > u16::MAX - 30 {
                    ip_id -= 30;
                }
                let mut last_response_ttl = 0;
                for ttl in 1..=30 {
                    let random_src_port = random_port_range(1000, 65535);
                    let dst_port = START_PORT + (ttl - 1) as u16;
                    let (hop_status, rtt) = udp::send_udp_trace_packet(
                        dst_ipv4,
                        dst_port,
                        src_ipv4,
                        random_src_port,
                        ip_id,
                        ttl,
                        timeout,
                    )?;
                    ip_id += 1;
                    match hop_status {
                        HopStatus::TimeExceeded(_) => last_response_ttl = ttl,
                        // icmpunreachable, means packet arrive the target machine
                        HopStatus::RecvReply(_) | HopStatus::Unreachable(_) => {
                            debug!("ttl: {}, {:?} - {:.2}s", ttl, hop_status, rtt.as_secs_f64());
                            return Ok(ttl);
                        }
                        _ => (),
                    }
                }
                debug!("last ttl: {}", last_response_ttl);
                Ok(last_response_ttl)
            } else {
                return Err(PistolError::AddressProtocolError { dst_addr, src_addr });
            }
        }
        IpAddr::V6(dst_ipv6) => {
            if let IpAddr::V6(src_ipv6) = src_addr {
                let mut last_response_hop_limit = 0;
                for hop_limit in 1..=30 {
                    let random_src_port = random_port_range(1000, 65535);
                    let dst_port = START_PORT + (hop_limit - 1) as u16;
                    let (hop_status, rtt) = udp6::send_udp_trace_packet(
                        dst_ipv6,
                        dst_port,
                        src_ipv6,
                        random_src_port,
                        hop_limit,
                        timeout,
                    )?;

                    match hop_status {
                        HopStatus::TimeExceeded(_) => last_response_hop_limit = hop_limit,
                        // icmpunreachable, means packet arrive the target machine
                        HopStatus::RecvReply(_) | HopStatus::Unreachable(_) => {
                            debug!(
                                "hop limit: {}, {:?} - {:.2}s",
                                hop_limit,
                                hop_status,
                                rtt.as_secs_f64()
                            );
                            return Ok(hop_limit);
                        }
                        _ => (),
                    }
                }
                debug!("last hop limit: {}", last_response_hop_limit);
                Ok(last_response_hop_limit)
            } else {
                return Err(PistolError::AddressProtocolError { dst_addr, src_addr });
            }
        }
    }
}

#[cfg(feature = "trace")]
#[cfg(test)]
mod tests {
    use super::*;
    use crate::PistolLogger;
    use crate::PistolRunner;
    use std::net::Ipv4Addr;
    use std::net::Ipv6Addr;
    use std::str::FromStr;
    #[test]
    fn test_get_hops_syn() {
        let _pr = PistolRunner::init(
            PistolLogger::Debug,
            None,
            None, // use default value
        )
        .unwrap();

        // let dst_ipv4 = Ipv4Addr::new(192, 168, 1, 3);
        // let dst_ipv4 = Ipv4Addr::new(192, 168, 5, 5);
        let dst_ipv4 = Ipv4Addr::new(182, 61, 244, 181);
        let src_ipv4 = Ipv4Addr::new(192, 168, 5, 3);
        let timeout = Some(Duration::from_secs_f64(5.0));
        let hops = syn_trace(dst_ipv4.into(), None, src_ipv4.into(), timeout).unwrap();
        println!("{}", hops);
    }
    #[test]
    fn test_get_hops_icmp() {
        let _pr = PistolRunner::init(
            PistolLogger::None,
            None,
            None, // use default value
        )
        .unwrap();

        // let dst_ipv4 = Ipv4Addr::new(192, 168, 1, 3);
        // let dst_ipv4 = Ipv4Addr::new(192, 168, 5, 5);
        let dst_ipv4 = Ipv4Addr::new(182, 61, 244, 181);
        let src_ipv4 = Ipv4Addr::new(192, 168, 5, 3);
        let timeout = Some(Duration::from_secs_f64(5.0));
        let hops = icmp_trace(dst_ipv4.into(), src_ipv4.into(), timeout).unwrap();
        println!("{}", hops);
    }
    #[test]
    fn test_get_hops_udp() {
        let _pr = PistolRunner::init(
            PistolLogger::Debug,
            None,
            None, // use default value
        )
        .unwrap();

        // let dst_ipv4 = Ipv4Addr::new(192, 168, 1, 3);
        let dst_ipv4 = Ipv4Addr::new(182, 61, 244, 181);
        let src_ipv4 = Ipv4Addr::new(192, 168, 5, 3);
        let timeout = Some(Duration::from_secs_f64(1.0));
        let hops = udp_trace(dst_ipv4.into(), src_ipv4.into(), timeout).unwrap();
        println!("{}", hops);
    }
    #[test]
    fn test_get_hops_udp6() {
        let _pr = PistolRunner::init(
            PistolLogger::Debug,
            None,
            None, // use default value
        )
        .unwrap();

        let src_ipv6 = Ipv6Addr::from_str("fe80::20c:29ff:fe5b:bd5c").unwrap();
        let dst_ipv6 = Ipv6Addr::from_str("fe80::20c:29ff:fe2c:9e4").unwrap();
        let timeout = Some(Duration::from_secs_f64(5.0));
        let hops = udp_trace(dst_ipv6.into(), src_ipv6.into(), timeout).unwrap();
        println!("{}", hops);
    }
}