tools/tools/netrate/udpecho/udpecho.c

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