lancer / programs / ping / src / main.rs

#![no_std]
#![no_main]

use lancer_core::packet_ring::{PacketRingReader, PacketRingWriter};
use lancer_user::io::write_bytes;
use lancer_user::net;
use lancer_user::syscall;

const MAX_PINGS: u8 = 8;
const DEFAULT_COUNT: u8 = 4;

enum Target {
    V4([u8; 4]),
    V6([u8; 16]),
}

struct DnsResult {
    v6: Option<[u8; 16]>,
    v4: Option<[u8; 4]>,
}

#[unsafe(no_mangle)]
pub extern "C" fn lancer_main() -> ! {
    let (rx, tx) = match net::init() {
        Some(pair) => pair,
        None => {
            lancer_user::show!(net, error, "netsock init failed");
            syscall::exit();
        }
    };

    if !net::has_relay() {
        write_bytes(b"networking not available (use SSH)\n");
        syscall::exit();
    }

    let mut args_buf = [0u8; 128];
    let args_len = net::recv_args(&rx, &mut args_buf);

    let args = &args_buf[..args_len];

    let (cmd_tok, rest) = net::next_token(args);
    let after_cmd = match net::bytes_eq(cmd_tok, b"ping") {
        true => rest,
        false => args,
    };

    let (host_tok, remainder) = net::next_token(after_cmd);
    if host_tok.is_empty() {
        write_bytes(b"usage: ping <host|ip> [count]\n");
        syscall::exit();
    }

    let (count_tok, _) = net::next_token(remainder);
    let count = match count_tok.is_empty() {
        true => DEFAULT_COUNT,
        false => match net::parse_decimal_u16(count_tok) {
            Some(n) if n >= 1 && n <= MAX_PINGS as u16 => n as u8,
            Some(_) => MAX_PINGS,
            None => DEFAULT_COUNT,
        },
    };

    let target = match net::parse_ip(host_tok) {
        Some(ip) => Target::V4(ip),
        None => match net::parse_ipv6(host_tok) {
            Some(addr) => Target::V6(addr),
            None => pick_target(&rx, &tx, host_tok),
        },
    };

    write_bytes(b"pinging...\n");
    send_ping(&tx, &target, count);
    recv_results(&rx, true);
    syscall::exit()
}

fn pick_target(rx: &PacketRingReader, tx: &PacketRingWriter, hostname: &[u8]) -> Target {
    let dns = resolve_both(rx, tx, hostname);
    match (dns.v6, dns.v4) {
        (None, None) => {
            write_bytes(b"DNS lookup failed\n");
            syscall::exit()
        }
        (Some(v6), None) => Target::V6(v6),
        (None, Some(v4)) => Target::V4(v4),
        (Some(v6), Some(v4)) => {
            send_ping(tx, &Target::V6(v6), 1);
            let (received, _) = recv_results(rx, false);
            match received > 0 {
                true => Target::V6(v6),
                false => Target::V4(v4),
            }
        }
    }
}

fn resolve_both(rx: &PacketRingReader, tx: &PacketRingWriter, hostname: &[u8]) -> DnsResult {
    let mut msg = [0u8; 62];
    let len = hostname.len().min(61);
    msg[1..1 + len].copy_from_slice(&hostname[..len]);
    let mut resp_buf = [0u8; 64];

    msg[0] = net::MSG_DNS_QUERY6;
    net::send_request(tx, &msg[..1 + len]);
    let n = net::recv_response_blocking(rx, &mut resp_buf);
    let v6 = match n >= 18 && resp_buf[0] == net::MSG_DNS_RESULT6 && resp_buf[1] == 0 {
        true => {
            let mut addr = [0u8; 16];
            addr.copy_from_slice(&resp_buf[2..18]);
            Some(addr)
        }
        false => None,
    };

    msg[0] = net::MSG_DNS_QUERY;
    net::send_request(tx, &msg[..1 + len]);
    let n = net::recv_response_blocking(rx, &mut resp_buf);
    let v4 = match n >= 6 && resp_buf[0] == net::MSG_DNS_RESULT && resp_buf[1] == 0 {
        true => Some([resp_buf[2], resp_buf[3], resp_buf[4], resp_buf[5]]),
        false => None,
    };

    DnsResult { v6, v4 }
}

fn send_ping(tx: &PacketRingWriter, target: &Target, count: u8) {
    match target {
        Target::V6(addr) => {
            let mut req = [0u8; 18];
            req[0] = net::MSG_PING_REQUEST6;
            req[1..17].copy_from_slice(addr);
            req[17] = count;
            net::send_request(tx, &req);
        }
        Target::V4(ip) => {
            let req = [net::MSG_PING_REQUEST, ip[0], ip[1], ip[2], ip[3], count];
            net::send_request(tx, &req);
        }
    }
}

fn recv_results(rx: &PacketRingReader, print: bool) -> (u8, u8) {
    let mut resp_buf = [0u8; 64];
    let mut result = (0u8, 0u8);
    let mut done = false;
    core::iter::from_fn(|| {
        if done {
            return None;
        }
        let n = net::recv_response_blocking(rx, &mut resp_buf);
        if n < 1 {
            return Some(());
        }
        match resp_buf[0] {
            net::MSG_PING_RESULT if n >= 4 && print => {
                print_ping_line(&resp_buf, n);
            }
            net::MSG_PING_DONE if n >= 3 => {
                let (received, total) = (resp_buf[1], resp_buf[2]);
                if print {
                    print_done_line(received, total);
                }
                result = (received, total);
                done = true;
            }
            _ => {}
        }
        Some(())
    })
    .take(64)
    .count();
    result
}

fn print_ping_line(resp_buf: &[u8], n: usize) {
    let seq = resp_buf[1];
    let status = resp_buf[2];
    let mut out = [0u8; 40];
    let mut pos = 0usize;
    let prefix = b"seq=";
    out[..prefix.len()].copy_from_slice(prefix);
    pos += prefix.len();
    pos += net::write_u8_decimal(&mut out[pos..], seq);
    match status {
        0 if n >= 5 => {
            let rtt = ((resp_buf[3] as u16) << 8) | resp_buf[4] as u16;
            let mid = b" time=";
            out[pos..pos + mid.len()].copy_from_slice(mid);
            pos += mid.len();
            pos += net::write_u16_decimal(&mut out[pos..], rtt);
            let suffix = b"ms\n";
            out[pos..pos + suffix.len()].copy_from_slice(suffix);
            pos += suffix.len();
        }
        _ => {
            let suffix = b" timeout\n";
            out[pos..pos + suffix.len()].copy_from_slice(suffix);
            pos += suffix.len();
        }
    }
    write_bytes(&out[..pos]);
}

fn print_done_line(received: u8, total: u8) {
    let mut out = [0u8; 24];
    let mut pos = 0usize;
    pos += net::write_u8_decimal(&mut out[pos..], received);
    out[pos] = b'/';
    pos += 1;
    pos += net::write_u8_decimal(&mut out[pos..], total);
    let suffix = b" received\n";
    out[pos..pos + suffix.len()].copy_from_slice(suffix);
    pos += suffix.len();
    write_bytes(&out[..pos]);
}