// SPDX-License-Identifier: GPL-2.0-only

/**
 * This is a reproducer for a regression introduced by:
 *   c82199061009 ("task_work: remove fifo ordering guarantee")
 * and proposed fixed by
 *   ............ ("tcp: reset late connection after listening socket close")
 *
 * This reproducer is written by Asbjørn Sloth Tønnesen <ast@fiberby.net>,
 * and is based on the original Perl-based reproducer written by
 * Kristian Nielsen, and distribued with the original report:
 *   https://lore.kernel.org/87sf0ldk41.fsf@urd.knielsen-hq.org
 *
 * Compile using:
 *   cc -o socket_teardown_test socket_teardown_test.c
 *
 * This reproducer has the following improvements:
 * - C-based
 * - Also tests IPv6
 * - More robust error handling
 */

#define _XOPEN_SOURCE 700
#define _DEFAULT_SOURCE

#include <stdlib.h>
#include <stdio.h>
#include <stdarg.h>
#include <stdbool.h>
#include <stdint.h>
#include <string.h>
#include <errno.h>
#include <fcntl.h>
#include <unistd.h>
#include <inttypes.h>
#include <arpa/inet.h>
#include <sys/select.h>
#include <sys/socket.h>
#include <sys/resource.h>
#include <sys/wait.h>
#include <stdlib.h>
#include <linux/in.h>
#include <time.h>

static struct __kernel_sockaddr_storage my_sockaddr_data;
static struct sockaddr *my_sockaddr = (struct sockaddr *) &my_sockaddr_data;
static size_t my_sockaddr_sz = 0;
static const int max_conns = 1000;

struct conn_profile {
	int min, max;
};

struct test_params {
	const struct conn_profile *conns_p;
	int sleep_ms_before_abort;
	bool skip_on_slow_connect;
	int iterations;
	int family;
};

struct test_context {
	const struct test_params *params;
	struct event_stats *stats;
	int *conns;
	int iter;
	int n;
};

struct event_stats {
	uint32_t iterations;
	uint32_t connections;
	uint32_t connection_failures;
	uint32_t extra_read_timeout_on_ipv4;
	uint32_t extra_read_timeout_on_ipv6;
	uint32_t extra_read_closed;
	uint32_t extra_connect_reset;
	uint32_t extra_connect_refused;
	uint32_t extra_connect_broken_pipe;
	uint32_t extra_connect_aborted;
	uint32_t extra_read_not_connected;
	uint32_t extra_read_reset;
	uint32_t extra_read_data;
	uint32_t listener_exited;
	uint32_t listener_dead_by_signal;
};

struct run_params {
	bool stop_after_first_bug;
	void (*run)(const struct run_params *rp, struct event_stats *stats);
};

static void print_stats(const struct event_stats *stats)
{
#define print_stat(key, extra) { \
	uint32_t val = stats->key; \
	if (val) \
		printf("  %-32s %6" PRIu32 "%s\n", \
		       #key ":", val, (val > 0 ? extra : "")); \
}
	printf("\nStats:\n");
	print_stat(iterations, "");
	print_stat(connections, "");
	print_stat(connection_failures, "");
	print_stat(extra_read_timeout_on_ipv4, " !!!");
	print_stat(extra_read_timeout_on_ipv6, " !!!");
	print_stat(extra_read_closed, "");
	print_stat(extra_connect_reset, "");
	print_stat(extra_connect_refused, "");
	print_stat(extra_connect_broken_pipe, "");
	print_stat(extra_connect_aborted, "");
	print_stat(extra_read_not_connected, "");
	print_stat(extra_read_reset, "");
	print_stat(extra_read_data, "");
	print_stat(listener_exited, "");
	print_stat(listener_dead_by_signal, "");
#undef print_stat
}

static void bail(const char *fmt, ...)
{
	va_list ap;

	va_start(ap, fmt);
	vfprintf(stderr, fmt, ap);
	va_end(ap);
	exit(EXIT_FAILURE);
}

#define bail_fn(fmt, ...) (bail("%s: " fmt "\n", __func__, __VA_ARGS__))

static void prep_ipv4_sockaddr(void)
{
	struct sockaddr_in *sa = (struct sockaddr_in *) my_sockaddr;
	size_t sz = sizeof(*sa);

	memset(sa, 0, sz);
	sa->sin_family = AF_INET;
	sa->sin_addr.s_addr = htonl((127 << 24) + 1); // 127.0.0.1
	sa->sin_port = htons(2345);
	my_sockaddr_sz = sz;
}

static void prep_ipv6_sockaddr(void)
{
	struct sockaddr_in6 *sa = (struct sockaddr_in6 *) my_sockaddr;
	size_t sz = sizeof(*sa);

	memset(sa, 0, sz);
	sa->sin6_family = AF_INET6;
	sa->sin6_addr.s6_addr[15] = 1; // ::1
	sa->sin6_port = htons(2345);
	my_sockaddr_sz = sz;
}

static void prep_sockaddr(const int family)
{
	switch (family) {
	case AF_INET:
		prep_ipv4_sockaddr();
		break;
	case AF_INET6:
		prep_ipv6_sockaddr();
		break;
	default:
		bail_fn("unknown family %d", family);
	}
}

static int setup_server_socket(const int test_fd)
{
	int sfd, ret;
	int one = 1;

	sfd = socket(my_sockaddr->sa_family, SOCK_STREAM | SOCK_CLOEXEC, 0);
	if (sfd == -1)
		bail_fn("select: %s", strerror(errno));

	ret = setsockopt(sfd, SOL_SOCKET, SO_REUSEADDR, &one, sizeof(one));
	if (ret)
		bail_fn("setsockopt: %s", strerror(errno));

	ret = bind(sfd, my_sockaddr, my_sockaddr_sz);
	if (ret)
		bail_fn("bind: %s", strerror(errno));

	ret = listen(sfd, SOMAXCONN);
	if (ret)
		bail_fn("listen: %s", strerror(errno));

	/* Signal our parent, that we are ready for connections */
	ret = write(test_fd, ".", 1);
	if (ret == -1)
		bail_fn("write: %s", strerror(errno));
	close(test_fd);

	return sfd;
}

static void run_server(struct test_context *ctx, const int test_fd)
{
	int n = ctx->n;
	int sfd, i, fd;
	int sleep_ms = ctx->params->sleep_ms_before_abort;

	sfd = setup_server_socket(test_fd);

	/* Accept n connections */
	for (i=0;i<n;i++) {
		fd = accept(sfd, NULL, NULL);
		if (fd == -1)
			bail_fn("accept: %s", strerror(errno));
		write(fd, "!", 1);
		/* Leak fd, as we abort abnormally before closing it */
	}

	/* Maybe wait some time */
	if (sleep_ms > 0)
		usleep(sleep_ms * 1e3);

	/* Something happended and we didn't do close(sfd) */
	abort();
}

static int spawn_server(struct test_context *ctx, pid_t *ret_pid)
{
	int fds[2], ret;
	pid_t pid;

	ret = pipe(fds);
	if (ret)
		bail_fn("pipe: %s", strerror(errno));

	pid = fork();
	if (pid == -1)
		bail_fn("fork: %s", strerror(errno));

	if (pid == 0) {
		close(fds[0]);
		run_server(ctx, fds[1]);
		exit(EXIT_SUCCESS);
	} else {
		close(fds[1]);
		*ret_pid = pid;
		return fds[0];
	}
}

static int connect_client_socket(void)
{
	int fd;

	fd = socket(my_sockaddr->sa_family, SOCK_STREAM | SOCK_CLOEXEC, 0);
	if (fd == -1)
		bail_fn("socket: %s", strerror(errno));

	return fd;
}

static int connect_client(void)
{
	int fd, ret;
	char buf[1];

	fd = connect_client_socket();

	ret = connect(fd, my_sockaddr, my_sockaddr_sz);
	if (ret) {
		if (errno == ECONNREFUSED)
			goto close_and_return;
		bail_fn("connect: %s", strerror(errno));
	}

	ret = read(fd, &buf, 1);
	if (ret == -1) {
		if (errno == ECONNRESET)
			goto close_and_return;
		bail_fn("read: %s", strerror(errno));
	}

	return fd;
close_and_return:
	close(fd);
	return -1;
}

static void wait_for_server_startup(const int child_fd)
{
	char buf[1];
	int ret;

	ret = read(child_fd, &buf, 1);
	if (ret == -1)
		bail_fn("read: %s", strerror(errno));
	close(child_fd);
}

static void establish_n_clients(struct test_context *ctx)
{
	struct event_stats *stats = ctx->stats;
	int *conns = ctx->conns;
	int n = ctx->n;
	int i;

	for (i = 0; i < n; i++)
		if ((conns[i] = connect_client()) == -1)
			stats->connection_failures++;
}

static void detect_server_teardown(struct test_context *ctx)
{
	int *conns = ctx->conns;
	struct timeval tv;
	fd_set rfds;
	int high_fd = 0;
	int n = ctx->n;
	int ret, i, fd;

	FD_ZERO(&rfds);
	tv.tv_sec = 5;
	tv.tv_usec = 0;

	for (i = 0; i < n; i++) {
		fd = conns[i];
		if (fd > 0) {
			FD_SET(fd, &rfds);
			if (fd > high_fd)
				high_fd = fd;
		}
	}

	ret = select(high_fd + 1, &rfds, NULL, NULL, &tv);
	if (ret == -1)
		bail_fn("select: %s", strerror(errno));
	if (!ret)
		bail_fn("select: failed to detect teardown (n=%d)", n);
}

static void set_blocking(const int fd, const bool blocking)
{
	int flags, new_flags;

	if ((flags = fcntl(fd, F_GETFL)) == -1)
		bail_fn("F_GETFL: %s", strerror(errno));

	if (blocking)
		new_flags = flags & (~O_NONBLOCK);
	else
		new_flags = flags | O_NONBLOCK;

	if (new_flags != flags && fcntl(fd, F_SETFL, new_flags) == -1)
		bail_fn("F_SETFL: %s", strerror(errno));
}

static bool connect_extra_client_sync(struct test_context *ctx, const int extra_fd)
{
	struct event_stats *stats = ctx->stats;
	int ret;

	ret = connect(extra_fd, my_sockaddr, my_sockaddr_sz);
	if (ret) {
		switch (errno) {
		case ECONNREFUSED:
			stats->extra_connect_refused++;
			break;
		case ECONNRESET:
			stats->extra_connect_reset++;
			break;
		default:
			bail_fn("connect: %s", strerror(errno));
		}
		return false;
	}
	return true;
}

static bool connect_extra_client_async(struct test_context *ctx, const int extra_fd)
{
	struct event_stats *stats = ctx->stats;
	struct timeval tv;
	int ret = false;
	fd_set wfds;

	while (true) {
do_connect:
		ret = connect(extra_fd, my_sockaddr, my_sockaddr_sz);
		if (!ret) {
			ret = true;
			goto do_return;
		}
		switch (errno) {
		case ECONNREFUSED:
			stats->extra_connect_refused++;
			break;
		case ECONNRESET:
			stats->extra_connect_reset++;
			break;
		case EPIPE:
			stats->extra_connect_broken_pipe++;
			break;
		case EINPROGRESS:
			goto do_select;
		default:
			bail_fn("connect: %s", strerror(errno));
		}
		goto do_return;
	}
do_select:
	FD_ZERO(&wfds);
	FD_SET(extra_fd, &wfds);
	tv.tv_sec = 0;
	tv.tv_usec = 2500;

	ret = select(extra_fd + 1, NULL, &wfds, NULL, &tv);
	if (ret == -1)
		bail_fn("select: %s", strerror(errno));
	if (ret) {
		goto do_connect;
	} else {
		/* timeout */
		stats->extra_connect_aborted++;
		goto do_return;
	}
do_return:
	set_blocking(extra_fd, true);
	return ret;
}

static bool connect_extra_client(struct test_context *ctx, const int extra_fd)
{
	if (ctx->params->skip_on_slow_connect)
		return connect_extra_client_async(ctx, extra_fd);
	else
		return connect_extra_client_sync(ctx, extra_fd);
}

static void try_read_from_extra_connection(struct test_context *ctx,
					   const int extra_fd)
{
	struct event_stats *stats = ctx->stats;
	struct timeval tv;
	fd_set rfds;
	char buf[1];
	int ret;

	FD_ZERO(&rfds);
	FD_SET(extra_fd, &rfds);
	tv.tv_sec = 2;
	tv.tv_usec = 0;

	/* Wait for extra_fd to get ECONNRESET */
	ret = select(extra_fd + 1, &rfds, NULL, NULL, &tv);
	if (ret == -1)
		bail_fn("select: %s", strerror(errno));

	if (ret) {
		if (!FD_ISSET(extra_fd, &rfds))
			bail_fn("extra_fd not set (fd = %d)", extra_fd);

		ret = read(extra_fd, &buf, 1);
		if (ret == -1) {
			switch (errno) {
			case ECONNRESET:
				stats->extra_read_reset++;
				break;
			case ENOTCONN:
				stats->extra_read_not_connected++;
				break;
			default:
				bail_fn("read: %s", strerror(errno));
			}
		} else if (ret == 0) {
			stats->extra_read_closed++;
		} else {
			stats->extra_read_data++;
		}
	} else {
		printf("BUG: read timeout on extra fd (n=%d, i=%d, fd=%d)\n", ctx->n, ctx->iter, extra_fd);
		switch (ctx->params->family) {
		case AF_INET:
			stats->extra_read_timeout_on_ipv4++;
			break;
		case AF_INET6:
			stats->extra_read_timeout_on_ipv6++;
			break;
		default:
			bail_fn("unknown family %d", ctx->params->family);
		}
	}
}

static void wait_for_server_exit(struct test_context *ctx, const pid_t child_pid)
{
	struct event_stats *stats = ctx->stats;
	int child_status;
	int ret;

	ret = waitpid(child_pid, &child_status, 0);
	if (ret == -1)
		bail_fn("waitpid: %s", strerror(errno));

	if (WIFEXITED(child_status))
		stats->listener_exited++;
	else if (WIFSIGNALED(child_status))
		stats->listener_dead_by_signal++;
}

static void cleanup_clients(struct test_context *ctx, const int extra_fd)
{
	struct event_stats *stats = ctx->stats;
	int *conns = ctx->conns;
	int n = ctx->n;
	int i, fd;

	for (i=0;i<n;i++) {
		fd = conns[i];
		if (fd > 0) {
			close(fd);
			stats->connections++;
		}
	}
	if (extra_fd > 0)
		close(extra_fd);
	stats->iterations++;
}

/* This is the main control flow for each test run */
static void run_test_iteration(const struct test_params *params,
			       const int iter, const int n_conns,
			       struct event_stats *stats)
{
	int conns[max_conns];
	struct test_context ctx = {
		.params = params,
		.n = n_conns,
		.iter = iter,
		.stats = stats,
		.conns = conns,
	};
	int child_fd, extra_fd;
	pid_t child_pid;

	if (n_conns >= max_conns)
		bail_fn("too many connections (%d > %d)", n_conns, max_conns);

	prep_sockaddr(params->family);

	/* Start server process */
	child_fd = spawn_server(&ctx, &child_pid);

	/* Wait for the server process to be listening */
	wait_for_server_startup(child_fd);

	/* Establish <n> client connections */
	establish_n_clients(&ctx);

	/* Open extra socket, use to simulate a client with bad timing */
	extra_fd = connect_client_socket();
	if (params->skip_on_slow_connect)
		set_blocking(extra_fd, false);

	/* Wait for first connection to be closed by server teardown */
	detect_server_teardown(&ctx);

	/* Attempt to connect the extra connection */
	if (connect_extra_client(&ctx, extra_fd))
		try_read_from_extra_connection(&ctx, extra_fd);

	/* Wait for server to exit */
	wait_for_server_exit(&ctx, child_pid);

	/* Cleanup file descriptors */
	cleanup_clients(&ctx, extra_fd);
}

static bool has_hit_bug(const struct event_stats *stats)
{
	int found = stats->extra_read_timeout_on_ipv4 +
		    stats->extra_read_timeout_on_ipv6;
	return found > 0;
}

static bool reached_goal(const struct run_params *rp,
			 const struct event_stats *stats)
{
	return rp->stop_after_first_bug && has_hit_bug(stats);
}


static void run_tests(const struct run_params *rp,
		      const struct test_params *params,
		      struct event_stats *stats)
{
	const struct conn_profile *cp = params->conns_p;
	const int conns_range = cp->max  - cp->min + 1;
	const int conns_offset = cp->min;
	int i, n;

	for (i = 0; i < params->iterations; i++) {
		n = (i % conns_range) + conns_offset;
		run_test_iteration(params, i, n, stats);

		if (reached_goal(rp, stats))
			break;
	}
}

/* Define connection profiles */
static const struct conn_profile conn_profiles[] = {
	{
		.min = 5,
		.max = 9,
	},
	{
		.min = 1,
		.max = 9,
	},
	{
		.min = 1,
		.max = 100,
	},
};
#define conn_profile_dflt (&conn_profiles[0])
#define conn_profiles_n \
	(sizeof(conn_profiles) / sizeof(conn_profiles[0]))

static const struct test_params full_test_series[] = {
	{
		.skip_on_slow_connect = true,
		.family = AF_INET,
		.sleep_ms_before_abort = 5,
		.iterations = 1000,
		.conns_p = conn_profile_dflt,
	},
	{
		.skip_on_slow_connect = true,
		.family = AF_INET6,
		.sleep_ms_before_abort = 5,
		.iterations = 1000,
		.conns_p = conn_profile_dflt,
	},
	{
		.skip_on_slow_connect = true,
		.family = AF_INET,
		.sleep_ms_before_abort = 0,
		.iterations = 1000,
		.conns_p = conn_profile_dflt,
	},
	{
		.skip_on_slow_connect = true,
		.family = AF_INET6,
		.sleep_ms_before_abort = 0,
		.iterations = 1000,
		.conns_p = conn_profile_dflt,
	},
	{
		.skip_on_slow_connect = false,
		.family = AF_INET,
		.sleep_ms_before_abort = 5,
		.iterations = 1000,
		.conns_p = conn_profile_dflt,
	},
	{
		.skip_on_slow_connect = false,
		.family = AF_INET6,
		.sleep_ms_before_abort = 5,
		.iterations = 1000,
		.conns_p = conn_profile_dflt,
	},
};
#define full_test_series_n \
	(sizeof(full_test_series) / sizeof(full_test_series[0]))

/* Run all tests with provided connection profile */
static void run_conn_profile(const struct run_params *rp,
			     const struct conn_profile *cp,
			     struct event_stats *stats)
{
	struct test_params params;
	int n = full_test_series_n;
	int i;

	for (i = 0; i < n; i++) {
		memcpy(&params, &full_test_series[i], sizeof(params));
		params.conns_p = cp;
		run_tests(rp, &params, stats);

		if (reached_goal(rp, stats))
			break;
	}
}

/* Short test: only use the first connection profile */
static void run_short_test(const struct run_params *rp, struct event_stats *stats)
{
	run_conn_profile(rp, conn_profile_dflt, stats);
}

/* Long test: run through all connection profiles */
static void run_long_test(const struct run_params *rp, struct event_stats *stats)
{
	int n = conn_profiles_n;
	int i;

	for (i = 0; i < n; i++) {
		run_conn_profile(rp, &conn_profiles[i], stats);
	}
}

static void print_usage(FILE *fp, const char *prog_name)
{
	fprintf(fp, "Usage: %s <keywords...>\n\n", prog_name);
	fprintf(fp, "This is a reproducer for a regression introduced by:\n");
	fprintf(fp, "  c82199061009 (\"task_work: remove fifo ordering guarantee\")\n");
	fprintf(fp, "and proposed fixed by\n");
	fprintf(fp, "  ............ (\"tcp: reset late connection after listening socket close\")\n\n");
	fprintf(fp, "Keywords:\n");
	fprintf(fp, "  short        Request the short test.\n");
	fprintf(fp, "  long         Request the long test.\n");
	fprintf(fp, "  stop-on-hit  Stop when the first bug is found.\n");
	fprintf(fp, "  quick        Run short test, but stop on hit.\n");
	fprintf(fp, "\nThe test is run, after processing all keywords\n");
}

/* Process command line keywords, modifying test parameters */
static void configure_test_run(struct run_params *rp,
			       const int argc, const char **argv)
{
	int i;

	for (i = 1; i < argc; i++) {
		const char *arg = argv[i];

		if (strcmp(arg, "help") == 0) {
			print_usage(stdout, argv[0]);
			exit(EXIT_SUCCESS);
		} else if (strcmp(arg, "short") == 0) {
			rp->run = run_short_test;
		} else if (strcmp(arg, "long") == 0) {
			rp->run = run_long_test;
		} else if (strcmp(arg, "stop-on-hit") == 0) {
			rp->stop_after_first_bug = true;
		} else if (strcmp(arg, "quick") == 0) {
			rp->run = run_short_test;
			rp->stop_after_first_bug = true;
		} else {
			bail("Unknown argument: %s\n", arg);
		}
	}
}

int main(const int argc, const char **argv)
{
	if (argc < 2) {
		print_usage(stderr, argv[0]);
		exit(EXIT_FAILURE);
	}

	struct run_params rp = {
		.run = run_short_test,
	};

	configure_test_run(&rp, argc, argv);

	struct event_stats stats = {0,};
	rp.run(&rp, &stats);
	print_stats(&stats);

	if (has_hit_bug(&stats)) {
		return EXIT_FAILURE;
	} else {
		printf("\nSuccess, the tests failed to trigger the bug!!!\n");
		return EXIT_SUCCESS;
	}
}