tools/tools/netrate/udpecho/udpecho.c

   1 #include <sys/types.h>
   2 #include <sys/socket.h>
   3 #include <sys/sysctl.h>
   4
   5 #include <arpa/inet.h>
   6 #include <netinet/in.h>
   7
   8 #include <err.h>
   9 #include <stdio.h>
  10 #include <stdint.h>
  11 #include <stdlib.h>
  12 #include <string.h>
  13 #include <unistd.h>
  14
  15 #define RCVBUF_SIZE     (256 * 1024)
  16 #define BUFLEN          2048
  17
  18 static void     mainloop(struct sockaddr_in *, int, int, int);
  19
  20 static void
  21 usage(const char *cmd)
  22 {
  23         fprintf(stderr, "%s -4 addr4 -p port [-i ninst] [-r rcvbuf] [-N]\n",
  24             cmd);
  25         exit(2);
  26 }
  27
  28 int
  29 main(int argc, char *argv[])
  30 {
  31         struct sockaddr_in in;
  32         int opt, ninst, i, s, rcvbuf, noreply;
  33         size_t prm_len;
  34
  35         prm_len = sizeof(ninst);
  36         if (sysctlbyname("hw.ncpu", &ninst, &prm_len, NULL, 0) != 0)
  37                 err(2, "sysctl hw.ncpu failed");
  38
  39         memset(&in, 0, sizeof(in));
  40         in.sin_family = AF_INET;
  41
  42         noreply = 0;
  43         rcvbuf = RCVBUF_SIZE;
  44
  45         while ((opt = getopt(argc, argv, "4:p:i:r:N")) != -1) {
  46                 switch (opt) {
  47                 case '4':
  48                         if (inet_pton(AF_INET, optarg, &in.sin_addr) <= 0)
  49                                 errx(1, "invalid -4 %s\n", optarg);
  50                         break;
  51
  52                 case 'p':
  53                         in.sin_port = strtoul(optarg, NULL, 10);
  54                         in.sin_port = htons(in.sin_port);
  55                         break;
  56
  57                 case 'i':
  58                         ninst = strtoul(optarg, NULL, 10);
  59                         break;
  60
  61                 case 'r':
  62                         rcvbuf = strtoul(optarg, NULL, 10);
  63                         rcvbuf *= 1024;
  64                         break;
  65
  66                 case 'N':
  67                         noreply = 1;
  68                         break;
  69
  70                 default:
  71                         usage(argv[0]);
  72                 }
  73         }
  74
  75         if (in.sin_port == 0 || in.sin_addr.s_addr == INADDR_ANY)
  76                 usage(argv[0]);
  77
  78         s = -1;
  79
  80         for (i = 0; i < ninst - 1; ++i) {
  81                 pid_t pid;
  82
  83                 pid = fork();
  84                 if (pid == 0)
  85                         mainloop(&in, s, rcvbuf, noreply);
  86                 else if (pid < 0)
  87                         err(1, "fork %d failed", i);
  88         }
  89
  90         mainloop(&in, s, rcvbuf, noreply);
  91         exit(0);
  92 }
  93
  94 static void
  95 mainloop(struct sockaddr_in *in, int s, int rcvbuf, int noreply)
  96 {
  97         int on;
  98         void *buf;
  99
 100         if (s < 0) {
 101                 s = socket(AF_INET, SOCK_DGRAM, 0);
 102                 if (s < 0)
 103                         err(1, "socket(INET, DGRAM) failed");
 104         }
 105
 106         buf = malloc(BUFLEN);
 107         if (buf == NULL)
 108                 err(1, "malloc buf failed");
 109
 110         on = 1;
 111         if (setsockopt(s, SOL_SOCKET, SO_REUSEPORT, &on, sizeof(on)) < 0)
 112                 err(1, "setsockopt(SOCK, REUSEPORT) failed");
 113
 114         if (setsockopt(s, SOL_SOCKET, SO_RCVBUF, &rcvbuf, sizeof(rcvbuf)) < 0)
 115                 err(1, "setsockopt(SOCK, RCVBUF) failed");
 116
 117         if (bind(s, (const struct sockaddr *)in, sizeof(*in)) < 0)
 118                 err(1, "bind failed");
 119
 120         for (;;) {
 121                 struct sockaddr_in cli;
 122                 socklen_t cli_len;
 123                 int n;
 124
 125                 cli_len = sizeof(cli);
 126                 n = recvfrom(s, buf, BUFLEN, 0,
 127                     (struct sockaddr *)&cli, &cli_len);
 128                 if (n > 0 && !noreply) {
 129                         sendto(s, buf, n, 0,
 130                            (const struct sockaddr *)&cli, cli_len);
 131                 }
 132         }
 133         exit(0);
 134 }