/*
 * Ouroboros - Copyright (C) 2016 - 2026
 *
 * Ouroboros ping application
 *
 *    Dimitri Staessens <dimitri@ouroboros.rocks>
 *    Sander Vrijders   <sander@ouroboros.rocks>
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 *
 * 1. Redistributions of source code must retain the above copyright
 * notice, this list of conditions and the following disclaimer.
 *
 * 2. Redistributions in binary form must reproduce the above
 * copyright notice, this list of conditions and the following
 * disclaimer in the documentation and/or other materials provided
 * with the distribution.
 *
 * 3. Neither the name of the copyright holder nor the names of its
 * contributors may be used to endorse or promote products derived
 * from this software without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
 * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
 * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
 * FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE
 * COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT,
 * INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
 * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
 * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
 * HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT,
 * STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
 * ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED
 * OF THE POSSIBILITY OF SUCH DAMAGE.
 */

volatile bool stop;

void shutdown_client(int signo, siginfo_t * info, void * c)
{
        (void) info;
        (void) c;

        switch(signo) {
        case SIGINT:
        case SIGTERM:
        case SIGHUP:
        case SIGALRM:
                stop = true;
        default:
                return;
        }
}

static void update_rtt_stats(double ms)
{
        double d;

        if (ms < client.rtt_min)
                client.rtt_min = ms;
        if (ms > client.rtt_max)
                client.rtt_max = ms;

        d = (ms - client.rtt_avg);
        client.rtt_avg += d / client.rcvd;
        client.rtt_m2 += d * (ms - client.rtt_avg);
}

static double rtt_val(double ms)
{
        return ms < 0.1 ? ms * 1000 : ms;
}

static const char * rtt_unit(double ms)
{
        return ms < 0.1 ? "µs" : "ms";
}

static void print_rtt(int len, int seq,
                      double ms, const char * suf)
{
        printf("%d bytes from %s: seq=%d "
               "time=%.3f %s%s\n",
               len, client.s_apn, seq,
               rtt_val(ms), rtt_unit(ms),
               suf != NULL ? suf : "");
}

void * reader(void * o)
{
        struct timespec timeout;
        struct timespec now = {0, 0};
        struct timespec sent;

        char               buf[OPING_BUF_SIZE];
        struct oping_msg * msg = (struct oping_msg *) buf;
        int                fd      = *((int *) o);
        int                msg_len = 0;
        double             ms      = 0;
        uint32_t           exp_id  = 0;

        timeout.tv_sec  = client.timeout / 1000;
        timeout.tv_nsec = (client.timeout % 1000) * MILLION;

        fccntl(fd, FLOWSRCVTIMEO, &timeout);

        while (!stop && client.rcvd != client.count) {
                uint32_t id;

                msg_len = flow_read(fd, buf, OPING_BUF_SIZE);
                if (msg_len == -ETIMEDOUT) {
                        printf("Server timed out.\n");
                        stop = true;
                        break;
                }

                if (msg_len < 0)
                        continue;

                if (ntohl(msg->type) != ECHO_REPLY) {
                        printf("Invalid message on fd %d.\n", fd);
                        continue;
                }

                id = (uint32_t) ntohl(msg->id);
                if (id >= client.count) {
                        printf("Invalid id.\n");
                        continue;
                }

                ++client.rcvd;

                clock_gettime(CLOCK_MONOTONIC, &now);

                sent.tv_sec = msg->tv_sec;
                sent.tv_nsec = msg->tv_nsec;

                ms = ts_diff_us(&now, &sent) / 1000.0;

                if (id < exp_id)
                        ++client.ooo;

                if (!client.quiet) {
                        if (client.timestamp) {
                                struct timespec rtc;
                                clock_gettime(CLOCK_REALTIME, &rtc);
                                printf("[%zd.%06zu] ",
                                       (ssize_t) rtc.tv_sec,
                                       (size_t) rtc.tv_nsec / 1000);
                        }

                        print_rtt(msg_len, ntohl(msg->id), ms,
                                  id < exp_id ?
                                  " [out-of-order]" : NULL);
                }

                update_rtt_stats(ms);

                if (id >= exp_id)
                        exp_id = id + 1;
        }

        return (void *) 0;
}

void * writer(void * o)
{
        int * fdp = (int *) o;
        struct timespec now;
        struct timespec wait = {client.interval / 1000,
                                (client.interval % 1000) * MILLION};
        struct oping_msg * msg;
        char * buf = malloc(client.size);

        if (buf == NULL)
                return (void *) -ENOMEM;

        if (fdp == NULL) {
                free(buf);
                return (void *) -EINVAL;
        }

        memset(buf, 0, client.size);

        msg = (struct oping_msg *) buf;

        if (!client.quiet)
                printf("Pinging %s with %d bytes of data (%u packets):\n\n",
                       client.s_apn, client.size, client.count);

        pthread_cleanup_push(free, buf);

        while (!stop && client.sent < client.count) {
                nanosleep(&wait, NULL);

                clock_gettime(CLOCK_MONOTONIC, &now);

                msg->type = htonl(ECHO_REQUEST);
                msg->id = htonl(client.sent++);
                msg->tv_sec = now.tv_sec;
                msg->tv_nsec = now.tv_nsec;

                if (flow_write(*fdp, buf, client.size) < 0) {
                        printf("Failed to send packet.\n");
                        stop = true;
                }
        }

        pthread_cleanup_pop(true);

        return (void *) 0;
}

static int client_init(void)
{
        stop = false;
        client.sent = 0;
        client.rcvd = 0;
        client.rtt_min = FLT_MAX;
        client.rtt_max = 0;
        client.rtt_avg = 0;
        client.rtt_m2 = 0;

        return 0;
}

static void client_fini(void)
{
        return;
}

static void print_stats(struct timespec * tic,
                        struct timespec * toc)
{
        printf("\n");
        printf("--- %s ping statistics ---\n", client.s_apn);
        printf("%d packets transmitted, ", client.sent);
        printf("%d received, ", client.rcvd);
        printf("%zd out-of-order, ", client.ooo);
        printf("%.0lf%% packet loss, ", client.sent == 0 ? 0 :
               ceil(100 - (100 * (client.rcvd / (float) client.sent))));
        printf("time: %.3f ms\n", ts_diff_us(toc, tic) / 1000.0);

        if (client.rcvd > 0) {
                double a = client.rtt_avg;
                double f = a < 0.1 ? 1000 : 1;
                printf("rtt min/avg/max/mdev = %.3f/%.3f/%.3f/",
                       client.rtt_min * f, client.rtt_avg * f,
                       client.rtt_max * f);
                if (client.rcvd > 1)
                        printf("%.3f %s\n",
                               sqrt(client.rtt_m2 /
                               (client.rcvd - 1)) * f,
                               rtt_unit(a));
                else
                        printf("NaN %s\n", rtt_unit(a));
        }
}

static int flood_busy_ping(int fd)
{
        char buf[OPING_BUF_SIZE];
        struct oping_msg * msg = (struct oping_msg *) buf;
        struct timespec sent;
        struct timespec rcvd;
        double ms;
        int n;

        memset(buf, 0, client.size);

        fccntl(fd, FLOWSFLAGS,
               FLOWFRDWR | FLOWFRNOPART | FLOWFRNOBLOCK);

        if (!client.quiet)
                printf("Pinging %s with %d bytes"
                       " of data (%u packets,"
                       " busy-poll):\n\n",
                       client.s_apn, client.size,
                       client.count);

        while (!stop && client.sent < client.count) {
                clock_gettime(CLOCK_MONOTONIC, &sent);

                msg->type = htonl(ECHO_REQUEST);
                msg->id = htonl(client.sent);
                msg->tv_sec = sent.tv_sec;
                msg->tv_nsec = sent.tv_nsec;

                if (flow_write(fd, buf, client.size) < 0) {
                        printf("Failed to send packet.\n");
                        break;
                }

                ++client.sent;

                do {
                        n = flow_read(fd, buf, OPING_BUF_SIZE);
                } while (n == -EAGAIN && !stop);

                if (n < 0)
                        break;

                clock_gettime(CLOCK_MONOTONIC, &rcvd);

                if (ntohl(msg->type) != ECHO_REPLY)
                        continue;

                ++client.rcvd;

                sent.tv_sec = msg->tv_sec;
                sent.tv_nsec = msg->tv_nsec;
                ms = ts_diff_us(&rcvd, &sent) / 1000.0;

                update_rtt_stats(ms);

                if (!client.quiet)
                        print_rtt(client.size, ntohl(msg->id), ms, NULL);
        }

        return 0;
}

static int flood_ping(int fd)
{
        char buf[OPING_BUF_SIZE];
        struct oping_msg * msg = (struct oping_msg *) buf;
        struct timespec sent;
        struct timespec rcvd;
        double ms;

        memset(buf, 0, client.size);

        if (!client.quiet)
                printf("Pinging %s with %d bytes of data (%u packets):\n\n",
                       client.s_apn, client.size, client.count);

        while (!stop && client.sent < client.count) {
                clock_gettime(CLOCK_MONOTONIC, &sent);

                msg->type = htonl(ECHO_REQUEST);
                msg->id = htonl(client.sent);
                msg->tv_sec = sent.tv_sec;
                msg->tv_nsec = sent.tv_nsec;

                if (flow_write(fd, buf, client.size) < 0) {
                        printf("Failed to send packet.\n");
                        break;
                }

                ++client.sent;

                if (flow_read(fd, buf, OPING_BUF_SIZE) < 0) {
                        printf("Failed to read packet.\n");
                        break;
                }

                clock_gettime(CLOCK_MONOTONIC, &rcvd);

                if (ntohl(msg->type) != ECHO_REPLY)
                        continue;

                ++client.rcvd;

                sent.tv_sec = msg->tv_sec;
                sent.tv_nsec = msg->tv_nsec;
                ms = ts_diff_us(&rcvd, &sent) / 1000.0;

                update_rtt_stats(ms);

                if (!client.quiet)
                        print_rtt(client.size, ntohl(msg->id), ms, NULL);
        }

        return 0;
}

static int threaded_ping(int fd)
{
        pthread_create(&client.reader_pt, NULL, reader, &fd);
        pthread_create(&client.writer_pt, NULL, writer, &fd);

        pthread_join(client.writer_pt, NULL);
        pthread_join(client.reader_pt, NULL);

        return 0;
}

static int client_main(void)
{
        struct sigaction sig_act;
        struct timespec tic;
        struct timespec toc;
        int fd;

        memset(&sig_act, 0, sizeof sig_act);
        sig_act.sa_sigaction = &shutdown_client;
        sig_act.sa_flags = 0;

        if (sigaction(SIGINT,  &sig_act, NULL) ||
            sigaction(SIGTERM, &sig_act, NULL) ||
            sigaction(SIGHUP,  &sig_act, NULL) ||
            sigaction(SIGPIPE, &sig_act, NULL) ||
            sigaction(SIGALRM, &sig_act, NULL)) {
                printf("Failed to install sighandler.\n");
                return 2;
        }

        if (client_init()) {
                printf("Failed to initialize client.\n");
                return 2;
        }

        fd = flow_alloc(client.s_apn, &client.qs, NULL);
        if (fd < 0) {
                printf("Failed to allocate flow: %d.\n", fd);
                client_fini();
                return 2;
        }

        fccntl(fd, FLOWSFLAGS, FLOWFRDWR | FLOWFRNOPART);

        if (client.duration > 0) {
                struct itimerval it;
                memset(&it, 0, sizeof(it));
                it.it_value.tv_sec  = client.duration / 1000;
                it.it_value.tv_usec = (client.duration % 1000) * 1000;
                setitimer(ITIMER_REAL, &it, NULL);
        }

        clock_gettime(CLOCK_REALTIME, &tic);

        if (client.flood_busy)
                flood_busy_ping(fd);
        else if (client.flood)
                flood_ping(fd);
        else
                threaded_ping(fd);

        clock_gettime(CLOCK_REALTIME, &toc);

        print_stats(&tic, &toc);

        flow_dealloc(fd);
        client_fini();

        return client.rcvd == client.sent ? 0 : 1;
}
