// SPDX-License-Identifier: GPL-2.0-or-later
// SPDX-FileCopyrightText: 2020-2021 Alyssa Ross <hi@alyssa.is>

#define _GNU_SOURCE

#include <errno.h>
#include <inttypes.h>
#include <poll.h>
#include <stdbool.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <stdnoreturn.h>
#include <sys/wait.h>
#include <sysexits.h>
#include <unistd.h>

#include "env.h"
#include "log.h"
#include "vsock.h"

noreturn static void ex_usage(void)
{
	if (verbosity)
		fprintf(stderr, "Usage: %s [ -1 ] [ -q | -Q | -v ] prog...\n",
		        program_invocation_short_name);
	exit(EX_USAGE);
}

int main(int argc, char *argv[])
{
	bool notify = false;
	int opt;
	pid_t child;
	uint32_t lcid, lport, rcid, rport;
	struct pollfd poll_fd
		= { .fd = STDIN_FILENO, .events = POLLIN, .revents = 0 };

	while ((opt = getopt(argc, argv, "+1qQv")) != -1) {
		switch (opt) {
		case '1':
			notify = true;
			break;
		case 'q':
		case 'Q':
		case 'v':
			set_verbosity(opt);
			break;
		default:
			ex_usage();
		}
	}

	// Check that there is something for us to exec into.
	if (optind > argc - 1)
		ex_usage();

	// Find out the CID and port of the socket.
	if (vsock_get_cid_and_port(STDIN_FILENO, &lcid, &lport)) {
		if (errno == ENOTSOCK || errno == EPROTOTYPE)
			die(EX_USAGE, "stdin is not an AF_VSOCK socket");
		diee(EX_OSERR, "getsockname");
	}

	// Notify the user of our port and indicate readiness.
	if (notify) {
		printf("%" PRIu32 "\n", lport);
		fclose(stdout);
	}

	// Set our CID and port as environment variables so `prog' can
	// know them.
	setenvf("VSOCKLOCALCID", 1, "%" PRIu32, lcid);
	setenvf("VSOCKLOCALPORT", 1, "%" PRIu32, lport);

	ilog("listening as %" PRIu32 " on port %" PRIu32, lcid, lport);

	// Wait for a connection.
	while (poll(&poll_fd, 1, -1) != -1) {
		// On Linux, conn will be blocking.  On other
		// platforms, this may not be the case.  If other
		// platforms are to be supported, we'd probably want
		// to set O_NONBLOCK here.
		int conn = vsock_accept(STDIN_FILENO, &rcid, &rport);
		if (conn == -1)
			diee(EX_OSERR, "accept");

		// Set the client's CID and port as environment
		// variables so `prog' can know them.
		setenvf("VSOCKREMOTECID", 1, "%" PRIu32, rcid);
		setenvf("VSOCKREMOTEPORT", 1, "%" PRIu32, rport);

		ilog("connection from %" PRIu32 " port %" PRIu32, rcid, rport);

		switch (child = fork()) {
		case -1:
			diee(EX_OSERR, "fork");
		case 0:
			// Set up the connection socket on prog's
			// stdin and stdout.  This has the happy side
			// effect of closing the listening socket in
			// the child, since it's also on stdin.
			if (dup2(conn, STDIN_FILENO) == -1)
				diee(EX_OSERR, "dup2");
			if (dup2(conn, STDOUT_FILENO) == -1)
				diee(EX_OSERR, "dup2");
			if (conn != STDIN_FILENO && conn != STDOUT_FILENO)
				close(conn);

			// exec into `prog'.
			execvp(argv[optind], &argv[optind]);
			diee(EX_OSERR, "exec");
		}

		if (waitpid(child, NULL, 0) == -1)
			diee(EX_OSERR, "waitpid");

		close(conn);
	}
	diee(EX_OSERR, "poll");
}