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

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

#include <err.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>

static int
test_sendsrcaddr(int s, const struct sockaddr_in *remote,
    const struct in_addr *src)
{
        struct msghdr msg;
        struct iovec iov;
        union {
                struct cmsghdr cm;
                uint8_t data[CMSG_SPACE(sizeof(struct in_addr))];
        } ctrl;
        struct cmsghdr *cm;
        int n;

        iov.iov_base = &n;
        iov.iov_len = sizeof(n);

        memset(&msg, 0, sizeof(msg));
        if (remote != NULL) {
                msg.msg_name = __DECONST(void *, remote);
                msg.msg_namelen = sizeof(*remote);
        }
        msg.msg_iov = &iov;
        msg.msg_iovlen = 1;
        msg.msg_control = ctrl.data;
        msg.msg_controllen = sizeof(ctrl.data);

        memset(&ctrl, 0, sizeof(ctrl));
        cm = CMSG_FIRSTHDR(&msg);
        cm->cmsg_len = CMSG_LEN(sizeof(int));
        cm->cmsg_level = IPPROTO_IP;
        cm->cmsg_type = IP_RECVDSTADDR;
        *((struct in_addr *)CMSG_DATA(cm)) = *src;

        return sendmsg(s, &msg, MSG_SYNC);
}

static void
usage(const char *cmd)
{
        fprintf(stderr, "%s -4 ip4 -p port [-s src]\n", cmd);
        exit(1);
}

int
main(int argc, char *argv[])
{
        struct sockaddr_in remote, local;
        struct in_addr src;
        socklen_t local_len;
        char local_str[INET_ADDRSTRLEN];
        int s, opt, s2, on, n;

        memset(&remote, 0, sizeof(remote));
        remote.sin_family = AF_INET;

        src.s_addr = INADDR_ANY;

        while ((opt = getopt(argc, argv, "4:p:s:")) != -1) {
                switch (opt) {
                case '4':
                        if (inet_pton(AF_INET, optarg, &remote.sin_addr) <= 0)
                                usage(argv[0]);
                        break;

                case 'p':
                        remote.sin_port = strtol(optarg, NULL, 10);
                        remote.sin_port = htons(remote.sin_port);
                        break;

                case 's':
                        if (inet_pton(AF_INET, optarg, &src) <= 0)
                                usage(argv[0]);
                        break;

                default:
                        usage(argv[0]);
                }
        }

        s = socket(AF_INET, SOCK_DGRAM, 0);
        if (s < 0)
                err(2, "socket faild");

        /*
         * inp_laddr == ANY && src valid --> pass.
         * inp_laddr == ANY && src invalid --> fail.
         */
        if (test_sendsrcaddr(s, &remote, &src) < 0)
                err(2, "sendsrcaddr failed");

        local_len = sizeof(local);
        if (getsockname(s, (struct sockaddr *)&local, &local_len) < 0)
                err(2, "getsockname failed");

        fprintf(stderr, "wildcard: laddr %s, lport %u\n",
            inet_ntop(AF_INET, &local.sin_addr, local_str, sizeof(local_str)),
            ntohs(local.sin_port));
        local.sin_addr = src;

        s2 = socket(AF_INET, SOCK_DGRAM, 0);
        if (s2 < 0)
                err(2, "socket 2 failed");

        on = 1;
        if (setsockopt(s2, SOL_SOCKET, SO_REUSEADDR, &on, sizeof(on)) < 0)
                err(2, "setsockopt(REUSEADDR) failed");

        if (bind(s2, (const struct sockaddr *)&local, sizeof(local)) < 0)
                err(2, "bind src failed");

        /*
         * inp_laddr != ANY && src == inp_laddr --> pass.
         */
        if (test_sendsrcaddr(s2, &remote, &src) < 0)
                err(2, "sendsrcaddr 2 failed");

        /*
         * Duplicated src/lport as s2 --> fail.
         */
        if (test_sendsrcaddr(s, &remote, &src) > 0)
                errx(2, "sendsrcaddr succeeded!?");

        close(s2);

        if (connect(s, (const struct sockaddr *)&remote, sizeof(remote)) < 0)
                err(2, "connect remote failed");

        local_len = sizeof(local);
        if (getsockname(s, (struct sockaddr *)&local, &local_len) < 0)
                err(2, "getsockname failed");
        fprintf(stderr, "connected: laddr %s, lport %u\n",
            inet_ntop(AF_INET, &local.sin_addr, local_str, sizeof(local_str)),
            ntohs(local.sin_port));

        /*
         * Connected socket:
         * if inp_laddr == src --> pass.
         * if inp_laddr != src --> fail.
         */
        n = test_sendsrcaddr(s, NULL, &src);
        if (local.sin_addr.s_addr == src.s_addr) {
                if (n < 0)
                        err(2, "sendsrcaddr 3 failed");
        } else {
                if (n > 0)
                        errx(2, "sendsrcaddr 2 succeeded!?");
        }

        s2 = socket(AF_INET, SOCK_DGRAM, 0);
        if (s2 < 0)
                err(2, "socket 3 failed");

        on = 1;
        if (setsockopt(s2, SOL_SOCKET, SO_REUSEADDR, &on, sizeof(on)) < 0)
                err(2, "setsockopt(REUSEADDR) failed");

        local.sin_addr.s_addr = INADDR_ANY;
        if (bind(s2, (const struct sockaddr *)&local, sizeof(local)) < 0)
                err(2, "bind * failed");

        /*
         * Connected socket, implied laddr/lport bound:
         * if inp_laddr == src --> pass.
         * if inp_laddr != src --> fail.
         *
         * Wildcard bound above should not matter.
         */
        local_len = sizeof(local);
        if (getsockname(s, (struct sockaddr *)&local, &local_len) < 0)
                err(2, "getsockname failed");
        n = test_sendsrcaddr(s, NULL, &src);
        if (local.sin_addr.s_addr == src.s_addr) {
                if (n < 0)
                        err(2, "sendsrcaddr 4 failed");
        } else {
                if (n > 0)
                        errx(2, "sendsrcaddr 3 succeeded!?");
        }

        close(s2);

        exit(0);
}