osmo-upf/contrib/udp_responder.c

/* UDP responder: listen on a UDP port, and respond to each received UDP packet back to the sender. */
/*
 * (C) 2024 by sysmocom - s.f.m.c. GmbH <info@sysmocom.de>
 * All Rights Reserved.
 *
 * Author: Neels Janosch Hofmeyr <nhofmeyr@sysmocom.de>
 *
 * SPDX-License-Identifier: GPL-2.0+
 *
 *  This program is free software; you can redistribute it and/or modify
 *  it under the terms of the GNU General Public License as published by
 *  the Free Software Foundation; either version 2 of the License, or
 *  (at your option) any later version.
 *
 *  This program is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *  GNU General Public License for more details.
 *
 *  You should have received a copy of the GNU General Public License
 *  along with this program.  If not, see <http://www.gnu.org/licenses/>.
 *
 */
#include "config.h"

#include <osmocom/core/application.h>
#include <osmocom/core/logging.h>
#include <osmocom/core/timer.h>

#if HAVE_URING

#define _GNU_SOURCE
#include <getopt.h>
#include <limits.h>

#include <liburing.h>

#include <osmocom/core/utils.h>
#include <osmocom/core/socket.h>
#include <osmocom/core/sockaddr_str.h>

struct cmdline_cmd {
	const char *short_option;
	const char *long_option;
	const char *arg_name;
	const char *doc;

	const char *value;
};

#define cmdline_foreach(ITER, CMDS) \
	for (const struct cmdline_cmd *ITER = (CMDS); \
	     ITER->short_option || ITER->long_option || ITER->arg_name; \
	     ITER++)

int cmdline_doc_str_buf(char *buf, size_t buflen, const struct cmdline_cmd *cmds)
{
	struct osmo_strbuf sb = { .buf = buf, .len = buflen };
	/* First find the longest options part */
	int w = 0;
	cmdline_foreach (cmd, cmds) {
		int cmd_w = 0;
		if (cmd->short_option)
			cmd_w += 2 + strlen(cmd->short_option);
		if (cmd->long_option)
			cmd_w += 3 + strlen(cmd->long_option);
		if (cmd->arg_name)
			cmd_w += 1 + strlen(cmd->arg_name);
		w = OSMO_MAX(w, cmd_w);
	}
	/* vertical gap */
	w += 2;

	OSMO_STRBUF_PRINTF(sb, "Options:\n");
	cmdline_foreach (cmd, cmds) {
		char *line_start = sb.pos;
		if (cmd->short_option)
			OSMO_STRBUF_PRINTF(sb, " -%s", cmd->short_option);
		if (cmd->long_option)
			OSMO_STRBUF_PRINTF(sb, " --%s", cmd->long_option);
		if (cmd->arg_name)
			OSMO_STRBUF_PRINTF(sb, " %s", cmd->arg_name);
		if (cmd->doc) {
			int have = sb.pos - line_start;
			int spaces = OSMO_MAX(1, w - have);
			OSMO_STRBUF_PRINTF(sb, "%*s", spaces, "");
			OSMO_STRBUF_PRINTF(sb, "%s", cmd->doc);
		}
		OSMO_STRBUF_PRINTF(sb, "\n");
	}
	return sb.chars_needed;
}

void cmdline_print_help(const struct cmdline_cmd *cmds)
{
	char buf[8192];
	cmdline_doc_str_buf(buf, sizeof(buf), cmds);
	printf("%s", buf);
}

void cmdline_cmd_store_optarg(struct cmdline_cmd *cmd)
{
	if (cmd->arg_name)
		cmd->value = optarg;
	else
		cmd->value = (cmd->short_option ? : cmd->long_option);
}

int cmdline_read(struct cmdline_cmd *cmds, int argc, char **argv)
{
	char short_options[256] = {};
	struct option long_options[128] = {};
	int long_options_i = 0;
	int long_option_val = 0;
	struct osmo_strbuf short_sb = { .buf = short_options, .len = sizeof(short_options) };

	cmdline_foreach (cmd, cmds) {
		if (cmd->short_option) {
			OSMO_STRBUF_PRINTF(short_sb, "%s", cmd->short_option);
			if (cmd->arg_name)
				OSMO_STRBUF_PRINTF(short_sb, ":");
		}
		if (cmd->long_option) {
			long_options[long_options_i] = (struct option){
				cmd->long_option,
				cmd->arg_name ? 1 : 0,
				&long_option_val,
				long_options_i,
			};
		}
	}

	while (1) {
		int option_index = 0;
		char c = getopt_long(argc, argv, short_options, long_options, &option_index);
		if (c == -1)
			break;
		if (c == 0) {
			struct cmdline_cmd *long_cmd = &cmds[long_option_val];
			cmdline_cmd_store_optarg(long_cmd);
		} else {
			bool found = false;
			cmdline_foreach (cc, cmds) {
				if (strchr(cc->short_option, c)) {
					cmdline_cmd_store_optarg((struct cmdline_cmd *)cc);
					found = true;
					break;
				}
			}
			if (!found) {
				fprintf(stderr, "%s: Error in command line options. Exiting.\n", argv[0]);
				return -1;
			}
		}
	}

	/* positional args */
	cmdline_foreach (cmd, cmds) {
		if (optind >= argc)
			break;
		if (cmd->short_option || cmd->long_option)
			continue;
		if (!cmd->arg_name)
			continue;
		((struct cmdline_cmd *)cmd)->value = argv[optind];
		optind++;
	}

	if (optind < argc) {
		cmdline_print_help(cmds);
		fprintf(stderr, "%s: Unsupported positional argument on command line\n", argv[optind]);
		return -1;
	}
	return 0;
}

const char *cmdline_get(const struct cmdline_cmd *cmds, const char *option_name, const char *default_val)
{
	cmdline_foreach (cmd, cmds) {
		if (cmd->long_option && !strcmp(cmd->long_option, option_name))
			return cmd->value;
		if (cmd->short_option && !strcmp(cmd->short_option, option_name))
			return cmd->value;
		if (cmd->arg_name && !strcmp(cmd->arg_name, option_name))
			return cmd->value;
	}
	return default_val;
}

bool cmdline_get_int(int *dst, int minval, int maxval, int default_val,
		     const struct cmdline_cmd *cmds, const char *option_name)
{
	const char *str = cmdline_get(cmds, option_name, NULL);
	if (!str) {
		*dst = default_val;
		return true;
	}
	if (osmo_str_to_int(dst, str, 10, minval, maxval)) {
		cmdline_print_help(cmds);
		printf("ERROR: invalid integer number: %s\n", str);
		return false;
	}
	if (*dst < minval || *dst > maxval) {
		cmdline_print_help(cmds);
		printf("ERROR: number out of range: %d <= %d <= %d\n", minval, *dst, maxval);
		return false;
	}
	return true;
}

struct udp_port {
	struct llist_head entry;

	/* IP address and UDP port from user input */
	struct osmo_sockaddr osa;

	/* locally bound socket */
	int fd;
};

enum data_io_type {
	IO_UNUSED = 0,
	IO_RECV,
	IO_SEND,
};

struct data_io {
	enum data_io_type type;
	struct osmo_sockaddr osa;
	struct iovec iov;
	struct msghdr msgh;
	uint8_t *data;
	size_t data_size;
	int n;
};

struct io_queue {
	size_t d_size;
	struct data_io d[0];
};

static void data_io_prep_recv(struct io_uring *ring, struct udp_port *port, struct data_io *d)
{
	struct io_uring_sqe *sqe;
	*d = (struct data_io){
		.type = IO_RECV,
		.iov = {
			.iov_base = d->data,
			.iov_len = d->data_size,
		},
		.msgh = {
			.msg_name = &d->osa,
			.msg_namelen = sizeof(d->osa),
			.msg_iov = &d->iov,
			.msg_iovlen = 1,
		},
		.data_size = d->data_size,
		.data = d->data,
	};

	sqe = io_uring_get_sqe(ring);
	OSMO_ASSERT(sqe);
	io_uring_prep_recvmsg(sqe, port->fd, &d->msgh, 0);
	io_uring_sqe_set_data(sqe, d);
}

static void data_io_prep_send(struct io_uring *ring, struct udp_port *port, struct data_io *d)
{
	struct io_uring_sqe *sqe;
	d->type = IO_SEND;
	sqe = io_uring_get_sqe(ring);
	OSMO_ASSERT(sqe);
	io_uring_prep_sendmsg(sqe, port->fd, &d->msgh, 0);
	io_uring_sqe_set_data(sqe, d);
}

uint32_t g_total_rx = 0;
uint32_t g_total_tx = 0;

static void data_io_handle_completion(struct io_uring *ring, struct udp_port *port, struct io_uring_cqe *cqe,
				      int response_size, int response_n)
{
	struct data_io *d;
	struct osmo_sockaddr *osa = NULL;
	int rc;

	d = io_uring_cqe_get_data(cqe);

	osa = &d->osa;
	rc = cqe->res;
	if (rc < 0) {
		LOGP(DLGLOBAL, LOGL_ERROR, "%s -> rx error rc=%d flags=0x%x\n",
		     osa ? osmo_sockaddr_to_str(osa) : "NULL",
		     rc, cqe->flags);
		return;
	}

	switch (d->type) {
	case IO_RECV:
		/* done reading */
		d->iov.iov_len = rc;
		LOGP(DLGLOBAL, LOGL_DEBUG, "%s -> rx rc=%d flags=0x%x: %s\n",
		     osa ? osmo_sockaddr_to_str(osa) : "NULL",
		     rc, cqe->flags,
		     osmo_quote_str(d->iov.iov_base, d->iov.iov_len));
		io_uring_cqe_seen(ring, cqe);
		g_total_rx++;

		if (response_n < 1)
			break;

		/* resubmit back to sender */
		if (response_size > 0)
			d->iov.iov_len = response_size;
		data_io_prep_send(ring, port, d);
		break;

	case IO_SEND:
		/* done writing. */
		LOGP(DLGLOBAL, LOGL_DEBUG, "%s <- tx rc=%d flags=0x%x: %s\n",
		     osa ? osmo_sockaddr_to_str(osa) : "NULL",
		     rc, cqe->flags,
		     osmo_quote_str(d->iov.iov_base, rc));
		io_uring_cqe_seen(ring, cqe);
		g_total_tx++;

		d->n++;

		/* Send again? If not, re-submit open slot for reading. */
		if (d->n < response_n)
			data_io_prep_send(ring, port, d);
		else
			data_io_prep_recv(ring, port, d);
		break;

	default:
		OSMO_ASSERT(0);
	}
}

struct cmdline_cmd cmds[] = {
	{
		.short_option = "h",
		.long_option = "help",
		.doc = "Show this help",
	},
	{
		.short_option = "l",
		.long_option = "local-addr",
		.arg_name = "IP-ADDR",
		.doc = "Listen on local IP address (default is 0.0.0.0).",
	},
	{
		.short_option = "p",
		.long_option = "port",
		.arg_name = "UDP-PORT",
		.doc = "Listen on local UDP port.",
	},
	/*
	{
		.short_option = "P",
		.long_option = "port-range-to",
		.arg_name = "UDP-PORT-TO",
		.doc = "Listen on a range of ports, from --port to --port-range-to, inclusive.",
	},
	*/
	{
		.short_option = "s",
		.long_option = "response-size",
		.arg_name = "BYTES",
		.doc = "When responding, enlarge or shorten the payload to this size.",
	},
	{
		.short_option = "n",
		.long_option = "response-repeat",
		.arg_name = "N",
		.doc = "Respond N times, i.e. multiply the returned traffic.",
	},
	{
		.long_option = "io-uring-queue",
		.arg_name = "SIZE",
		.doc = "I/O tuning: queue size to use for io_uring, default is 4000.",
	},
	{
		.long_option = "io-uring-buf",
		.arg_name = "SIZE",
		.doc = "I/O tuning: maximum payload size, default is 2048.",
	},
	{}
};

static const struct log_info_cat categories[] = {
};

const struct log_info udp_responder_log_info = {
	.cat = categories,
	.num_cat = ARRAY_SIZE(categories),
};

int main(int argc, char **argv)
{
	struct udp_port port = {};
	struct osmo_sockaddr_str addr = {};
	struct osmo_sockaddr osa = {};

	osmo_init_logging2(OTC_GLOBAL, &udp_responder_log_info);
	log_set_log_level(osmo_stderr_target, LOGL_ERROR);

	if (cmdline_read(cmds, argc, argv)
	    || cmdline_get(cmds, "help", NULL)) {
		cmdline_print_help(cmds);
		return -1;
	}

	int port_nr;
	if (!cmdline_get_int(&port_nr, 1, 65525, 23000, cmds, "port"))
		return -1;

	const char *local_addr = cmdline_get(cmds, "local-addr", "0.0.0.0");
	if (osmo_sockaddr_str_from_str(&addr, local_addr, port_nr)
	    || osmo_sockaddr_str_to_osa(&addr, &osa)) {
		printf("ERROR: invalid interface or port number: %s:%d\n", local_addr, port_nr);
		return -1;
	}

	int queue_size;
	if (!cmdline_get_int(&queue_size, 1, 65536, 4000, cmds, "io-uring-queue"))
		return -1;
	int buf_size;
	if (!cmdline_get_int(&buf_size, 1, 65536, 2048, cmds, "io-uring-buf"))
		return -1;

	int response_size = 0;
	if (!cmdline_get_int(&response_size, 0, buf_size, 0, cmds, "response-size"))
		return -1;

	int response_n = 1;
	if (!cmdline_get_int(&response_n, 0, INT_MAX, 1, cmds, "response-repeat"))
		return -1;


	port.osa = osa;

	/* create and bind socket */
	int rc;
	rc = osmo_sock_init_osa(SOCK_DGRAM, IPPROTO_UDP, &port.osa, NULL, OSMO_SOCK_F_BIND);
	/* (logging of errors already happens in osmo_sock_init_osa() */
	if (rc < 0)
		return -1;
	port.fd = rc;
	LOGP(DLGLOBAL, LOGL_NOTICE, "bound UDP %s fd=%d\n", osmo_sock_get_name2(port.fd), port.fd);

	struct io_uring ring = {};
	rc = io_uring_queue_init(queue_size, &ring, 0);
	OSMO_ASSERT(rc >= 0);

	struct io_queue *q = talloc_size(OTC_GLOBAL, sizeof(struct io_queue) + queue_size * sizeof(struct data_io));
	OSMO_ASSERT(q);
	*q = (struct io_queue){
		.d_size = queue_size,
	};

	for (int i = 0; i < q->d_size; i++) {
		struct data_io *d = &q->d[i];
		*d = (struct data_io){
			.data = talloc_size(q, buf_size),
			.data_size = buf_size,
		};

		/* fill once with random printable data */
		for (int j = 0; j < d->data_size; j++)
			d->data[j] = 32 + random() % (126 - 32 + 1);
	}

	for (int i = 0; i < q->d_size; i++) {
		data_io_prep_recv(&ring, &port, &q->d[i]);
	}

	struct __kernel_timespec ts = {};
	struct timespec next = {};

	while (1) {
		uint32_t submitted;
		uint32_t completed = 0;
		struct io_uring_cqe *cqe;
		bool show = false;

		/* submit any requests from previous loop */
		submitted = io_uring_submit(&ring);

		/* process all pending completions */
		while (io_uring_wait_cqe_timeout(&ring, &cqe, &ts) == 0) {
			data_io_handle_completion(&ring, &port, cqe, response_size, response_n);
			completed++;
		}

		{
			struct timespec now;
			osmo_clock_gettime(CLOCK_MONOTONIC, &now);
			if (timespeccmp(&now, &next, >=)) {
				next = now;
				next.tv_sec++;
				show = true;
			}
		}

		if (show || (!submitted && !completed)) {
			printf("%6u rx %6u tx   \r", g_total_rx, g_total_tx);
			fflush(stdout);
		}

		if (!submitted && !completed) {
			io_uring_wait_cqe(&ring, &cqe);
			data_io_handle_completion(&ring, &port, cqe, response_size, response_n);
		}
	}

	talloc_free(q);
	return 0;
}

#endif /* HAVE_URING */