// nettest -- prove the TCP/IP stack carries real traffic.
//
// Not a ping and not a status readout: both of those can pass on a stack that
// cannot actually move bytes. This opens a TCP connection to a host outside the
// guest, sends a request, reads the reply and reports what came back -- which
// exercises ARP, IP, the TCP handshake, the send and receive paths, the window
// and the close, in one exchange.
//
//   nettest [-u] <a.b.c.d> <port>
//
// Under QEMU's user-mode networking the host is 10.0.2.2, so pointing this at a
// listener there is the intended use:
//
//   host$  printf 'HELLO-FROM-HOST\n' | nc -l -p 9000
//   guest# nettest 10.0.2.2 9000
//
// -u does the same exchange over UDP, and checks the one thing TCP cannot: that
// recvfrom reports who sent the datagram. On a connected stream the source is
// known before the first byte arrives, so a kernel that filled the address in
// with anything at all would look correct there.
//
//   host$  nc -u -l -p 9000
//   guest# nettest -u 10.0.2.2 9000

#include "core/types.h"
#include "uapi/user.h"
#include "net/socket.h"

static void
puts_(const char *s)
{
    int n = 0;
    while (s[n]) n++;
    write(1, (void *)s, n);
}

static void
putn_(int v)
{
    char buf[12];
    int i = 0;
    if (v == 0) { write(1, "0", 1); return; }
    while (v > 0 && i < (int)sizeof(buf)) { buf[i++] = '0' + (v % 10); v /= 10; }
    while (i--) write(1, &buf[i], 1);
}

// Host byte order to network. Written out rather than taken from a header
// because the userland has no <arpa/inet.h> and this is all of what it would
// provide here.
static ushort
hton16(ushort v)
{
    return (ushort)((v << 8) | (v >> 8));
}

// "10.0.2.2" -> the 32-bit address, already in network order. Returns 0 on
// anything malformed, which doubles as "0.0.0.0 is not a destination".
static uint
parse_ip(const char *s)
{
    uint out = 0;
    for (int part = 0; part < 4; part++) {
        int v = 0, digits = 0;
        while (*s >= '0' && *s <= '9') {
            v = v * 10 + (*s++ - '0');
            digits++;
        }
        if (digits == 0 || v > 255)
            return 0;
        // Network order is most-significant byte first, and the first part of
        // a dotted quad is the most significant.
        out |= ((uint)v) << (8 * part);
        if (part < 3) {
            if (*s != '.')
                return 0;
            s++;
        }
    }
    return *s == '\0' ? out : 0;
}

static int
parse_port(const char *s)
{
    int v = 0;
    while (*s >= '0' && *s <= '9')
        v = v * 10 + (*s++ - '0');
    return (*s == '\0' && v > 0 && v < 65536) ? v : -1;
}

// The UDP half. One datagram out, one back, and the source of the reply
// printed -- which is the part that has no equivalent over TCP.
static void
udp_exchange(uint ip, int port, const char *ipstr)
{
    int fd = socket(AF_INET, SOCK_DGRAM, 0);
    if (fd < 0) {
        puts_("nettest: udp socket failed\n");
        return;
    }

    struct sockaddr_in to;
    for (int i = 0; i < (int)sizeof(to); i++)
        ((char *)&to)[i] = 0;
    to.sin_family = AF_INET;
    to.sin_port   = hton16((ushort)port);
    to.sin_addr   = ip;

    const char *msg = "HELLO-UDP-FROM-SCHOS\n";
    int len = 0;
    while (msg[len]) len++;

    puts_("nettest: udp sending to "); puts_(ipstr);
    puts_(":"); putn_(port); puts_("\n");

    if (sendto(fd, msg, len, 0, (const struct sockaddr *)&to, sizeof(to)) != len) {
        puts_("nettest: udp send failed\n");
        close(fd);
        return;
    }

    // Deliberately not zeroed first. If the kernel does not fill this in, the
    // address printed is whatever was on the stack, and a test that asserts on
    // 10.0.2.2 catches that -- where a zeroed struct would quietly read as
    // 0.0.0.0 and look like a mere formatting problem.
    struct sockaddr_in from;
    int fromlen = sizeof(from);
    char buf[512];

    int n = recvfrom(fd, buf, sizeof(buf) - 1, 0,
                     (struct sockaddr *)&from, &fromlen);
    if (n <= 0) {
        puts_("nettest: udp FAILED, nothing came back\n");
        close(fd);
        return;
    }
    buf[n] = '\0';

    puts_("nettest: udp received "); putn_(n); puts_(" bytes: ");
    puts_(buf);
    if (buf[0] && buf[n - 1] != '\n')
        puts_("\n");

    if (fromlen < (int)sizeof(from)) {
        puts_("nettest: udp FAILED, no source address returned\n");
        close(fd);
        return;
    }

    // Network order, most significant byte first, which is how parse_ip built
    // it -- so the same shifts unpack it.
    uint a = from.sin_addr;
    puts_("nettest: udp source ");
    putn_((int)(a & 0xFF));        puts_(".");
    putn_((int)((a >> 8) & 0xFF)); puts_(".");
    putn_((int)((a >> 16) & 0xFF));puts_(".");
    putn_((int)((a >> 24) & 0xFF));
    puts_(":");
    putn_((int)hton16(from.sin_port));
    puts_("\n");

    // A second datagram, read separately, because message boundaries are the
    // whole difference between UDP and TCP. The responder sends two; if the
    // kernel kept them as a byte stream the first read would have swallowed
    // both, and the byte counts printed above and below are what says which
    // happened.
    int m = recvfrom(fd, buf, sizeof(buf) - 1, 0, (struct sockaddr *)&from,
                     &fromlen);
    close(fd);
    if (m <= 0) {
        puts_("nettest: udp FAILED, no second datagram\n");
        return;
    }
    buf[m] = '\0';
    puts_("nettest: udp second "); putn_(m); puts_(" bytes: ");
    puts_(buf);
    if (buf[0] && buf[m - 1] != '\n')
        puts_("\n");

    puts_("nettest: udp OK\n");
}

int
main(int argc, char *argv[])
{
    int udp = 0;

    if (argc == 4 && argv[1][0] == '-' && argv[1][1] == 'u' && argv[1][2] == '\0') {
        udp = 1;
        argv++;
        argc--;
    }
    if (argc != 3) {
        puts_("usage: nettest [-u] <a.b.c.d> <port>\n");
        exit();
    }

    uint ip = parse_ip(argv[1]);
    int port = parse_port(argv[2]);
    if (ip == 0 || port < 0) {
        puts_("nettest: bad address or port\n");
        exit();
    }

    if (udp) {
        udp_exchange(ip, port, argv[1]);
        exit();
    }

    int fd = socket(AF_INET, SOCK_STREAM, 0);
    if (fd < 0) {
        puts_("nettest: socket failed\n");
        exit();
    }

    // Zeroed by hand: this userland has no memset, and sin_zero must be clear
    // because the kernel compares the whole struct's length against sizeof.
    struct sockaddr_in sa;
    for (int i = 0; i < (int)sizeof(sa); i++)
        ((char *)&sa)[i] = 0;
    sa.sin_family = AF_INET;
    sa.sin_port   = hton16((ushort)port);
    sa.sin_addr   = ip;

    puts_("nettest: connecting to "); puts_(argv[1]);
    puts_(":"); putn_(port); puts_("\n");

    if (connect(fd, (const struct sockaddr *)&sa, sizeof(sa)) < 0) {
        puts_("nettest: connect failed\n");
        close(fd);
        exit();
    }
    puts_("nettest: connected\n");

    const char *msg = "HELLO-FROM-SCHOS\n";
    int len = 0;
    while (msg[len]) len++;
    if (write(fd, (void *)msg, len) != len) {
        puts_("nettest: short write\n");
        close(fd);
        exit();
    }
    puts_("nettest: sent "); putn_(len); puts_(" bytes\n");

    // Read until the peer closes. One read is not enough to claim the receive
    // path works: TCP is a stream and may well deliver in pieces.
    char buf[512];
    int total = 0;
    for (;;) {
        int n = read(fd, buf + total, (int)sizeof(buf) - 1 - total);
        if (n <= 0)
            break;
        total += n;
        if (total >= (int)sizeof(buf) - 1)
            break;
    }
    buf[total] = '\0';
    close(fd);

    puts_("nettest: received "); putn_(total); puts_(" bytes: ");
    puts_(buf);
    if (buf[0] && buf[total ? total - 1 : 0] != '\n')
        puts_("\n");

    if (total == 0) {
        puts_("nettest: FAILED, nothing came back\n");
        exit();
    }
    puts_("nettest: OK\n");
    exit();
}