statsd-bridge.cpp

#include <sys/types.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <unistd.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <netdb.h>
#include <getopt.h>
#include <algorithm>
#include <vector>

static char msg[65535] = {0};

sockaddr_in mktarget(const char* ipaddr, const char* port) {
  sockaddr_in sin;
  memset(&sin, 0, sizeof(sin));
  sin.sin_family = AF_INET;
  sin.sin_port = htons(atoi(port));
  inet_aton(ipaddr, &sin.sin_addr);

  return sin;
}

sockaddr_in mktarget(const char* target) {
  char* p = strchr(target, ':');
  if (p == nullptr) {
    printf("invalid ipaddr:port format for target '%s'.\n", target);
    exit(1);
  }

  char ipaddr[64];
  size_t n = std::min(sizeof(ipaddr), static_cast<size_t>(p - target));
  memcpy(ipaddr, target, n);
  ipaddr[n] = '\0';
  return mktarget(ipaddr, p + 1);
}

void help(const char* errmsg) {
  if (errmsg) {
    fprintf(stderr, "Error: %s\n\n", errmsg);
  }

  fprintf(stderr,
      "statsd-bridge -b BIND_IP:PORT [-t TARGET:PORT]*\n"
      "\n"
      " -b,--bind=IP:PORT    Binds UDP listener to given IP:PORT [0.0.0.0:8125]\n"
      " -t,--target=IP:PORT  Adds given UDP target to proxy to (repeat me)\n"
      "\n");

  exit(1);
}

int main(int argc, char* argv[]) {
  sockaddr_in saddr = mktarget("0.0.0.0:8125");
  std::vector<sockaddr_in> targets;

  static struct option opts[] = {
    {"bind", required_argument, 0, 'b'},
    {"target", optional_argument, 0, 't'},
    {"help", no_argument, 0, 'h'},
    {0, 0, 0, 0}
  };

  for (;;) {
    int c = getopt_long(argc, argv, "b:p:t:h", opts, nullptr);
    if (c < 0) {
      if (optind == argc)
        break;

      help("superfluous arguments");
      return 1;
    }

    switch (c) {
      case 'b':
        saddr = mktarget(optarg);
        break;
      case 't':
        targets.push_back(mktarget(optarg));
        break;
      case 'h':
        help(nullptr);
        return 0;
      default:
        help("Invalid argument");
        return 1;
    }
  }

  if (targets.empty()) {
    help("No targets given.");
    return 1;
  }

  int sfd = socket(AF_INET, SOCK_DGRAM, 0);
  if (sfd < 0) {
    perror("socket");
    return 1;
  }

  if (bind(sfd, (sockaddr *) &saddr, sizeof(saddr)) < 0) {
    perror("bind");
    return 1;
  }

  for (;;) {
    sockaddr_in caddr;
    socklen_t len = sizeof(caddr);
    int n = recvfrom(sfd, msg, sizeof(msg) - 1, 0, (sockaddr *) &caddr, &len);

    for (const auto& target: targets) {
      sendto(sfd, msg, n, MSG_NOSIGNAL,
             (struct sockaddr *)&target, sizeof(target));
    }
  }

  close(sfd);
  return 0;
}