/*-
 * SPDX-License-Identifier: BSD-2-Clause
 *
 * Copyright (c) 2023 Igor Ostapenko <pm@igoro.pro>
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 * 1. Redistributions of source code must retain the above copyright
 *    notice, this list of conditions and the following disclaimer.
 * 2. Redistributions in binary form must reproduce the above copyright
 *    notice, this list of conditions and the following disclaimer in the
 *    documentation and/or other materials provided with the distribution.
 *
 * THIS SOFTWARE IS PROVIDED BY THE PROJECT AND CONTRIBUTORS ``AS IS'' AND
 * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
 * ARE DISCLAIMED.  IN NO EVENT SHALL THE PROJECT OR CONTRIBUTORS BE LIABLE
 * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
 * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS
 * OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
 * HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
 * LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY
 * OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
 * SUCH DAMAGE.
 */

/* Used by divert(4) related tests */

#include <errno.h>
#include <stdlib.h>
#include <stdbool.h>
#include <err.h>
#include <sysexits.h>
#include <string.h>

#include <sys/types.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <netinet/ip.h>


struct context {
	unsigned short divert_port;
	bool divert_back;

	int fd;
	struct sockaddr_in sin;
	socklen_t sin_len;
	char pkt[IP_MAXPACKET];
	ssize_t pkt_n;
};

static void
init(struct context *c)
{
	c->fd = socket(PF_DIVERT, SOCK_RAW, 0);
	if (c->fd == -1)
		errx(EX_OSERR, "init: Cannot create divert socket.");

	memset(&c->sin, 0, sizeof(c->sin));
	c->sin.sin_family = AF_INET;
	c->sin.sin_port = htons(c->divert_port);
	c->sin.sin_addr.s_addr = INADDR_ANY;
	c->sin_len = sizeof(struct sockaddr_in);

	if (bind(c->fd, (struct sockaddr *) &c->sin, c->sin_len) != 0)
		errx(EX_OSERR, "init: Cannot bind divert socket.");
}

static ssize_t
recv_pkt(struct context *c)
{
	fd_set readfds;
	struct timeval timeout;
	int s;

	FD_ZERO(&readfds);
	FD_SET(c->fd, &readfds);
	timeout.tv_sec = 3;
	timeout.tv_usec = 0;

	s = select(c->fd + 1, &readfds, 0, 0, &timeout);
	if (s == -1)
		errx(EX_IOERR, "recv_pkt: select() errors.");
	if (s != 1) /* timeout */
		return (-1);

	c->pkt_n = recvfrom(c->fd, c->pkt, sizeof(c->pkt), 0,
	    (struct sockaddr *) &c->sin, &c->sin_len);
	if (c->pkt_n == -1)
		errx(EX_IOERR, "recv_pkt: recvfrom() errors.");

	return (c->pkt_n);
}

static void
send_pkt(struct context *c)
{
	ssize_t n;

	n = sendto(c->fd, c->pkt, c->pkt_n, 0,
	    (struct sockaddr *) &c->sin, c->sin_len);
	if (n == -1)
		err(EX_IOERR, "send_pkt: sendto() errors");
	if (n != c->pkt_n)
		errx(EX_IOERR, "send_pkt: sendto() sent %zd of %zd bytes.",
		    n, c->pkt_n);
}

int
main(int argc, char *argv[])
{
	struct context c;
	int npkt;

	if (argc < 2)
		errx(EX_USAGE,
		    "Usage: %s <divert-port> [divert-back]", argv[0]);

	memset(&c, 0, sizeof(struct context));

	c.divert_port = (unsigned short) strtol(argv[1], NULL, 10);
	if (c.divert_port == 0)
		errx(EX_USAGE, "divert port is not defined.");

	if (argc >= 3 && strcmp(argv[2], "divert-back") == 0)
		c.divert_back = true;


	init(&c);

	npkt = 0;
	while (recv_pkt(&c) > 0) {
		if (c.divert_back)
			send_pkt(&c);
		npkt++;
		if (npkt >= 20)
			break;
	}

	if (npkt != 1)
		errx(EXIT_FAILURE, "%d: npkt=%d.", c.divert_port, npkt);

	return (EXIT_SUCCESS);
}