// SPDX-License-Identifier: GPL-2.0-or-later

/* fuzz-server.c - Host-side test peer for AFL++ fuzzing of passt
 *
 * Accepts TCP connections from passt and sends AFL++-controlled
 * payload read directly from AFL++'s shared memory (region c of
 * the testcase buffer).
 *
 * Started by passt (fork+exec before __AFL_INIT), inherits the
 * __AFL_SHM_FUZZ_ID env var and attaches directly.
 *
 * Runs in a network namespace with AnyIP routing, listening on
 * TCP ports 1 through FUZZ_LISTEN_MAX on both IPv4 and IPv6.
 * The top FUZZ_RESERVED_PORTS ports are left free for passt's
 * own connect() and bind() calls.
 *
 * Build:  make fuzz-server
 * Run:    started automatically by passt when built with FUZZING
 *
 * Copyright Red Hat
 * Author: Anshu Kumari <anskuma@redhat.com>
 */

#ifndef _GNU_SOURCE
#define _GNU_SOURCE
#endif

#include <stdio.h>
#include <stdint.h>
#include <string.h>
#include <unistd.h>
#include <errno.h>
#include <signal.h>
#include <poll.h>
#include <stdlib.h>
#include <sys/shm.h>
#include <sys/socket.h>
#include <sys/epoll.h>
#include <sys/prctl.h>
#include <sys/resource.h>
#include <netinet/in.h>

#include "fuzz-testbuf.h"

#define FUZZ_MAX_PORT		65535
#define FUZZ_RESERVED_PORTS	10000
#define FUZZ_LISTEN_MAX		(FUZZ_MAX_PORT - FUZZ_RESERVED_PORTS)
#define MAX_EVENTS		64

#define TYPE_LISTENER		1
#define TYPE_CONNECTION		2

static uint8_t *afl_shmem;
static int epfd = -1;

/**
 * map_afl_shmem() - Attach to AFL++'s shared memory segment
 *
 * Reads __AFL_SHM_FUZZ_ID env var (inherited when passt fork+execs).
 *
 * Return: 0 on success, -1 on failure
 */
static int map_afl_shmem(void)
{
	const char *env;
	char *endptr;
	void *ptr;
	long id;

	env = getenv("__AFL_SHM_FUZZ_ID");
	if (!env)
		return -1;

	errno = 0;
	id = strtol(env, &endptr, 10);
	if (errno || *endptr || endptr == env)
		return -1;

	ptr = shmat((int)id, NULL, SHM_RDONLY);
	if (ptr == (void *)-1)
		return -1;

	afl_shmem = ptr;
	return 0;
}

/**
 * listen_on() - Create one non-blocking listening socket
 * @af:		Address family, AF_INET or AF_INET6
 * @port:	TCP port to bind
 *
 * Return: listening socket, or -1 if it can't be created or bound
 */
static int listen_on(int af, int port)
{
	struct sockaddr_in6 sa6 = {
		.sin6_family = AF_INET6,
		.sin6_addr = in6addr_any,
		.sin6_port = htons(port),
	};
	struct sockaddr_in sa4 = {
		.sin_family = AF_INET,
		.sin_addr.s_addr = htonl(INADDR_ANY),
		.sin_port = htons(port),
	};
	const struct sockaddr *sa;
	int fd, opt = 1;
	socklen_t sl;

	fd = socket(af, SOCK_STREAM | SOCK_NONBLOCK | SOCK_CLOEXEC, 0);
	if (fd < 0)
		return -1;

	setsockopt(fd, SOL_SOCKET, SO_REUSEADDR, &opt, sizeof(opt));
	setsockopt(fd, SOL_SOCKET, SO_REUSEPORT, &opt, sizeof(opt));

	if (af == AF_INET6) {
		setsockopt(fd, IPPROTO_IPV6, IPV6_V6ONLY, &opt, sizeof(opt));
		sa = (const struct sockaddr *)&sa6;
		sl = sizeof(sa6);
	} else {
		sa = (const struct sockaddr *)&sa4;
		sl = sizeof(sa4);
	}

	if (bind(fd, sa, sl) || listen(fd, 128)) {
		close(fd);
		return -1;
	}

	return fd;
}

/**
 * epoll_add() - Register @fd in the event loop, tagged with @type
 * @fd:		File descriptor to watch
 * @type:	TYPE_LISTENER or TYPE_CONNECTION
 * @events:	epoll event mask
 */
static void epoll_add(int fd, uint32_t type, uint32_t events)
{
	struct epoll_event ev = {
		.events = events,
		.data.u64 = ((uint64_t)type << 32) | (uint32_t)fd,
	};

	epoll_ctl(epfd, EPOLL_CTL_ADD, fd, &ev);
}

/**
 * server_init() - Create TCP listeners on every port, IPv4 and IPv6
 *
 * Binds to 0.0.0.0 and [::] on ports 1 through FUZZ_LISTEN_MAX.
 * Ports that fail to bind are skipped.
 *
 * Return: 0 on success, -1 on failure
 */
static int server_init(void)
{
	struct rlimit rl;
	int port;

	if (getrlimit(RLIMIT_NOFILE, &rl))
		return -1;

	rl.rlim_cur = rl.rlim_max;
	if (setrlimit(RLIMIT_NOFILE, &rl))
		return -1;

	epfd = epoll_create1(EPOLL_CLOEXEC);
	if (epfd < 0)
		return -1;

	for (port = 1; port <= FUZZ_LISTEN_MAX; port++) {
		int fd;

		if ((fd = listen_on(AF_INET, port)) >= 0)
			epoll_add(fd, TYPE_LISTENER, EPOLLIN);

		if ((fd = listen_on(AF_INET6, port)) >= 0)
			epoll_add(fd, TYPE_LISTENER, EPOLLIN);
	}

	return 0;
}

/**
 * send_fuzz_payload() - Send region (c) from AFL++ shared memory
 * @fd:		Connected non-blocking socket
 *
 * Reads the testcase length and event count from AFL++'s shared
 * memory, computes the region (c) offset and length, and sends
 * the payload to @fd, tolerating short writes.
 */
static void send_fuzz_payload(int fd)
{
	const uint8_t *testcase, *data;
	struct fuzz_layout fl;
	uint32_t total_len;
	size_t off = 0;
	size_t len;

	if (!afl_shmem)
		return;

	memcpy(&total_len, afl_shmem, sizeof(total_len));
	testcase = afl_shmem + sizeof(uint32_t);
	fl = fuzz_parse_layout(testcase, total_len);

	if (!fl.testbuf_len)
		return;

	data = testcase + fl.testbuf_off;
	len = (size_t)fl.testbuf_len;

	while (off < len) {
		ssize_t n = send(fd, data + off, len - off, MSG_NOSIGNAL);

		if (n > 0) {
			off += (size_t)n;
			continue;
		}

		if (n < 0 && (errno == EAGAIN || errno == EWOULDBLOCK)) {
			struct pollfd pfd = { .fd = fd, .events = POLLOUT };

			if (poll(&pfd, 1, 50) <= 0)
				return;
			continue;
		}

		if (n < 0 && errno == EINTR)
			continue;

		return;
	}
}

/**
 * handle_accept() - Accept connections and send test payload
 * @lfd:	Listening socket file descriptor
 *
 * For each accepted connection, sends the current region (c) data
 * from AFL++'s shared memory and registers the connection for
 * further I/O events.
 */
static void handle_accept(int lfd)
{
	int fd;

	while ((fd = accept4(lfd, NULL, NULL,
			     SOCK_NONBLOCK | SOCK_CLOEXEC)) >= 0) {
		send_fuzz_payload(fd);
		epoll_add(fd, TYPE_CONNECTION,
			  EPOLLIN | EPOLLRDHUP | EPOLLHUP | EPOLLERR);
	}
}

/**
 * handle_data() - Read data from passt and reply with fuzz payload
 * @fd:		Connected TCP socket
 */
static void handle_data(int fd)
{
	uint8_t buf[4096];
	ssize_t n;

	n = read(fd, buf, sizeof(buf));
	if (n <= 0) {
		epoll_ctl(epfd, EPOLL_CTL_DEL, fd, NULL);
		close(fd);
		return;
	}

	send_fuzz_payload(fd);
}

/**
 * main() - Server entry point
 *
 * Attaches to AFL++ shared memory via env var,
 * binds all ports, then enters the epoll loop. Dies automatically
 * when the parent (passt forkserver) exits.
 *
 * Return: 0 on success, 1 on initialization failure
 */
int main(void)
{
	struct epoll_event events[MAX_EVENTS];
	int nfds, i;

	signal(SIGPIPE, SIG_IGN);
	prctl(PR_SET_PDEATHSIG, SIGTERM);

	if (map_afl_shmem() < 0) {
		(void)fprintf(stderr,
			"fuzz-server: __AFL_SHM_FUZZ_ID not set, "
			"running without AFL++ shared memory\n");
	} else {
		(void)fprintf(stderr,
			"fuzz-server: attached to AFL++ shared memory\n");
	}

	if (server_init() < 0) {
		perror("fuzz-server: server_init");
		return 1;
	}

	(void)fprintf(stderr,
		"fuzz-server: listening on ports 1-%d, "
		"%d-%d reserved for passt, IPv4+IPv6\n",
		FUZZ_LISTEN_MAX, FUZZ_LISTEN_MAX + 1, FUZZ_MAX_PORT);

	while (1) {
		nfds = epoll_wait(epfd, events, MAX_EVENTS, -1);
		if (nfds < 0) {
			if (errno == EINTR)
				continue;
			perror("fuzz-server: epoll_wait");
			return 1;
		}

		for (i = 0; i < nfds; i++) {
			uint32_t type = events[i].data.u64 >> 32;
			int fd = (int)(events[i].data.u64 & 0xFFFFFFFF);

			if (type == TYPE_LISTENER) {
				handle_accept(fd);
			} else if (type == TYPE_CONNECTION) {
				if (events[i].events &
				    (EPOLLHUP | EPOLLERR | EPOLLRDHUP)) {
					epoll_ctl(epfd, EPOLL_CTL_DEL,
						  fd, NULL);
					close(fd);
				} else if (events[i].events & EPOLLIN) {
					handle_data(fd);
				}
			}
		}
	}

	return 0;
}