/*
 * Copyright (c) 2016 Vincent Gross <vincent.gross@kilob.yt>
 * Copyright (c) 2017 Alexander Bluhm <bluhm@openbsd.org>
 *
 * Permission to use, copy, modify, and distribute this software for any
 * purpose with or without fee is hereby granted, provided that the above
 * copyright notice and this permission notice appear in all copies.
 *
 * THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
 * WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
 * MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
 * ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
 * WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
 * ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
 * OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
 */

#include <sys/types.h>
#include <sys/socket.h>

#include <arpa/inet.h>
#include <netinet/in.h>

#include <err.h>
#include <errno.h>
#include <getopt.h>
#include <netdb.h>
#include <poll.h>
#include <stdlib.h>
#include <string.h>
#include <stdio.h>
#include <unistd.h>

#define PAYLOAD "payload"

int fuzzit;

void __dead usage(const char *);
int udp_bind(struct sockaddr_in *);
int udp_send(int, struct sockaddr_in *, struct sockaddr_in *);
struct sockaddr_in * udp_recv(int s, struct sockaddr_in *);

void __dead
usage(const char *msg)
{
        if (msg != NULL)
                fprintf(stderr, "%s\n", msg);
        fprintf(stderr, "runtest [-f] -D destination -B bind [-C cmesg] "
            "[-E error] -R reserved -W wire\n");
        exit(1);
}

int
main(int argc, char *argv[])
{
        int ch, error, errexpect, bind_sock, dest_sock, resv_sock;
        char addr[16];
        const char *errstr;
        struct addrinfo hints, *res;
        struct sockaddr_in *bind_sin, *cmsg_sin, *dest_sin, *resv_sin,
            *wire_sin, *from_sin;

        bzero(&hints, sizeof(hints));
        hints.ai_family = AF_INET;
        hints.ai_socktype = SOCK_DGRAM;

        bind_sin = cmsg_sin = dest_sin = resv_sin = wire_sin = NULL;
        errexpect = 0;

        while ((ch = getopt(argc, argv, "B:C:D:E:fR:W:")) != -1) {
                switch (ch) {
                case 'B':
                        error = getaddrinfo(optarg, NULL, &hints, &res);
                        if (error)
                                errx(1, "-B: %s", gai_strerror(error));
                        bind_sin = (struct sockaddr_in *)res->ai_addr;
                        break;
                case 'C':
                        error = getaddrinfo(optarg, NULL, &hints, &res);
                        if (error)
                                errx(1, "-C: %s", gai_strerror(error));
                        cmsg_sin = (struct sockaddr_in *)res->ai_addr;
                        break;
                case 'D':
                        error = getaddrinfo(optarg, NULL, &hints, &res);
                        if (error)
                                errx(1, "-D: %s", gai_strerror(error));
                        dest_sin = (struct sockaddr_in *)res->ai_addr;
                        break;
                case 'E':
                        errexpect = strtonum(optarg, 1, 255, &errstr);
                        if (errstr != NULL)
                                errx(1, "error number is %s: %s",
                                    errstr, optarg);
                        break;
                case 'f':
                        fuzzit = 1;
                        break;
                case 'R':
                        error = getaddrinfo(optarg, NULL, &hints, &res);
                        if (error)
                                errx(1, "-R: %s", gai_strerror(error));
                        resv_sin = (struct sockaddr_in *)res->ai_addr;
                        break;
                case 'W':
                        error = getaddrinfo(optarg, NULL, &hints, &res);
                        if (error)
                                errx(1, "-W: %s", gai_strerror(error));
                        wire_sin = (struct sockaddr_in *)res->ai_addr;
                        break;
                default:
                        usage(NULL);
                }
        }
        argc -= optind;
        argv += optind;

        if (argc > 0)
                usage("too many arguments");

        if (bind_sin == NULL)
                usage("no bind addr");

        if (dest_sin == NULL)
                usage("no destination addr");

        if (resv_sin == NULL)
                usage("no reserved addr");

        /* bind on address that cannot be used */
        resv_sock = udp_bind(resv_sin);

        /* bind socket that should receive the packet */
        dest_sock = udp_bind(dest_sin);

        /* bind socket that is used to send the packet */
        bind_sin->sin_port = resv_sin->sin_port;
        bind_sock = udp_bind(bind_sin);
        error = udp_send(bind_sock, cmsg_sin, dest_sin);

        if (errexpect && !error) {
                errno = errexpect;
                err(2, "udp send succeeded, but expected error");
        }
        if (!errexpect && error) {
                errno = error;
                err(2, "no error expected, but udp send failed");
        }
        if (errexpect != error) {
                errno = error;
                err(2, "expected error %d, but udp send failed", errexpect);
        }

        if (wire_sin != NULL) {
                from_sin = udp_recv(dest_sock, dest_sin);
                if (from_sin == NULL)
                        errx(2, "receive timeout");
                inet_ntop(from_sin->sin_family, &from_sin->sin_addr,
                    addr, sizeof(addr));
                if (from_sin->sin_addr.s_addr != wire_sin->sin_addr.s_addr)
                        errx(2, "receive addr %s", addr);
                if (from_sin->sin_port != bind_sin->sin_port)
                        errx(2, "receive port %d", ntohs(from_sin->sin_port));
        }

        return 0;
}

int
udp_bind(struct sockaddr_in *src)
{
        int s, reuse, salen;
        char addr[16];

        inet_ntop(src->sin_family, &src->sin_addr, addr, sizeof(addr));

        if ((s = socket(src->sin_family, SOCK_DGRAM, 0)) == -1)
                err(1, "socket %s", addr);
        reuse = ntohl(src->sin_addr.s_addr) == INADDR_ANY ? 1 : 0;
        if (setsockopt(s, SOL_SOCKET, SO_REUSEADDR, &reuse, sizeof(reuse))
            == -1)
                err(1, "setsockopt %s", addr);
        if (bind(s, (struct sockaddr *)src, src->sin_len) == -1)
                err(1, "bind %s", addr);
        /* fill out port */
        salen = sizeof(*src);
        if (getsockname(s, (struct sockaddr *)src, &salen))
                err(1, "getsockname %s", addr);

        return s;
}

int
udp_send(int s, struct sockaddr_in *src, struct sockaddr_in *dst)
{
        struct msghdr            msg;
        struct iovec             iov;
        struct cmsghdr          *cmsg;
        struct in_addr          *sendopt;
        int                     *hopopt;
#define CMSGSP_SADDR    CMSG_SPACE(sizeof(u_int32_t))
#define CMSGSP_HOPLIM   CMSG_SPACE(sizeof(int))
#define CMSGSP_BOGUS    CMSG_SPACE(12)
#define CMSGBUF_SP      CMSGSP_SADDR + CMSGSP_HOPLIM + CMSGSP_BOGUS + 3
        unsigned char            cmsgbuf[CMSGBUF_SP];

        iov.iov_base = PAYLOAD;
        iov.iov_len = strlen(PAYLOAD) + 1;
        bzero(&msg, sizeof(msg));
        msg.msg_name = dst;
        msg.msg_namelen = dst->sin_len;
        msg.msg_iov = &iov;
        msg.msg_iovlen = 1;

        if (src) {
                bzero(&cmsgbuf, sizeof(cmsgbuf));
                msg.msg_control = &cmsgbuf;
                msg.msg_controllen = CMSGSP_SADDR;
                cmsg = CMSG_FIRSTHDR(&msg);
                cmsg->cmsg_len = CMSG_LEN(sizeof(struct in_addr));
                cmsg->cmsg_level = IPPROTO_IP;
                cmsg->cmsg_type = IP_SENDSRCADDR;
                sendopt = (struct in_addr *)CMSG_DATA(cmsg);
                memcpy(sendopt, &src->sin_addr, sizeof(*sendopt));
                if (fuzzit) {
                        msg.msg_controllen = CMSGBUF_SP;
                        cmsg = CMSG_NXTHDR(&msg, cmsg);
                        cmsg->cmsg_len = CMSG_LEN(sizeof(int));
                        cmsg->cmsg_level = IPPROTO_IPV6;
                        cmsg->cmsg_type = IPV6_UNICAST_HOPS;
                        hopopt = (int *)CMSG_DATA(cmsg);
                        *hopopt = 8;

                        cmsg = CMSG_NXTHDR(&msg, cmsg);
                        cmsg->cmsg_len = CMSG_LEN(sizeof(int)) + 15;
                        cmsg->cmsg_level = IPPROTO_IPV6;
                        cmsg->cmsg_type = IPV6_UNICAST_HOPS;
                }
        }

        if (sendmsg(s, &msg, 0) == -1)
                return errno;

        return 0;
}

struct sockaddr_in *
udp_recv(int s, struct sockaddr_in *dst)
{
        struct sockaddr_in *src;
        struct pollfd pfd[1];
        char addr[16], buf[256];
        int nready, len, salen;

        inet_ntop(dst->sin_family, &dst->sin_addr, addr, sizeof(addr));

        pfd[0].fd = s;
        pfd[0].events = POLLIN;
        nready = poll(pfd, 1, 2 * 1000);
        if (nready == -1)
                err(1, "poll");
        if (nready == 0)
                return NULL;
        if ((pfd[0].revents & POLLIN) == 0)
                errx(1, "event %d %s", pfd[0].revents, addr);

        if ((src = malloc(sizeof(*src))) == NULL)
                err(1, "malloc");
        salen = sizeof(*src);
        if ((len = recvfrom(s, buf, sizeof(buf), 0, (struct sockaddr *)src,
            &salen)) == -1)
                err(1, "recvfrom %s", addr);

        if (len != strlen(PAYLOAD) + 1)
                errx(1, "recvfrom %s len %d", addr, len);
        if (strcmp(buf, PAYLOAD) != 0)
                errx(1, "recvfrom %s payload", addr);

        return src;
}