/*
 * This file and its contents are supplied under the terms of the
 * Common Development and Distribution License ("CDDL"), version 1.0.
 * You may only use this file in accordance with the terms of version
 * 1.0 of the CDDL.
 *
 * A full copy of the text of the CDDL should have accompanied this
 * source.  A copy of the CDDL is also available via the Internet at
 * http://www.illumos.org/license/CDDL.
 */

/*
 * Copyright 2020 OmniOS Community Edition (OmniOSce) Association.
 * Copyright 2020 Joyent, Inc.
 */

/*
 * Test file descriptor passing via SCM_RIGHTS, and in particular what happens
 * on message truncation in terms of the represented size of the data in the
 * control message. Ensure that no file descriptors are leaked - the kernel
 * must close any that would not fit in the available buffer space and the
 * userland application must close the rest.
 */

#include <stdio.h>
#include <errno.h>
#include <fcntl.h>
#include <signal.h>
#include <stdlib.h>
#include <string.h>
#include <strings.h>
#include <unistd.h>
#include <libproc.h>

#include <sys/types.h>
#include <sys/param.h>
#include <sys/socket.h>
#include <sys/stat.h>
#include <sys/un.h>
#include <sys/wait.h>
#include <assert.h>
#include <alloca.h>
#include <err.h>

static boolean_t debug;

typedef struct cmsg_test {
        char *name;             /* Name of the test */
        uint_t send;            /* Number of FDs to send */
        uint_t recv;            /* Size receive buffer for this number of FDs */
        size_t predata;         /* Prepend dummy cmsg of this size */
        int bufsize;            /* Explicitly set receive buffer size. */
                                /* Overrides 'recv' if non-zero */
        uint_t x_controllen;    /* Expected received msg_controllen */
        uint_t x_cmsg_datalen;  /* Expected received cmsg data length */
        uint32_t x_flags;       /* Expected received msf_flags */
} cmsg_test_t;

static cmsg_test_t tests[] = {
        {
                .name = "send 1, recv 1",
                .send = 1,
                .recv = 1,
                .predata = 0,
                .bufsize = 0,
                .x_controllen = 16,
                .x_cmsg_datalen = 4,
                .x_flags = 0,
        },
        {
                .name = "send 10, recv 10",
                .send = 10,
                .recv = 10,
                .predata = 0,
                .bufsize = 0,
                .x_controllen = 52,
                .x_cmsg_datalen = 40,
                .x_flags = 0,
        },
        {
                .name = "send 2, recv 1",
                .send = 2,
                .recv = 1,
                .predata = 0,
                .bufsize = 0,
                .x_controllen = 16,
                .x_cmsg_datalen = 4,
                .x_flags = MSG_CTRUNC,
        },
        {
                .name = "send 2, recv 1, buffer 5",
                .send = 2,
                .recv = 1,
                .predata = 0,
                .bufsize = sizeof (int) * 2 - 3,
                .x_controllen = 17,
                .x_cmsg_datalen = 5,
                .x_flags = MSG_CTRUNC,
        },
        {
                .name = "send 2, recv 1, buffer 6",
                .send = 2,
                .recv = 1,
                .predata = 0,
                .bufsize = sizeof (int) * 2 - 2,
                .x_controllen = 18,
                .x_cmsg_datalen = 6,
                .x_flags = MSG_CTRUNC,
        },
        {
                .name = "send 2, recv 1, buffer 7",
                .send = 2,
                .recv = 1,
                .predata = 0,
                .bufsize = sizeof (int) * 2 - 1,
                .x_controllen = 19,
                .x_cmsg_datalen = 7,
                .x_flags = MSG_CTRUNC,
        },

        /* Tests where there is no room allowed for data */

        {
                .name = "send 2, recv 0, hdronly",
                .send = 2,
                .recv = 0,
                .predata = 0,
                .bufsize = 0,
                .x_controllen = 12,
                .x_cmsg_datalen = 0,
                .x_flags = MSG_CTRUNC,
        },

        {
                .name = "send 2, recv 0, hdr - 1",
                .send = 2,
                .recv = 0,
                .predata = 0,
                .bufsize = -1,
                .x_controllen = 11,
                .x_cmsg_datalen = 0,
                .x_flags = MSG_CTRUNC,
        },

        {
                .name = "send 2, recv 0, hdr - 5",
                .send = 2,
                .recv = 0,
                .predata = 0,
                .bufsize = -5,
                .x_controllen = 7,
                .x_cmsg_datalen = 0,
                .x_flags = MSG_CTRUNC,
        },

        /* Tests where SCM_RIGHTS is not the first message */

        {
                .name = "send 1, recv 1, pre 8",
                .send = 1,
                .recv = 1,
                .predata = 8,
                .bufsize = 0,
                .x_controllen = 36,
                .x_cmsg_datalen = 4,
                .x_flags = 0,
        },
        {
                .name = "send 1, recv 1, pre 7",
                .send = 1,
                .recv = 1,
                .predata = 7,
                .bufsize = 0,
                .x_controllen = 35,
                .x_cmsg_datalen = 4,
                .x_flags = 0,
        },
        {
                .name = "send 1, recv 1, pre 6",
                .send = 1,
                .recv = 1,
                .predata = 6,
                .bufsize = 0,
                .x_controllen = 34,
                .x_cmsg_datalen = 4,
                .x_flags = 0,
        },
        {
                .name = "send 1, recv 1, pre 5",
                .send = 1,
                .recv = 1,
                .predata = 5,
                .bufsize = 0,
                .x_controllen = 33,
                .x_cmsg_datalen = 4,
                .x_flags = 0,
        },

        {
                .name = "send 2, recv 1, pre 8",
                .send = 2,
                .recv = 1,
                .predata = 8,
                .bufsize = 0,
                .x_controllen = 36,
                .x_cmsg_datalen = 8,
                .x_flags = MSG_CTRUNC,
        },
        {
                .name = "send 2, recv 1, pre 7",
                .send = 2,
                .recv = 1,
                .predata = 7,
                .bufsize = 0,
                .x_controllen = 36,
                .x_cmsg_datalen = 8,
                .x_flags = MSG_CTRUNC,
        },
        {
                .name = "send 2, recv 1, pre 6",
                .send = 2,
                .recv = 1,
                .predata = 6,
                .bufsize = 0,
                .x_controllen = 36,
                .x_cmsg_datalen = 8,
                .x_flags = MSG_CTRUNC,
        },
        {
                .name = "send 2, recv 1, pre 5",
                .send = 2,
                .recv = 1,
                .predata = 5,
                .bufsize = 0,
                .x_controllen = 36,
                .x_cmsg_datalen = 8,
                .x_flags = MSG_CTRUNC,
        },
        {
                .name = "send 2, recv 1, pre 4",
                .send = 2,
                .recv = 1,
                .predata = 4,
                .bufsize = 0,
                .x_controllen = 32,
                .x_cmsg_datalen = 8,
                .x_flags = MSG_CTRUNC,
        },
        {
                .name = "send 2, recv 1, pre 3",
                .send = 2,
                .recv = 1,
                .predata = 3,
                .bufsize = 0,
                .x_controllen = 32,
                .x_cmsg_datalen = 8,
                .x_flags = MSG_CTRUNC,
        },
        {
                .name = "send 2, recv 1, pre 2",
                .send = 2,
                .recv = 1,
                .predata = 2,
                .bufsize = 0,
                .x_controllen = 32,
                .x_cmsg_datalen = 8,
                .x_flags = MSG_CTRUNC,
        },
        {
                .name = "send 2, recv 1, pre 1",
                .send = 2,
                .recv = 1,
                .predata = 1,
                .bufsize = 0,
                .x_controllen = 32,
                .x_cmsg_datalen = 8,
                .x_flags = MSG_CTRUNC,
        },

        {
                .name = "send 2, recv 1, pre 8, buffer 5",
                .send = 2,
                .recv = 1,
                .predata = 8,
                .bufsize = sizeof (int) * 2 - 3,
                .x_controllen = 37,
                .x_cmsg_datalen = 8,
                .x_flags = MSG_CTRUNC,
        },
        {
                .name = "send 2, recv 1, pre 8, buffer 6",
                .send = 2,
                .recv = 1,
                .predata = 8,
                .bufsize = sizeof (int) * 2 - 2,
                .x_controllen = 38,
                .x_cmsg_datalen = 8,
                .x_flags = MSG_CTRUNC,
        },
        {
                .name = "send 2, recv 1, pre 8, buffer 7",
                .send = 2,
                .recv = 1,
                .predata = 8,
                .bufsize = sizeof (int) * 2 - 1,
                .x_controllen = 39,
                .x_cmsg_datalen = 8,
                .x_flags = MSG_CTRUNC,
        },
        {
                .name = "send 10, recv 1, pre 8",
                .send = 10,
                .recv = 1,
                .predata = 8,
                .bufsize = 0,
                .x_controllen = 36,
                .x_cmsg_datalen = 24,
                .x_flags = MSG_CTRUNC,
        },

        /* End of tests */

        {
                .name = NULL
        }
};

static int sock = -1, testfd = -1, cfd = -1;
static int fdcount;

static int
fdwalkcb(const prfdinfo_t *info, void *arg)
{
        if (!S_ISDIR(info->pr_mode) && info->pr_fd > 2 &&
            info->pr_fd != sock && info->pr_fd != testfd &&
            info->pr_fd != cfd) {
                if (debug) {
                        fprintf(stderr, "%s: unexpected fd: %d\n",
                            (char *)arg, info->pr_fd);
                }
                fdcount++;
        }

        return (0);

}

static void
check_fds(char *tag)
{
        fdcount = 0;
        proc_fdwalk(getpid(), fdwalkcb, tag);
}

static void
send_and_wait(pid_t pid, sigset_t *set, int osig, int isig)
{
        int sig;

        if (osig > 0)
                kill(pid, osig);

        if (isig > 0) {
                if (sigwait(set, &sig) != 0) {
                        err(EXIT_FAILURE,
                            "sigwait failed waiting for %d", isig);
                }
                if (sig == SIGINT) {
                        exit(1);
                }
                if (sig != isig) {
                        err(EXIT_FAILURE,
                            "sigwait returned unexpected signal %d", sig);
                }
        }
}

static void
sendtest(cmsg_test_t *t)
{
        struct msghdr msg;
        struct cmsghdr *cm;
        struct iovec iov;
        ssize_t nbytes;
        char c = '*';
        int i, *p;

        bzero(&msg, sizeof (msg));

        msg.msg_name = NULL;
        msg.msg_namelen = 0;

        iov.iov_base = &c;
        iov.iov_len = sizeof (c);
        msg.msg_iov = &iov;
        msg.msg_iovlen = 1;

        msg.msg_flags = 0;

        msg.msg_controllen = CMSG_SPACE(sizeof (int) * t->send);

        if (t->predata > 0) {
                /* A dummy cmsg will be inserted at the head of the data */
                msg.msg_controllen += CMSG_SPACE(t->predata);
        }

        msg.msg_control = alloca(msg.msg_controllen);
        bzero(msg.msg_control, msg.msg_controllen);

        cm = CMSG_FIRSTHDR(&msg);

        if (t->predata > 0) {
                /* Insert the dummy cmsg */
                cm->cmsg_len = CMSG_LEN(t->predata);
                cm->cmsg_level = SOL_SOCKET;
                cm->cmsg_type = 0;
                cm = CMSG_NXTHDR(&msg, cm);
        }

        cm->cmsg_len = CMSG_LEN(sizeof (int) * t->send);
        cm->cmsg_level = SOL_SOCKET;
        cm->cmsg_type = SCM_RIGHTS;

        p = (int *)CMSG_DATA(cm);
        for (i = 0; i < t->send; i++) {
                int s = dup(testfd);
                if (s == -1)
                        err(EXIT_FAILURE, "dup()");
                *p++ = s;
        }

        if (debug)
                printf("Sending: controllen=%u\n", msg.msg_controllen);

        nbytes = sendmsg(cfd, &msg, 0);
        if (nbytes == -1)
                err(EXIT_FAILURE, "sendmsg()");

        p = (int *)CMSG_DATA(cm);
        for (i = 0; i < t->send; i++)
                (void) close(*p++);
}

static int
server(const char *sockpath, pid_t pid)
{
        struct sockaddr_un addr;
        sigset_t set;
        cmsg_test_t *t;

        sigemptyset(&set);
        sigaddset(&set, SIGUSR2);
        sigaddset(&set, SIGINT);

        sock = socket(PF_LOCAL, SOCK_STREAM, 0);
        if (sock == -1)
                err(EXIT_FAILURE, "failed to create socket");
        addr.sun_family = AF_UNIX;
        strlcpy(addr.sun_path, sockpath, sizeof (addr.sun_path));
        if (bind(sock, (struct sockaddr *)&addr, sizeof (addr)) == -1)
                err(EXIT_FAILURE, "bind failed");
        if (listen(sock, 0) == -1)
                err(EXIT_FAILURE, "listen failed");

        if ((testfd = open("/dev/null", O_RDONLY)) == -1)
                err(EXIT_FAILURE, "/dev/null");

        check_fds("server");

        /* Signal the child to connect to the socket */
        send_and_wait(pid, &set, SIGUSR1, SIGUSR2);

        if ((cfd = accept(sock, NULL, 0)) == -1)
                err(EXIT_FAILURE, "accept failed");

        for (t = tests; t->name != NULL; t++) {
                if (debug)
                        printf("\n>>> Starting test %s\n", t->name);

                sendtest(t);
                check_fds("server");

                send_and_wait(pid, &set, SIGUSR1, SIGUSR2);
        }

        close(cfd);
        close(testfd);
        close(sock);

        return (0);
}

static boolean_t pass;

static void
check(uint_t actual, uint_t expected, char *tag)
{
        if (actual != expected) {
                fprintf(stderr, "    !!!: "
                    "%1$s = %2$u(%2$#x) (expected %3$u(%3$#x))\n",
                    tag, actual, expected);
                pass = _B_FALSE;
        } else if (debug) {
                fprintf(stderr, "       : "
                    "%1$s = %2$u(%2$#x)\n",
                    tag, actual);
        }
}

static boolean_t
recvtest(cmsg_test_t *t)
{
        struct msghdr msg;
        struct cmsghdr *cm;
        struct iovec iov;
        size_t bufsize;
        ssize_t nbytes;
        char c = '*';

        bzero(&msg, sizeof (msg));

        msg.msg_name = NULL;
        msg.msg_namelen = 0;

        iov.iov_base = &c;
        iov.iov_len = sizeof (c);
        msg.msg_iov = &iov;
        msg.msg_iovlen = 1;

        msg.msg_flags = 0;

        /*
         * If the test does not specify a receive buffer size, calculate one
         * from the number of file descriptors to receive.
         */
        if (t->bufsize == 0) {
                bufsize = sizeof (int) * t->recv;
                bufsize = CMSG_SPACE(bufsize);
        } else {
                /*
                 * Use the specific buffer size provided but add in
                 * space for the header
                 */
                bufsize = t->bufsize + CMSG_LEN(0);
        }

        if (t->predata > 0) {
                /* A dummy cmsg will be found at the head of the data */
                bufsize += CMSG_SPACE(t->predata);
        }

        msg.msg_controllen = bufsize;
        msg.msg_control = alloca(bufsize);
        bzero(msg.msg_control, msg.msg_controllen);

        pass = _B_TRUE;

        if (debug)
                printf("Receiving: controllen=%u, \n", msg.msg_controllen);

        nbytes = recvmsg(sock, &msg, 0);

        if (nbytes == -1) {
                pass = _B_FALSE;
                fprintf(stderr, "recvmsg() failed: %s\n", strerror(errno));
                goto out;
        }

        if (debug) {
                printf("Received: controllen=%u, flags=%#x\n",
                    msg.msg_controllen, msg.msg_flags);
        }

        check(msg.msg_flags, t->x_flags, "msg_flags");
        check(msg.msg_controllen, t->x_controllen, "msg_controllen");

        for (cm = CMSG_FIRSTHDR(&msg); cm; cm = CMSG_NXTHDR(&msg, cm)) {
                void *data, *end;

                if (debug) {
                        printf("    >> : Got cmsg %x/%x - %u\n", cm->cmsg_level,
                            cm->cmsg_type, cm->cmsg_len);
                }

                if (cm->cmsg_type != SCM_RIGHTS) {
                        if (debug)
                                printf("       : skipping cmsg\n");
                        continue;
                }

                check(cm->cmsg_len - CMSG_LEN(0),
                    t->x_cmsg_datalen, "cmsg_len");

                /* Close any received file descriptors */
                data = CMSG_DATA(cm);

                if ((msg.msg_flags & MSG_CTRUNC) &&
                    CMSG_NXTHDR(&msg, cm) == NULL) {
                        /*
                         * illumos did not previously adjust cmsg_len on
                         * truncation. This is the last cmsg, derive the
                         * length from msg_controllen
                         */
                        end = msg.msg_control + msg.msg_controllen;
                } else {
                        end = data + cm->cmsg_len - CMSG_LEN(0);
                }

                while (data <= end - sizeof (int)) {
                        int *a = (int *)data;
                        if (debug)
                                printf("       : close(%d)\n", *a);
                        if (close(*a) == -1) {
                                pass = _B_FALSE;
                                fprintf(stderr, "    !!!: "
                                    "failed to close fd %d - %s\n", *a,
                                    strerror(errno));
                        }
                        data += sizeof (int);
                }
        }

out:

        check_fds("client");
        check(fdcount, 0, "client descriptors");
        printf("     + : %s %s\n", pass ? "PASS" : "FAIL", t->name);

        return (pass);
}

static int
client(const char *sockpath, pid_t pid)
{
        struct sockaddr_un addr;
        sigset_t set;
        cmsg_test_t *t;
        int ret = 0;

        sigemptyset(&set);
        sigaddset(&set, SIGUSR1);
        sigaddset(&set, SIGINT);

        send_and_wait(pid, &set, 0, SIGUSR1);

        sock = socket(PF_LOCAL, SOCK_STREAM, 0);
        if (sock == -1)
                err(EXIT_FAILURE, "failed to create socket");
        addr.sun_family = AF_UNIX;
        strlcpy(addr.sun_path, sockpath, sizeof (addr.sun_path));
        if (connect(sock, (struct sockaddr *)&addr, sizeof (addr)) == -1)
                err(EXIT_FAILURE, "could not connect to server socket");

        for (t = tests; t->name != NULL; t++) {
                send_and_wait(pid, &set, SIGUSR2, SIGUSR1);
                if (!recvtest(t))
                        ret = 1;
        }

        close(sock);

        return (ret);
}

int
main(int argc, const char **argv)
{
        char sockpath[] = "/tmp/cmsg.testsock.XXXXXX";
        pid_t pid, ppid;
        sigset_t set;
        int ret = 0;

        /*
         * The tests make assumptions about the number of open file descriptors
         * present. In case we are invoked with more than just STDIN_FILENO,
         * STDOUT_FILENO, and STDERR_FILENO open, close any other open
         * descriptors that might exist. Otherwise their presence will violate
         * the assumptions of the test and cause an erroneous failure.
         */
        closefrom(STDERR_FILENO + 1);

        if (argc > 1 && strcmp(argv[1], "-d") == 0)
                debug = _B_TRUE;

        sigfillset(&set);
        sigdelset(&set, SIGINT);
        sigdelset(&set, SIGTSTP);
        sigprocmask(SIG_BLOCK, &set, NULL);

        if (mktemp(sockpath) == NULL)
                err(EXIT_FAILURE, "Failed to make temporary socket path");

        ppid = getpid();
        pid = fork();
        switch (pid) {
        case -1:
                err(EXIT_FAILURE, "fork failed");
        case 0:
                return (server(sockpath, ppid));
        default:
                break;
        }

        ret = client(sockpath, pid);
        kill(pid, SIGINT);

        unlink(sockpath);

        return (ret);
}