1 /*-
2 * Copyright (c) 2008-2009 Robert N. M. Watson
3 * All rights reserved.
4 *
5 * Redistribution and use in source and binary forms, with or without
6 * modification, are permitted provided that the following conditions
7 * are met:
8 * 1. Redistributions of source code must retain the above copyright
9 * notice, this list of conditions and the following disclaimer.
10 * 2. Redistributions in binary form must reproduce the above copyright
11 * notice, this list of conditions and the following disclaimer in the
12 * documentation and/or other materials provided with the distribution.
13 *
14 * THIS SOFTWARE IS PROVIDED BY THE AUTHOR AND CONTRIBUTORS ``AS IS'' AND
15 * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
16 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
17 * ARE DISCLAIMED. IN NO EVENT SHALL THE AUTHOR OR CONTRIBUTORS BE LIABLE
18 * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
19 * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS
20 * OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
21 * HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
22 * LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY
23 * OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
24 * SUCH DAMAGE.
25 *
26 * $FreeBSD$
27 */
28
29 #include <sys/types.h>
30 #include <sys/endian.h>
31 #include <sys/event.h>
32 #include <sys/resource.h>
33 #include <sys/sched.h>
34 #include <sys/socket.h>
35 #include <sys/sysctl.h>
36 #include <sys/time.h>
37 #include <sys/wait.h>
38
39 #include <netinet/in.h>
40 #include <netinet/tcp.h>
41
42 #include <err.h>
43 #include <fcntl.h>
44 #include <inttypes.h>
45 #include <signal.h>
46 #include <stdio.h>
47 #include <stdlib.h>
48 #include <string.h>
49 #include <unistd.h>
50
51 #include "tcpp.h"
52
53 /*
54 * Server side -- create a pool of processes, each listening on its own TCP
55 * port number for new connections. The first 8 bytes of each connection
56 * will be a network byte order length, then there will be that number of
57 * bytes of data. We use non-blocking sockets with kqueue to to avoid the
58 * overhead of threading or more than one process per processor, which makes
59 * things a bit awkward when dealing with data we care about. As such, we
60 * read into a small character buffer which we then convert to a length once
61 * we have all the data.
62 */
63 #define CONNECTION_MAGIC 0x6392af27
64 struct connection {
65 uint32_t conn_magic; /* Just magic. */
66 int conn_fd;
67 struct tcpp_header conn_header; /* Header buffer. */
68 u_int conn_header_len; /* Bytes so far. */
69 u_int64_t conn_data_len; /* How much to sink. */
70 u_int64_t conn_data_received; /* How much so far. */
71 };
72
73 static pid_t *pid_list;
74 static int kq;
75
76 static struct connection *
tcpp_server_newconn(int listen_fd)77 tcpp_server_newconn(int listen_fd)
78 {
79 struct connection *conn;
80 struct kevent kev;
81 int fd;
82
83 fd = accept(listen_fd, NULL, NULL);
84 if (fd < 0) {
85 warn("accept");
86 return (NULL);
87 }
88
89 if (fcntl(fd, F_SETFL, O_NONBLOCK) < 0)
90 err(-1, "fcntl");
91
92 conn = malloc(sizeof(*conn));
93 if (conn == NULL)
94 return (NULL);
95 bzero(conn, sizeof(*conn));
96 conn->conn_magic = CONNECTION_MAGIC;
97 conn->conn_fd = fd;
98
99 /*
100 * Register to read on the socket, and set our conn pointer as the
101 * udata so we can find it quickly in the future.
102 */
103 EV_SET(&kev, fd, EVFILT_READ, EV_ADD, 0, 0, conn);
104 if (kevent(kq, &kev, 1, NULL, 0, NULL) < 0)
105 err(-1, "kevent");
106
107 return (conn);
108 }
109
110 static void
tcpp_server_closeconn(struct connection * conn)111 tcpp_server_closeconn(struct connection *conn)
112 {
113
114 /*
115 * Kqueue cleans up after itself once we close the socket, and since
116 * we are processing only one kevent at a time, we don't need to
117 * worry about watching out for future kevents referring to it.
118 *
119 * ... right?
120 */
121 close(conn->conn_fd);
122 bzero(conn, sizeof(*conn));
123 free(conn);
124 }
125
126 static u_char buffer[256*1024]; /* Buffer in which to sink data. */
127 static void
tcpp_server_handleconn(struct kevent * kev)128 tcpp_server_handleconn(struct kevent *kev)
129 {
130 struct connection *conn;
131 ssize_t len;
132
133 conn = kev->udata;
134 if (conn->conn_magic != CONNECTION_MAGIC)
135 errx(-1, "tcpp_server_handleconn: magic");
136
137 if (conn->conn_header_len < sizeof(conn->conn_header)) {
138 len = read(conn->conn_fd,
139 ((u_char *)&conn->conn_header) + conn->conn_header_len,
140 sizeof(conn->conn_header) - conn->conn_header_len);
141 if (len < 0) {
142 warn("tcpp_server_handleconn: header read");
143 tcpp_server_closeconn(conn);
144 return;
145 }
146 if (len == 0) {
147 warnx("tcpp_server_handleconn: header premature eof");
148 tcpp_server_closeconn(conn);
149 return;
150 }
151 conn->conn_header_len += len;
152 if (conn->conn_header_len == sizeof(conn->conn_header)) {
153 tcpp_header_decode(&conn->conn_header);
154 if (conn->conn_header.th_magic != TCPP_MAGIC) {
155 warnx("tcpp_server_handleconn: bad magic");
156 tcpp_server_closeconn(conn);
157 return;
158 }
159 }
160 } else {
161 /*
162 * Drain up to a buffer from the connection, so that we pay
163 * attention to other connections too.
164 */
165 len = read(conn->conn_fd, buffer, sizeof(buffer));
166 if (len < 0) {
167 warn("tcpp_server_handleconn: data bad read");
168 tcpp_server_closeconn(conn);
169 return;
170 }
171 if (len == 0 && conn->conn_data_received <
172 conn->conn_header.th_len) {
173 warnx("tcpp_server_handleconn: data premature eof");
174 tcpp_server_closeconn(conn);
175 return;
176 }
177 conn->conn_data_received += len;
178 if (conn->conn_data_received > conn->conn_header.th_len) {
179 warnx("tcpp_server_handleconn: too much data");
180 tcpp_server_closeconn(conn);
181 return;
182 }
183 if (conn->conn_data_received == conn->conn_header.th_len) {
184 /*
185 * All is well.
186 */
187 tcpp_server_closeconn(conn);
188 return;
189 }
190 }
191 }
192
193 static void
tcpp_server_worker(int workernum)194 tcpp_server_worker(int workernum)
195 {
196 int i, listen_sock, numevents;
197 struct kevent kev, *kev_array;
198 int kev_bytes;
199 #if defined(CPU_SETSIZE) && 0
200 cpu_set_t mask;
201 int ncpus;
202 ssize_t len;
203
204 if (Pflag) {
205 len = sizeof(ncpus);
206 if (sysctlbyname(SYSCTLNAME_CPUS, &ncpus, &len, NULL, 0) < 0)
207 err(-1, "sysctlbyname: %s", SYSCTLNAME_CPUS);
208 if (len != sizeof(ncpus))
209 errx(-1, "sysctlbyname: %s: len %jd", SYSCTLNAME_CPUS,
210 (intmax_t)len);
211
212 CPU_ZERO(&mask);
213 CPU_SET(workernum % ncpus, &mask);
214 if (sched_setaffinity(0, CPU_SETSIZE, &mask) < 0)
215 err(-1, "sched_setaffinity");
216 }
217 #endif
218 setproctitle("tcpp_server %d", workernum);
219
220 /* Allow an extra kevent for the listen socket. */
221 kev_bytes = sizeof(*kev_array) * (mflag + 1);
222 kev_array = malloc(kev_bytes);
223 if (kev_array == NULL)
224 err(-1, "malloc");
225 bzero(kev_array, kev_bytes);
226
227 /* XXXRW: Want to set and pin the CPU here. */
228
229 /*
230 * Add the worker number to the local port.
231 */
232 localipbase.sin_port = htons(rflag + workernum);
233
234 listen_sock = socket(PF_INET, SOCK_STREAM, 0);
235 if (listen_sock < 0)
236 err(-1, "socket");
237 i = 1;
238 if (setsockopt(listen_sock, SOL_SOCKET, SO_NOSIGPIPE, &i, sizeof(i))
239 < 0)
240 err(-1, "setsockopt");
241 i = 1;
242 if (setsockopt(listen_sock, SOL_SOCKET, SO_REUSEPORT, &i, sizeof(i))
243 < 0)
244 err(-1, "setsockopt");
245 i = 1;
246 if (setsockopt(listen_sock, IPPROTO_TCP, TCP_NODELAY, &i, sizeof(i))
247 < 0)
248 err(-1, "setsockopt");
249 if (bind(listen_sock, (struct sockaddr *)&localipbase,
250 sizeof(localipbase)) < 0)
251 err(-1, "bind");
252 if (listen(listen_sock, 16384))
253 err(-1, "listen");
254 if (fcntl(listen_sock, F_SETFL, O_NONBLOCK) < 0)
255 err(-1, "fcntl");
256
257 kq = kqueue();
258 if (kq < 0)
259 err(-1, "kqueue");
260
261 EV_SET(&kev, listen_sock, EVFILT_READ, EV_ADD, 0, 0, NULL);
262 if (kevent(kq, &kev, 1, NULL, 0, NULL) < 0)
263 err(-1, "kevent");
264
265 while ((numevents = kevent(kq, NULL, 0, kev_array, mflag + 1, NULL))
266 > 0) {
267 for (i = 0; i < numevents; i++) {
268 if (kev_array[i].ident == (u_int)listen_sock)
269 (void)tcpp_server_newconn(listen_sock);
270 else
271 tcpp_server_handleconn(&kev_array[i]);
272 }
273 }
274 printf("Worker %d done\n", workernum);
275 }
276
277 void
tcpp_server(void)278 tcpp_server(void)
279 {
280 #if 0
281 long cp_time_last[CPUSTATES], cp_time_now[CPUSTATES], ticks;
282 size_t size;
283 #endif
284 pid_t pid;
285 int i;
286
287 pid_list = malloc(sizeof(*pid_list) * pflag);
288 if (pid_list == NULL)
289 err(-1, "malloc pid_list");
290 bzero(pid_list, sizeof(*pid_list) * pflag);
291
292 /*
293 * Start workers.
294 */
295 for (i = 0; i < pflag; i++) {
296 pid = fork();
297 if (pid < 0) {
298 warn("fork");
299 for (i = 0; i < pflag; i++) {
300 if (pid_list[i] != 0)
301 (void)kill(pid_list[i], SIGKILL);
302 }
303 exit(-1);
304 }
305 if (pid == 0) {
306 tcpp_server_worker(i);
307 exit(0);
308 }
309 pid_list[i] = pid;
310 }
311
312 #if 0
313 size = sizeof(cp_time_last);
314 if (sysctlbyname(SYSCTLNAME_CPTIME, &cp_time_last, &size,
315 NULL, 0) < 0)
316 err(-1, "sysctlbyname: %s", SYSCTLNAME_CPTIME);
317 while (1) {
318 sleep(10);
319 size = sizeof(cp_time_last);
320 if (sysctlbyname(SYSCTLNAME_CPTIME, &cp_time_now,
321 &size, NULL, 0) < 0)
322 err(-1, "sysctlbyname: %s",
323 SYSCTLNAME_CPTIME);
324 ticks = 0;
325 for (i = 0; i < CPUSTATES; i++) {
326 cp_time_last[i] = cp_time_now[i] -
327 cp_time_last[i];
328 ticks += cp_time_last[i];
329 }
330 printf("user%% %lu nice%% %lu sys%% %lu intr%% %lu "
331 "idle%% %lu\n",
332 (100 * cp_time_last[CP_USER]) / ticks,
333 (100 * cp_time_last[CP_NICE]) / ticks,
334 (100 * cp_time_last[CP_SYS]) / ticks,
335 (100 * cp_time_last[CP_INTR]) / ticks,
336 (100 * cp_time_last[CP_IDLE]) / ticks);
337 bcopy(cp_time_now, cp_time_last, sizeof(cp_time_last));
338 }
339 #endif
340
341 /*
342 * GC workers.
343 */
344 for (i = 0; i < pflag; i++) {
345 if (pid_list[i] != 0) {
346 while (waitpid(pid_list[i], NULL, 0) != pid_list[i]);
347 }
348 }
349 }
350