/* UDP server displaying incoming ECN (see RFC 10055) */

#include <sys/socket.h>
#include <netinet/in.h>
#include <netinet/udp.h>
#include <getopt.h>
#include <stdio.h>
#include <stdbool.h>
#include <string.h>
#include <errno.h>
#include <stdlib.h>
#include <arpa/inet.h>

#define DEFAULT_PORT 4242

#define MAXHOSTNAMELEN 256

/* Need to be global */
char            recvbuffer[2048];
static int      port = 0;

static struct option long_options[] = {
    {"port", required_argument, NULL, 0},
    {0, 0, 0, 0}
};

static char     progname[MAXHOSTNAMELEN + 1];

void
usage()
{
    printf("Usage: %s [-p port]\n", progname);
}

int
main(int argc, char **argv)
{
    char            c;
    int             socketfd;
    struct sockaddr_storage recvaddr;
    socklen_t       struct_sockaddr_size;
    struct sockaddr_in6 *me6;
    int             enable = 1;
    int             result;
    ssize_t         numbytes;
    struct sockaddr_storage me;
    struct sockaddr_storage *me_p;
    int             family;
    int            *tclass = NULL;
    char            cmbuf[0x100];
    struct iovec    iov;
    struct msghdr   mh = {
        .msg_name = &recvaddr,
        .msg_namelen = sizeof(recvaddr),
        .msg_iov = &iov,
        .msg_iovlen = 1,
        .msg_control = cmbuf,
        .msg_controllen = sizeof(cmbuf),
    };
    iov.iov_len = sizeof(recvbuffer);
    iov.iov_base = malloc(sizeof(recvbuffer));
    while (true) {
        /* getopt_long stores the option index here.  */
        int             option_index = 0;

        c = getopt_long(argc, argv, "p:", long_options, &option_index);

        /* Detect the end of the options.  */
        if (c == -1)
            break;

        switch (c) {
        case 0:
            /* If this option set a flag, do nothing else now.  */
            if (long_options[option_index].flag != 0)
                break;
            else if (strcmp(long_options[option_index].name, "port") == 0) {
                port = strtol(optarg, NULL, 10);        /* Assume the port number is 
                                                         * in base 10 */
                if (port == 0) {
                    fprintf(stderr, "Invalid port number %s\n", optarg);
                    exit(1);
                }
            } else {
                fprintf(stderr, "Internal error, option %s not known\n",
                        long_options[option_index].name);
            }
            break;

        case 'p':
            port = strtol(optarg, NULL, 10);    /* Assume the port number is in base 
                                                 * 10 */
            if (port == 0) {
                fprintf(stderr, "Invalid port number %s\n", optarg);
                exit(1);
            }
            break;

        case '?':
            usage();
            break;

        default:
            exit(1);
        }
    }

    /* Print any remaining command line arguments (not options).  */
    if (argc - optind != 0) {
        usage();
        exit(1);
    }
    if (port == 0)
        port = DEFAULT_PORT;

    socketfd = socket(AF_INET6, SOCK_DGRAM, IPPROTO_UDP);
    if (socketfd < 0) {
        perror("Cannot create socket");
        exit(1);
    }
    result =
        setsockopt(socketfd, IPPROTO_IPV6, IPV6_RECVTCLASS, &enable, sizeof(enable));
    if (result != 0) {
        perror("Cannot setsockopt");
        exit(1);
    }
    family = AF_INET6;
    struct_sockaddr_size = (socklen_t) sizeof(struct sockaddr_in6);
    me6 = (struct sockaddr_in6 *) &me;
    me6->sin6_family = family;
    me6->sin6_addr = in6addr_any;
    me6->sin6_port = htons(port);
    me_p = memcpy(&me, me6, sizeof(*me6));
    result = bind(socketfd, (struct sockaddr *) me_p, struct_sockaddr_size);
    if (result != 0) {
        perror("Cannot bind");
        exit(1);
    }
    while (true) {
        /* We don't use recvfrom because we need the ECN, so we have to rely on
         * recvmsg and control messages. */
        numbytes = recvmsg(socketfd, &mh, 0);
        if (numbytes < 0) {
            perror("recvmsg");
            exit(1);
        }
        if (numbytes < 0) {
            perror("Cannot receive data");
            continue;
        }
        if (numbytes == 0) {
            fprintf(stderr, "Zero data\n");
            continue;
        }
        fprintf(stderr, "Message received, %li bytes.\n", numbytes);

        for (struct cmsghdr * cmsg = CMSG_FIRSTHDR(&mh);
             cmsg != NULL; cmsg = CMSG_NXTHDR(&mh, cmsg)) {
            if (cmsg->cmsg_level != IPPROTO_IPV6 || cmsg->cmsg_type != IPV6_TCLASS) {
                continue;
            }
            tclass = (int *) CMSG_DATA(cmsg);
            break;
        }
        if (tclass == NULL) {
            fprintf(stderr, "No info on DSCP\n");
            exit(1);
        } else {
            fprintf(stdout, "Received DSCP %i\n", *tclass);
        }
    }                           /* End of the main loop */
    return 0;
}