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 22 #include <net/route.h> 23 24 #include <netinet/in.h> 25 #include <netinet/in_systm.h> 26 #include <netinet/ip.h> 27 #include <netinet/tcp.h> 28 #include <netinet/tcp_timer.h> 29 #include <netinet/tcp_fsm.h> 30 #include <netinet/in_pcb.h> 31 #include <netinet/tcp_var.h> 32 33 #include <arpa/inet.h> 34 35 #include <unistd.h> 36 #include <limits.h> 37 #include <stdlib.h> 38 #include <stdio.h> 39 #include <string.h> 40 #include <errno.h> 41 #include <netdb.h> 42 #include <signal.h> 43 #include <err.h> 44 #include <fcntl.h> 45 #include <poll.h> 46 47 #include <kvm.h> 48 #include <nlist.h> 49 50 #define DEFAULT_PORT "12345" 51 #define DEFAULT_STATS_INTERVAL 1000 /* ms */ 52 #define DEFAULT_BUF 256 * 1024 53 54 sig_atomic_t done = 0; 55 sig_atomic_t print_stats = 0; 56 57 struct statctx { 58 struct timeval t_start, t_last, t_cur; 59 unsigned long long bytes; 60 pid_t pid; 61 u_long tcbaddr; 62 kvm_t *kh; 63 char **kvars; 64 }; 65 66 /* When adding variables, also add to stats_display() */ 67 static const char *allowed_kvars[] = { 68 "inpcb.inp_flags", 69 "sockb.so_rcv.sb_cc", 70 "sockb.so_rcv.sb_hiwat", 71 "sockb.so_snd.sb_cc", 72 "sockb.so_snd.sb_hiwat", 73 "tcpcb.snd_una", 74 "tcpcb.snd_nxt", 75 "tcpcb.snd_wl1", 76 "tcpcb.snd_wl2", 77 "tcpcb.snd_wnd", 78 "tcpcb.rcv_wnd", 79 "tcpcb.rcv_nxt", 80 "tcpcb.rcv_adv", 81 "tcpcb.snd_max", 82 "tcpcb.snd_cwnd", 83 "tcpcb.snd_ssthresh", 84 "tcpcb.t_rcvtime", 85 "tcpcb.t_rtttime", 86 "tcpcb.t_rtseq", 87 "tcpcb.t_srtt", 88 "tcpcb.t_rttvar", 89 "tcpcb.t_rttmin", 90 "tcpcb.max_sndwnd", 91 "tcpcb.snd_scale", 92 "tcpcb.rcv_scale", 93 "tcpcb.last_ack_sent", 94 NULL 95 }; 96 97 static void 98 exitsighand(int signo) 99 { 100 done = signo; 101 } 102 103 static void 104 alarmhandler(int signo) 105 { 106 print_stats = 1; 107 signal(signo, alarmhandler); 108 } 109 110 static void __dead 111 usage(void) 112 { 113 fprintf(stderr, 114 "usage: tcpbench -l\n" 115 " tcpbench [-v] [-B buf] [-k kvars] [-n connections]" 116 " [-p port] [-r rate]\n" 117 " [-S space] hostname\n" 118 " tcpbench -s [-v] [-B buf] [-k kvars] [-p port] [-r rate]" 119 " [-S space]\n"); 120 exit(1); 121 } 122 123 static void 124 saddr_ntop(const struct sockaddr *addr, socklen_t alen, char *buf, size_t len) 125 { 126 char hbuf[NI_MAXHOST], pbuf[NI_MAXSERV]; 127 int herr; 128 129 if ((herr = getnameinfo(addr, alen, hbuf, sizeof(hbuf), 130 pbuf, sizeof(pbuf), NI_NUMERICHOST|NI_NUMERICSERV)) != 0) { 131 if (herr == EAI_SYSTEM) 132 err(1, "getnameinfo"); 133 else 134 errx(1, "getnameinfo: %s", gai_strerror(herr)); 135 } 136 snprintf(buf, len, "[%s]:%s", hbuf, pbuf); 137 } 138 139 static void 140 kget(kvm_t *kh, u_long addr, void *buf, int size) 141 { 142 if (kvm_read(kh, addr, buf, size) != size) 143 errx(1, "kvm_read: %s", kvm_geterr(kh)); 144 } 145 146 static u_long 147 kfind_tcb(kvm_t *kh, u_long ktcbtab, int sock, int vflag) 148 { 149 struct inpcbtable tcbtab; 150 struct inpcb *head, *next, *prev; 151 struct inpcb inpcb; 152 struct tcpcb tcpcb; 153 154 struct sockaddr_storage me, them; 155 socklen_t melen, themlen; 156 struct sockaddr_in *in4; 157 struct sockaddr_in6 *in6; 158 char tmp1[64], tmp2[64]; 159 160 melen = themlen = sizeof(struct sockaddr_storage); 161 if (getsockname(sock, (struct sockaddr *)&me, &melen) == -1) 162 err(1, "getsockname"); 163 if (getpeername(sock, (struct sockaddr *)&them, &themlen) == -1) 164 err(1, "getpeername"); 165 if (me.ss_family != them.ss_family) 166 errx(1, "%s: me.ss_family != them.ss_family", __func__); 167 if (me.ss_family != AF_INET && me.ss_family != AF_INET6) 168 errx(1, "%s: unknown socket family", __func__); 169 if (vflag >= 2) { 170 saddr_ntop((struct sockaddr *)&me, me.ss_len, 171 tmp1, sizeof(tmp1)); 172 saddr_ntop((struct sockaddr *)&them, them.ss_len, 173 tmp2, sizeof(tmp2)); 174 fprintf(stderr, "Our socket local %s remote %s\n", tmp1, tmp2); 175 } 176 if (vflag >= 2) 177 fprintf(stderr, "Using PCB table at %lu\n", ktcbtab); 178 179 kget(kh, ktcbtab, &tcbtab, sizeof(tcbtab)); 180 prev = head = (struct inpcb *)&CIRCLEQ_FIRST( 181 &((struct inpcbtable *)ktcbtab)->inpt_queue); 182 next = CIRCLEQ_FIRST(&tcbtab.inpt_queue); 183 184 if (vflag >= 2) 185 fprintf(stderr, "PCB head at %p\n", head); 186 while (next != head) { 187 if (vflag >= 2) 188 fprintf(stderr, "Checking PCB %p\n", next); 189 kget(kh, (u_long)next, &inpcb, sizeof(inpcb)); 190 if (CIRCLEQ_PREV(&inpcb, inp_queue) != prev) 191 errx(1, "pcb prev pointer insane"); 192 prev = next; 193 next = CIRCLEQ_NEXT(&inpcb, inp_queue); 194 195 if (me.ss_family == AF_INET) { 196 if ((inpcb.inp_flags & INP_IPV6) != 0) { 197 if (vflag >= 2) 198 fprintf(stderr, "Skip: INP_IPV6"); 199 continue; 200 } 201 if (vflag >= 2) { 202 inet_ntop(AF_INET, &inpcb.inp_laddr, 203 tmp1, sizeof(tmp1)); 204 inet_ntop(AF_INET, &inpcb.inp_faddr, 205 tmp2, sizeof(tmp2)); 206 fprintf(stderr, "PCB %p local: [%s]:%d " 207 "remote: [%s]:%d\n", prev, 208 tmp1, inpcb.inp_lport, 209 tmp2, inpcb.inp_fport); 210 } 211 in4 = (struct sockaddr_in *)&me; 212 if (memcmp(&in4->sin_addr, &inpcb.inp_laddr, 213 sizeof(struct in_addr)) != 0 || 214 in4->sin_port != inpcb.inp_lport) 215 continue; 216 in4 = (struct sockaddr_in *)&them; 217 if (memcmp(&in4->sin_addr, &inpcb.inp_faddr, 218 sizeof(struct in_addr)) != 0 || 219 in4->sin_port != inpcb.inp_fport) 220 continue; 221 } else { 222 if ((inpcb.inp_flags & INP_IPV6) == 0) 223 continue; 224 if (vflag >= 2) { 225 inet_ntop(AF_INET6, &inpcb.inp_laddr6, 226 tmp1, sizeof(tmp1)); 227 inet_ntop(AF_INET6, &inpcb.inp_faddr6, 228 tmp2, sizeof(tmp2)); 229 fprintf(stderr, "PCB %p local: [%s]:%d " 230 "remote: [%s]:%d\n", prev, 231 tmp1, inpcb.inp_lport, 232 tmp2, inpcb.inp_fport); 233 } 234 in6 = (struct sockaddr_in6 *)&me; 235 if (memcmp(&in6->sin6_addr, &inpcb.inp_laddr6, 236 sizeof(struct in6_addr)) != 0 || 237 in6->sin6_port != inpcb.inp_lport) 238 continue; 239 in6 = (struct sockaddr_in6 *)&them; 240 if (memcmp(&in6->sin6_addr, &inpcb.inp_faddr6, 241 sizeof(struct in6_addr)) != 0 || 242 in6->sin6_port != inpcb.inp_fport) 243 continue; 244 } 245 kget(kh, (u_long)inpcb.inp_ppcb, &tcpcb, sizeof(tcpcb)); 246 if (tcpcb.t_state != TCPS_ESTABLISHED) { 247 if (vflag >= 2) 248 fprintf(stderr, "Not established\n"); 249 continue; 250 } 251 if (vflag >= 2) 252 fprintf(stderr, "Found PCB at %p\n", prev); 253 return (u_long)prev; 254 } 255 256 errx(1, "No matching PCB found"); 257 } 258 259 static void 260 kupdate_stats(kvm_t *kh, u_long tcbaddr, 261 struct inpcb *inpcb, struct tcpcb *tcpcb, struct socket *sockb) 262 { 263 kget(kh, tcbaddr, inpcb, sizeof(*inpcb)); 264 kget(kh, (u_long)inpcb->inp_ppcb, tcpcb, sizeof(*tcpcb)); 265 kget(kh, (u_long)inpcb->inp_socket, sockb, sizeof(*sockb)); 266 } 267 268 static void 269 check_kvar(const char *var) 270 { 271 size_t i; 272 273 for (i = 0; allowed_kvars[i] != NULL; i++) 274 if (strcmp(allowed_kvars[i], var) == 0) 275 return; 276 errx(1, "Unrecognised kvar: %s", var); 277 } 278 279 static void 280 list_kvars(void) 281 { 282 size_t i; 283 284 fprintf(stderr, "Supported kernel variables:\n"); 285 for (i = 0; allowed_kvars[i] != NULL; i++) 286 fprintf(stderr, "\t%s\n", allowed_kvars[i]); 287 } 288 289 static char ** 290 check_prepare_kvars(char *list) 291 { 292 char *item, **ret = NULL; 293 size_t n = 0; 294 295 while ((item = strsep(&list, ", \t\n")) != NULL) { 296 check_kvar(item); 297 if ((ret = realloc(ret, sizeof(*ret) * (++n + 1))) == NULL) 298 errx(1, "realloc(kvars)"); 299 if ((ret[n - 1] = strdup(item)) == NULL) 300 errx(1, "strdup"); 301 ret[n] = NULL; 302 } 303 return ret; 304 } 305 306 static void 307 stats_prepare(struct statctx *sc, int fd, kvm_t *kh, u_long ktcbtab, 308 int rflag, int vflag, char **kflag) 309 { 310 struct itimerval itv; 311 int i; 312 313 if (rflag <= 0) 314 return; 315 sc->kh = kh; 316 sc->kvars = kflag; 317 if (kflag) 318 sc->tcbaddr = kfind_tcb(kh, ktcbtab, fd, vflag); 319 gettimeofday(&sc->t_start, NULL); 320 sc->t_last = sc->t_start; 321 signal(SIGALRM, alarmhandler); 322 itv.it_interval.tv_sec = rflag / 1000; 323 itv.it_interval.tv_usec = (rflag % 1000) * 1000; 324 itv.it_value = itv.it_interval; 325 setitimer(ITIMER_REAL, &itv, NULL); 326 sc->bytes = 0; 327 sc->pid = getpid(); 328 329 printf("%8s %12s %14s %12s ", "pid", "elapsed_ms", "bytes", "Mbps"); 330 if (sc->kvars != NULL) { 331 for (i = 0; sc->kvars[i] != NULL; i++) 332 printf("%s%s", i > 0 ? "," : "", sc->kvars[i]); 333 } 334 printf("\n"); 335 fflush(stdout); 336 } 337 338 static void 339 stats_update(struct statctx *sc, ssize_t n) 340 { 341 sc->bytes += n; 342 } 343 344 static void 345 stats_display(struct statctx *sc) 346 { 347 struct timeval t_diff; 348 unsigned long long total_elapsed, since_last; 349 size_t i; 350 struct inpcb inpcb; 351 struct tcpcb tcpcb; 352 struct socket sockb; 353 354 gettimeofday(&sc->t_cur, NULL); 355 timersub(&sc->t_cur, &sc->t_start, &t_diff); 356 total_elapsed = t_diff.tv_sec * 1000 + t_diff.tv_usec / 1000; 357 timersub(&sc->t_cur, &sc->t_last, &t_diff); 358 since_last = t_diff.tv_sec * 1000 + t_diff.tv_usec / 1000; 359 printf("%8ld %12llu %14llu %12.3Lf ", (long)sc->pid, 360 total_elapsed, sc->bytes, 361 (long double)(sc->bytes * 8) / (since_last * 1000.0)); 362 sc->t_last = sc->t_cur; 363 sc->bytes = 0; 364 365 if (sc->kvars != NULL) { 366 kupdate_stats(sc->kh, sc->tcbaddr, &inpcb, &tcpcb, &sockb); 367 for (i = 0; sc->kvars[i] != NULL; i++) { 368 #define P(v, f) \ 369 if (strcmp(sc->kvars[i], #v) == 0) { \ 370 printf("%s"f, i > 0 ? "," : "", v); \ 371 continue; \ 372 } 373 P(inpcb.inp_flags, "0x%08x") 374 P(sockb.so_rcv.sb_cc, "%lu") 375 P(sockb.so_rcv.sb_hiwat, "%lu") 376 P(sockb.so_snd.sb_cc, "%lu") 377 P(sockb.so_snd.sb_hiwat, "%lu") 378 P(tcpcb.snd_una, "%u") 379 P(tcpcb.snd_nxt, "%u") 380 P(tcpcb.snd_wl1, "%u") 381 P(tcpcb.snd_wl2, "%u") 382 P(tcpcb.snd_wnd, "%lu") 383 P(tcpcb.rcv_wnd, "%lu") 384 P(tcpcb.rcv_nxt, "%u") 385 P(tcpcb.rcv_adv, "%u") 386 P(tcpcb.snd_max, "%u") 387 P(tcpcb.snd_cwnd, "%lu") 388 P(tcpcb.snd_ssthresh, "%lu") 389 P(tcpcb.t_rcvtime, "%u") 390 P(tcpcb.t_rtttime, "%u") 391 P(tcpcb.t_rtseq, "%u") 392 P(tcpcb.t_srtt, "%hu") 393 P(tcpcb.t_rttvar, "%hu") 394 P(tcpcb.t_rttmin, "%hu") 395 P(tcpcb.max_sndwnd, "%lu") 396 P(tcpcb.snd_scale, "%u") 397 P(tcpcb.rcv_scale, "%u") 398 P(tcpcb.last_ack_sent, "%u") 399 #undef P 400 } 401 } 402 printf("\n"); 403 fflush(stdout); 404 } 405 406 static void 407 stats_finish(struct statctx *sc) 408 { 409 struct itimerval itv; 410 411 signal(SIGALRM, SIG_DFL); 412 bzero(&itv, sizeof(itv)); 413 setitimer(ITIMER_REAL, &itv, NULL); 414 } 415 416 static void __dead 417 handle_connection(kvm_t *kvmh, u_long ktcbtab, int sock, int vflag, 418 int rflag, char **kflag, int Bflag) 419 { 420 char *buf; 421 struct pollfd pfd; 422 ssize_t n; 423 int r; 424 struct statctx sc; 425 426 if ((buf = malloc(Bflag)) == NULL) 427 err(1, "malloc"); 428 if ((r = fcntl(sock, F_GETFL, 0)) == -1) 429 err(1, "fcntl(F_GETFL)"); 430 r |= O_NONBLOCK; 431 if (fcntl(sock, F_SETFL, r) == -1) 432 err(1, "fcntl(F_SETFL, O_NONBLOCK)"); 433 434 signal(SIGINT, exitsighand); 435 signal(SIGTERM, exitsighand); 436 signal(SIGHUP, exitsighand); 437 signal(SIGPIPE, SIG_IGN); 438 439 bzero(&pfd, sizeof(pfd)); 440 pfd.fd = sock; 441 pfd.events = POLLIN; 442 443 stats_prepare(&sc, sock, kvmh, ktcbtab, rflag, vflag, kflag); 444 445 while (!done) { 446 if (print_stats) { 447 stats_display(&sc); 448 print_stats = 0; 449 } 450 if (poll(&pfd, 1, INFTIM) == -1) { 451 if (errno == EINTR) 452 continue; 453 err(1, "poll"); 454 } 455 if ((n = read(pfd.fd, buf, Bflag)) == -1) { 456 if (errno == EINTR || errno == EAGAIN) 457 continue; 458 err(1, "read"); 459 } 460 if (n == 0) { 461 fprintf(stderr, "%8ld closed by remote end\n", 462 (long)getpid()); 463 done = -1; 464 break; 465 } 466 if (vflag >= 3) 467 fprintf(stderr, "read: %zd bytes\n", n); 468 stats_update(&sc, n); 469 } 470 stats_finish(&sc); 471 472 free(buf); 473 close(sock); 474 exit(1); 475 } 476 477 static void __dead 478 serverloop(kvm_t *kvmh, u_long ktcbtab, struct addrinfo *aitop, 479 int vflag, int rflag, char **kflag, int Sflag, int Bflag) 480 { 481 char tmp[128]; 482 int r, sock, client_id, on = 1; 483 struct addrinfo *ai; 484 struct pollfd *pfd; 485 struct sockaddr_storage ss; 486 socklen_t sslen; 487 size_t nfds, i, j; 488 489 pfd = NULL; 490 nfds = 0; 491 for (ai = aitop; ai != NULL; ai = ai->ai_next) { 492 saddr_ntop(ai->ai_addr, ai->ai_addrlen, tmp, sizeof(tmp)); 493 if (vflag) 494 fprintf(stderr, "Try listen on %s\n", tmp); 495 if ((sock = socket(ai->ai_family, ai->ai_socktype, 496 ai->ai_protocol)) == -1) { 497 if (ai->ai_next == NULL) 498 err(1, "socket"); 499 if (vflag) 500 warn("socket"); 501 continue; 502 } 503 if (setsockopt(sock, SOL_SOCKET, SO_REUSEADDR, 504 &on, sizeof(on)) == -1) 505 warn("reuse port"); 506 if (bind(sock, ai->ai_addr, ai->ai_addrlen) != 0) { 507 if (ai->ai_next == NULL) 508 err(1, "bind"); 509 if (vflag) 510 warn("bind"); 511 close(sock); 512 continue; 513 } 514 if (Sflag) { 515 if (setsockopt(sock, SOL_SOCKET, SO_RCVBUF, 516 &Sflag, sizeof(Sflag)) == -1) 517 warn("set TCP receive buffer size"); 518 } 519 if (listen(sock, 64) == -1) { 520 if (ai->ai_next == NULL) 521 err(1, "listen"); 522 if (vflag) 523 warn("listen"); 524 close(sock); 525 continue; 526 } 527 if (nfds > 128) 528 break; 529 if ((pfd = realloc(pfd, ++nfds * sizeof(*pfd))) == NULL) 530 errx(1, "realloc(pfd * %zu)", nfds); 531 pfd[nfds - 1].fd = sock; 532 pfd[nfds - 1].events = POLLIN; 533 } 534 freeaddrinfo(aitop); 535 if (nfds == 0) 536 errx(1, "No working listen addresses found"); 537 538 signal(SIGINT, exitsighand); 539 signal(SIGTERM, exitsighand); 540 signal(SIGHUP, exitsighand); 541 signal(SIGPIPE, SIG_IGN); 542 signal(SIGCHLD, SIG_IGN); 543 544 if (setpgid(0, 0) == -1) 545 err(1, "setpgid"); 546 547 client_id = 0; 548 while (!done) { 549 if ((r = poll(pfd, nfds, INFTIM)) == -1) { 550 if (errno == EINTR) 551 continue; 552 warn("poll"); 553 break; 554 } 555 if (vflag >= 3) 556 fprintf(stderr, "poll: %d\n", r); 557 for (i = 0 ; r > 0 && i < nfds; i++) { 558 if ((pfd[i].revents & POLLIN) == 0) 559 continue; 560 if (vflag >= 3) 561 fprintf(stderr, "fd %d active\n", pfd[i].fd); 562 r--; 563 sslen = sizeof(ss); 564 if ((sock = accept(pfd[i].fd, (struct sockaddr *)&ss, 565 &sslen)) == -1) { 566 if (errno == EINTR) 567 continue; 568 warn("accept"); 569 break; 570 } 571 saddr_ntop((struct sockaddr *)&ss, sslen, 572 tmp, sizeof(tmp)); 573 if (vflag) 574 fprintf(stderr, "Accepted connection %d from " 575 "%s, fd = %d\n", client_id++, tmp, sock); 576 switch (fork()) { 577 case -1: 578 warn("fork"); 579 done = -1; 580 break; 581 case 0: 582 for (j = 0; j < nfds; j++) 583 if (j != i) 584 close(pfd[j].fd); 585 handle_connection(kvmh, ktcbtab, sock, 586 vflag, rflag, kflag, Bflag); 587 /* NOTREACHED */ 588 _exit(1); 589 default: 590 close(sock); 591 break; 592 } 593 if (done == -1) 594 break; 595 } 596 } 597 for (i = 0; i < nfds; i++) 598 close(pfd[i].fd); 599 if (done > 0) 600 warnx("Terminated by signal %d", done); 601 signal(SIGTERM, SIG_IGN); 602 killpg(0, SIGTERM); 603 exit(1); 604 } 605 606 static void __dead 607 clientloop(kvm_t *kvmh, u_long ktcbtab, const char *host, const char *port, 608 int vflag, int rflag, char **kflag, int Sflag, int Bflag, int nconn) 609 { 610 char tmp[128]; 611 char *buf; 612 int r, sock, herr; 613 struct addrinfo *aitop, *ai, hints; 614 struct pollfd *pfd; 615 ssize_t n; 616 struct statctx sc; 617 u_int i, scnt = 0; 618 619 if ((buf = malloc(Bflag)) == NULL) 620 err(1, "malloc"); 621 622 if ((pfd = calloc(nconn, sizeof(struct pollfd))) == NULL) 623 err(1, "clientloop pfd calloc"); 624 625 for (i = 0; i < nconn; i++) { 626 bzero(&hints, sizeof(hints)); 627 hints.ai_socktype = SOCK_STREAM; 628 hints.ai_flags = 0; 629 if ((herr = getaddrinfo(host, port, &hints, &aitop)) != 0) { 630 if (herr == EAI_SYSTEM) 631 err(1, "getaddrinfo"); 632 else 633 errx(1, "c getaddrinfo: %s", gai_strerror(herr)); 634 } 635 636 for (sock = -1, ai = aitop; ai != NULL; ai = ai->ai_next) { 637 saddr_ntop(ai->ai_addr, ai->ai_addrlen, tmp, 638 sizeof(tmp)); 639 if (vflag && scnt == 0) 640 fprintf(stderr, "Trying %s\n", tmp); 641 if ((sock = socket(ai->ai_family, ai->ai_socktype, 642 ai->ai_protocol)) == -1) { 643 if (ai->ai_next == NULL) 644 err(1, "socket"); 645 if (vflag) 646 warn("socket"); 647 continue; 648 } 649 if (Sflag) { 650 if (setsockopt(sock, SOL_SOCKET, SO_SNDBUF, 651 &Sflag, sizeof(Sflag)) == -1) 652 warn("set TCP send buffer size"); 653 } 654 if (connect(sock, ai->ai_addr, ai->ai_addrlen) != 0) { 655 if (ai->ai_next == NULL) 656 err(1, "connect"); 657 if (vflag) 658 warn("connect"); 659 close(sock); 660 sock = -1; 661 continue; 662 } 663 break; 664 } 665 freeaddrinfo(aitop); 666 if (sock == -1) 667 errx(1, "No host found"); 668 669 if ((r = fcntl(sock, F_GETFL, 0)) == -1) 670 err(1, "fcntl(F_GETFL)"); 671 r |= O_NONBLOCK; 672 if (fcntl(sock, F_SETFL, r) == -1) 673 err(1, "fcntl(F_SETFL, O_NONBLOCK)"); 674 675 pfd[i].fd = sock; 676 pfd[i].events = POLLOUT; 677 scnt++; 678 } 679 680 if (vflag && scnt > 1) 681 fprintf(stderr, "%u connections established\n", scnt); 682 arc4random_buf(buf, Bflag); 683 684 signal(SIGINT, exitsighand); 685 signal(SIGTERM, exitsighand); 686 signal(SIGHUP, exitsighand); 687 signal(SIGPIPE, SIG_IGN); 688 689 stats_prepare(&sc, sock, kvmh, ktcbtab, rflag, vflag, kflag); 690 691 while (!done) { 692 if (print_stats) { 693 stats_display(&sc); 694 print_stats = 0; 695 } 696 if (poll(pfd, nconn, INFTIM) == -1) { 697 if (errno == EINTR) 698 continue; 699 err(1, "poll"); 700 } 701 for (i = 0; i < nconn; i++) { 702 if (pfd[i].revents & POLLOUT) { 703 if ((n = write(pfd[i].fd, buf, Bflag)) == -1) { 704 if (errno == EINTR || errno == EAGAIN) 705 continue; 706 err(1, "write"); 707 } 708 if (n == 0) { 709 warnx("Remote end closed connection"); 710 done = -1; 711 break; 712 } 713 if (vflag >= 3) 714 fprintf(stderr, "write: %zd bytes\n", 715 n); 716 stats_update(&sc, n); 717 } 718 } 719 } 720 stats_finish(&sc); 721 722 if (done > 0) 723 warnx("Terminated by signal %d", done); 724 725 free(buf); 726 close(sock); 727 exit(0); 728 } 729 730 static void 731 drop_gid(void) 732 { 733 gid_t gid; 734 735 gid = getgid(); 736 if (setresgid(gid, gid, gid) == -1) 737 err(1, "setresgid"); 738 } 739 740 int 741 main(int argc, char **argv) 742 { 743 extern int optind; 744 extern char *optarg; 745 746 char kerr[_POSIX2_LINE_MAX], *tmp; 747 const char *errstr; 748 int ch, herr; 749 struct addrinfo *aitop, hints; 750 kvm_t *kvmh = NULL; 751 752 const char *host = NULL, *port = DEFAULT_PORT; 753 char **kflag = NULL; 754 int sflag = 0, vflag = 0, rflag = DEFAULT_STATS_INTERVAL, Sflag = 0; 755 int Bflag = DEFAULT_BUF; 756 int nconn = 1; 757 758 struct nlist nl[] = { { "_tcbtable" }, { "" } }; 759 760 while ((ch = getopt(argc, argv, "B:hlk:n:p:r:sS:v")) != -1) { 761 switch (ch) { 762 case 'l': 763 list_kvars(); 764 exit(0); 765 case 'k': 766 if ((tmp = strdup(optarg)) == NULL) 767 errx(1, "strdup"); 768 kflag = check_prepare_kvars(tmp); 769 free(tmp); 770 break; 771 case 'r': 772 rflag = strtonum(optarg, 0, 60 * 60 * 24 * 1000, 773 &errstr); 774 if (errstr != NULL) 775 errx(1, "statistics interval is %s: %s", 776 errstr, optarg); 777 break; 778 case 'p': 779 port = optarg; 780 break; 781 case 's': 782 sflag = 1; 783 break; 784 case 'S': 785 Sflag = strtonum(optarg, 0, 1024*1024*1024, 786 &errstr); 787 if (errstr != NULL) 788 errx(1, "receive space interval is %s: %s", 789 errstr, optarg); 790 break; 791 case 'B': 792 Bflag = strtonum(optarg, 0, 1024*1024*1024, 793 &errstr); 794 if (errstr != NULL) 795 errx(1, "read/write buffer size is %s: %s", 796 errstr, optarg); 797 break; 798 case 'v': 799 vflag++; 800 break; 801 case 'n': 802 nconn = strtonum(optarg, 0, 65535, &errstr); 803 if (errstr != NULL) 804 errx(1, "number of connections is %s: %s", 805 errstr, optarg); 806 break; 807 case 'h': 808 default: 809 usage(); 810 } 811 } 812 813 argv += optind; 814 argc -= optind; 815 if (argc != (sflag ? 0 : 1)) 816 usage(); 817 818 if (kflag != NULL && nconn > 1) 819 errx(1, "-k currently only works with a single tcp connection"); 820 821 if (!sflag) 822 host = argv[0]; 823 824 if (sflag) { 825 bzero(&hints, sizeof(hints)); 826 hints.ai_socktype = SOCK_STREAM; 827 hints.ai_flags = AI_PASSIVE; 828 if ((herr = getaddrinfo(host, port, &hints, &aitop)) != 0) { 829 if (herr == EAI_SYSTEM) 830 err(1, "getaddrinfo"); 831 else 832 errx(1, "s getaddrinfo: %s", gai_strerror(herr)); 833 } 834 } 835 836 if (kflag) { 837 if ((kvmh = kvm_openfiles(NULL, NULL, NULL, 838 O_RDONLY, kerr)) == NULL) 839 errx(1, "kvm_open: %s", kerr); 840 drop_gid(); 841 if (kvm_nlist(kvmh, nl) < 0 || nl[0].n_type == 0) 842 errx(1, "kvm: no namelist"); 843 } else 844 drop_gid(); 845 846 if (sflag) 847 serverloop(kvmh, nl[0].n_value, aitop, vflag, rflag, kflag, 848 Sflag, Bflag); 849 else 850 clientloop(kvmh, nl[0].n_value, host, port, vflag, rflag, kflag, 851 Sflag, Bflag, nconn); 852 853 return 0; 854 } 855