1 /* 2 * Copyright (c) 2008 Damien Miller <djm@mindrot.org> 3 * 4 * Permission to use, copy, modify, and distribute this software for any 5 * purpose with or without fee is hereby granted, provided that the above 6 * copyright notice and this permission notice appear in all copies. 7 * 8 * THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES 9 * WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF 10 * MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR 11 * ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES 12 * WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN 13 * ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF 14 * OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE. 15 */ 16 17 #include <sys/types.h> 18 #include <sys/time.h> 19 #include <sys/socket.h> 20 #include <sys/socketvar.h> 21 #include <sys/resource.h> 22 23 #include <net/route.h> 24 25 #include <netinet/in.h> 26 #include <netinet/in_systm.h> 27 #include <netinet/ip.h> 28 #include <netinet/tcp.h> 29 #include <netinet/tcp_timer.h> 30 #include <netinet/tcp_fsm.h> 31 #include <netinet/in_pcb.h> 32 #include <netinet/tcp_var.h> 33 34 #include <arpa/inet.h> 35 36 #include <unistd.h> 37 #include <limits.h> 38 #include <stdlib.h> 39 #include <stdio.h> 40 #include <string.h> 41 #include <errno.h> 42 #include <netdb.h> 43 #include <signal.h> 44 #include <err.h> 45 #include <fcntl.h> 46 #include <poll.h> 47 48 #include <kvm.h> 49 #include <nlist.h> 50 51 #define DEFAULT_PORT "12345" 52 #define DEFAULT_STATS_INTERVAL 1000 /* ms */ 53 #define DEFAULT_BUF 256 * 1024 54 #define MAX_FD 1024 55 56 sig_atomic_t done = 0; 57 sig_atomic_t proc_slice = 0; 58 59 static u_int rdomain; 60 static char **kflag; 61 static size_t Bflag; 62 static int Sflag; 63 static int rflag; 64 static int sflag; 65 static int vflag; 66 67 /* stats for a single connection */ 68 struct statctx { 69 struct timeval t_start, t_last; 70 unsigned long long bytes; 71 u_long tcbaddr; 72 char **kvars; 73 kvm_t *kh; 74 }; 75 76 /* 77 * We account the mainstats here, that is the stats 78 * for all connections, all variables starting with slice 79 * are used to account information for the timeslice 80 * between each output. Peak variables record the highest 81 * between all slices so far. 82 */ 83 static struct { 84 unsigned long long slice_bytes; /* bytes for last slice */ 85 struct timeval t_start; /* when we started counting */ 86 long double peak_mbps; /* peak mbps so far */ 87 int nconns; /* connected clients */ 88 } mainstats; 89 90 /* When adding variables, also add to stats_display() */ 91 static const char *allowed_kvars[] = { 92 "inpcb.inp_flags", 93 "sockb.so_rcv.sb_cc", 94 "sockb.so_rcv.sb_hiwat", 95 "sockb.so_snd.sb_cc", 96 "sockb.so_snd.sb_hiwat", 97 "tcpcb.snd_una", 98 "tcpcb.snd_nxt", 99 "tcpcb.snd_wl1", 100 "tcpcb.snd_wl2", 101 "tcpcb.snd_wnd", 102 "tcpcb.rcv_wnd", 103 "tcpcb.rcv_nxt", 104 "tcpcb.rcv_adv", 105 "tcpcb.snd_max", 106 "tcpcb.snd_cwnd", 107 "tcpcb.snd_ssthresh", 108 "tcpcb.t_rcvtime", 109 "tcpcb.t_rtttime", 110 "tcpcb.t_rtseq", 111 "tcpcb.t_srtt", 112 "tcpcb.t_rttvar", 113 "tcpcb.t_rttmin", 114 "tcpcb.max_sndwnd", 115 "tcpcb.snd_scale", 116 "tcpcb.rcv_scale", 117 "tcpcb.last_ack_sent", 118 NULL 119 }; 120 121 static void 122 exitsighand(int signo) 123 { 124 done = signo; 125 } 126 127 static void 128 alarmhandler(int signo) 129 { 130 proc_slice = 1; 131 signal(signo, alarmhandler); 132 } 133 134 static void __dead 135 usage(void) 136 { 137 fprintf(stderr, 138 "usage: tcpbench -l\n" 139 " tcpbench [-v] [-B buf] [-k kvars] [-n connections] [-p port]\n" 140 " [-r rate] [-S space] [-V rdomain] hostname\n" 141 " tcpbench -s [-v] [-B buf] [-k kvars] [-p port]\n" 142 " [-r rate] [-S space] [-V rdomain]\n"); 143 exit(1); 144 } 145 146 static void 147 saddr_ntop(const struct sockaddr *addr, socklen_t alen, char *buf, size_t len) 148 { 149 char hbuf[NI_MAXHOST], pbuf[NI_MAXSERV]; 150 int herr; 151 152 if ((herr = getnameinfo(addr, alen, hbuf, sizeof(hbuf), 153 pbuf, sizeof(pbuf), NI_NUMERICHOST|NI_NUMERICSERV)) != 0) { 154 if (herr == EAI_SYSTEM) 155 err(1, "getnameinfo"); 156 else 157 errx(1, "getnameinfo: %s", gai_strerror(herr)); 158 } 159 snprintf(buf, len, "[%s]:%s", hbuf, pbuf); 160 } 161 162 static void 163 set_timer(int toggle) 164 { 165 struct itimerval itv; 166 167 if (rflag <= 0) 168 return; 169 170 if (toggle) { 171 itv.it_interval.tv_sec = rflag / 1000; 172 itv.it_interval.tv_usec = (rflag % 1000) * 1000; 173 itv.it_value = itv.it_interval; 174 } 175 else 176 bzero(&itv, sizeof(itv)); 177 178 setitimer(ITIMER_REAL, &itv, NULL); 179 } 180 181 static void 182 print_header(void) 183 { 184 char **kv; 185 186 printf("%12s %14s %12s %8s ", "elapsed_ms", "bytes", "mbps", 187 "bwidth"); 188 189 for (kv = kflag; kflag != NULL && *kv != NULL; kv++) 190 printf("%s%s", kv != kflag ? "," : "", *kv); 191 192 printf("\n"); 193 } 194 195 static void 196 kget(kvm_t *kh, u_long addr, void *buf, int size) 197 { 198 if (kvm_read(kh, addr, buf, size) != size) 199 errx(1, "kvm_read: %s", kvm_geterr(kh)); 200 } 201 202 static u_long 203 kfind_tcb(kvm_t *kh, u_long ktcbtab, int sock) 204 { 205 struct inpcbtable tcbtab; 206 struct inpcb *head, *next, *prev; 207 struct inpcb inpcb; 208 struct tcpcb tcpcb; 209 210 struct sockaddr_storage me, them; 211 socklen_t melen, themlen; 212 struct sockaddr_in *in4; 213 struct sockaddr_in6 *in6; 214 char tmp1[64], tmp2[64]; 215 int nretry; 216 217 nretry = 10; 218 melen = themlen = sizeof(struct sockaddr_storage); 219 if (getsockname(sock, (struct sockaddr *)&me, &melen) == -1) 220 err(1, "getsockname"); 221 if (getpeername(sock, (struct sockaddr *)&them, &themlen) == -1) 222 err(1, "getpeername"); 223 if (me.ss_family != them.ss_family) 224 errx(1, "%s: me.ss_family != them.ss_family", __func__); 225 if (me.ss_family != AF_INET && me.ss_family != AF_INET6) 226 errx(1, "%s: unknown socket family", __func__); 227 if (vflag >= 2) { 228 saddr_ntop((struct sockaddr *)&me, me.ss_len, 229 tmp1, sizeof(tmp1)); 230 saddr_ntop((struct sockaddr *)&them, them.ss_len, 231 tmp2, sizeof(tmp2)); 232 fprintf(stderr, "Our socket local %s remote %s\n", tmp1, tmp2); 233 } 234 if (vflag >= 2) 235 fprintf(stderr, "Using PCB table at %lu\n", ktcbtab); 236 retry: 237 kget(kh, ktcbtab, &tcbtab, sizeof(tcbtab)); 238 prev = head = (struct inpcb *)&CIRCLEQ_FIRST( 239 &((struct inpcbtable *)ktcbtab)->inpt_queue); 240 next = CIRCLEQ_FIRST(&tcbtab.inpt_queue); 241 242 if (vflag >= 2) 243 fprintf(stderr, "PCB head at %p\n", head); 244 while (next != head) { 245 if (vflag >= 2) 246 fprintf(stderr, "Checking PCB %p\n", next); 247 kget(kh, (u_long)next, &inpcb, sizeof(inpcb)); 248 if (CIRCLEQ_PREV(&inpcb, inp_queue) != prev) { 249 if (nretry--) { 250 warnx("pcb prev pointer insane"); 251 goto retry; 252 } 253 else 254 errx(1, "pcb prev pointer insane," 255 " all attempts exausted"); 256 } 257 prev = next; 258 next = CIRCLEQ_NEXT(&inpcb, inp_queue); 259 260 if (me.ss_family == AF_INET) { 261 if ((inpcb.inp_flags & INP_IPV6) != 0) { 262 if (vflag >= 2) 263 fprintf(stderr, "Skip: INP_IPV6"); 264 continue; 265 } 266 if (vflag >= 2) { 267 inet_ntop(AF_INET, &inpcb.inp_laddr, 268 tmp1, sizeof(tmp1)); 269 inet_ntop(AF_INET, &inpcb.inp_faddr, 270 tmp2, sizeof(tmp2)); 271 fprintf(stderr, "PCB %p local: [%s]:%d " 272 "remote: [%s]:%d\n", prev, 273 tmp1, inpcb.inp_lport, 274 tmp2, inpcb.inp_fport); 275 } 276 in4 = (struct sockaddr_in *)&me; 277 if (memcmp(&in4->sin_addr, &inpcb.inp_laddr, 278 sizeof(struct in_addr)) != 0 || 279 in4->sin_port != inpcb.inp_lport) 280 continue; 281 in4 = (struct sockaddr_in *)&them; 282 if (memcmp(&in4->sin_addr, &inpcb.inp_faddr, 283 sizeof(struct in_addr)) != 0 || 284 in4->sin_port != inpcb.inp_fport) 285 continue; 286 } else { 287 if ((inpcb.inp_flags & INP_IPV6) == 0) 288 continue; 289 if (vflag >= 2) { 290 inet_ntop(AF_INET6, &inpcb.inp_laddr6, 291 tmp1, sizeof(tmp1)); 292 inet_ntop(AF_INET6, &inpcb.inp_faddr6, 293 tmp2, sizeof(tmp2)); 294 fprintf(stderr, "PCB %p local: [%s]:%d " 295 "remote: [%s]:%d\n", prev, 296 tmp1, inpcb.inp_lport, 297 tmp2, inpcb.inp_fport); 298 } 299 in6 = (struct sockaddr_in6 *)&me; 300 if (memcmp(&in6->sin6_addr, &inpcb.inp_laddr6, 301 sizeof(struct in6_addr)) != 0 || 302 in6->sin6_port != inpcb.inp_lport) 303 continue; 304 in6 = (struct sockaddr_in6 *)&them; 305 if (memcmp(&in6->sin6_addr, &inpcb.inp_faddr6, 306 sizeof(struct in6_addr)) != 0 || 307 in6->sin6_port != inpcb.inp_fport) 308 continue; 309 } 310 kget(kh, (u_long)inpcb.inp_ppcb, &tcpcb, sizeof(tcpcb)); 311 if (tcpcb.t_state != TCPS_ESTABLISHED) { 312 if (vflag >= 2) 313 fprintf(stderr, "Not established\n"); 314 continue; 315 } 316 if (vflag >= 2) 317 fprintf(stderr, "Found PCB at %p\n", prev); 318 return (u_long)prev; 319 } 320 321 errx(1, "No matching PCB found"); 322 } 323 324 static void 325 kupdate_stats(kvm_t *kh, u_long tcbaddr, 326 struct inpcb *inpcb, struct tcpcb *tcpcb, struct socket *sockb) 327 { 328 kget(kh, tcbaddr, inpcb, sizeof(*inpcb)); 329 kget(kh, (u_long)inpcb->inp_ppcb, tcpcb, sizeof(*tcpcb)); 330 kget(kh, (u_long)inpcb->inp_socket, sockb, sizeof(*sockb)); 331 } 332 333 static void 334 check_kvar(const char *var) 335 { 336 size_t i; 337 338 for (i = 0; allowed_kvars[i] != NULL; i++) 339 if (strcmp(allowed_kvars[i], var) == 0) 340 return; 341 errx(1, "Unrecognised kvar: %s", var); 342 } 343 344 static void 345 list_kvars(void) 346 { 347 size_t i; 348 349 fprintf(stderr, "Supported kernel variables:\n"); 350 for (i = 0; allowed_kvars[i] != NULL; i++) 351 fprintf(stderr, "\t%s\n", allowed_kvars[i]); 352 } 353 354 static char ** 355 check_prepare_kvars(char *list) 356 { 357 char *item, **ret = NULL; 358 size_t n = 0; 359 360 while ((item = strsep(&list, ", \t\n")) != NULL) { 361 check_kvar(item); 362 if ((ret = realloc(ret, sizeof(*ret) * (++n + 1))) == NULL) 363 errx(1, "realloc(kvars)"); 364 if ((ret[n - 1] = strdup(item)) == NULL) 365 errx(1, "strdup"); 366 ret[n] = NULL; 367 } 368 return ret; 369 } 370 371 static void 372 stats_prepare(struct statctx *sc, int fd, kvm_t *kh, u_long ktcbtab) 373 { 374 if (rflag <= 0) 375 return; 376 sc->kh = kh; 377 sc->kvars = kflag; 378 if (kflag) 379 sc->tcbaddr = kfind_tcb(kh, ktcbtab, fd); 380 if (gettimeofday(&sc->t_start, NULL) == -1) 381 err(1, "gettimeofday"); 382 sc->t_last = sc->t_start; 383 sc->bytes = 0; 384 } 385 386 static void 387 stats_update(struct statctx *sc, ssize_t n) 388 { 389 sc->bytes += n; 390 mainstats.slice_bytes += n; 391 } 392 393 static void 394 stats_cleanslice(void) 395 { 396 mainstats.slice_bytes = 0; 397 } 398 399 static void 400 stats_display(unsigned long long total_elapsed, long double mbps, 401 float bwperc, struct statctx *sc, struct inpcb *inpcb, 402 struct tcpcb *tcpcb, struct socket *sockb) 403 { 404 int j; 405 406 printf("%12llu %14llu %12.3Lf %7.2f%% ", total_elapsed, sc->bytes, 407 mbps, bwperc); 408 409 if (sc->kvars != NULL) { 410 kupdate_stats(sc->kh, sc->tcbaddr, inpcb, tcpcb, 411 sockb); 412 413 for (j = 0; sc->kvars[j] != NULL; j++) { 414 #define S(a) #a 415 #define P(b, v, f) \ 416 if (strcmp(sc->kvars[j], S(b.v)) == 0) { \ 417 printf("%s"f, j > 0 ? "," : "", b->v); \ 418 continue; \ 419 } 420 P(inpcb, inp_flags, "0x%08x") 421 P(sockb, so_rcv.sb_cc, "%lu") 422 P(sockb, so_rcv.sb_hiwat, "%lu") 423 P(sockb, so_snd.sb_cc, "%lu") 424 P(sockb, so_snd.sb_hiwat, "%lu") 425 P(tcpcb, snd_una, "%u") 426 P(tcpcb, snd_nxt, "%u") 427 P(tcpcb, snd_wl1, "%u") 428 P(tcpcb, snd_wl2, "%u") 429 P(tcpcb, snd_wnd, "%lu") 430 P(tcpcb, rcv_wnd, "%lu") 431 P(tcpcb, rcv_nxt, "%u") 432 P(tcpcb, rcv_adv, "%u") 433 P(tcpcb, snd_max, "%u") 434 P(tcpcb, snd_cwnd, "%lu") 435 P(tcpcb, snd_ssthresh, "%lu") 436 P(tcpcb, t_rcvtime, "%u") 437 P(tcpcb, t_rtttime, "%u") 438 P(tcpcb, t_rtseq, "%u") 439 P(tcpcb, t_srtt, "%hu") 440 P(tcpcb, t_rttvar, "%hu") 441 P(tcpcb, t_rttmin, "%hu") 442 P(tcpcb, max_sndwnd, "%lu") 443 P(tcpcb, snd_scale, "%u") 444 P(tcpcb, rcv_scale, "%u") 445 P(tcpcb, last_ack_sent, "%u") 446 #undef S 447 #undef P 448 } 449 } 450 printf("\n"); 451 } 452 453 static void 454 mainstats_display(long double slice_mbps, long double avg_mbps) 455 { 456 printf("Conn: %3d Mbps: %12.3Lf Peak Mbps: %12.3Lf Avg Mbps: %12.3Lf\n", 457 mainstats.nconns, slice_mbps, mainstats.peak_mbps, avg_mbps); 458 } 459 460 static void 461 process_slice(struct statctx *sc, size_t nsc) 462 { 463 unsigned long long total_elapsed, since_last; 464 long double mbps, slice_mbps = 0; 465 float bwperc; 466 nfds_t i; 467 struct timeval t_cur, t_diff; 468 struct inpcb inpcb; 469 struct tcpcb tcpcb; 470 struct socket sockb; 471 472 for (i = 0; i < nsc; i++, sc++) { 473 if (gettimeofday(&t_cur, NULL) == -1) 474 err(1, "gettimeofday"); 475 if (sc->kvars != NULL) /* process kernel stats */ 476 kupdate_stats(sc->kh, sc->tcbaddr, &inpcb, &tcpcb, 477 &sockb); 478 timersub(&t_cur, &sc->t_start, &t_diff); 479 total_elapsed = t_diff.tv_sec * 1000 + t_diff.tv_usec / 1000; 480 timersub(&t_cur, &sc->t_last, &t_diff); 481 since_last = t_diff.tv_sec * 1000 + t_diff.tv_usec / 1000; 482 bwperc = (sc->bytes * 100.0) / mainstats.slice_bytes; 483 mbps = (sc->bytes * 8) / (since_last * 1000.0); 484 slice_mbps += mbps; 485 486 stats_display(total_elapsed, mbps, bwperc, sc, 487 &inpcb, &tcpcb, &sockb); 488 489 sc->t_last = t_cur; 490 sc->bytes = 0; 491 492 } 493 494 /* process stats for this slice */ 495 if (slice_mbps > mainstats.peak_mbps) 496 mainstats.peak_mbps = slice_mbps; 497 mainstats_display(slice_mbps, slice_mbps / mainstats.nconns); 498 } 499 500 static int 501 handle_connection(struct statctx *sc, int fd, char *buf, size_t buflen) 502 { 503 ssize_t n; 504 505 again: 506 n = read(fd, buf, buflen); 507 if (n == -1) { 508 if (errno == EINTR) 509 goto again; 510 else if (errno == EWOULDBLOCK) 511 return 0; 512 warn("fd %d read error", fd); 513 514 return -1; 515 } 516 else if (n == 0) { 517 if (vflag) 518 fprintf(stderr, "%8d closed by remote end\n", fd); 519 close(fd); 520 return -1; 521 } 522 if (vflag >= 3) 523 fprintf(stderr, "read: %zd bytes\n", n); 524 525 stats_update(sc, n); 526 return 0; 527 } 528 529 static nfds_t 530 serverbind(struct pollfd *pfd, nfds_t max_nfds, struct addrinfo *aitop) 531 { 532 char tmp[128]; 533 int sock, on = 1; 534 struct addrinfo *ai; 535 nfds_t lnfds; 536 537 lnfds = 0; 538 for (ai = aitop; ai != NULL; ai = ai->ai_next) { 539 if (lnfds == max_nfds) { 540 fprintf(stderr, 541 "maximum number of listening fds reached\n"); 542 break; 543 } 544 saddr_ntop(ai->ai_addr, ai->ai_addrlen, tmp, sizeof(tmp)); 545 if (vflag) 546 fprintf(stderr, "Try to listen on %s\n", tmp); 547 if ((sock = socket(ai->ai_family, ai->ai_socktype, 548 ai->ai_protocol)) == -1) { 549 if (ai->ai_next == NULL) 550 err(1, "socket"); 551 if (vflag) 552 warn("socket"); 553 continue; 554 } 555 if (rdomain && ai->ai_family == AF_INET) { 556 if (setsockopt(sock, IPPROTO_IP, SO_RDOMAIN, 557 &rdomain, sizeof(rdomain)) == -1) 558 err(1, "setsockopt SO_RDOMAIN"); 559 } else if (rdomain) 560 warnx("rdomain only supported on AF_INET"); 561 if (setsockopt(sock, SOL_SOCKET, SO_REUSEADDR, 562 &on, sizeof(on)) == -1) 563 warn("reuse port"); 564 if (bind(sock, ai->ai_addr, ai->ai_addrlen) != 0) { 565 if (ai->ai_next == NULL) 566 err(1, "bind"); 567 if (vflag) 568 warn("bind"); 569 close(sock); 570 continue; 571 } 572 if (Sflag) { 573 if (setsockopt(sock, SOL_SOCKET, SO_RCVBUF, 574 &Sflag, sizeof(Sflag)) == -1) 575 warn("set TCP receive buffer size"); 576 } 577 if (listen(sock, 64) == -1) { 578 if (ai->ai_next == NULL) 579 err(1, "listen"); 580 if (vflag) 581 warn("listen"); 582 close(sock); 583 continue; 584 } 585 if (vflag >= 3) 586 fprintf(stderr, "listening on fd %d\n", sock); 587 lnfds++; 588 pfd[lnfds - 1].fd = sock; 589 pfd[lnfds - 1].events = POLLIN; 590 591 } 592 freeaddrinfo(aitop); 593 if (lnfds == 0) 594 errx(1, "No working listen addresses found"); 595 596 return lnfds; 597 } 598 599 static void 600 set_listening(struct pollfd *pfd, nfds_t lfds, int toggle) { 601 int i; 602 603 for (i = 0; i < (int)lfds; i++) { 604 if (toggle) 605 pfd[i].events = POLLIN; 606 else 607 pfd[i].events = 0; 608 } 609 610 } 611 static void __dead 612 serverloop(kvm_t *kvmh, u_long ktcbtab, struct addrinfo *aitop) 613 { 614 socklen_t sslen; 615 struct pollfd *pfd; 616 char tmp[128], *buf; 617 struct statctx *psc; 618 struct sockaddr_storage ss; 619 nfds_t i, nfds, lfds; 620 size_t nalloc; 621 int r, sock, client_id; 622 623 sslen = sizeof(ss); 624 nalloc = 128; 625 if ((pfd = calloc(sizeof(*pfd), nalloc)) == NULL) 626 err(1, "calloc"); 627 if ((psc = calloc(sizeof(*psc), nalloc)) == NULL) 628 err(1, "calloc"); 629 if ((buf = malloc(Bflag)) == NULL) 630 err(1, "malloc"); 631 lfds = nfds = serverbind(pfd, nalloc - 1, aitop); 632 if (vflag >= 3) 633 fprintf(stderr, "listening on %d fds\n", lfds); 634 if (setpgid(0, 0) == -1) 635 err(1, "setpgid"); 636 637 print_header(); 638 639 client_id = 0; 640 while (!done) { 641 if (proc_slice) { 642 process_slice(psc + lfds, nfds - lfds); 643 stats_cleanslice(); 644 proc_slice = 0; 645 } 646 if (vflag >= 3) 647 fprintf(stderr, "mainstats.nconns = %u\n", 648 mainstats.nconns); 649 if ((r = poll(pfd, nfds, INFTIM)) == -1) { 650 if (errno == EINTR) 651 continue; 652 warn("poll"); 653 break; 654 } 655 656 if (vflag >= 3) 657 fprintf(stderr, "poll: %d\n", r); 658 for (i = 0 ; r > 0 && i < nfds; i++) { 659 if ((pfd[i].revents & POLLIN) == 0) 660 continue; 661 if (pfd[i].fd == -1) 662 errx(1, "pfd insane"); 663 r--; 664 if (vflag >= 3) 665 fprintf(stderr, "fd %d active i = %d\n", 666 pfd[i].fd, i); 667 /* new connection */ 668 if (i < lfds) { 669 if ((sock = accept(pfd[i].fd, 670 (struct sockaddr *)&ss, 671 &sslen)) == -1) { 672 if (errno == EINTR) 673 continue; 674 else if (errno == EMFILE || 675 errno == ENFILE) 676 set_listening(pfd, lfds, 0); 677 warn("accept"); 678 continue; 679 } 680 if ((r = fcntl(sock, F_GETFL, 0)) == -1) 681 err(1, "fcntl(F_GETFL)"); 682 r |= O_NONBLOCK; 683 if (fcntl(sock, F_SETFL, r) == -1) 684 err(1, "fcntl(F_SETFL, O_NONBLOCK)"); 685 saddr_ntop((struct sockaddr *)&ss, sslen, 686 tmp, sizeof(tmp)); 687 if (vflag) 688 fprintf(stderr, 689 "Accepted connection %d from " 690 "%s, fd = %d\n", client_id++, tmp, 691 sock); 692 /* alloc more space if we're full */ 693 if (nfds == nalloc) { 694 nalloc *= 2; 695 if ((pfd = realloc(pfd, 696 sizeof(*pfd) * nalloc)) == NULL) 697 err(1, "realloc"); 698 if ((psc = realloc(psc, 699 sizeof(*psc) * nalloc)) == NULL) 700 err(1, "realloc"); 701 } 702 pfd[nfds].fd = sock; 703 pfd[nfds].events = POLLIN; 704 stats_prepare(&psc[nfds], sock, kvmh, ktcbtab); 705 nfds++; 706 if (!mainstats.nconns++) 707 set_timer(1); 708 continue; 709 } 710 /* event in fd */ 711 if (vflag >= 3) 712 fprintf(stderr, 713 "fd %d active", pfd[i].fd); 714 while (handle_connection(&psc[i], pfd[i].fd, 715 buf, Bflag) == -1) { 716 pfd[i] = pfd[nfds - 1]; 717 pfd[nfds - 1].fd = -1; 718 psc[i] = psc[nfds - 1]; 719 mainstats.nconns--; 720 nfds--; 721 /* stop display if no clients */ 722 if (!mainstats.nconns) { 723 proc_slice = 1; 724 set_timer(0); 725 } 726 /* if we were full */ 727 set_listening(pfd, lfds, 1); 728 729 /* is there an event pending on the last fd? */ 730 if (pfd[i].fd == -1 || 731 (pfd[i].revents & POLLIN) == 0) 732 break; 733 } 734 } 735 } 736 exit(1); 737 } 738 739 void 740 clientconnect(struct addrinfo *aitop, struct pollfd *pfd, int nconn) 741 { 742 char tmp[128]; 743 struct addrinfo *ai; 744 int i, r, sock; 745 746 for (i = 0; i < nconn; i++) { 747 for (sock = -1, ai = aitop; ai != NULL; ai = ai->ai_next) { 748 saddr_ntop(ai->ai_addr, ai->ai_addrlen, tmp, 749 sizeof(tmp)); 750 if (vflag && i == 0) 751 fprintf(stderr, "Trying %s\n", tmp); 752 if ((sock = socket(ai->ai_family, ai->ai_socktype, 753 ai->ai_protocol)) == -1) { 754 if (ai->ai_next == NULL) 755 err(1, "socket"); 756 if (vflag) 757 warn("socket"); 758 continue; 759 } 760 if (rdomain && ai->ai_family == AF_INET) { 761 if (setsockopt(sock, IPPROTO_IP, SO_RDOMAIN, 762 &rdomain, sizeof(rdomain)) == -1) 763 err(1, "setsockopt SO_RDOMAIN"); 764 } else if (rdomain) 765 warnx("rdomain only supported on AF_INET"); 766 if (Sflag) { 767 if (setsockopt(sock, SOL_SOCKET, SO_SNDBUF, 768 &Sflag, sizeof(Sflag)) == -1) 769 warn("set TCP send buffer size"); 770 } 771 if (connect(sock, ai->ai_addr, ai->ai_addrlen) != 0) { 772 if (ai->ai_next == NULL) 773 err(1, "connect"); 774 if (vflag) 775 warn("connect"); 776 close(sock); 777 sock = -1; 778 continue; 779 } 780 break; 781 } 782 if (sock == -1) 783 errx(1, "No host found"); 784 785 if ((r = fcntl(sock, F_GETFL, 0)) == -1) 786 err(1, "fcntl(F_GETFL)"); 787 r |= O_NONBLOCK; 788 if (fcntl(sock, F_SETFL, r) == -1) 789 err(1, "fcntl(F_SETFL, O_NONBLOCK)"); 790 791 pfd[i].fd = sock; 792 pfd[i].events = POLLOUT; 793 } 794 freeaddrinfo(aitop); 795 796 if (vflag && nconn > 1) 797 fprintf(stderr, "%u connections established\n", nconn); 798 } 799 800 static void __dead 801 clientloop(kvm_t *kvmh, u_long ktcbtab, struct addrinfo *aitop, int nconn) 802 { 803 struct statctx *psc; 804 struct pollfd *pfd; 805 char *buf; 806 int i, sock = -1; 807 ssize_t n; 808 809 if ((pfd = calloc(nconn, sizeof(*pfd))) == NULL) 810 err(1, "clientloop pfd calloc"); 811 if ((psc = calloc(nconn, sizeof(*psc))) == NULL) 812 err(1, "clientloop psc calloc"); 813 814 clientconnect(aitop, pfd, nconn); 815 816 for (i = 0; i < nconn; i++) { 817 stats_prepare(psc + i, sock, kvmh, ktcbtab); 818 mainstats.nconns++; 819 } 820 821 if ((buf = malloc(Bflag)) == NULL) 822 err(1, "malloc"); 823 arc4random_buf(buf, Bflag); 824 825 print_header(); 826 set_timer(1); 827 828 while (!done) { 829 if (proc_slice) { 830 process_slice(psc, nconn); 831 stats_cleanslice(); 832 proc_slice = 0; 833 } 834 if (poll(pfd, nconn, INFTIM) == -1) { 835 if (errno == EINTR) 836 continue; 837 err(1, "poll"); 838 } 839 for (i = 0; i < nconn; i++) { 840 if (pfd[i].revents & POLLOUT) { 841 if ((n = write(pfd[i].fd, buf, Bflag)) == -1) { 842 if (errno == EINTR || errno == EAGAIN) 843 continue; 844 err(1, "write"); 845 } 846 if (n == 0) { 847 warnx("Remote end closed connection"); 848 done = -1; 849 break; 850 } 851 if (vflag >= 3) 852 fprintf(stderr, "write: %zd bytes\n", 853 n); 854 stats_update(psc + i, n); 855 } 856 } 857 } 858 859 if (done > 0) 860 warnx("Terminated by signal %d", done); 861 862 free(buf); 863 close(sock); 864 exit(0); 865 } 866 867 static void 868 drop_gid(void) 869 { 870 gid_t gid; 871 872 gid = getgid(); 873 if (setresgid(gid, gid, gid) == -1) 874 err(1, "setresgid"); 875 } 876 877 int 878 main(int argc, char **argv) 879 { 880 extern int optind; 881 extern char *optarg; 882 883 char kerr[_POSIX2_LINE_MAX], *tmp; 884 struct addrinfo *aitop, hints; 885 const char *errstr; 886 kvm_t *kvmh = NULL; 887 struct rlimit rl; 888 int ch, herr; 889 struct nlist nl[] = { { "_tcbtable" }, { "" } }; 890 const char *host = NULL, *port = DEFAULT_PORT; 891 int nconn = 1; 892 893 Bflag = DEFAULT_BUF; 894 Sflag = sflag = vflag = rdomain = 0; 895 kflag = NULL; 896 rflag = DEFAULT_STATS_INTERVAL; 897 898 while ((ch = getopt(argc, argv, "B:hlk:n:p:r:sS:vV:")) != -1) { 899 switch (ch) { 900 case 'l': 901 list_kvars(); 902 exit(0); 903 case 'k': 904 if ((tmp = strdup(optarg)) == NULL) 905 errx(1, "strdup"); 906 kflag = check_prepare_kvars(tmp); 907 free(tmp); 908 break; 909 case 'r': 910 rflag = strtonum(optarg, 0, 60 * 60 * 24 * 1000, 911 &errstr); 912 if (errstr != NULL) 913 errx(1, "statistics interval is %s: %s", 914 errstr, optarg); 915 break; 916 case 'p': 917 port = optarg; 918 break; 919 case 's': 920 sflag = 1; 921 break; 922 case 'S': 923 Sflag = strtonum(optarg, 0, 1024*1024*1024, 924 &errstr); 925 if (errstr != NULL) 926 errx(1, "receive space interval is %s: %s", 927 errstr, optarg); 928 break; 929 case 'B': 930 Bflag = strtonum(optarg, 0, 1024*1024*1024, 931 &errstr); 932 if (errstr != NULL) 933 errx(1, "read/write buffer size is %s: %s", 934 errstr, optarg); 935 break; 936 case 'v': 937 vflag++; 938 break; 939 case 'V': 940 rdomain = (unsigned int)strtonum(optarg, 0, 941 RT_TABLEID_MAX, &errstr); 942 if (errstr) 943 errx(1, "rdomain value is %s: %s", 944 errstr, optarg); 945 break; 946 case 'n': 947 nconn = strtonum(optarg, 0, 65535, &errstr); 948 if (errstr != NULL) 949 errx(1, "number of connections is %s: %s", 950 errstr, optarg); 951 break; 952 case 'h': 953 default: 954 usage(); 955 } 956 } 957 958 argv += optind; 959 argc -= optind; 960 if (argc != (sflag ? 0 : 1)) 961 usage(); 962 963 if (!sflag) 964 host = argv[0]; 965 966 bzero(&hints, sizeof(hints)); 967 hints.ai_socktype = SOCK_STREAM; 968 if (sflag) 969 hints.ai_flags = AI_PASSIVE; 970 if ((herr = getaddrinfo(host, port, &hints, &aitop)) != 0) { 971 if (herr == EAI_SYSTEM) 972 err(1, "getaddrinfo"); 973 else 974 errx(1, "getaddrinfo: %s", gai_strerror(herr)); 975 } 976 977 if (kflag) { 978 if ((kvmh = kvm_openfiles(NULL, NULL, NULL, 979 O_RDONLY, kerr)) == NULL) 980 errx(1, "kvm_open: %s", kerr); 981 drop_gid(); 982 if (kvm_nlist(kvmh, nl) < 0 || nl[0].n_type == 0) 983 errx(1, "kvm: no namelist"); 984 } else 985 drop_gid(); 986 987 signal(SIGINT, exitsighand); 988 signal(SIGTERM, exitsighand); 989 signal(SIGHUP, exitsighand); 990 signal(SIGPIPE, SIG_IGN); 991 signal(SIGALRM, alarmhandler); 992 993 if (getrlimit(RLIMIT_NOFILE, &rl) == -1) 994 err(1, "getrlimit"); 995 if (rl.rlim_cur < MAX_FD) 996 rl.rlim_cur = MAX_FD; 997 if (setrlimit(RLIMIT_NOFILE, &rl)) 998 err(1, "setrlimit"); 999 if (getrlimit(RLIMIT_NOFILE, &rl) == -1) 1000 err(1, "getrlimit"); 1001 1002 if (sflag) 1003 serverloop(kvmh, nl[0].n_value, aitop); 1004 else 1005 clientloop(kvmh, nl[0].n_value, aitop, nconn); 1006 1007 return 0; 1008 } 1009