1 #include <sys/types.h>
2 #include <sys/socket.h>
3 #include <sys/sysctl.h>
4 #include <sys/usched.h>
5 
6 #include <arpa/inet.h>
7 #include <netinet/in.h>
8 
9 #include <err.h>
10 #include <stdio.h>
11 #include <stdint.h>
12 #include <stdlib.h>
13 #include <string.h>
14 #include <unistd.h>
15 
16 #define RCVBUF_SIZE	(256 * 1024)
17 #define BUFLEN		2048
18 
19 #define FLAG_NOREPLY	0x0001
20 #define FLAG_BINDCPU	0x0002
21 #define FLAG_NOREUSE	0x0004
22 
23 static void	mainloop(struct sockaddr_in *, int, int, uint32_t);
24 
25 static void
26 usage(const char *cmd)
27 {
28 	fprintf(stderr, "%s [-4 addr4] -p port [-i ninst] [-r rcvbuf] "
29 	    "[-N] [-B] [-U]\n", cmd);
30 	exit(2);
31 }
32 
33 int
34 main(int argc, char *argv[])
35 {
36 	struct sockaddr_in in;
37 	int opt, ninst, i, s, rcvbuf, flags;
38 	size_t prm_len;
39 
40 	prm_len = sizeof(ninst);
41 	if (sysctlbyname("hw.ncpu", &ninst, &prm_len, NULL, 0) != 0)
42 		err(2, "sysctl hw.ncpu failed");
43 
44 	memset(&in, 0, sizeof(in));
45 	in.sin_family = AF_INET;
46 
47 	flags = 0;
48 	rcvbuf = RCVBUF_SIZE;
49 
50 	while ((opt = getopt(argc, argv, "4:p:i:r:NBU")) != -1) {
51 		switch (opt) {
52 		case '4':
53 			if (inet_pton(AF_INET, optarg, &in.sin_addr) <= 0)
54 				errx(1, "invalid -4 %s\n", optarg);
55 			break;
56 
57 		case 'p':
58 			in.sin_port = strtoul(optarg, NULL, 10);
59 			in.sin_port = htons(in.sin_port);
60 			break;
61 
62 		case 'i':
63 			ninst = strtoul(optarg, NULL, 10);
64 			break;
65 
66 		case 'r':
67 			rcvbuf = strtoul(optarg, NULL, 10);
68 			rcvbuf *= 1024;
69 			break;
70 
71 		case 'N':
72 			flags |= FLAG_NOREPLY;
73 			break;
74 
75 		case 'B':
76 			flags |= FLAG_BINDCPU;
77 			break;
78 
79 		case 'U':
80 			flags |= FLAG_NOREUSE;
81 			break;
82 
83 		default:
84 			usage(argv[0]);
85 		}
86 	}
87 
88 	if (in.sin_port == 0)
89 		usage(argv[0]);
90 
91 	s = -1;
92 
93 	for (i = 0; i < ninst - 1; ++i) {
94 		pid_t pid;
95 
96 		pid = fork();
97 		if (pid == 0)
98 			mainloop(&in, s, rcvbuf, flags);
99 		else if (pid < 0)
100 			err(1, "fork %d failed", i);
101 	}
102 
103 	mainloop(&in, s, rcvbuf, flags);
104 	exit(0);
105 }
106 
107 static void
108 mainloop(struct sockaddr_in *in, int s, int rcvbuf, uint32_t flags)
109 {
110 	void *buf;
111 
112 	if (s < 0) {
113 		s = socket(AF_INET, SOCK_DGRAM, 0);
114 		if (s < 0)
115 			err(1, "socket(INET, DGRAM) failed");
116 	}
117 
118 	buf = malloc(BUFLEN);
119 	if (buf == NULL)
120 		err(1, "malloc buf failed");
121 
122 	if ((flags & FLAG_NOREUSE) == 0) {
123 		int on = 1;
124 
125 		if (setsockopt(s, SOL_SOCKET, SO_REUSEPORT, &on, sizeof(on)) < 0)
126 			err(1, "setsockopt(SOCK, REUSEPORT) failed");
127 	}
128 
129 	if (setsockopt(s, SOL_SOCKET, SO_RCVBUF, &rcvbuf, sizeof(rcvbuf)) < 0)
130 		err(1, "setsockopt(SOCK, RCVBUF) failed");
131 
132 	if (bind(s, (const struct sockaddr *)in, sizeof(*in)) < 0)
133 		err(1, "bind failed");
134 
135 	if (flags & FLAG_BINDCPU) {
136 		socklen_t cpu_len;
137 		int cpu;
138 
139 		cpu_len = sizeof(cpu);
140 		if (getsockopt(s, SOL_SOCKET, SO_CPUHINT, &cpu, &cpu_len) < 0)
141 			err(1, "getsockopt(SOCK, CPUHINT) failed");
142 		if (cpu >= 0)
143 			usched_set(getpid(), USCHED_SET_CPU, &cpu, sizeof(cpu));
144 	}
145 
146 	for (;;) {
147 		struct sockaddr_in cli;
148 		socklen_t cli_len;
149 		int n;
150 
151 		cli_len = sizeof(cli);
152 		n = recvfrom(s, buf, BUFLEN, 0,
153 		    (struct sockaddr *)&cli, &cli_len);
154 		if (n > 0 && (flags & FLAG_NOREPLY) == 0) {
155 			sendto(s, buf, n, 0,
156 			   (const struct sockaddr *)&cli, cli_len);
157 		}
158 	}
159 	exit(0);
160 }
161