1 /* vi: set sw=4 ts=4: */
2 /*
3  * Generic non-forking server infrastructure.
4  * Intended to make writing telnetd-type servers easier.
5  *
6  * Copyright (C) 2007 Denys Vlasenko
7  *
8  * Licensed under GPLv2, see file LICENSE in this source tree.
9  */
10 
11 #include "libbb.h"
12 #include "isrv.h"
13 
14 #define DEBUG 0
15 
16 #if DEBUG
17 #define DPRINTF(args...) bb_error_msg(args)
18 #else
19 #define DPRINTF(args...) ((void)0)
20 #endif
21 
22 /* Helpers */
23 
24 /* Opaque structure */
25 
26 struct isrv_state_t {
27 	short  *fd2peer; /* one per registered fd */
28 	void  **param_tbl; /* one per registered peer */
29 	/* one per registered peer; doesn't exist if !timeout */
30 	time_t *timeo_tbl;
31 	int   (*new_peer)(isrv_state_t *state, int fd);
32 	time_t  curtime;
33 	int     timeout;
34 	int     fd_count;
35 	int     peer_count;
36 	int     wr_count;
37 	fd_set  rd;
38 	fd_set  wr;
39 };
40 #define FD2PEER    (state->fd2peer)
41 #define PARAM_TBL  (state->param_tbl)
42 #define TIMEO_TBL  (state->timeo_tbl)
43 #define CURTIME    (state->curtime)
44 #define TIMEOUT    (state->timeout)
45 #define FD_COUNT   (state->fd_count)
46 #define PEER_COUNT (state->peer_count)
47 #define WR_COUNT   (state->wr_count)
48 
49 /* callback */
isrv_want_rd(isrv_state_t * state,int fd)50 void isrv_want_rd(isrv_state_t *state, int fd)
51 {
52 	FD_SET(fd, &state->rd);
53 }
54 
55 /* callback */
isrv_want_wr(isrv_state_t * state,int fd)56 void isrv_want_wr(isrv_state_t *state, int fd)
57 {
58 	if (!FD_ISSET(fd, &state->wr)) {
59 		WR_COUNT++;
60 		FD_SET(fd, &state->wr);
61 	}
62 }
63 
64 /* callback */
isrv_dont_want_rd(isrv_state_t * state,int fd)65 void isrv_dont_want_rd(isrv_state_t *state, int fd)
66 {
67 	FD_CLR(fd, &state->rd);
68 }
69 
70 /* callback */
isrv_dont_want_wr(isrv_state_t * state,int fd)71 void isrv_dont_want_wr(isrv_state_t *state, int fd)
72 {
73 	if (FD_ISSET(fd, &state->wr)) {
74 		WR_COUNT--;
75 		FD_CLR(fd, &state->wr);
76 	}
77 }
78 
79 /* callback */
isrv_register_fd(isrv_state_t * state,int peer,int fd)80 int isrv_register_fd(isrv_state_t *state, int peer, int fd)
81 {
82 	int n;
83 
84 	DPRINTF("register_fd(peer:%d,fd:%d)", peer, fd);
85 
86 	if (FD_COUNT >= FD_SETSIZE) return -1;
87 	if (FD_COUNT <= fd) {
88 		n = FD_COUNT;
89 		FD_COUNT = fd + 1;
90 
91 		DPRINTF("register_fd: FD_COUNT %d", FD_COUNT);
92 
93 		FD2PEER = xrealloc(FD2PEER, FD_COUNT * sizeof(FD2PEER[0]));
94 		while (n < fd) FD2PEER[n++] = -1;
95 	}
96 
97 	DPRINTF("register_fd: FD2PEER[%d] = %d", fd, peer);
98 
99 	FD2PEER[fd] = peer;
100 	return 0;
101 }
102 
103 /* callback */
isrv_close_fd(isrv_state_t * state,int fd)104 void isrv_close_fd(isrv_state_t *state, int fd)
105 {
106 	DPRINTF("close_fd(%d)", fd);
107 
108 	close(fd);
109 	isrv_dont_want_rd(state, fd);
110 	if (WR_COUNT) isrv_dont_want_wr(state, fd);
111 
112 	FD2PEER[fd] = -1;
113 	if (fd == FD_COUNT-1) {
114 		do fd--; while (fd >= 0 && FD2PEER[fd] == -1);
115 		FD_COUNT = fd + 1;
116 
117 		DPRINTF("close_fd: FD_COUNT %d", FD_COUNT);
118 
119 		FD2PEER = xrealloc(FD2PEER, FD_COUNT * sizeof(FD2PEER[0]));
120 	}
121 }
122 
123 /* callback */
isrv_register_peer(isrv_state_t * state,void * param)124 int isrv_register_peer(isrv_state_t *state, void *param)
125 {
126 	int n;
127 
128 	if (PEER_COUNT >= FD_SETSIZE) return -1;
129 	n = PEER_COUNT++;
130 
131 	DPRINTF("register_peer: PEER_COUNT %d", PEER_COUNT);
132 
133 	PARAM_TBL = xrealloc(PARAM_TBL, PEER_COUNT * sizeof(PARAM_TBL[0]));
134 	PARAM_TBL[n] = param;
135 	if (TIMEOUT) {
136 		TIMEO_TBL = xrealloc(TIMEO_TBL, PEER_COUNT * sizeof(TIMEO_TBL[0]));
137 		TIMEO_TBL[n] = CURTIME;
138 	}
139 	return n;
140 }
141 
remove_peer(isrv_state_t * state,int peer)142 static void remove_peer(isrv_state_t *state, int peer)
143 {
144 	int movesize;
145 	int fd;
146 
147 	DPRINTF("remove_peer(%d)", peer);
148 
149 	fd = FD_COUNT - 1;
150 	while (fd >= 0) {
151 		if (FD2PEER[fd] == peer) {
152 			isrv_close_fd(state, fd);
153 			fd--;
154 			continue;
155 		}
156 		if (FD2PEER[fd] > peer)
157 			FD2PEER[fd]--;
158 		fd--;
159 	}
160 
161 	PEER_COUNT--;
162 	DPRINTF("remove_peer: PEER_COUNT %d", PEER_COUNT);
163 
164 	movesize = (PEER_COUNT - peer) * sizeof(void*);
165 	if (movesize > 0) {
166 		memcpy(&PARAM_TBL[peer], &PARAM_TBL[peer+1], movesize);
167 		if (TIMEOUT)
168 			memcpy(&TIMEO_TBL[peer], &TIMEO_TBL[peer+1], movesize);
169 	}
170 	PARAM_TBL = xrealloc(PARAM_TBL, PEER_COUNT * sizeof(PARAM_TBL[0]));
171 	if (TIMEOUT)
172 		TIMEO_TBL = xrealloc(TIMEO_TBL, PEER_COUNT * sizeof(TIMEO_TBL[0]));
173 }
174 
handle_accept(isrv_state_t * state,int fd)175 static void handle_accept(isrv_state_t *state, int fd)
176 {
177 	int n, newfd;
178 
179 	/* suppress gcc warning "cast from ptr to int of different size" */
180 	fcntl(fd, F_SETFL, (int)(ptrdiff_t)(PARAM_TBL[0]) | O_NONBLOCK);
181 	newfd = accept(fd, NULL, 0);
182 	fcntl(fd, F_SETFL, (int)(ptrdiff_t)(PARAM_TBL[0]));
183 	if (newfd < 0) {
184 		if (errno == EAGAIN) return;
185 		/* Most probably someone gave us wrong fd type
186 		 * (for example, non-socket). Don't want
187 		 * to loop forever. */
188 		bb_simple_perror_msg_and_die("accept");
189 	}
190 
191 	DPRINTF("new_peer(%d)", newfd);
192 	n = state->new_peer(state, newfd);
193 	if (n)
194 		remove_peer(state, n); /* unsuccessful peer start */
195 }
196 
handle_fd_set(isrv_state_t * state,fd_set * fds,int (* h)(int,void **))197 static void handle_fd_set(isrv_state_t *state, fd_set *fds, int (*h)(int, void **))
198 {
199 	enum { LONG_CNT = sizeof(fd_set) / sizeof(long) };
200 	int fds_pos;
201 	int fd, peer;
202 	/* need to know value at _the beginning_ of this routine */
203 	int fd_cnt = FD_COUNT;
204 
205 	BUILD_BUG_ON(LONG_CNT * sizeof(long) != sizeof(fd_set));
206 
207 	fds_pos = 0;
208 	while (1) {
209 		/* Find next nonzero bit */
210 		while (fds_pos < LONG_CNT) {
211 			if (((long*)fds)[fds_pos] == 0) {
212 				fds_pos++;
213 				continue;
214 			}
215 			/* Found non-zero word */
216 			fd = fds_pos * sizeof(long)*8; /* word# -> bit# */
217 			while (1) {
218 				if (FD_ISSET(fd, fds)) {
219 					FD_CLR(fd, fds);
220 					goto found_fd;
221 				}
222 				fd++;
223 			}
224 		}
225 		break; /* all words are zero */
226  found_fd:
227 		if (fd >= fd_cnt) { /* paranoia */
228 			DPRINTF("handle_fd_set: fd > fd_cnt?? (%d > %d)",
229 					fd, fd_cnt);
230 			break;
231 		}
232 		DPRINTF("handle_fd_set: fd %d is active", fd);
233 		peer = FD2PEER[fd];
234 		if (peer < 0)
235 			continue; /* peer is already gone */
236 		if (peer == 0) {
237 			handle_accept(state, fd);
238 			continue;
239 		}
240 		DPRINTF("h(fd:%d)", fd);
241 		if (h(fd, &PARAM_TBL[peer])) {
242 			/* this peer is gone */
243 			remove_peer(state, peer);
244 		} else if (TIMEOUT) {
245 			TIMEO_TBL[peer] = monotonic_sec();
246 		}
247 	}
248 }
249 
handle_timeout(isrv_state_t * state,int (* do_timeout)(void **))250 static void handle_timeout(isrv_state_t *state, int (*do_timeout)(void **))
251 {
252 	int n, peer;
253 	peer = PEER_COUNT-1;
254 	/* peer 0 is not checked */
255 	while (peer > 0) {
256 		DPRINTF("peer %d: time diff %d", peer,
257 				(int)(CURTIME - TIMEO_TBL[peer]));
258 		if ((CURTIME - TIMEO_TBL[peer]) >= TIMEOUT) {
259 			DPRINTF("peer %d: do_timeout()", peer);
260 			n = do_timeout(&PARAM_TBL[peer]);
261 			if (n)
262 				remove_peer(state, peer);
263 		}
264 		peer--;
265 	}
266 }
267 
268 /* Driver */
isrv_run(int listen_fd,int (* new_peer)(isrv_state_t * state,int fd),int (* do_rd)(int fd,void **),int (* do_wr)(int fd,void **),int (* do_timeout)(void **),int timeout,int linger_timeout)269 void isrv_run(
270 	int listen_fd,
271 	int (*new_peer)(isrv_state_t *state, int fd),
272 	int (*do_rd)(int fd, void **),
273 	int (*do_wr)(int fd, void **),
274 	int (*do_timeout)(void **),
275 	int timeout,
276 	int linger_timeout)
277 {
278 	isrv_state_t *state = xzalloc(sizeof(*state));
279 	state->new_peer = new_peer;
280 	state->timeout  = timeout;
281 
282 	/* register "peer" #0 - it will accept new connections */
283 	isrv_register_peer(state, NULL);
284 	isrv_register_fd(state, /*peer:*/ 0, listen_fd);
285 	isrv_want_rd(state, listen_fd);
286 	/* remember flags to make blocking<->nonblocking switch faster */
287 	/* (suppress gcc warning "cast from ptr to int of different size") */
288 	PARAM_TBL[0] = (void*)(ptrdiff_t)(fcntl(listen_fd, F_GETFL));
289 
290 	while (1) {
291 		struct timeval tv;
292 		fd_set rd;
293 		fd_set wr;
294 		fd_set *wrp = NULL;
295 		int n;
296 
297 		tv.tv_sec = timeout;
298 		if (PEER_COUNT <= 1)
299 			tv.tv_sec = linger_timeout;
300 		tv.tv_usec = 0;
301 		rd = state->rd;
302 		if (WR_COUNT) {
303 			wr = state->wr;
304 			wrp = &wr;
305 		}
306 
307 		DPRINTF("run: select(FD_COUNT:%d,timeout:%d)...",
308 				FD_COUNT, (int)tv.tv_sec);
309 		n = select(FD_COUNT, &rd, wrp, NULL, tv.tv_sec ? &tv : NULL);
310 		DPRINTF("run: ...select:%d", n);
311 
312 		if (n < 0) {
313 			if (errno != EINTR)
314 				bb_simple_perror_msg("select");
315 			continue;
316 		}
317 
318 		if (n == 0 && linger_timeout && PEER_COUNT <= 1)
319 			break;
320 
321 		if (timeout) {
322 			time_t t = monotonic_sec();
323 			if (t != CURTIME) {
324 				CURTIME = t;
325 				handle_timeout(state, do_timeout);
326 			}
327 		}
328 		if (n > 0) {
329 			handle_fd_set(state, &rd, do_rd);
330 			if (wrp)
331 				handle_fd_set(state, wrp, do_wr);
332 		}
333 	}
334 	DPRINTF("run: bailout");
335 	/* NB: accept socket is not closed. Caller is to decide what to do */
336 }
337