Skip to content

Instantly share code, notes, and snippets.

@MosakujiHokuto
Created December 29, 2020 17:07
Show Gist options
  • Save MosakujiHokuto/26024f16ad87717402980482b0b46943 to your computer and use it in GitHub Desktop.
Save MosakujiHokuto/26024f16ad87717402980482b0b46943 to your computer and use it in GitHub Desktop.
SO_REUSEPORT with UDP multicast
#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");
}
}
}
#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);
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment