SO_REUSEPORT with UDP multicast

reuseport_test_udp_mc_client.c

#include <stdio.h>
#include <stdint.h>
#include <stdlib.h>
#include <string.h>

#include <unistd.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <netinet/ip.h>
#include <arpa/inet.h>

struct {
	uint32_t port;
	int count;
	uint32_t multicast;
} cmdargs;

static _Bool
opt_have_arg(char c, const char* optstr)
{
	for (size_t i = 0; i < strlen(optstr); ++i) {
		if (optstr[i] == c)
			return optstr[i+1] == ':';
	}
	return 0;
}

static void
usage() {
	printf("Usage: \n");
	printf("reuseport_cnt [-p port] [-c count]\n");
	printf("\t-p port:\tDestination port\n");
	printf("\t-c count:\tTotal connections to make\n");
}

static void
usage_and_exit() {
	usage();
	exit(-1);
}

static void
parse_arg(int argc, char *argv[]) {
	cmdargs.port = 12345;
	cmdargs.count = 1;
	cmdargs.multicast = 1;

	const char* optstr = "p:c:";
	char *arg = NULL;
	char optchar = '\0';

	for (int i = 1; i < argc; ++i) {
		char *opt = argv[i];
		if (opt[0] != '-')
			usage_and_exit();
		
		optchar = opt[1];
		if (optchar == '\0')
			usage_and_exit();

		if (!opt_have_arg(optchar, optstr)) {
			arg = NULL;
		} else if (opt[2] == '\0') {
			if (++i == argc)
				usage_and_exit();
			arg = argv[i];
		} else {
			arg = &opt[2];
		}

		switch (optchar) {
		case 'p':
			cmdargs.port = strtoul(arg, NULL, 10);
			if (cmdargs.port == 0 || cmdargs.port >= 65536)
				usage_and_exit();
			break;
		case 'c':
			cmdargs.count = strtoul(arg, NULL, 10);
			if (cmdargs.count == 0)
				usage_and_exit();
			break;
		case 'm':
			cmdargs.multicast = 0;
			break;
		default:
			usage_and_exit();
		}
	}
}

int
bind_socket() {
	int fd = socket(AF_INET, SOCK_DGRAM, 0);
	if (fd < 0) {
		perror("socket");
		return -1;
	}

	if (cmdargs.multicast) {
	int val = 1;
	if (setsockopt(fd, SOL_SOCKET, SO_BROADCAST, &val, sizeof(val)) < 0) {
		perror("setsockopt");
		return -1;
	}
	}

	return fd;
}

int
main(int argc, char** argv) {
	parse_arg(argc, argv);

	int s[100];
	char buf[100];

	memset(s, 0, sizeof(s));

	struct sockaddr_in addr;
	memset(&addr, 0, sizeof(addr));
	addr.sin_family = AF_INET;
	addr.sin_addr.s_addr = cmdargs.multicast ? (INADDR_BROADCAST) : inet_addr("127.0.0.1");
	addr.sin_port = htons(cmdargs.port);

	for (uint i = 0; i < cmdargs.count; ++i) {
		int fd = bind_socket();
		if (fd < 0) {
			return -1;
		}

		char msg[] = "TESTPING";
		if (sendto(fd, msg, sizeof(msg), 0, (void*)&addr, sizeof(addr)) < 0) {
			perror("sendto");
		}
	}
}

reuseport_udp_mc.c

#include <stdio.h>
#include <stdint.h>
#include <stdlib.h>
#include <string.h>

#include <unistd.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <netinet/ip.h>
#include <arpa/inet.h>

struct {
	uint32_t port;
	_Bool reuseport;
	int workers;
} cmdargs;

static _Bool
opt_have_arg(char c, const char* optstr)
{
	for (size_t i = 0; i < strlen(optstr); ++i) {
		if (optstr[i] == c)
			return optstr[i+1] == ':';
	}
	return 0;
}

static void
usage() {
	printf("Usage: \n");
	printf("reuseport [-p port] [-n workers] [-r]\n");
	printf("\t-p port:\tListen port\n");
	printf("\t-n workers:\ttotal worker processes to spawn\n");
	printf("\t-r:\t\tDO NOT use SO_REUSEPORT; To test if bind() will fail correctly\n");
}

static void
usage_and_exit() {
	usage();
	exit(-1);
}

static void
parse_arg(int argc, char *argv[]) {
	cmdargs.port = 12345;
	cmdargs.reuseport = 1;
	cmdargs.workers = 4;

	const char* optstr = "p:n:r";
	char *arg = NULL;
	char optchar = '\0';

	for (int i = 1; i < argc; ++i) {
		char *opt = argv[i];
		if (opt[0] != '-')
			usage_and_exit();
		
		optchar = opt[1];
		if (optchar == '\0')
			usage_and_exit();

		if (!opt_have_arg(optchar, optstr)) {
			arg = NULL;
		} else if (opt[2] == '\0') {
			if (++i == argc)
				usage_and_exit();
			arg = argv[i];
		} else {
			arg = &opt[2];
		}

		switch (optchar) {
		case 'p':
			cmdargs.port = strtoul(arg, NULL, 10);
			if (cmdargs.port == 0 || cmdargs.port >= 65536)
				usage_and_exit();
			break;
		case 'n':
			cmdargs.workers = strtoul(arg, NULL, 10);
			if (cmdargs.workers == 0)
				usage_and_exit();
			break;
		case 'r':
			cmdargs.reuseport = 0;
			break;
		default:
			usage_and_exit();
		}
	}
}

int
bind_sock() {
	int fd = socket(AF_INET, SOCK_DGRAM, 0);
	if (fd < 0) {
		perror("socket");
		return -1;
	}

	int val = 1;

	int option = cmdargs.reuseport ? SO_REUSEPORT : SO_REUSEADDR;
	if (setsockopt(fd, SOL_SOCKET, option, &val, sizeof(val)) < 0) {
		perror("setsockopt");
		return -1;
	}

	struct sockaddr_in addr;
	memset(&addr, 0, sizeof(addr));
	addr.sin_family = AF_INET;
	addr.sin_addr.s_addr = INADDR_ANY;
	addr.sin_port = htons(cmdargs.port);
	
	if (bind(fd, (void*)&addr, sizeof(addr)) < 0) {
		perror("bind");
		return -1;
	}

	return fd;
}

void
worker_main(int n)
{
	printf("worker %d started\n", n);

	int fd = bind_sock();
	if (fd < 0)
		exit(-1);
	
	struct sockaddr c_addr;
	socklen_t addrlen = sizeof(c_addr);
	ssize_t size;
	char rbuf[128];

	while ((size = recvfrom(fd, rbuf, sizeof(rbuf), 0, &c_addr, &addrlen)) > 0) {
		printf("worker %d received data\r\n", n);
	}

	perror("accept");
	exit(-1);
}

int
main(int argc, char** argv) {
	parse_arg(argc, argv);

	int n = 0;
	for (int i = 1; i < cmdargs.workers; ++i) {
		if (fork()) {
			n = i;
			break;
		}
	}

	worker_main(n);
}