// SPDX-License-Identifier: GPL-2.0-only
/* MSG_ZEROCOPY feature tests for vsock
 *
 * Copyright (C) 2023 SberDevices.
 *
 * Author: Arseniy Krasnov <avkrasnov@salutedevices.com>
 */

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/ioctl.h>
#include <sys/mman.h>
#include <unistd.h>
#include <poll.h>
#include <linux/errqueue.h>
#include <linux/kernel.h>
#include <linux/sockios.h>
#include <linux/time64.h>
#include <errno.h>

#include "control.h"
#include "timeout.h"
#include "vsock_test_zerocopy.h"
#include "msg_zerocopy_common.h"

#ifndef PAGE_SIZE
#define PAGE_SIZE               4096
#endif

#define VSOCK_TEST_DATA_MAX_IOV 3

struct vsock_test_data {
        /* This test case if for SOCK_STREAM only. */
        bool stream_only;
        /* Data must be zerocopied. This field is checked against
         * field 'ee_code' of the 'struct sock_extended_err', which
         * contains bit to detect that zerocopy transmission was
         * fallbacked to copy mode.
         */
        bool zerocopied;
        /* Enable SO_ZEROCOPY option on the socket. Without enabled
         * SO_ZEROCOPY, every MSG_ZEROCOPY transmission will behave
         * like without MSG_ZEROCOPY flag.
         */
        bool so_zerocopy;
        /* 'errno' after 'sendmsg()' call. */
        int sendmsg_errno;
        /* Number of valid elements in 'vecs'. */
        int vecs_cnt;
        struct iovec vecs[VSOCK_TEST_DATA_MAX_IOV];
};

static struct vsock_test_data test_data_array[] = {
        /* Last element has non-page aligned size. */
        {
                .zerocopied = true,
                .so_zerocopy = true,
                .sendmsg_errno = 0,
                .vecs_cnt = 3,
                {
                        { NULL, PAGE_SIZE },
                        { NULL, PAGE_SIZE },
                        { NULL, 200 }
                }
        },
        /* All elements have page aligned base and size. */
        {
                .zerocopied = true,
                .so_zerocopy = true,
                .sendmsg_errno = 0,
                .vecs_cnt = 3,
                {
                        { NULL, PAGE_SIZE },
                        { NULL, PAGE_SIZE * 2 },
                        { NULL, PAGE_SIZE * 3 }
                }
        },
        /* All elements have page aligned base and size. But
         * data length is bigger than 64Kb.
         */
        {
                .zerocopied = true,
                .so_zerocopy = true,
                .sendmsg_errno = 0,
                .vecs_cnt = 3,
                {
                        { NULL, PAGE_SIZE * 16 },
                        { NULL, PAGE_SIZE * 16 },
                        { NULL, PAGE_SIZE * 16 }
                }
        },
        /* Middle element has both non-page aligned base and size. */
        {
                .zerocopied = true,
                .so_zerocopy = true,
                .sendmsg_errno = 0,
                .vecs_cnt = 3,
                {
                        { NULL, PAGE_SIZE },
                        { (void *)1, 100 },
                        { NULL, PAGE_SIZE }
                }
        },
        /* Middle element is unmapped. */
        {
                .zerocopied = false,
                .so_zerocopy = true,
                .sendmsg_errno = ENOMEM,
                .vecs_cnt = 3,
                {
                        { NULL, PAGE_SIZE },
                        { MAP_FAILED, PAGE_SIZE },
                        { NULL, PAGE_SIZE }
                }
        },
        /* Valid data, but SO_ZEROCOPY is off. This
         * will trigger fallback to copy.
         */
        {
                .zerocopied = false,
                .so_zerocopy = false,
                .sendmsg_errno = 0,
                .vecs_cnt = 1,
                {
                        { NULL, PAGE_SIZE }
                }
        },
        /* Valid data, but message is bigger than peer's
         * buffer, so this will trigger fallback to copy.
         * This test is for SOCK_STREAM only, because
         * for SOCK_SEQPACKET, 'sendmsg()' returns EMSGSIZE.
         */
        {
                .stream_only = true,
                .zerocopied = false,
                .so_zerocopy = true,
                .sendmsg_errno = 0,
                .vecs_cnt = 1,
                {
                        { NULL, 100 * PAGE_SIZE }
                }
        },
};

#define POLL_TIMEOUT_MS         100

static void test_client(const struct test_opts *opts,
                        const struct vsock_test_data *test_data,
                        bool sock_seqpacket)
{
        struct pollfd fds = { 0 };
        struct msghdr msg = { 0 };
        ssize_t sendmsg_res;
        struct iovec *iovec;
        int fd;

        if (sock_seqpacket)
                fd = vsock_seqpacket_connect(opts->peer_cid, opts->peer_port);
        else
                fd = vsock_stream_connect(opts->peer_cid, opts->peer_port);

        if (fd < 0) {
                perror("connect");
                exit(EXIT_FAILURE);
        }

        if (test_data->so_zerocopy)
                enable_so_zerocopy_check(fd);

        iovec = alloc_test_iovec(test_data->vecs, test_data->vecs_cnt);

        msg.msg_iov = iovec;
        msg.msg_iovlen = test_data->vecs_cnt;

        errno = 0;

        sendmsg_res = sendmsg(fd, &msg, MSG_ZEROCOPY);
        if (errno != test_data->sendmsg_errno) {
                fprintf(stderr, "expected 'errno' == %i, got %i\n",
                        test_data->sendmsg_errno, errno);
                exit(EXIT_FAILURE);
        }

        if (!errno) {
                if (sendmsg_res != iovec_bytes(iovec, test_data->vecs_cnt)) {
                        fprintf(stderr, "expected 'sendmsg()' == %li, got %li\n",
                                iovec_bytes(iovec, test_data->vecs_cnt),
                                sendmsg_res);
                        exit(EXIT_FAILURE);
                }
        }

        fds.fd = fd;
        fds.events = 0;

        if (poll(&fds, 1, POLL_TIMEOUT_MS) < 0) {
                perror("poll");
                exit(EXIT_FAILURE);
        }

        if (fds.revents & POLLERR) {
                vsock_recv_completion(fd, &test_data->zerocopied);
        } else if (test_data->so_zerocopy && !test_data->sendmsg_errno) {
                /* If we don't have data in the error queue, but
                 * SO_ZEROCOPY was enabled and 'sendmsg()' was
                 * successful - this is an error.
                 */
                fprintf(stderr, "POLLERR expected\n");
                exit(EXIT_FAILURE);
        }

        if (!test_data->sendmsg_errno)
                control_writeulong(iovec_hash_djb2(iovec, test_data->vecs_cnt));
        else
                control_writeulong(0);

        control_writeln("DONE");
        free_test_iovec(test_data->vecs, iovec, test_data->vecs_cnt);
        close(fd);
}

void test_stream_msgzcopy_client(const struct test_opts *opts)
{
        int i;

        for (i = 0; i < ARRAY_SIZE(test_data_array); i++)
                test_client(opts, &test_data_array[i], false);
}

void test_seqpacket_msgzcopy_client(const struct test_opts *opts)
{
        int i;

        for (i = 0; i < ARRAY_SIZE(test_data_array); i++) {
                if (test_data_array[i].stream_only)
                        continue;

                test_client(opts, &test_data_array[i], true);
        }
}

static void test_server(const struct test_opts *opts,
                        const struct vsock_test_data *test_data,
                        bool sock_seqpacket)
{
        unsigned long remote_hash;
        unsigned long local_hash;
        ssize_t total_bytes_rec;
        unsigned char *data;
        size_t data_len;
        int fd;

        if (sock_seqpacket)
                fd = vsock_seqpacket_accept(VMADDR_CID_ANY, opts->peer_port, NULL);
        else
                fd = vsock_stream_accept(VMADDR_CID_ANY, opts->peer_port, NULL);

        if (fd < 0) {
                perror("accept");
                exit(EXIT_FAILURE);
        }

        data_len = iovec_bytes(test_data->vecs, test_data->vecs_cnt);

        data = malloc(data_len);
        if (!data) {
                perror("malloc");
                exit(EXIT_FAILURE);
        }

        total_bytes_rec = 0;

        while (total_bytes_rec != data_len) {
                ssize_t bytes_rec;

                bytes_rec = read(fd, data + total_bytes_rec,
                                 data_len - total_bytes_rec);
                if (bytes_rec <= 0)
                        break;

                total_bytes_rec += bytes_rec;
        }

        if (test_data->sendmsg_errno == 0)
                local_hash = hash_djb2(data, data_len);
        else
                local_hash = 0;

        free(data);

        /* Waiting for some result. */
        remote_hash = control_readulong();
        if (remote_hash != local_hash) {
                fprintf(stderr, "hash mismatch\n");
                exit(EXIT_FAILURE);
        }

        control_expectln("DONE");
        close(fd);
}

void test_stream_msgzcopy_server(const struct test_opts *opts)
{
        int i;

        for (i = 0; i < ARRAY_SIZE(test_data_array); i++)
                test_server(opts, &test_data_array[i], false);
}

void test_seqpacket_msgzcopy_server(const struct test_opts *opts)
{
        int i;

        for (i = 0; i < ARRAY_SIZE(test_data_array); i++) {
                if (test_data_array[i].stream_only)
                        continue;

                test_server(opts, &test_data_array[i], true);
        }
}

void test_stream_msgzcopy_empty_errq_client(const struct test_opts *opts)
{
        struct msghdr msg = { 0 };
        char cmsg_data[128];
        ssize_t res;
        int fd;

        fd = vsock_stream_connect(opts->peer_cid, opts->peer_port);
        if (fd < 0) {
                perror("connect");
                exit(EXIT_FAILURE);
        }

        msg.msg_control = cmsg_data;
        msg.msg_controllen = sizeof(cmsg_data);

        res = recvmsg(fd, &msg, MSG_ERRQUEUE);
        if (res != -1) {
                fprintf(stderr, "expected 'recvmsg(2)' failure, got %zi\n",
                        res);
                exit(EXIT_FAILURE);
        }

        control_writeln("DONE");
        close(fd);
}

void test_stream_msgzcopy_empty_errq_server(const struct test_opts *opts)
{
        int fd;

        fd = vsock_stream_accept(VMADDR_CID_ANY, opts->peer_port, NULL);
        if (fd < 0) {
                perror("accept");
                exit(EXIT_FAILURE);
        }

        control_expectln("DONE");
        close(fd);
}

#define GOOD_COPY_LEN   128     /* net/vmw_vsock/virtio_transport_common.c */

void test_stream_msgzcopy_mangle_client(const struct test_opts *opts)
{
        char sbuf1[PAGE_SIZE + 1], sbuf2[GOOD_COPY_LEN];
        unsigned long hash;
        struct pollfd fds;
        int fd, i;

        fd = vsock_stream_connect(opts->peer_cid, opts->peer_port);
        if (fd < 0) {
                perror("connect");
                exit(EXIT_FAILURE);
        }

        enable_so_zerocopy_check(fd);

        memset(sbuf1, 'x', sizeof(sbuf1));
        send_buf(fd, sbuf1, sizeof(sbuf1), 0, sizeof(sbuf1));

        for (i = 0; i < sizeof(sbuf2); i++)
                sbuf2[i] = rand() & 0xff;

        send_buf(fd, sbuf2, sizeof(sbuf2), MSG_ZEROCOPY, sizeof(sbuf2));

        hash = hash_djb2(sbuf2, sizeof(sbuf2));
        control_writeulong(hash);

        fds.fd = fd;
        fds.events = 0;

        if (poll(&fds, 1, TIMEOUT * MSEC_PER_SEC) != 1 ||
            !(fds.revents & POLLERR)) {
                perror("poll");
                exit(EXIT_FAILURE);
        }

        close(fd);
}

void test_stream_msgzcopy_mangle_server(const struct test_opts *opts)
{
        unsigned long local_hash, remote_hash;
        char rbuf[PAGE_SIZE + 1];
        int fd;

        fd = vsock_stream_accept(VMADDR_CID_ANY, opts->peer_port, NULL);
        if (fd < 0) {
                perror("accept");
                exit(EXIT_FAILURE);
        }

        /* Wait, don't race the (buggy) skbs coalescence. */
        vsock_ioctl_int(fd, SIOCINQ, PAGE_SIZE + 1 + GOOD_COPY_LEN);

        /* Discard the first packet. */
        recv_buf(fd, rbuf, PAGE_SIZE + 1, 0, PAGE_SIZE + 1);

        recv_buf(fd, rbuf, GOOD_COPY_LEN, 0, GOOD_COPY_LEN);
        remote_hash = control_readulong();
        local_hash = hash_djb2(rbuf, GOOD_COPY_LEN);

        if (local_hash != remote_hash) {
                fprintf(stderr, "Data received corrupted\n");
                exit(EXIT_FAILURE);
        }

        close(fd);
}