/* SPDX-License-Identifier: MIT */
/*
 * Test case for socket read/write through IORING_OP_READV and
 * IORING_OP_WRITEV, using both TCP and sockets and blocking and
 * non-blocking IO.
 *
 * Heavily based on a test case from Hrvoje Zeba <zeba.hrvoje@gmail.com>
 */
#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>
#include <assert.h>

#include <pthread.h>
#include <errno.h>
#include <fcntl.h>
#include <unistd.h>
#include <sys/socket.h>
#include <sys/un.h>
#include <netinet/tcp.h>
#include <netinet/in.h>

#include "liburing.h"

#define RECV_BUFF_SIZE 2
#define SEND_BUFF_SIZE 3

#define PORT	0x1235

struct params {
	int tcp;
	int non_blocking;
};

pthread_mutex_t mutex = PTHREAD_MUTEX_INITIALIZER;
pthread_cond_t cond = PTHREAD_COND_INITIALIZER;
int rcv_ready = 0;

static void set_rcv_ready(void)
{
	pthread_mutex_lock(&mutex);

	rcv_ready = 1;
	pthread_cond_signal(&cond);

	pthread_mutex_unlock(&mutex);
}

static void wait_for_rcv_ready(void)
{
	pthread_mutex_lock(&mutex);

	while (!rcv_ready)
		pthread_cond_wait(&cond, &mutex);

	pthread_mutex_unlock(&mutex);
}

static void *rcv(void *arg)
{
	struct params *p = arg;
	int s0;

	if (p->tcp) {
		int val = 1;

		s0 = socket(AF_INET, SOCK_STREAM | SOCK_CLOEXEC, IPPROTO_TCP);
		assert(setsockopt(s0, SOL_SOCKET, SO_REUSEPORT, &val, sizeof(val)) != -1);
		assert(setsockopt(s0, SOL_SOCKET, SO_REUSEADDR, &val, sizeof(val)) != -1);

		struct sockaddr_in addr;

		addr.sin_family = AF_INET;
		addr.sin_port = PORT;
		addr.sin_addr.s_addr = 0x0100007fU;
		assert(bind(s0, (struct sockaddr *) &addr, sizeof(addr)) != -1);
	} else {
		s0 = socket(AF_UNIX, SOCK_STREAM | SOCK_CLOEXEC, 0);
		assert(s0 != -1);

		struct sockaddr_un addr;
		memset(&addr, 0, sizeof(addr));

		addr.sun_family = AF_UNIX;
		memcpy(addr.sun_path, "\0sock", 6);
		assert(bind(s0, (struct sockaddr *) &addr, sizeof(addr)) != -1);
	}

	assert(listen(s0, 128) != -1);

	set_rcv_ready();

	int s1 = accept(s0, NULL, NULL);
	assert(s1 != -1);

	if (p->non_blocking) {
		int flags = fcntl(s1, F_GETFL, 0);
		assert(flags != -1);

		flags |= O_NONBLOCK;
		assert(fcntl(s1, F_SETFL, flags) != -1);
	}

	struct io_uring m_io_uring;
	void *ret = NULL;

	assert(io_uring_queue_init(32, &m_io_uring, 0) >= 0);

	int bytes_read = 0;
	int expected_byte = 0;
	int done = 0;

	while (!done && bytes_read != 33) {
		char buff[RECV_BUFF_SIZE];
		struct iovec iov;

		iov.iov_base = buff;
		iov.iov_len = sizeof(buff);

		struct io_uring_sqe *sqe = io_uring_get_sqe(&m_io_uring);
		assert(sqe != NULL);

		io_uring_prep_readv(sqe, s1, &iov, 1, 0);

		assert(io_uring_submit(&m_io_uring) != -1);

		struct io_uring_cqe *cqe;
		unsigned head;
		unsigned count = 0;

		while (!done && count != 1) {
			io_uring_for_each_cqe(&m_io_uring, head, cqe) {
				if (cqe->res < 0)
					assert(cqe->res == -EAGAIN);
				else {
					int i;

					for (i = 0; i < cqe->res; i++) {
						if (buff[i] != expected_byte) {
							fprintf(stderr,
								"Received %d, wanted %d\n",
								buff[i], expected_byte);
							ret++;
							done = 1;
						 }
						 expected_byte++;
					}
					bytes_read += cqe->res;
				}

				count++;
			}

			assert(count <= 1);
			io_uring_cq_advance(&m_io_uring, count);
		}
	}

	shutdown(s1, SHUT_RDWR);
	close(s1);
	close(s0);
	io_uring_queue_exit(&m_io_uring);
	return ret;
}

static void *snd(void *arg)
{
	struct params *p = arg;
	int s0;

	wait_for_rcv_ready();

	if (p->tcp) {
		int val = 1;

		s0 = socket(AF_INET, SOCK_STREAM | SOCK_CLOEXEC, IPPROTO_TCP);
		assert(setsockopt(s0, IPPROTO_TCP, TCP_NODELAY, &val, sizeof(val)) != -1);

		struct sockaddr_in addr;

		addr.sin_family = AF_INET;
		addr.sin_port = PORT;
		addr.sin_addr.s_addr = 0x0100007fU;
		assert(connect(s0, (struct sockaddr*) &addr, sizeof(addr)) != -1);
	} else {
		s0 = socket(AF_UNIX, SOCK_STREAM | SOCK_CLOEXEC, 0);
		assert(s0 != -1);

		struct sockaddr_un addr;
		memset(&addr, 0, sizeof(addr));

		addr.sun_family = AF_UNIX;
		memcpy(addr.sun_path, "\0sock", 6);
		assert(connect(s0, (struct sockaddr*) &addr, sizeof(addr)) != -1);
	}

	if (p->non_blocking) {
		int flags = fcntl(s0, F_GETFL, 0);
		assert(flags != -1);

		flags |= O_NONBLOCK;
		assert(fcntl(s0, F_SETFL, flags) != -1);
	}

	struct io_uring m_io_uring;

	assert(io_uring_queue_init(32, &m_io_uring, 0) >= 0);

	int bytes_written = 0;
	int done = 0;

	while (!done && bytes_written != 33) {
		char buff[SEND_BUFF_SIZE];
		int i;

		for (i = 0; i < SEND_BUFF_SIZE; i++)
			buff[i] = i + bytes_written;

		struct iovec iov;

		iov.iov_base = buff;
		iov.iov_len = sizeof(buff);

		struct io_uring_sqe *sqe = io_uring_get_sqe(&m_io_uring);
		assert(sqe != NULL);

		io_uring_prep_writev(sqe, s0, &iov, 1, 0);

		assert(io_uring_submit(&m_io_uring) != -1);

		struct io_uring_cqe *cqe;
		unsigned head;
		unsigned count = 0;

		while (!done && count != 1) {
			io_uring_for_each_cqe(&m_io_uring, head, cqe) {
				if (cqe->res < 0) {
					if (cqe->res == -EPIPE) {
						done = 1;
						break;
					}
					assert(cqe->res == -EAGAIN);
				} else {
					bytes_written += cqe->res;
				}

				count++;
			}

			assert(count <= 1);
			io_uring_cq_advance(&m_io_uring, count);
		}
		usleep(100000);
	}

	shutdown(s0, SHUT_RDWR);
	close(s0);
	io_uring_queue_exit(&m_io_uring);
	return NULL;
}

int main(int argc, char *argv[])
{
	struct params p;
	pthread_t t1, t2;
	void *res1, *res2;
	int i, exit_val = 0;

	for (i = 0; i < 4; i++) {
		p.tcp = i & 1;
		p.non_blocking = (i & 2) >> 1;

		rcv_ready = 0;

		pthread_create(&t1, NULL, rcv, &p);
		pthread_create(&t2, NULL, snd, &p);
		pthread_join(t1, &res1);
		pthread_join(t2, &res2);
		if (res1 || res2) {
			fprintf(stderr, "Failed tcp=%d, non_blocking=%d\n", p.tcp, p.non_blocking);
			exit_val = 1;
		}
	}

	return exit_val;
}