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