c_cpp BrickstorOS的快速单文件网络统计实用程序

Posted 2021-05-14

#include <errno.h>
#include <ctype.h>
#include <stdbool.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>

#include <kstat.h>
#include <sys/sysinfo.h>
#include <sys/time.h>
// cc[ flag… ] file… -lkstat [ library…]

#define GET_NAMED(ksp) ((kstat_named_t *)(ksp)->ks_data)
#define MKNAME(ptr, a, b) (asprintf(&ptr, "%s%d", a, b))

#define LABEL(ptr, module, index) asprintf(&ptr, "%s%d", module, index)
typedef struct {
  int instance;
  char module[31];
  hrtime_t snaptime;
  uint64_t inbound_err;
  uint64_t outbound_err;
  uint64_t inbound_pkts;
  uint64_t outbound_pkts;
  uint64_t inbound_bytes;
  uint64_t outbound_bytes;
  struct timeval *t;
} netstat_t;

#define TABLESZ 137
int hash(char *str) {
  int h = 0, a = 127;
  for (; *str != '\0'; str++)
    h = (a * h + *str) % TABLESZ;
  return h;
}
// is_relevant returns true if keys of interest are found in given named kstat
// kstat_named_t knp.
boolean_t is_relevant(kstat_named_t *knp, char **keys) {
  for (size_t i = 0; *keys != NULL; keys++) {
    if (strcmp(knp->name, *keys) == 0) {
      return true;
    }
  }
  return false;
}

// make_statistic copies data from kstats into a netstat_t structure, which
// we store in a table indexed by hashed name of interface which is a
// combination of module and instance, i.e. e1000g + 1 == e1000g1.
netstat_t make_statistic(kstat_t *ksp, char **keys, struct timeval *tstamp) {
  kstat_named_t *fields = GET_NAMED(ksp);
  netstat_t s = {0};
  strncpy(s.module, ksp->ks_module, 31);
  s.instance = ksp->ks_instance;
  s.snaptime = ksp->ks_snaptime;
  for (size_t i = 0; i < ksp->ks_ndata; i++, fields++) {
    if (!is_relevant(fields, keys))
      continue;
    if (strncmp(fields->name, keys[0], 32) == 0) {
      s.inbound_bytes = fields->value.ui64;
    } else if (strncmp(fields->name, keys[1], 32) == 0) {
      s.outbound_bytes = fields->value.ui64;
    } else if (strncmp(fields->name, keys[2], 32) == 0) {
      s.inbound_pkts = fields->value.ui64;
    } else if (strncmp(fields->name, keys[3], 32) == 0) {
      s.outbound_pkts = fields->value.ui64;
    } else if (strncmp(fields->name, keys[4], 32) == 0) {
      s.inbound_err = fields->value.ui64;
    } else if (strncmp(fields->name, keys[5], 32) == 0) {
      s.outbound_err = fields->value.ui64;
    }
    s.t = tstamp;
  }
  return s;
}

// output_delta converts previous and latest stats into a delta as a
// per second rate.
void output_delta(netstat_t *prev, netstat_t *latest) {
  char *ptr;
  LABEL(ptr, prev->module, prev->instance);
  double snap_delta = (latest->snaptime - prev->snaptime) * 1e-9;
  printf("%lu.%lu,%s,%f,%f,%f,%f,%f,%f\n", latest->t->tv_sec,
         latest->t->tv_usec, ptr,
         (double)(latest->inbound_bytes - prev->inbound_bytes) / snap_delta,
         (double)(latest->outbound_bytes - prev->outbound_bytes) / snap_delta,
         (double)(latest->inbound_pkts - prev->inbound_pkts) / snap_delta,
         (double)(latest->outbound_pkts - prev->outbound_pkts) / snap_delta,
         (double)(latest->inbound_err - prev->inbound_err) / snap_delta,
         (double)(latest->outbound_err - prev->outbound_err) / snap_delta);
  free(ptr);
}

// module_matches checks that a module for associated kstat actually matches
// name of module(s) that we passed in as arguments on the command line.
boolean_t module_matches(kstat_t *ksp, char **modules) {
  for (size_t i = 0; *modules != NULL; modules++) {
    if (strcmp(ksp->ks_module, *modules) == 0) {
      return true;
    }
  }
  return false;
}

// make_lower mutates string array in-place, replacing anything upper-case with
// lower-case version of that letter.
char *make_lower(char *w) {
  size_t cnt = 0 ;
  for (; *w != '\0' ; w++ , cnt++) {
    *w = isupper (*w) ? *w - 'A' + 'a' : *w;
  }
  return w-cnt;
}

// trim_suffix removes any numbers that follow name of module. In other words
// if argument is e1000g0, trailing zero (0) is removed with resulting string
// mutated in place to e1000g.
char *trim_suffix(char *w) {
  while (isdigit(w[strlen(w)-1])) {
    w[strlen(w)-1] = '\0';
  }
  return w;
}

int main(int argc, char **argv) {
  if (argc < 2) {
    fprintf(stderr, "At least one network module name is required\n");
    return 1;
  } else {
    // Convert any mixed case or upper-case module names to lower-case
    for (size_t i = 1; i < argc ; i++) {
      trim_suffix(make_lower(argv[i]));
    }
  }

  struct timeval time_now = {0}; // Timestamp to pin onto each delta
  kstat_ctl_t *kc = kstat_open();

  char **modules = (char **)calloc(argc, sizeof(char *));
  for (size_t i = 0; *(modules + i) == NULL; i++) {
    *(modules + i) = (char *)malloc(31 * sizeof(char));
    if (*(modules + i) && argv[i + 1] != NULL) {
      strcpy(modules[i], argv[i + 1]);
    } else {
      modules[i] = NULL;
    }
  }

  if (!kc) {
    perror("what just happened?");
    exit(1);
  }

  char *wanted_names[] = {
      "rbytes64", "obytes64", "ipackets64", "opackets64",
      "ierrors",  "oerrors",  NULL,
  };
  kstat_t *ksp = NULL;
  bool first_run = true, found = false;
  netstat_t *lastp = NULL, *curp = NULL, current[TABLESZ], last[TABLESZ];
  char *name;
  do {
    kstat_chain_update(kc);
    gettimeofday(&time_now, NULL);
    for (ksp = kc->kc_chain; ksp != NULL; ksp = ksp->ks_next) {
      kstat_read(kc, ksp, NULL);
      if (strcmp(ksp->ks_class, "net") == 0 &&
          strcmp(ksp->ks_name, "mac") == 0 && module_matches(ksp, modules)) {
        MKNAME(name, ksp->ks_module, ksp->ks_instance);
        // First time around only collect data for future delta calculations
        if (first_run) {
          // If we are here, at least one module matched command line arguments
          if (!found)
            found = true;
          current[hash(name)] = make_statistic(ksp, wanted_names, &time_now);
          continue;
        }
        last[hash(name)] = current[hash(name)];
        lastp = &last[hash(name)];
        current[hash(name)] = make_statistic(ksp, wanted_names, &time_now);
        curp = &current[hash(name)];
        output_delta(lastp, curp);
      }
    }

    // If no arguments matched to valid module names, don't bother printing
    // header, instead inform user that they need o
    if (first_run) {
      if (found) {
        printf("timestamp,interface,inBytes_sec,outBytes_sec,"
               "inPkts_sec,outPkts_sec,inErrors_sec,outErrors_sec\n");
      } else {
        fprintf(stderr,
                "No module names matched this request; check name(s)\n");
        return 1;
      }
    }
    first_run = false;
    fflush(stdout);
  } while (sleep(1) == 0); // sleep, rinse, repeat, eternally ;)

  free(name); // right now it does not matter, never reached
  return kstat_close(kc); // Never reached, so far, due to unbounded loop
}