anti-ping 0.1.2

//! UDP ping implementation
//!
//! This module provides UDP-based connectivity testing by sending UDP packets to
//! high-numbered ports and listening for ICMP Port Unreachable responses.
//! This technique is commonly used by tools like mtr and traceroute.

use anti_common::{icmp, ports, PingConfig, PingError, PingReply, PingResult};
use socket2::{Domain, Protocol, Socket, Type};
use std::io::Read;
use std::net::{IpAddr, Ipv4Addr, SocketAddr, UdpSocket};
use std::time::{Duration, Instant};

/// UDP pinger that sends packets to high ports and listens for ICMP responses
pub struct UdpPinger {
    config: PingConfig,
    udp_socket: UdpSocket,
    icmp_socket: Socket,
    base_port: u16,
}

impl UdpPinger {
    /// Create a new UDP pinger
    pub fn new(config: PingConfig) -> PingResult<Self> {
        // Create UDP socket for sending
        let udp_socket =
            UdpSocket::bind("0.0.0.0:0").map_err(|e| PingError::SocketCreation(e.to_string()))?;

        // Create ICMP socket for receiving Port Unreachable responses
        let icmp_socket = Socket::new(Domain::IPV4, Type::DGRAM, Some(Protocol::ICMPV4))
            .map_err(|e| {
                if e.kind() == std::io::ErrorKind::PermissionDenied {
                    PingError::PermissionDenied {
                        context: "UDP ping requires ICMP socket access. Try running with elevated privileges or use TCP ping instead.".to_string(),
                    }
                } else {
                    PingError::SocketCreation(e.to_string())
                }
            })?;

        icmp_socket
            .set_read_timeout(Some(config.timeout))
            .map_err(|e| PingError::SocketCreation(e.to_string()))?;

        // Use traceroute-style base port
        let base_port = ports::TRACEROUTE_BASE;

        Ok(Self {
            config,
            udp_socket,
            icmp_socket,
            base_port,
        })
    }

    /// Create a new UDP pinger with custom port
    pub fn new_with_port(config: PingConfig, port: u16) -> PingResult<Self> {
        let mut pinger = Self::new(config)?;
        pinger.base_port = port;
        Ok(pinger)
    }

    /// Send a single UDP ping and wait for ICMP Port Unreachable response
    pub fn ping(&self, sequence: u16) -> PingResult<PingReply> {
        // Calculate target port (like traceroute: base + sequence)
        let target_port = self.base_port.wrapping_add(sequence);
        let target_addr = SocketAddr::new(IpAddr::V4(self.config.target), target_port);

        let start_time = Instant::now();

        // Send UDP packet to likely closed port
        let test_data = format!(
            "UDP_PING_SEQ_{:04}_{}",
            sequence,
            start_time.elapsed().as_nanos()
        );
        self.udp_socket
            .send_to(test_data.as_bytes(), target_addr)
            .map_err(|e| PingError::SocketCreation(e.to_string()))?;

        // Listen for ICMP Port Unreachable response
        self.wait_for_icmp_response(sequence, target_port, start_time)
    }

    /// Wait for and parse ICMP Port Unreachable response
    fn wait_for_icmp_response(
        &self,
        sequence: u16,
        target_port: u16,
        start_time: Instant,
    ) -> PingResult<PingReply> {
        let mut buffer = [0u8; 1500]; // MTU-sized buffer
        let deadline = start_time + self.config.timeout;

        loop {
            let now = Instant::now();
            if now >= deadline {
                return Err(PingError::Timeout {
                    duration: self.config.timeout,
                });
            }

            // Update timeout for remaining time
            let remaining = deadline - now;
            self.icmp_socket
                .set_read_timeout(Some(remaining))
                .map_err(|e| PingError::SocketCreation(e.to_string()))?;

            match (&self.icmp_socket).read(&mut buffer) {
                Ok(bytes_received) => {
                    let recv_time = Instant::now();

                    if let Some(reply) = self.parse_icmp_response(
                        &buffer[..bytes_received],
                        target_port,
                        sequence,
                        recv_time - start_time,
                    ) {
                        return Ok(reply);
                    }
                    // Continue listening for the right packet
                }
                Err(e) => {
                    if e.kind() == std::io::ErrorKind::TimedOut {
                        return Err(PingError::Timeout {
                            duration: self.config.timeout,
                        });
                    } else if e.kind() == std::io::ErrorKind::WouldBlock {
                        continue;
                    } else {
                        return Err(PingError::SocketCreation(e.to_string()));
                    }
                }
            }
        }
    }

    /// Parse ICMP response to check for Port Unreachable
    fn parse_icmp_response(
        &self,
        buffer: &[u8],
        expected_port: u16,
        sequence: u16,
        rtt: Duration,
    ) -> Option<PingReply> {
        // Minimum size for IP header + ICMP header + original IP header + UDP header
        if buffer.len() < 48 {
            return None;
        }

        // Parse IP header
        let ip_header_len = ((buffer[0] & 0x0F) * 4) as usize;
        if buffer.len() < ip_header_len + 8 {
            return None;
        }

        // Extract source IP
        let source_ip = if buffer.len() >= 20 {
            Ipv4Addr::new(buffer[12], buffer[13], buffer[14], buffer[15])
        } else {
            self.config.target
        };

        // Get ICMP data
        let icmp_data = &buffer[ip_header_len..];
        if icmp_data.len() < 8 {
            return None;
        }

        // Check for ICMP Destination Unreachable, Port Unreachable
        if icmp_data[0] != icmp::DEST_UNREACHABLE || icmp_data[1] != icmp::PORT_UNREACHABLE {
            return None;
        }

        // Parse the original packet embedded in ICMP payload
        if icmp_data.len() < 28 {
            return None; // Not enough data for original IP + UDP headers
        }

        let orig_ip_data = &icmp_data[8..]; // Skip ICMP header
        if orig_ip_data.len() < 20 {
            return None;
        }

        // Check if original packet was UDP (protocol 17)
        if orig_ip_data[9] != 17 {
            return None;
        }

        let orig_ip_header_len = ((orig_ip_data[0] & 0x0F) * 4) as usize;
        if orig_ip_data.len() < orig_ip_header_len + 8 {
            return None;
        }

        // Parse original UDP header
        let orig_udp_data = &orig_ip_data[orig_ip_header_len..];
        if orig_udp_data.len() < 8 {
            return None;
        }

        let orig_dest_port = u16::from_be_bytes([orig_udp_data[2], orig_udp_data[3]]);

        // Verify this is our packet
        if orig_dest_port == expected_port {
            Some(PingReply {
                sequence,
                rtt,
                bytes_received: buffer.len(),
                from: source_ip,
                ttl: None, // Could extract from IP header if needed
            })
        } else {
            None
        }
    }

    /// Get the target configuration
    pub fn config(&self) -> &PingConfig {
        &self.config
    }

    /// Get the base port being used
    pub fn base_port(&self) -> u16 {
        self.base_port
    }
}

/// Utility function to check if a UDP port is likely closed
pub fn is_port_likely_closed(port: u16) -> bool {
    // Ports typically closed or used for system services
    match port {
        // Traceroute range
        33434..=33534 => true,
        // High ephemeral ports
        49152..=65535 => true,
        // Other commonly closed ports
        1234 | 5678 | 9999 => true,
        _ => false,
    }
}

/// Get a suitable UDP port for pinging based on sequence
pub fn get_udp_port(sequence: u16, base_port: Option<u16>) -> u16 {
    let base = base_port.unwrap_or(ports::TRACEROUTE_BASE);
    base.wrapping_add(sequence % 100) // Wrap to avoid going too high
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::time::Duration;

    #[test]
    fn test_port_selection() {
        assert!(is_port_likely_closed(33434));
        assert!(is_port_likely_closed(50000));
        assert!(!is_port_likely_closed(80));
        assert!(!is_port_likely_closed(443));
    }

    #[test]
    fn test_get_udp_port() {
        let port1 = get_udp_port(1, Some(10000));
        let port2 = get_udp_port(2, Some(10000));
        assert_eq!(port1, 10001);
        assert_eq!(port2, 10002);

        // Test default base port
        let port3 = get_udp_port(1, None);
        assert_eq!(port3, ports::TRACEROUTE_BASE + 1);
    }

    #[test]
    fn test_config_creation() {
        let config = PingConfig {
            target: Ipv4Addr::new(8, 8, 8, 8),
            count: 1,
            timeout: Duration::from_secs(2),
            ..Default::default()
        };

        // This will fail in test environment due to socket permissions,
        // but we can at least verify the config is accepted
        assert_eq!(config.target, Ipv4Addr::new(8, 8, 8, 8));
        assert_eq!(config.count, 1);
    }

    #[test]
    fn test_parse_icmp_response_insufficient_data() {
        let _config = PingConfig::default();
        // This will fail due to socket creation, but we're testing parsing logic
        let buffer = vec![0u8; 10]; // Too small

        // We can't easily test the parsing without a real UdpPinger instance,
        // but this tests that small buffers are handled correctly
        assert!(buffer.len() < 48);
    }
}