async-icmp 0.2.1

use async_icmp::{
    message::echo::{EchoId, EchoSeq},
    ping::PingMultiplexer,
    socket::SocketConfig,
};
use log::{debug, warn};
use rand::seq::SliceRandom;
use std::{collections, net, time};

/// Allow using smaller params on macOS in a const-friendly way. MacOS chokes on IPv6 localhost
/// pings, and I haven't found away around that, so we just do less on macOS.
#[cfg(target_os = "linux")]
macro_rules! platform_param {
    ($linux_val:expr, $macos_val:expr) => {
        $linux_val
    };
}
#[cfg(target_os = "macos")]
macro_rules! platform_param {
    ($linux_val:expr, $macos_val:expr) => {
        $macos_val
    };
}

#[tokio::test]
async fn single_session_ipv4_ping() -> anyhow::Result<()> {
    multi_ping_test(&[(net::Ipv4Addr::LOCALHOST.into(), 1000)]).await
}

#[tokio::test]
async fn single_session_ipv6_ping() -> anyhow::Result<()> {
    multi_ping_test(&[(net::Ipv6Addr::LOCALHOST.into(), platform_param!(1000, 100))]).await
}

#[tokio::test]
async fn multi_session_ipv4_ping() -> anyhow::Result<()> {
    multi_ping_test(
        &[(net::Ipv4Addr::LOCALHOST.into(), platform_param!(50, 5)); platform_param!(50, 5)],
    )
    .await
}

#[tokio::test]
async fn multi_session_ipv6_ping() -> anyhow::Result<()> {
    multi_ping_test(
        &[(net::Ipv6Addr::LOCALHOST.into(), platform_param!(50, 5)); platform_param!(50, 5)],
    )
    .await
}

#[tokio::test]
async fn multi_session_multi_protocol_ping() -> anyhow::Result<()> {
    multi_ping_test(
        &[
            [(net::Ipv4Addr::LOCALHOST.into(), platform_param!(25, 3)); platform_param!(50, 5)]
                .as_slice(),
            [(net::Ipv6Addr::LOCALHOST.into(), platform_param!(25, 3)); platform_param!(50, 5)]
                .as_slice(),
        ]
        .concat(),
    )
    .await
}

/// random id if we can, otherwise local port as id
fn local_port_or_random_id(multiplexer: &PingMultiplexer, ip: net::IpAddr) -> EchoId {
    multiplexer
        .platform_echo_id(ip.into())
        .unwrap_or_else(rand::random)
}

/// For each ip and number of pings, send pings with a suitable id and 32-byte random data and
/// validate that all replies are present and correct.
async fn multi_ping_test(targets: &[(net::IpAddr, u16)]) -> anyhow::Result<()> {
    let _ = env_logger::builder()
        .is_test(true)
        .format_timestamp_millis()
        .try_init();

    let multiplexer = PingMultiplexer::new(SocketConfig::default(), SocketConfig::default())?;
    let mut rng = rand::thread_rng();
    // macos falls over if you send too quickly
    let (rate_limiter_tx, rate_limiter_rx) = flume::bounded(1);
    let rate_limiter_handle = tokio::spawn({
        async move {
            loop {
                if rate_limiter_tx.send_async(()).await.is_err() {
                    // receivers dropped
                    debug!("Rate limiter stopping");
                    break;
                }

                tokio::time::sleep(time::Duration::from_millis(1)).await;
            }
        }
    });

    let mut send_sessions = vec![];
    let mut recv_tasks = tokio::task::JoinSet::new();
    for (ip, num_pings) in targets.iter().cloned() {
        let id = local_port_or_random_id(&multiplexer, ip);
        let data = rand::random::<[u8; 32]>().to_vec();
        // if this succeeds we know the id & data are unique
        let (handle, mut rx) = multiplexer.add_session(ip, id, data.clone()).await?;

        // collect all timestamps into a vec

        let hash_key = (id, data);
        let recv_hash_key = hash_key.clone();
        recv_tasks.spawn(async move {
            let mut recv_timestamps = vec![];

            while recv_timestamps.len() < num_pings.into() {
                if let Some(recv_ts) = rx.recv().await {
                    recv_timestamps.push(recv_ts);
                    debug!("len now {}", recv_timestamps.len());
                } else {
                    warn!("Recv session closed?");
                    break;
                }
            }

            (recv_hash_key, recv_timestamps)
        });

        send_sessions.push((handle, num_pings, hash_key));
    }

    // scramble things a bit so we're not so dependent on scheduler ordering
    send_sessions.shuffle(&mut rng);

    // send everything
    let send_handles = send_sessions.iter().cloned().fold(
        tokio::task::JoinSet::<Result<_, anyhow::Error>>::new(),
        |mut set, (handle, num_pings, hash_key)| {
            let multiplexer = multiplexer.clone();
            let rate_limiter_clone = rate_limiter_rx.clone();
            set.spawn(async move {
                let mut sent_timestamps = collections::HashMap::new();
                for seq in (0..num_pings).map(EchoSeq::from_be) {
                    rate_limiter_clone.recv_async().await?;
                    let sent_at = multiplexer.send_ping(handle, seq).await?;
                    sent_timestamps.insert(seq, sent_at);
                }
                Ok((hash_key, sent_timestamps))
            });
            set
        },
    );

    // don't keep tx alive once sends are done
    drop(rate_limiter_rx);

    // wait for all the sends
    let send_tasks_data = send_handles
        .join_all()
        .await
        .into_iter()
        .collect::<Result<collections::HashMap<_, _>, _>>()?;

    // wait for all the receives
    let mut recv_tasks_data =
        tokio::time::timeout(time::Duration::from_secs(5), recv_tasks.join_all())
            .await?
            .into_iter()
            .collect::<collections::HashMap<_, _>>();

    // close sessions
    for (handle, _, _) in send_sessions {
        multiplexer.close_session(handle).await?;
    }

    // match sends to receives
    for (key, sent_at_by_seq) in send_tasks_data {
        // convert to map
        let recv_ts_vec = recv_tasks_data.remove(&key).unwrap();
        assert_eq!(sent_at_by_seq.len(), recv_ts_vec.len());

        let recv_at_by_seq = recv_ts_vec
            .into_iter()
            .map(|ts| (ts.seq, ts.received_at))
            .collect::<collections::HashMap<_, _>>();

        // make sure we didn't have any duplicates
        assert_eq!(sent_at_by_seq.len(), recv_at_by_seq.len());

        for (seq, sent_at) in sent_at_by_seq.into_iter() {
            let recv_at = recv_at_by_seq.get(&seq).unwrap();
            assert!(*recv_at >= sent_at);
        }
    }

    multiplexer.shutdown().await;
    rate_limiter_handle.await?;
    Ok(())
}