//! Stage 12: networking, end to end through QEMU's user-mode network.
//! Everything here is deterministic: it only talks to things QEMU itself
//! provides (see test-args in Cargo.toml), never to the internet.
//! - DHCP: QEMU's built-in server gives us 10.0.2.15/24, gateway 10.0.2.2
//! - ICMP: ping the gateway (10.0.2.2)
//! - UDP: read a file from QEMU's TFTP server (10.0.2.2:69, serves disk/)
//! - TCP: a 64 KiB round trip through an echo server (guestfwd -> `cat`),
//!   an HTTP GET of 108894 bytes (guestfwd -> scripts/http-reply.sh), and a
//!   refused connection
//! - blocking: recv with SO_RCVTIMEO times out on time
//! - user programs: ping and fetch (saving to a disk)
//! - closed sockets are really freed
//! - DNS through 10.0.2.3 if the host has DNS (skipped otherwise)
#![no_std]
#![no_main]
#![feature(custom_test_frameworks)]
#![test_runner(rustbelt::test_runner)]
#![reexport_test_harness_main = "test_main"]

extern crate alloc;

use alloc::{string::String, vec::Vec};
use bootloader::{BootInfo, entry_point};
use core::panic::PanicInfo;
use rustbelt::errno::Errno;
use rustbelt::net::{self, socket::{Kind, Socket}};
use rustbelt::process::{self, ExitStatus, SpawnArgs};
use rustbelt::{fs, memory, serial_println, thread, time};

entry_point!(main);

fn main(boot_info: &'static BootInfo) -> ! {
    rustbelt::init();
    rustbelt::init_memory(boot_info);
    rustbelt::init_smp();
    fs::init();
    rustbelt::init_storage();
    rustbelt::init_net();
    assert!(net::present(), "no virtio-net card: test-args need -device virtio-net-pci");
    test_main();
    rustbelt::hlt_loop();
}

const GATEWAY: [u8; 4] = [10, 0, 2, 2];
const TEST_SERVER: [u8; 4] = [10, 0, 2, 100]; // guestfwd: :80 HTTP, :7 echo

fn run(name: &str, args: &[&str]) -> ExitStatus {
    let mut spawn = SpawnArgs::kernel(name);
    spawn.args.extend(args.iter().map(|a| String::from(*a)));
    let p = process::spawn_path(&alloc::format!("/mnt/disk0/bin/{}", name), spawn).expect("spawn failed");
    process::wait(&p)
}

/// Send all of `data` on a TCP socket.
fn send_all(s: &Socket, mut data: &[u8]) {
    while !data.is_empty() {
        let n = s.send(data, None).expect("send failed");
        data = &data[n..];
    }
}

/// Read a TCP socket until the peer closes it.
fn recv_to_end(s: &Socket) -> Vec<u8> {
    let mut all = Vec::new();
    let mut buf = [0u8; 4096];
    loop {
        match s.recv(&mut buf).expect("recv failed").0 {
            0 => return all,
            n => all.extend_from_slice(&buf[..n]),
        }
    }
}

#[test_case]
fn dhcp_gives_us_the_slirp_address() {
    let config = net::wait_configured(10_000_000).expect("no DHCP lease within 10 s");
    assert_eq!(config.address, [10, 0, 2, 15]);
    assert_eq!(config.prefix_len, 24);
    assert_eq!(config.gateway, Some(GATEWAY));
    assert_eq!(config.dns.first(), Some(&[10, 0, 2, 3]));
    let info = net::info();
    assert_eq!(info.mac, [0x52, 0x54, 0, 0x12, 0x34, 0x56], "QEMU's default MAC");
    assert!(info.rx_packets >= 2 && info.tx_packets >= 2, "DHCP is 4 packets");
}

#[test_case]
fn ping_the_gateway() {
    let s = Socket::new(Kind::Icmp).unwrap();
    s.set_recv_timeout(Some(2_000_000));
    for seq in 1..=3u16 {
        // Echo request: type 8, code 0, checksum + id (kernel's job), seq, data.
        let mut packet = [0u8; 8 + 32];
        packet[0] = 8;
        packet[6..8].copy_from_slice(&seq.to_be_bytes());
        packet[8..].fill(seq as u8);
        s.send(&packet, Some((GATEWAY, 0))).unwrap();
        let mut reply = [0u8; 128];
        let (n, from) = s.recv(&mut reply).expect("no echo reply within 2 s");
        assert_eq!(from.unwrap().0, GATEWAY);
        assert_eq!(reply[0], 0, "type 0 = echo reply");
        assert_eq!(&reply[6..8], &seq.to_be_bytes());
        assert_eq!(&reply[8..n], &packet[8..], "payload comes back unchanged");
    }
}

#[test_case]
fn udp_tftp_read_from_slirp() {
    // TFTP (RFC 1350) over UDP: read request = opcode 1, file name, 0,
    // mode, 0. The server answers with DATA = opcode 3, block number, up to
    // 512 bytes. (A real server answers from a fresh port; slirp's built-in
    // one simply keeps using 69, so we ACK to wherever it came from.)
    let s = Socket::new(Kind::Udp).unwrap();
    s.set_recv_timeout(Some(3_000_000));
    s.send(b"\x00\x01readme.txt\x00octet\x00", Some((GATEWAY, 69))).unwrap();
    let mut buf = [0u8; 600];
    let (n, from) = s.recv(&mut buf).expect("no TFTP answer within 3 s");
    let (ip, port) = from.unwrap();
    assert_eq!(ip, GATEWAY);
    assert_eq!(&buf[..4], &[0, 3, 0, 1], "DATA, block 1");
    assert!(buf[4..n].starts_with(b"Hello from the FAT32 disk!"), "got {:?}", core::str::from_utf8(&buf[4..n.min(40)]));
    s.send(&[0, 4, 0, 1], Some((ip, port))).unwrap(); // ACK block 1
}

#[test_case]
fn tcp_echo_round_trip_64k() {
    let s = Socket::new(Kind::Tcp).unwrap();
    s.connect((TEST_SERVER, 7)).expect("connect to the echo server");
    let data: Vec<u8> = (0..64 * 1024u32).map(|i| (i * 7 + i / 256) as u8).collect();
    // Send and receive in turns, so neither side's buffers fill up.
    let mut echoed = Vec::new();
    let mut buf = [0u8; 4096];
    for chunk in data.chunks(4096) {
        send_all(&s, chunk);
        let want = echoed.len() + chunk.len();
        while echoed.len() < want {
            let (n, _) = s.recv(&mut buf).unwrap();
            assert!(n > 0, "echo server hung up early");
            echoed.extend_from_slice(&buf[..n]);
        }
    }
    assert!(echoed == data, "echoed bytes differ");
}

#[test_case]
fn tcp_http_get_108k() {
    let s = Socket::new(Kind::Tcp).unwrap();
    s.connect((TEST_SERVER, 80)).unwrap();
    send_all(&s, b"GET /big HTTP/1.0\r\nHost: test\r\n\r\n");
    let response = recv_to_end(&s);
    let body_at = response.windows(4).position(|w| w == b"\r\n\r\n").expect("no end of headers") + 4;
    assert!(response.starts_with(b"HTTP/1.0 200 OK\r\n"));
    let body = &response[body_at..];
    assert_eq!(body.len(), 108894, "seq 1 20000 is 108894 bytes");
    assert!(body.starts_with(b"1\n2\n3\n") && body.ends_with(b"19999\n20000\n"));
}

#[test_case]
fn connect_to_a_closed_port_is_refused() {
    let s = Socket::new(Kind::Tcp).unwrap();
    // slirp forwards to the host's 127.0.0.1:1, where nothing listens.
    assert_eq!(s.connect((GATEWAY, 1)), Err(Errno::ECONNREFUSED));
}

#[test_case]
fn recv_timeout_blocks_then_gives_up() {
    let s = Socket::new(Kind::Udp).unwrap();
    s.bind(40000).unwrap();
    s.set_recv_timeout(Some(300_000));
    let start = time::now_us();
    let mut buf = [0u8; 16];
    assert_eq!(s.recv(&mut buf).err(), Some(Errno::EAGAIN));
    let waited = time::now_us() - start;
    assert!(waited >= 250_000 && waited < 1_000_000, "waited {} us", waited);
}

#[test_case]
fn ping_and_fetch_programs() {
    assert_eq!(run("ping", &["-c", "2", "-i", "100", "10.0.2.2"]), ExitStatus::Exited(0));
    let _ = fs::remove("/tmp/big.txt", false);
    assert_eq!(run("fetch", &["-o", "/tmp/big.txt", "http://10.0.2.100/big"]), ExitStatus::Exited(0));
    assert_eq!(fs::lookup("/tmp/big.txt").unwrap().size(), 108894);
    assert_eq!(run("fetch", &["http://10.0.2.100/nope"]), ExitStatus::Exited(1), "404 -> exit 1");
}

#[test_case]
fn closed_sockets_are_freed() {
    let wait_for_count = |want: usize| {
        let deadline = rustbelt::interrupts::ticks() + 18 * 5;
        while net::socket_count() != want && rustbelt::interrupts::ticks() < deadline {
            thread::sleep_ticks(1);
        }
        net::socket_count()
    };
    let baseline = wait_for_count(2); // just DHCP + DNS once earlier tests are done
    assert_eq!(baseline, 2);
    let frames_before = memory::frames_in_use();
    for i in 0..20u8 {
        let s = Socket::new(Kind::Tcp).unwrap();
        s.connect((TEST_SERVER, 7)).unwrap();
        send_all(&s, &[i; 100]);
        let mut buf = [0u8; 100];
        let mut got = 0;
        while got < 100 {
            got += s.recv(&mut buf[got..]).unwrap().0;
        }
        assert_eq!(buf, [i; 100]);
        let _ = Socket::new(Kind::Udp).unwrap(); // and a UDP one, dropped at once
    }
    assert_eq!(wait_for_count(baseline), baseline, "closed TCP sockets linger");
    assert_eq!(memory::frames_in_use(), frames_before, "frames leaked");
}

#[test_case]
fn dns_lookup_or_skip() {
    match net::resolve("example.com", 3_000_000) {
        Ok(addrs) => {
            assert!(!addrs.is_empty());
            serial_println!("    (example.com -> {:?}) ", addrs[0]);
        }
        Err(e) => serial_println!("    [skipped: the host has no working DNS ({:?})] ", e),
    }
}

#[panic_handler]
fn panic(info: &PanicInfo) -> ! {
    rustbelt::test_panic_handler(info)
}