nbeng

a non-blocking client/server engine
git clone git://git.2f30.org/nbeng.git
Log | Files | Refs | README

nbeng.c (9422B)


      1 /*
      2  * Non-blocking network engine for tcp/udp/ssl connections
      3  * The design uses the same socket for sending and receiving
      4  * data.  The relevant data are stored in a connection context
      5  * which is passed around the functions for dealing with it.
      6  * Also provided are sample reading and writing functions.
      7  * Copyleft:
      8  *    DsP <dsp@2f30.org>
      9  *    Lazaros Koromilas <lostd@2f30.org>
     10  * Cheers go to PhotoRec for helping recover this file :)
     11  */
     12 
     13 #include <sys/types.h>
     14 #include <sys/socket.h>
     15 
     16 #include <netdb.h>
     17 
     18 #include <err.h>
     19 #include <fcntl.h>
     20 #include <stdio.h>
     21 #include <stdlib.h>
     22 #include <string.h>
     23 #include <unistd.h>
     24 #include <errno.h>
     25 
     26 /* global flags */
     27 unsigned int tflag = 0;
     28 unsigned int sflag = 0;
     29 unsigned int qflag = 0;
     30 
     31 /* connection context */
     32 typedef struct concontxt_t {
     33 #define UDPCON 1
     34 #define TCPCON 2
     35 #define SSLCON 3
     36 	unsigned int contype; /* connection type */
     37 	int confd; /* net connection fd */
     38 	int clifd;
     39 	int clifds[FD_SETSIZE]; /* for tcp connected clients */
     40 	fd_set lset; /* the set that we select() on */
     41 	struct addrinfo *clinfo; /* info on where we connect */
     42 	struct addrinfo *srvinfo; /* info for the local part */
     43 #define INBUFSIZ 512
     44 	char *buf; /* buffer used for input/output */
     45 	unsigned int buflen;
     46 	int wfd, lfd; /* local file descriptors */
     47 } concontxt;
     48 
     49 void
     50 usage(int ret)
     51 {
     52         fprintf(stderr, "usage: nbeng [-ht] rhost rport\n"
     53 		"\t-t	Use TCP for transport\n"
     54 		"\t-h	This help screen\n");
     55 	if (ret)
     56 		exit(1);
     57 }
     58 
     59 static void
     60 set_nonblocking(int fd)
     61 {
     62 	int opts;
     63 	opts = fcntl(fd, F_GETFL);
     64 	if (opts < 0)
     65 		err(1, "fcntl");
     66 	opts = (opts | O_NONBLOCK);
     67 	if (fcntl(fd, F_SETFL, opts) < 0)
     68 		err(1, "fcntl");
     69 }
     70 
     71 static void
     72 set_blocking(int fd)
     73 {
     74 	int opts;
     75 	opts = fcntl(fd, F_GETFL);
     76 	if (opts < 0)
     77 		err(1, "fcntl");
     78 	opts &= (~O_NONBLOCK);
     79 	if (fcntl(fd, F_SETFL, opts) < 0)
     80 		err(1, "fcntl");
     81 }
     82 
     83 /*
     84  * This is the main function that prepares a socket for:
     85  *   a) connection to the other endpoint.
     86  *   b) being able to receive data.
     87  * When it is ready it sets the fields in the context.
     88  * The flag for the connection type lives in the context.
     89  */
     90 static void
     91 prepare_socket(concontxt *con, char *host, char *port, int lflag)
     92 {
     93 	struct addrinfo cli_hints, *cli_servinfo, srv_hints, *srv_servinfo, *p0;
     94 	int rv, cli_sockfd, optval = 1;
     95 
     96 	if ((con == NULL) || (host == NULL) || (port == NULL))
     97 		errx(1, "prepare_socket was passed a null arg");
     98 	memset(&cli_hints, 0, sizeof (cli_hints));
     99 	memset(&srv_hints, 0, sizeof (srv_hints));
    100 	cli_hints.ai_family = srv_hints.ai_family = AF_INET;
    101 	srv_hints.ai_flags = AI_PASSIVE;
    102 	cli_hints.ai_socktype = srv_hints.ai_socktype =
    103 		((con->contype == TCPCON) || (con->contype == SSLCON))
    104 		? SOCK_STREAM : SOCK_DGRAM;
    105 	rv = getaddrinfo(host, port, &cli_hints, &cli_servinfo);
    106 	if (rv)
    107 		errx(1, "getaddrinfo: %s", gai_strerror(rv));
    108 	rv = getaddrinfo(NULL, port, &srv_hints, &srv_servinfo);
    109 	if (rv)
    110 		errx(1, "getaddrinfo: %s", gai_strerror(rv));
    111 	/* getaddrinfo returns a list of results so we iterate on it */
    112 	for (p0 = cli_servinfo; p0; p0 = p0->ai_next) {
    113 		cli_sockfd = socket(p0->ai_family, p0->ai_socktype,
    114 				    p0->ai_protocol);
    115 		if (cli_sockfd < 0)
    116 			continue; /* until the socket is ready */
    117 		rv = setsockopt(cli_sockfd, SOL_SOCKET,
    118 				SO_REUSEADDR, &optval, sizeof (optval));
    119 		if (rv < 0) {
    120 			close(cli_sockfd);
    121 			warn("setsockopt");
    122 			continue;
    123 		}
    124 		break;
    125 	}
    126 	if (!p0)
    127 		err(1, "socket");
    128 	/* the same for our local part */
    129 	for (p0 = srv_servinfo; p0; p0 = p0->ai_next) {
    130 		if (bind(cli_sockfd, p0->ai_addr, p0->ai_addrlen) < 0) {
    131 			close(cli_sockfd);
    132 			warn("bind");
    133 			continue;
    134 		}
    135 		break;
    136 	}
    137 	if (!p0)
    138 		err(1, "bind");
    139 	/* if it's a tcp connection we have to listen() */
    140 	if (con->contype != UDPCON && lflag)
    141 		listen(cli_sockfd, 5);
    142 	/* all was ok, so we register the socket to the context */
    143 	con->confd = cli_sockfd;
    144 	con->clinfo = cli_servinfo;
    145 }
    146 
    147 /*
    148  * Reads data from standard input and saves them in the context.
    149  * Sets the qflag on "Q".
    150  */
    151 static void
    152 readstdin(concontxt *con)
    153 {
    154 	char buf[INBUFSIZ];
    155 	ssize_t n;
    156 
    157 	/* handle input data */
    158 	do {
    159 		n = read(con->wfd, buf, INBUFSIZ);
    160 		if (n < 0) {
    161 			if (errno != EAGAIN)
    162 				err(1, "read");
    163 			continue;
    164 		}
    165 		break;
    166 	} while (1);
    167 	buf[n] = '\0';
    168 
    169 	/* user quits? */
    170 	if (n == 0 || strncmp(buf, "Q\n", 3) == 0) {
    171 		qflag = 1;
    172 		return;
    173 	}
    174        	printf("n=%zd buf=%s", n, buf);
    175 
    176 	/* copy to context */
    177 	if ((con->buf = malloc(n)) == NULL) {
    178 		warn("malloc");
    179 		goto freex;
    180 	}
    181 	memcpy(con->buf, buf, n);
    182 	con->buflen = n;
    183 	return;
    184 
    185 freex:
    186 	if (con->buf != NULL)
    187                 free(con->buf);
    188 	con->buf = NULL;
    189 	con->buflen = 0;
    190 }
    191 
    192 /*
    193  * Sends data available in the context to all peers.
    194  */
    195 static void
    196 writesock(concontxt *con)
    197 {
    198 	ssize_t n;
    199 
    200 	/* nothing to do */
    201 	if (con->buf == NULL)
    202 		return;
    203 
    204 	if (con->contype == UDPCON) {
    205         	n = sendto(con->confd, con->buf, con->buflen, 0,
    206 			   con->clinfo->ai_addr,
    207 			   con->clinfo->ai_addrlen);
    208 		if (n < 0)
    209 			warn("sendto");
    210 	} else {
    211 		n = write(con->clifd, con->buf, con->buflen);
    212 		if (n < 0) {
    213 			warn("write");
    214 			close(con->clifd);
    215 			con->clifd = -1;
    216 			goto freex;
    217 		}
    218 	}
    219 
    220 freex:
    221 	if (con->buf != NULL)
    222                 free(con->buf);
    223 	con->buf = NULL;
    224 	con->buflen = 0;
    225 }
    226 
    227 /*
    228  * Reads data from the socket and saves them in the context.
    229  */
    230 static void
    231 readsock(concontxt *con)
    232 {
    233 	char buf[INBUFSIZ];
    234 	ssize_t n;
    235 	int r;
    236 	char host[NI_MAXHOST];
    237 	socklen_t addr_len;
    238 	struct sockaddr_storage their_addr;
    239 
    240 	addr_len = sizeof (their_addr);
    241 	if (con->contype == UDPCON) {
    242 		n = recvfrom(con->confd, buf, sizeof (buf), MSG_DONTWAIT,
    243 			     (struct sockaddr *)&their_addr, &addr_len);
    244 		if (n < 0) {
    245 			warn("recvfrom");
    246 			goto freex;
    247 		}
    248 
    249 		/* resolv */
    250 		r = getnameinfo((struct sockaddr *)&their_addr, addr_len,
    251 				host, sizeof (host), NULL, 0, 0);
    252 		if (r < 0) {
    253 			warn("getnameinfo");
    254 			snprintf(host, sizeof (host), "unknown");
    255 		}
    256 		printf("host=%s\n", host);
    257 	} else {
    258 		do {
    259 			n = read(con->clifd, buf, sizeof (buf));
    260 			if (n < 0) {
    261 				if (errno != EAGAIN)
    262 					err(1, "read");
    263 				continue;
    264 			}
    265 			if (n == 0) {
    266 				qflag = 1;
    267 				close(con->clifd);
    268 				con->clifd = -1;
    269 				goto freex;
    270 			}
    271 			break;
    272 		} while (1);
    273 	}
    274 	buf[n] = '\0';
    275 
    276 	/* copy to context */
    277 	if ((con->buf = malloc(n)) == NULL) {
    278 		warn("malloc");
    279 		goto freex;
    280 	}
    281 	memcpy(con->buf, buf, n);
    282 	con->buflen = n;
    283 	printf("n=%zd buf=%s", n, buf);
    284 	return;
    285 
    286 freex:
    287 	if (con->buf != NULL)
    288                 free(con->buf);
    289 	con->buf = NULL;
    290 	con->buflen = 0;
    291 }
    292 
    293 /*
    294  * Writes data available in the context to standard output.
    295  */
    296 static void
    297 writestdout(concontxt *con)
    298 {
    299 	write(con->lfd, con->buf, con->buflen);
    300 
    301 	if (con->buf != NULL)
    302                 free(con->buf);
    303 	con->buf = NULL;
    304 	con->buflen = 0;
    305 }
    306 
    307 static void
    308 myaccept(concontxt *con)
    309 {
    310 	int r;
    311 	char host[NI_MAXHOST];
    312 	int newfd;
    313 	struct sockaddr_storage sa;
    314 	socklen_t salen = sizeof (sa);
    315 
    316 	newfd = accept(con->confd, (struct sockaddr *)&sa, &salen);
    317 	printf("newfd=%d\n", newfd);
    318 	if (newfd < 0)
    319 		err(1, "accept");
    320 	con->clifd = newfd;
    321 
    322 	/* resolv */
    323 	r = getnameinfo((struct sockaddr *)&sa, salen,
    324 			host, sizeof (host), NULL, 0, 0);
    325 	if (r < 0) {
    326 		warn("getnameinfo");
    327 		snprintf(host, sizeof (host), "unknown");
    328 	}
    329 	printf("host=%s\n", host);
    330 }
    331 
    332 static void
    333 myconnect(concontxt *con)
    334 {
    335 	int r;
    336 
    337 	r = connect(con->confd,
    338 		    con->clinfo->ai_addr,
    339 		    con->clinfo->ai_addrlen);
    340 	if (r < 0)
    341 		warn("connect");
    342 	con->clifd = con->confd;
    343 	con->confd = -1;
    344 }
    345 
    346 int
    347 main(int argc, char *argv[])
    348 {
    349 	int c;
    350 	int highsock, readsocks;
    351 	struct timeval timeout;
    352 	char *host, *port;
    353 	concontxt *con;
    354 
    355 	while ((c = getopt(argc, argv, "ht")) != -1) {
    356 		switch (c) {
    357 		case 'h':
    358 			usage(1);
    359 			break;
    360 		case 't':
    361 			tflag = 1;
    362 			break;
    363 		default:
    364 			usage(1);
    365 		}
    366 	}
    367 	argc -= optind;
    368 	argv += optind;
    369 
    370 	if (argc != 2)
    371 		usage(1);
    372 	host = argv[0];
    373 	port = argv[1];
    374 
    375 	con = malloc(sizeof (concontxt));
    376 	if (con == NULL)
    377 		err(1, "malloc");
    378 	/* defaults */
    379 	con->wfd = STDIN_FILENO;
    380 	con->lfd = STDOUT_FILENO;
    381 	con->contype = (tflag || sflag) ? TCPCON : UDPCON;
    382 	con->clifd = -1;
    383 
    384 	prepare_socket(con, host, port, 1);
    385 
    386 	/* file descriptors for reading are stdin and confd */
    387 	while (1) {
    388 		timeout.tv_sec = 1;
    389 		timeout.tv_usec = 0;
    390 		FD_ZERO(&(con->lset));
    391 		FD_SET(con->wfd, &con->lset);
    392 		highsock = con->wfd;
    393 		if (con->confd != -1) {
    394 			FD_SET(con->confd, &con->lset);
    395 			highsock = con->confd;
    396 		}
    397 		if (con->clifd != -1) {
    398 			FD_SET(con->clifd, &con->lset);
    399 			highsock = con->clifd;
    400 		}
    401 		if (qflag)
    402 			goto freex;
    403 		readsocks = select(highsock + 1, &con->lset, NULL,
    404 				   NULL, &timeout);
    405 		if (readsocks < 0) {
    406 			warn("select");
    407 			goto freex;
    408 		}
    409 		/* change this to == 0 and put default printout if necessary */
    410 		if (readsocks != 0) {
    411 			if (FD_ISSET(con->wfd, &con->lset)) {
    412 				if (con->clifd == -1
    413 				    && con->contype != UDPCON) {
    414 					close(con->confd);
    415 					prepare_socket(con, host, port, 0);
    416 					myconnect(con);
    417 				}
    418 				readstdin(con);
    419 				writesock(con);
    420 			}
    421 			if (con->confd != -1 &&
    422 			    FD_ISSET(con->confd, &con->lset)) {
    423 				if (con->contype == UDPCON) {
    424 					readsock(con);
    425 					writestdout(con);
    426 				} else {
    427 					myaccept(con);
    428 				}
    429 			}
    430 			if (con->clifd != -1 &&
    431 			    FD_ISSET(con->clifd, &con->lset)) {
    432 				readsock(con);
    433 				writestdout(con);
    434 			}
    435 		}
    436 	}
    437 
    438 freex:
    439 	printf("exiting\n");
    440 	if (con != NULL) {
    441 		close(con->confd);
    442 		close(con->clifd);
    443 		free(con);
    444 	}
    445 	return 0;
    446 }