Finish verion 0.1.0

crates/wtfnet-probe/src/lib.rs

@@ -0,0 +1,520 @@
+#[cfg(unix)]
+use pnet::packet::icmp::{IcmpPacket, IcmpTypes};
+#[cfg(unix)]
+use pnet::packet::icmpv6::{Icmpv6Packet, Icmpv6Types};
+#[cfg(unix)]
+use pnet::packet::ip::IpNextHeaderProtocols;
+#[cfg(unix)]
+use pnet::transport::{
+    icmp_packet_iter, icmpv6_packet_iter, transport_channel, TransportChannelType,
+    TransportProtocol,
+};
+use serde::{Deserialize, Serialize};
+use socket2::{Domain, Protocol, Socket, Type};
+use std::net::{IpAddr, SocketAddr};
+use std::time::{Duration, Instant};
+use thiserror::Error;
+use tokio::net::{lookup_host, TcpStream};
+use tokio::time::timeout;
+use wtfnet_geoip::GeoIpRecord;
+
+#[derive(Debug, Error)]
+pub enum ProbeError {
+    #[error("resolution failed: {0}")]
+    Resolve(String),
+    #[error("io error: {0}")]
+    Io(String),
+    #[error("timeout")]
+    Timeout,
+    #[error("ping error: {0}")]
+    Ping(String),
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct PingResult {
+    pub seq: u16,
+    pub rtt_ms: Option<u128>,
+    pub error: Option<String>,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct PingSummary {
+    pub sent: u32,
+    pub received: u32,
+    pub loss_pct: f64,
+    pub min_ms: Option<u128>,
+    pub avg_ms: Option<f64>,
+    pub max_ms: Option<u128>,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct PingReport {
+    pub target: String,
+    pub ip: Option<String>,
+    pub geoip: Option<GeoIpRecord>,
+    pub timeout_ms: u64,
+    pub count: u32,
+    pub results: Vec<PingResult>,
+    pub summary: PingSummary,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct TcpPingResult {
+    pub seq: u16,
+    pub rtt_ms: Option<u128>,
+    pub error: Option<String>,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct TcpPingReport {
+    pub target: String,
+    pub ip: Option<String>,
+    pub geoip: Option<GeoIpRecord>,
+    pub port: u16,
+    pub timeout_ms: u64,
+    pub count: u32,
+    pub results: Vec<TcpPingResult>,
+    pub summary: PingSummary,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct TraceHop {
+    pub ttl: u8,
+    pub addr: Option<String>,
+    pub rtt_ms: Option<u128>,
+    pub note: Option<String>,
+    pub geoip: Option<GeoIpRecord>,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct TraceReport {
+    pub target: String,
+    pub ip: Option<String>,
+    pub geoip: Option<GeoIpRecord>,
+    pub port: u16,
+    pub max_hops: u8,
+    pub timeout_ms: u64,
+    pub protocol: String,
+    pub hops: Vec<TraceHop>,
+}
+
+pub async fn ping(
+    target: &str,
+    count: u32,
+    timeout_ms: u64,
+    interval_ms: u64,
+) -> Result<PingReport, ProbeError> {
+    let addr = resolve_one(target).await?;
+    let mut results = Vec::new();
+    let mut received = 0u32;
+    let mut min = None;
+    let mut max = None;
+    let mut sum = 0u128;
+
+    let config = match addr {
+        IpAddr::V4(_) => surge_ping::Config::default(),
+        IpAddr::V6(_) => surge_ping::Config::builder()
+            .kind(surge_ping::ICMP::V6)
+            .build(),
+    };
+    let client = surge_ping::Client::new(&config)
+        .map_err(|err| ProbeError::Ping(err.to_string()))?;
+    let mut pinger = client
+        .pinger(addr, surge_ping::PingIdentifier(0))
+        .await;
+    let timeout_dur = Duration::from_millis(timeout_ms);
+
+    for seq in 0..count {
+        let seq = seq as u16;
+        let start = Instant::now();
+        let response =
+            timeout(timeout_dur, pinger.ping(surge_ping::PingSequence(seq), &[0; 8])).await;
+        match response {
+            Ok(Ok((_packet, _))) => {
+                let rtt = start.elapsed().as_millis();
+                received += 1;
+                min = Some(min.map_or(rtt, |value: u128| value.min(rtt)));
+                max = Some(max.map_or(rtt, |value: u128| value.max(rtt)));
+                sum += rtt;
+                results.push(PingResult {
+                    seq,
+                    rtt_ms: Some(rtt),
+                    error: None,
+                });
+            }
+            Ok(Err(err)) => {
+                results.push(PingResult {
+                    seq,
+                    rtt_ms: None,
+                    error: Some(err.to_string()),
+                });
+            }
+            Err(_) => {
+                results.push(PingResult {
+                    seq,
+                    rtt_ms: None,
+                    error: Some("timeout".to_string()),
+                });
+            }
+        }
+        if interval_ms > 0 {
+            tokio::time::sleep(Duration::from_millis(interval_ms)).await;
+        }
+    }
+
+    let summary = build_summary(count, received, min, max, sum);
+    Ok(PingReport {
+        target: target.to_string(),
+        ip: Some(addr.to_string()),
+        geoip: None,
+        timeout_ms,
+        count,
+        results,
+        summary,
+    })
+}
+
+pub async fn tcp_ping(
+    target: &str,
+    port: u16,
+    count: u32,
+    timeout_ms: u64,
+) -> Result<TcpPingReport, ProbeError> {
+    let addr = resolve_one(target).await?;
+    let socket_addr = SocketAddr::new(addr, port);
+    let timeout_dur = Duration::from_millis(timeout_ms);
+    let mut results = Vec::new();
+    let mut received = 0u32;
+    let mut min = None;
+    let mut max = None;
+    let mut sum = 0u128;
+
+    for seq in 0..count {
+        let seq = seq as u16;
+        let start = Instant::now();
+        let attempt = timeout(timeout_dur, TcpStream::connect(socket_addr)).await;
+        match attempt {
+            Ok(Ok(_stream)) => {
+                let rtt = start.elapsed().as_millis();
+                received += 1;
+                min = Some(min.map_or(rtt, |value: u128| value.min(rtt)));
+                max = Some(max.map_or(rtt, |value: u128| value.max(rtt)));
+                sum += rtt;
+                results.push(TcpPingResult {
+                    seq,
+                    rtt_ms: Some(rtt),
+                    error: None,
+                });
+            }
+            Ok(Err(err)) => {
+                results.push(TcpPingResult {
+                    seq,
+                    rtt_ms: None,
+                    error: Some(err.to_string()),
+                });
+            }
+            Err(_) => {
+                results.push(TcpPingResult {
+                    seq,
+                    rtt_ms: None,
+                    error: Some("timeout".to_string()),
+                });
+            }
+        }
+    }
+
+    let summary = build_summary(count, received, min, max, sum);
+    Ok(TcpPingReport {
+        target: target.to_string(),
+        ip: Some(addr.to_string()),
+        geoip: None,
+        port,
+        timeout_ms,
+        count,
+        results,
+        summary,
+    })
+}
+
+pub async fn tcp_trace(
+    target: &str,
+    port: u16,
+    max_hops: u8,
+    timeout_ms: u64,
+) -> Result<TraceReport, ProbeError> {
+    let addr = resolve_one(target).await?;
+    let socket_addr = SocketAddr::new(addr, port);
+    let timeout_dur = Duration::from_millis(timeout_ms);
+    let mut hops = Vec::new();
+
+    for ttl in 1..=max_hops {
+        let addr = socket_addr;
+        let start = Instant::now();
+        let result = tokio::task::spawn_blocking(move || {
+            tcp_connect_with_ttl(addr, ttl, timeout_dur)
+        })
+        .await
+        .map_err(|err| ProbeError::Io(err.to_string()))?;
+
+        match result {
+            Ok(()) => {
+                let rtt = start.elapsed().as_millis();
+                hops.push(TraceHop {
+                    ttl,
+                    addr: Some(socket_addr.ip().to_string()),
+                    rtt_ms: Some(rtt),
+                    note: None,
+                    geoip: None,
+                });
+                break;
+            }
+            Err(err) => {
+                let rtt = start.elapsed().as_millis();
+                hops.push(TraceHop {
+                    ttl,
+                    addr: None,
+                    rtt_ms: Some(rtt),
+                    note: Some(err.to_string()),
+                    geoip: None,
+                });
+            }
+        }
+    }
+
+    Ok(TraceReport {
+        target: target.to_string(),
+        ip: Some(addr.to_string()),
+        geoip: None,
+        port,
+        max_hops,
+        timeout_ms,
+        protocol: "tcp".to_string(),
+        hops,
+    })
+}
+
+pub async fn udp_trace(
+    target: &str,
+    port: u16,
+    max_hops: u8,
+    timeout_ms: u64,
+) -> Result<TraceReport, ProbeError> {
+    let addr = resolve_one(target).await?;
+
+    let timeout_dur = Duration::from_millis(timeout_ms);
+    let mut hops = Vec::new();
+
+    for ttl in 1..=max_hops {
+        let addr = SocketAddr::new(addr, port);
+        let start = Instant::now();
+        let result = tokio::task::spawn_blocking(move || udp_trace_hop(addr, ttl, timeout_dur))
+        .await
+        .map_err(|err| ProbeError::Io(err.to_string()))?;
+
+        match result {
+            Ok((hop_addr, reached)) => {
+                let rtt = start.elapsed().as_millis();
+                hops.push(TraceHop {
+                    ttl,
+                    addr: hop_addr.map(|ip| ip.to_string()),
+                    rtt_ms: Some(rtt),
+                    note: None,
+                    geoip: None,
+                });
+                if reached {
+                    break;
+                }
+            }
+            Err(err) => {
+                hops.push(TraceHop {
+                    ttl,
+                    addr: None,
+                    rtt_ms: None,
+                    note: Some(err.to_string()),
+                    geoip: None,
+                });
+            }
+        }
+    }
+
+    Ok(TraceReport {
+        target: target.to_string(),
+        ip: Some(addr.to_string()),
+        geoip: None,
+        port,
+        max_hops,
+        timeout_ms,
+        protocol: "udp".to_string(),
+        hops,
+    })
+}
+
+fn build_summary(
+    sent: u32,
+    received: u32,
+    min: Option<u128>,
+    max: Option<u128>,
+    sum: u128,
+) -> PingSummary {
+    let loss_pct = if sent == 0 {
+        0.0
+    } else {
+        ((sent - received) as f64 / sent as f64) * 100.0
+    };
+    let avg_ms = if received == 0 {
+        None
+    } else {
+        Some(sum as f64 / received as f64)
+    };
+    PingSummary {
+        sent,
+        received,
+        loss_pct,
+        min_ms: min,
+        avg_ms,
+        max_ms: max,
+    }
+}
+
+async fn resolve_one(target: &str) -> Result<IpAddr, ProbeError> {
+    let mut iter = lookup_host((target, 0))
+        .await
+        .map_err(|err| ProbeError::Resolve(err.to_string()))?;
+    iter.next()
+        .map(|addr| addr.ip())
+        .ok_or_else(|| ProbeError::Resolve("no address found".to_string()))
+}
+
+fn tcp_connect_with_ttl(
+    addr: SocketAddr,
+    ttl: u8,
+    timeout: Duration,
+) -> Result<(), ProbeError> {
+    let domain = match addr.ip() {
+        IpAddr::V4(_) => Domain::IPV4,
+        IpAddr::V6(_) => Domain::IPV6,
+    };
+    let socket = Socket::new(domain, Type::STREAM, Some(Protocol::TCP))
+        .map_err(|err| ProbeError::Io(err.to_string()))?;
+    match addr.ip() {
+        IpAddr::V4(_) => socket
+            .set_ttl_v4(u32::from(ttl))
+            .map_err(|err| ProbeError::Io(err.to_string()))?,
+        IpAddr::V6(_) => socket
+            .set_unicast_hops_v6(u32::from(ttl))
+            .map_err(|err| ProbeError::Io(err.to_string()))?,
+    }
+    socket
+        .connect_timeout(&addr.into(), timeout)
+        .map_err(|err| ProbeError::Io(err.to_string()))?;
+    Ok(())
+}
+
+#[cfg(unix)]
+fn udp_trace_hop(
+    addr: SocketAddr,
+    ttl: u8,
+    timeout: Duration,
+) -> Result<(Option<IpAddr>, bool), ProbeError> {
+    match addr.ip() {
+        IpAddr::V4(_) => udp_trace_hop_v4(addr, ttl, timeout),
+        IpAddr::V6(_) => udp_trace_hop_v6(addr, ttl, timeout),
+    }
+}
+
+#[cfg(not(unix))]
+fn udp_trace_hop(
+    _addr: SocketAddr,
+    _ttl: u8,
+    _timeout: Duration,
+) -> Result<(Option<IpAddr>, bool), ProbeError> {
+    Err(ProbeError::Io(
+        "udp trace not supported on this platform".to_string(),
+    ))
+}
+
+#[cfg(unix)]
+fn udp_trace_hop_v4(
+    addr: SocketAddr,
+    ttl: u8,
+    timeout: Duration,
+) -> Result<(Option<IpAddr>, bool), ProbeError> {
+    let protocol =
+        TransportChannelType::Layer4(TransportProtocol::Ipv4(IpNextHeaderProtocols::Icmp));
+    let (_tx, mut rx) = transport_channel(4096, protocol)
+        .map_err(|err| ProbeError::Io(err.to_string()))?;
+
+    let socket = std::net::UdpSocket::bind("0.0.0.0:0")
+        .map_err(|err| ProbeError::Io(err.to_string()))?;
+    socket
+        .set_ttl(u32::from(ttl))
+        .map_err(|err| ProbeError::Io(err.to_string()))?;
+    let _ = socket.send_to(&[0u8; 4], addr);
+
+    let mut iter = icmp_packet_iter(&mut rx);
+    match iter.next_with_timeout(timeout) {
+        Ok(Some((packet, addr))) => {
+            if let Some(result) = interpret_icmp_v4(&packet) {
+                return Ok((Some(addr), result));
+            }
+            Ok((Some(addr), false))
+        }
+        Ok(None) => Err(ProbeError::Timeout),
+        Err(err) => Err(ProbeError::Io(err.to_string())),
+    }
+}
+
+#[cfg(unix)]
+fn udp_trace_hop_v6(
+    addr: SocketAddr,
+    ttl: u8,
+    timeout: Duration,
+) -> Result<(Option<IpAddr>, bool), ProbeError> {
+    let protocol =
+        TransportChannelType::Layer4(TransportProtocol::Ipv6(IpNextHeaderProtocols::Icmpv6));
+    let (_tx, mut rx) = transport_channel(4096, protocol)
+        .map_err(|err| ProbeError::Io(err.to_string()))?;
+
+    let socket = std::net::UdpSocket::bind("[::]:0")
+        .map_err(|err| ProbeError::Io(err.to_string()))?;
+    socket
+        .set_unicast_hops_v6(u32::from(ttl))
+        .map_err(|err| ProbeError::Io(err.to_string()))?;
+    let _ = socket.send_to(&[0u8; 4], addr);
+
+    let mut iter = icmpv6_packet_iter(&mut rx);
+    match iter.next_with_timeout(timeout) {
+        Ok(Some((packet, addr))) => {
+            if let Some(result) = interpret_icmp_v6(&packet) {
+                return Ok((Some(addr), result));
+            }
+            Ok((Some(addr), false))
+        }
+        Ok(None) => Err(ProbeError::Timeout),
+        Err(err) => Err(ProbeError::Io(err.to_string())),
+    }
+}
+
+#[cfg(unix)]
+fn interpret_icmp_v4(packet: &IcmpPacket) -> Option<bool> {
+    let icmp_type = packet.get_icmp_type();
+    if icmp_type == IcmpTypes::TimeExceeded {
+        return Some(false);
+    }
+    if icmp_type == IcmpTypes::DestinationUnreachable {
+        return Some(true);
+    }
+    None
+}
+
+#[cfg(unix)]
+fn interpret_icmp_v6(packet: &Icmpv6Packet) -> Option<bool> {
+    let icmp_type = packet.get_icmpv6_type();
+    if icmp_type == Icmpv6Types::TimeExceeded {
+        return Some(false);
+    }
+    if icmp_type == Icmpv6Types::DestinationUnreachable {
+        return Some(true);
+    }
+    None
+}