#define _GNU_SOURCE
#include <assert.h>

#include "test-runner.h"
#include "virtio_wl.c"

TEST(iov_len_fit)
{
	int f1[] = { 1, 2 };
	int f2[] = { 3 };

	struct iovec iov[] = {
		{ .iov_base = f1, .iov_len = sizeof(f1) },
		{ .iov_base = NULL, .iov_len = 0 },
		{ .iov_base = f2, .iov_len = sizeof(f2) },
	};

	assert(iov_len(iov, sizeof(iov) / sizeof(*iov)) == 3 * sizeof(int));
}

TEST(iov_len_overflow)
{
	struct iovec iov[] = {
		{ .iov_base = NULL, .iov_len = SIZE_MAX },
		{ .iov_base = NULL, .iov_len = SIZE_MAX },
	};

	assert(iov_len(iov, sizeof(iov) / sizeof(*iov)) == -1);
}

TEST(iov_flatten_test)
{
	int f1[] = { 1 };
	int f2[] = { 2, 3 };

	struct iovec iov[] = {
		{ .iov_base = f1, .iov_len = sizeof(f1) },
		{ .iov_base = NULL, .iov_len = 0 },
		{ .iov_base = f2, .iov_len = sizeof(f2) },
	};

	int buf[2];
	iov_flatten(buf, sizeof(buf), iov, sizeof(iov) / sizeof(*iov));
	assert(buf[0] == 1);
	assert(buf[1] == 2);
}

TEST(iov_flatten_zero_iov)
{
	unsigned char buf[] = { 0xFF };
	assert(iov_flatten(buf, 1, NULL, 0) == 0);
}

TEST(iov_flatten_null_iov)
{
	unsigned char buf[] = { 0xFF };
	struct iovec iov[] = { { .iov_base = NULL, .iov_len = 0 } };

	assert(iov_flatten(buf, sizeof buf, iov, sizeof(iov) / sizeof(*iov)) == 0);
}

TEST(iov_flatten_zero_buf)
{
	unsigned char f1[] = { 0 };
	struct iovec iov[] = { { .iov_base = f1, .iov_len = sizeof f1 } };

	assert(iov_flatten(NULL, 0, iov, sizeof(iov) / sizeof(*iov)) == 0);
}

TEST(iov_flatten_short_iov)
{
	int f1[] = { 1 };
	int f2[] = { 2, 3 };

	struct iovec iov[] = {
		{ .iov_base = f1, .iov_len = sizeof f1 },
		{ .iov_base = NULL, .iov_len = 0 },
		{ .iov_base = f2, .iov_len = sizeof f2 },
	};

	int buf[] = { 9, 8, 7, 6 };

	assert(iov_flatten(buf, sizeof buf, iov, sizeof(iov) / sizeof(*iov)) == 12);
	assert(buf[0] == 1);
	assert(buf[1] == 2);
	assert(buf[2] == 3);
	assert(buf[3] == 6);
}

TEST(iov_flatten_exact_iov)
{
	int f1[] = { 1 };
	int f2[] = { 2, 3 };

	struct iovec iov[] = {
		{ .iov_base = f1, .iov_len = sizeof f1 },
		{ .iov_base = NULL, .iov_len = 0 },
		{ .iov_base = f2, .iov_len = sizeof f2 },
	};

	int buf[] = { 9, 8, 7 };

	assert(iov_flatten(buf, sizeof buf, iov, sizeof(iov) / sizeof(*iov)) == 12);
	assert(f1[0] == 1);
	assert(f2[0] == 2);
	assert(f2[1] == 3);
}

TEST(iov_flatten_long_iov)
{
	int f1[] = { 1 };
	int f2[] = { 2, 3, 4 };

	struct iovec iov[] = {
		{ .iov_base = f1, .iov_len = sizeof f1 },
		{ .iov_base = NULL, .iov_len = 0 },
		{ .iov_base = f2, .iov_len = sizeof f2 },
	};

	int buf[] = { 9, 8, 7 };

	assert(iov_flatten(buf, sizeof buf, iov, sizeof(iov) / sizeof(*iov)) == 12);
	assert(buf[0] == 1);
	assert(buf[1] == 2);
	assert(buf[2] == 3);
}

TEST(iov_fill_zero_iov)
{
	unsigned char buf[] = { 0xFF };
	assert(iov_fill(NULL, 0, buf, 1) == 0);
}

TEST(iov_fill_null_iov)
{
	unsigned char buf[] = { 0xFF };
	struct iovec iov[] = { { .iov_base = NULL, .iov_len = 0 } };

	assert(iov_fill(iov, sizeof(iov) / sizeof(*iov), buf, sizeof buf) == 0);
}

TEST(iov_fill_zero_buf)
{
	unsigned char f1[] = { 0 };
	struct iovec iov[] = { { .iov_base = f1, .iov_len = sizeof f1 } };

	assert(iov_fill(iov, sizeof(iov) / sizeof(*iov), NULL, 0) == 0);
}

TEST(iov_fill_short_iov)
{
	int f1[] = { 1 };
	int f2[] = { 2, 3 };

	struct iovec iov[] = {
		{ .iov_base = f1, .iov_len = sizeof f1 },
		{ .iov_base = NULL, .iov_len = 0 },
		{ .iov_base = f2, .iov_len = sizeof f2 },
	};

	int buf[] = { 9, 8, 7, 6 };

	assert(iov_fill(iov, sizeof(iov) / sizeof(*iov), buf, sizeof buf) == 12);
	assert(f1[0] == 9);
	assert(f2[0] == 8);
	assert(f2[1] == 7);
}

TEST(iov_fill_exact)
{
	int f1[] = { 1 };
	int f2[] = { 2, 3 };

	struct iovec iov[] = {
		{ .iov_base = f1, .iov_len = sizeof f1 },
		{ .iov_base = NULL, .iov_len = 0 },
		{ .iov_base = f2, .iov_len = sizeof f2 },
	};

	int buf[] = { 9, 8, 7 };

	assert(iov_fill(iov, sizeof(iov) / sizeof(*iov), buf, sizeof buf) == 12);
	assert(f1[0] == 9);
	assert(f2[0] == 8);
	assert(f2[1] == 7);
}

TEST(iov_fill_long_iov)
{
	int f1[] = { 1 };
	int f2[] = { 2, 3, 4 };

	struct iovec iov[] = {
		{ .iov_base = f1, .iov_len = sizeof f1 },
		{ .iov_base = NULL, .iov_len = 0 },
		{ .iov_base = f2, .iov_len = sizeof f2 },
	};

	int buf[] = { 9, 8, 7 };

	assert(iov_fill(iov, sizeof(iov) / sizeof(*iov), buf, sizeof buf) == 12);
	assert(f1[0] == 9);
	assert(f2[0] == 8);
	assert(f2[1] == 7);
	assert(f2[2] == 4);
}

TEST(cmsg_to_fdbuf_zero_buf)
{
	int fds[] = { 0 };

	union {
		char buf[CMSG_SPACE(sizeof fds)];
		struct cmsghdr align;
	} u;

	struct msghdr msg = { 0 };
	msg.msg_control = u.buf;
	msg.msg_controllen = sizeof u.buf;

	struct cmsghdr *cmsg = CMSG_FIRSTHDR(&msg);
	cmsg->cmsg_level = SOL_SOCKET;
	cmsg->cmsg_type = SCM_RIGHTS;
	cmsg->cmsg_len = CMSG_LEN(sizeof fds);
	memcpy(CMSG_DATA(cmsg), fds, sizeof fds);

	assert(cmsg_to_fdbuf(NULL, 0, &msg) == 0);
}

TEST(cmsg_to_fdbuf_zero_cmsg)
{
	int fds[2] = { 0 };
	struct msghdr msg = { 0 };
	assert(cmsg_to_fdbuf(fds, sizeof fds, &msg) == 0);
}

TEST(cmsg_to_fdbuf_small_buf)
{
	int fds[] = { 0, 1 };

	union {
		char buf[CMSG_SPACE(sizeof fds)];
		struct cmsghdr align;
	} u;

	struct msghdr msg = { 0 };
	msg.msg_control = u.buf;
	msg.msg_controllen = sizeof u.buf;

	struct cmsghdr *cmsg = CMSG_FIRSTHDR(&msg);
	cmsg->cmsg_level = SOL_SOCKET;
	cmsg->cmsg_type = SCM_RIGHTS;
	cmsg->cmsg_len = CMSG_LEN(sizeof fds);
	memcpy(CMSG_DATA(cmsg), fds, sizeof fds);

	int buf[] = { 0xFF };
	assert(cmsg_to_fdbuf(buf, sizeof(buf) / sizeof(*buf), &msg) == 1);
	assert(buf[0] == 0);
}

TEST(cmsg_to_fdbuf_exact)
{
	int fds[] = { 0, 1 };

	union {
		char buf[CMSG_SPACE(sizeof fds)];
		struct cmsghdr align;
	} u;

	struct msghdr msg = { 0 };
	msg.msg_control = u.buf;
	msg.msg_controllen = sizeof u.buf;

	struct cmsghdr *cmsg = CMSG_FIRSTHDR(&msg);
	cmsg->cmsg_level = SOL_SOCKET;
	cmsg->cmsg_type = SCM_RIGHTS;
	cmsg->cmsg_len = CMSG_LEN(sizeof fds);
	memcpy(CMSG_DATA(cmsg), fds, sizeof fds);

	int buf[] = { 0xFF, 0xFE };
	assert(cmsg_to_fdbuf(buf, sizeof(buf) / sizeof(*buf), &msg) == 2);
	assert(buf[0] == 0);
	assert(buf[1] == 1);
}

TEST(cmsg_to_fdbuf_small_cmsg)
{
	int fds[] = { 0 };

	union {
		char buf[CMSG_SPACE(sizeof fds)];
		struct cmsghdr align;
	} u;

	struct msghdr msg = { 0 };
	msg.msg_control = u.buf;
	msg.msg_controllen = sizeof u.buf;

	struct cmsghdr *cmsg = CMSG_FIRSTHDR(&msg);
	cmsg->cmsg_level = SOL_SOCKET;
	cmsg->cmsg_type = SCM_RIGHTS;
	cmsg->cmsg_len = CMSG_LEN(sizeof fds);
	memcpy(CMSG_DATA(cmsg), fds, sizeof fds);

	int buf[] = { 0xFF, 0xFE };
	assert(cmsg_to_fdbuf(buf, sizeof(buf) / sizeof(*buf), &msg) == 1);
	assert(buf[0] == 0);
	assert(buf[1] == 0xFE);
}

TEST(cmsg_to_fdbuf_multiple_cmsg)
{
	int f1[] = { 0 };
	int f2[] = { 1 };

	union {
		char buf[CMSG_SPACE(sizeof f1) + CMSG_SPACE(sizeof f2)];
		struct cmsghdr align;
	} u;

	struct msghdr msg = { 0 };
	msg.msg_control = u.buf;
	msg.msg_controllen = sizeof u.buf;

	struct cmsghdr *cmsg1 = CMSG_FIRSTHDR(&msg);
	cmsg1->cmsg_level = SOL_SOCKET;
	cmsg1->cmsg_type = SCM_RIGHTS;
	cmsg1->cmsg_len = CMSG_LEN(sizeof f1);
	memcpy(CMSG_DATA(cmsg1), f1, sizeof f1);

	struct cmsghdr *cmsg2 = CMSG_NXTHDR(&msg, cmsg1);
	cmsg2->cmsg_level = SOL_SOCKET;
	cmsg2->cmsg_type = SCM_RIGHTS;
	cmsg2->cmsg_len = CMSG_LEN(sizeof f2);
	memcpy(CMSG_DATA(cmsg2), f2, sizeof f2);

	int buf[] = { 0xFF, 0xFE };
	assert(cmsg_to_fdbuf(buf, sizeof(buf) / sizeof(*buf), &msg) == 2);
	assert(buf[0] == 0);
	assert(buf[1] == 1);
}

TEST(cmsg_to_fdbuf_other_cmsg)
{
	struct ucred cred = {
		.pid = getpid(),
		.uid = getuid(),
		.gid = getgid(),
	};

	int fds[] = { 0, 1 };

	union {
		char buf[CMSG_SPACE(sizeof cred) + CMSG_SPACE(sizeof fds)];
		struct cmsghdr align;
	} u;

	struct msghdr msg = { 0 };
	msg.msg_control = u.buf;
	msg.msg_controllen = sizeof u.buf;

	struct cmsghdr *cmsg1 = CMSG_FIRSTHDR(&msg);
	cmsg1->cmsg_level = SOL_SOCKET;
	cmsg1->cmsg_type = SCM_CREDENTIALS;
	cmsg1->cmsg_len = CMSG_LEN(sizeof cred);
	memcpy(CMSG_DATA(cmsg1), &cred, sizeof cred);

	struct cmsghdr *cmsg2 = CMSG_NXTHDR(&msg, cmsg1);
	cmsg2->cmsg_level = SOL_SOCKET;
	cmsg2->cmsg_type = SCM_RIGHTS;
	cmsg2->cmsg_len = CMSG_LEN(sizeof fds);
	memcpy(CMSG_DATA(cmsg2), fds, sizeof fds);

	int buf[] = { 0xFF, 0xFE };
	assert(cmsg_to_fdbuf(buf, sizeof(buf) / sizeof(*buf), &msg) == -1);
	assert(errno == EINVAL);
}

TEST(cmsg_to_fdbuf_misaligned)
{
	int fds[] = { 0, 1 };

	union {
		char buf[CMSG_SPACE(sizeof fds + 1)];
		struct cmsghdr align;
	} u;

	struct msghdr msg = { 0 };
	msg.msg_control = u.buf;
	msg.msg_controllen = sizeof u.buf;

	struct cmsghdr *cmsg = CMSG_FIRSTHDR(&msg);
	cmsg->cmsg_level = SOL_SOCKET;
	cmsg->cmsg_type = SCM_RIGHTS;
	cmsg->cmsg_len = CMSG_LEN(sizeof(fds) + 1);
	memcpy(CMSG_DATA(cmsg), fds, sizeof fds);
	CMSG_DATA(cmsg)[sizeof(fds)] = 0;

	int buf[] = { 0xFF, 0xFE };
	assert(cmsg_to_fdbuf(buf, sizeof(buf) / sizeof(*buf), &msg) == -1);
	assert(errno == EINVAL);
}

TEST(fdbuf_to_cmsg_null_cmsg)
{
	struct msghdr msg = { 0 };
	int fds[] = { 1 };
	assert(fdbuf_to_cmsg(&msg, fds, sizeof(fds) / sizeof(*fds)) == 0);
	assert(msg.msg_flags & MSG_CTRUNC);
}

TEST(fdbuf_to_cmsg_empty_both)
{
	struct msghdr msg = { 0 };
	int fds[] = { -1 };
	assert(fdbuf_to_cmsg(&msg, fds, sizeof(fds) / sizeof(*fds)) == 0);
	assert(!(msg.msg_flags & MSG_CTRUNC));
}

TEST(fdbuf_to_cmsg_null_buf)
{
	union {
		char buf[CMSG_SPACE(sizeof(int))];
		struct cmsghdr align;
	} u = { 0 };

	struct msghdr msg = { 0 };
	msg.msg_control = u.buf;
	msg.msg_controllen = sizeof u.buf;

	assert(fdbuf_to_cmsg(&msg, NULL, 0) == 0);
	assert(!(msg.msg_flags & MSG_CTRUNC));
}

TEST(fdbuf_to_cmsg_tiny_msg)
{
	union {
		char buf[1];
		struct cmsghdr align;
	} u = { 0 };

	struct msghdr msg = { 0 };
	msg.msg_control = u.buf;
	msg.msg_controllen = sizeof u.buf;

	int fds[] = { 1 };
	assert(fdbuf_to_cmsg(&msg, fds, sizeof(fds) / sizeof(*fds)) == 0);
	assert(msg.msg_flags & MSG_CTRUNC);
}

TEST(fdbuf_to_cmsg_small_msg)
{
	union {
		char buf[CMSG_SPACE(sizeof(int))];
		struct cmsghdr align;
	} u = { 0 };

	struct msghdr msg = { 0 };
	msg.msg_control = u.buf;
	msg.msg_controllen = sizeof u.buf;

	int fds[] = { 1, 2, 3 };
	assert(CMSG_SPACE(sizeof fds) != sizeof u.buf);

	// CMSG_SPACE(sizeof(int)) can equal CMSG_SPACE(sizeof(int) * 2)
	size_t n = fdbuf_to_cmsg(&msg, fds, sizeof(fds) / sizeof(*fds));
	assert(CMSG_SPACE(n * sizeof(int)) == sizeof u.buf);
	assert(msg.msg_flags & MSG_CTRUNC);
}

TEST(fdbuf_to_cmsg_equal)
{
	int fds[] = { 1, 2 };

	union {
		char buf[CMSG_SPACE(sizeof fds)];
		struct cmsghdr align;
	} u = { 0 };

	struct msghdr msg = { 0 };
	msg.msg_control = u.buf;
	msg.msg_controllen = sizeof u.buf;

	assert(fdbuf_to_cmsg(&msg, fds, sizeof(fds) / sizeof(*fds)) == 2);
	assert(!memcmp(CMSG_DATA(CMSG_FIRSTHDR(&msg)), fds, sizeof fds));
	assert(!(msg.msg_flags & MSG_CTRUNC));
}

TEST(fdbuf_to_cmsg_negative_one)
{
	int fds[] = { 1, -1, -1, 2 };

	union {
		char buf[CMSG_SPACE(2 * sizeof(int))];
		struct cmsghdr align;
	} u = { 0 };

	struct msghdr msg = { 0 };
	msg.msg_control = u.buf;
	msg.msg_controllen = sizeof u.buf;

	assert(fdbuf_to_cmsg(&msg, fds, sizeof(fds) / sizeof(*fds)) == 2);

	unsigned char *data = CMSG_DATA(CMSG_FIRSTHDR(&msg));
	assert(!memcmp(data, &fds[0], sizeof(int)));
	assert(!memcmp(data + sizeof(int), &fds[3], sizeof(int)));
	assert(!(msg.msg_flags & MSG_CTRUNC));
}

TEST(fdbuf_to_cmsg_small_buf)
{
	int fds[] = { 1 };

	union {
		char buf[CMSG_SPACE(2 * sizeof fds)];
		struct cmsghdr align;
	} u = { 0 };

	struct msghdr msg = { 0 };
	msg.msg_control = u.buf;
	msg.msg_controllen = sizeof u.buf;

	assert(fdbuf_to_cmsg(&msg, fds, sizeof(fds) / sizeof(*fds)) == 1);
	assert(!memcmp(CMSG_DATA(CMSG_FIRSTHDR(&msg)), fds, sizeof fds));
	assert(!(msg.msg_flags & MSG_CTRUNC));
}

TEST(msghdr_to_txn_test)
{
	struct virtwl_ioctl_txn *txn;

	int f1[] = { 1 };
	int f2[] = { 2, 3 };

	const int fds[] = { 0, 1, 2 };

	struct iovec iov[] = {
		{ .iov_base = f1, .iov_len = sizeof f1 },
		{ .iov_base = f2, .iov_len = sizeof f2 },
	};

	union {
		char buf[CMSG_SPACE(sizeof fds)];
		struct cmsghdr align;
	} u;

	struct msghdr msg = { 0 };
	msg.msg_iov = iov;
	msg.msg_iovlen = sizeof(iov) / sizeof(*iov);
	msg.msg_control = u.buf;
	msg.msg_controllen = sizeof u.buf;

	struct cmsghdr *cmsg = CMSG_FIRSTHDR(&msg);
	cmsg->cmsg_level = SOL_SOCKET;
	cmsg->cmsg_type = SCM_RIGHTS;
	cmsg->cmsg_len = CMSG_LEN(sizeof fds);
	memcpy(CMSG_DATA(cmsg), fds, sizeof fds);

	assert(!msghdr_to_txn(&msg, &txn));

	assert(txn->len == sizeof(f1) + sizeof(f2));
	assert(!memcmp(txn->fds, fds, sizeof fds));
	for (size_t i = sizeof(fds) / sizeof(*fds);
	     i < VIRTWL_SEND_MAX_ALLOCS; i++)
		assert(txn->fds[i] == -1);
	assert(!memcmp(txn->data, f1, sizeof f1));
	assert(!memcmp(txn->data + sizeof f1, f2, sizeof f2));
}

TEST(txn_to_msghdr_test)
{
	const int data[] = { 1, 2, 3 };

	size_t len = sizeof(data);
	struct virtwl_ioctl_txn *txn = malloc(sizeof(*txn) + len);
	assert(txn);

	txn->len = len;

	const int fds[] = { 0, 1, 2 };
	memcpy(txn->fds, fds, sizeof fds);
	for (int i = sizeof(fds) / sizeof(*fds); i < VIRTWL_SEND_MAX_ALLOCS; i++)
		txn->fds[i] = -1;

	memcpy(txn->data, data, sizeof data);

	int buf1[1], buf2[2];

	struct iovec iov[] = {
		{ .iov_base = buf1, .iov_len = sizeof buf1 },
		{ .iov_base = buf2, .iov_len = sizeof buf2 },
	};

	unsigned char cmsg_buf[VIRTWL_SEND_MAX_ALLOCS];

	struct msghdr msg = { 0 };
	msg.msg_iov = iov;
	msg.msg_iovlen = sizeof(iov) / sizeof(*iov);
	msg.msg_control = cmsg_buf;
	msg.msg_controllen = sizeof cmsg_buf;

	txn_to_msghdr(&msg, txn);

	assert(!memcmp(buf1, data, sizeof buf1));
	assert(!memcmp(buf2, (unsigned char *)data + sizeof buf1, sizeof buf2));

	struct cmsghdr *cmsg = CMSG_FIRSTHDR(&msg);
	assert(cmsg->cmsg_level == SOL_SOCKET);
	assert(cmsg->cmsg_type == SCM_RIGHTS);
	assert(cmsg->cmsg_len == CMSG_LEN(3 * sizeof(int)));
	for (int i = 0; i < 3; i++)
		assert(!memcmp(CMSG_DATA(cmsg), fds, sizeof fds));

	assert(!CMSG_NXTHDR(&msg, cmsg));
}