sqm-autorate-rust/src/pinger.rs

use crate::SHUTDOWN;
use crate::MeasurementType;
use crate::metrics::{Metric, MetricsSender};
use crate::util::RwLockExt;
use icmp_socket2::socket::IcmpSocket;
use icmp_socket2::{IcmpSocket4, Icmpv4Packet};
use log::{debug, info, warn};
use std::net::{IpAddr, Ipv4Addr};
use std::sync::atomic::Ordering;
use std::sync::mpsc::Sender;
use std::sync::{Arc, RwLock};
use std::time::{Duration, Instant};
use std::{io, thread};
use thiserror::Error;

#[derive(Error, Debug)]
pub enum PingError {
    #[error("Couldn't parse number")]
    InvalidNumber(#[from] io::Error),
    #[error("Error parsing packet")]
    InvalidPacket(String),
    #[error("Invalid packet type")]
    InvalidType(String),
    #[error("Wrong ICMP identifier (expected {expected:?}, found {found:?})")]
    WrongID { expected: u16, found: u16 },
}

pub struct PingReply {
    pub reflector: IpAddr,
    pub measurement_type: MeasurementType,
    pub seq: u16,
    pub rtt: f64,
    pub current_time: i64,
    pub down_time: f64,
    pub up_time: f64,
    pub originate_timestamp: i64,
    pub receive_timestamp: i64,
    pub transmit_timestamp: i64,
    pub last_receive_time_s: Instant,
}

fn open_socket(type_: MeasurementType) -> io::Result<IcmpSocket4> {
    match type_ {
        MeasurementType::Icmp | MeasurementType::IcmpTimestamps => IcmpSocket4::new(),
        _ => {
            unimplemented!()
        }
    }
}

pub trait PingListener {
    fn listen(
        &mut self,
        id: u16,
        type_: MeasurementType,
        reflectors_lock: Arc<RwLock<Vec<IpAddr>>>,
        stats_tx: Sender<PingReply>,
        ping_metrics: MetricsSender,
    ) -> anyhow::Result<()> {
        let socket = &mut open_socket(type_)?;

        loop {
            if SHUTDOWN.load(Ordering::Relaxed) {
                info!("Ping listener shutting down");
                return Ok(());
            }
            let (pkt, sender) = match socket.rcv_from() {
                Ok(val) => val,
                Err(_) => continue,
            };

            let addr: IpAddr = sender.as_socket().unwrap().ip();

            let reflectors = reflectors_lock.read_anyhow()?;
            if !reflectors.contains(&addr) {
                continue;
            }

            let reply = match self.parse_packet(id, addr, type_, pkt) {
                Ok(val) => val,
                Err(_) => {
                    // parse_packet will throw an error if it's an unknown protocol etc.
                    // so just quietly move on
                    continue;
                }
            };

            ping_metrics.send(Metric::Ping {
                reflector: reply.reflector,
                measurement_type: reply.measurement_type,
                rtt: reply.rtt,
                up_time: reply.up_time,
                down_time: reply.down_time,
            });

            debug!(
                "Type: {:4}  | Reflector IP: {:>15}  | Seq: {:5}  | Current time: {:8}  |  Originate: {:8}  |  Received time: {:8}  |  Transmit time : {:8}  |  RTT: {:8}  | UL time: {:5}  | DL time: {:5}",
                "ICMP",
                addr.to_string(),
                reply.seq,
                reply.current_time,
                reply.originate_timestamp,
                reply.receive_timestamp,
                reply.transmit_timestamp,
                reply.rtt,
                reply.up_time,
                reply.down_time
            );
            if let Err(e) = stats_tx.send(reply) {
                warn!("Stats channel closed, stopping listener: {}", e);
                break Ok(());
            }
        }
    }

    fn parse_packet(
        &self,
        id: u16,
        reflector: IpAddr,
        measurement_type: MeasurementType,
        pkt: Icmpv4Packet,
    ) -> Result<PingReply, PingError>;
}

pub trait PingSender {
    fn send(
        &mut self,
        id: u16,
        type_: MeasurementType,
        reflectors_lock: Arc<RwLock<Vec<IpAddr>>>,
        tick_interval: f64,
    ) -> anyhow::Result<()> {
        let mut socket = open_socket(type_)?;

        let mut seq: u16 = 0;
        let tick_duration_ms: u16 = (tick_interval * 1000.0) as u16;

        loop {
            if SHUTDOWN.load(Ordering::Relaxed) {
                info!("Ping sender shutting down");
                return Ok(());
            }
            // Clone the reflectors vec and drop the read lock immediately —
            // holding it for the entire tick would starve the reflector selector's write lock.
            let reflectors_unlocked = reflectors_lock.read_anyhow()?;
            let reflectors = reflectors_unlocked.clone();
            drop(reflectors_unlocked);

            if reflectors.is_empty() {
                thread::sleep(Duration::from_millis(tick_duration_ms as u64));
                continue;
            }

            let sleep_duration =
                Duration::from_millis((tick_duration_ms / reflectors.len() as u16) as u64);

            for reflector in reflectors.iter() {
                let addr: Ipv4Addr = match reflector {
                    IpAddr::V4(ipv4) => *ipv4,
                    IpAddr::V6(_) => unimplemented!(),
                };

                socket.send_to(addr, self.craft_packet(id, seq))?;
                thread::sleep(sleep_duration);
            }

            if seq == u16::MAX {
                seq = 0;
            } else {
                seq += 1;
            }
        }
    }

    fn craft_packet(&self, id: u16, seq: u16) -> icmp_socket2::packet::Icmpv4Packet;
}