nfdump_reader_test.cpp

/**
 * \file nfdump_reader_test.cpp
 * \brief Special version of nfdump reader for throughput testing - it reads
 *  whole file into memory before sending flows to the output.
 * \author Vaclav Bartos <ibartosv@fit.vutbr.cz>
 * \date 2013
 * \date 2014
 */
/*
 * Copyright (C) 2013,2014 CESNET
 *
 * LICENSE TERMS
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 * 1. Redistributions of source code must retain the above copyright
 *    notice, this list of conditions and the following disclaimer.
 * 2. Redistributions in binary form must reproduce the above copyright
 *    notice, this list of conditions and the following disclaimer in
 *    the documentation and/or other materials provided with the
 *    distribution.
 * 3. Neither the name of the Company nor the names of its contributors
 *    may be used to endorse or promote products derived from this
 *    software without specific prior written permission.
 *
 * ALTERNATIVELY, provided that this notice is retained in full, this
 * product may be distributed under the terms of the GNU General Public
 * License (GPL) version 2 or later, in which case the provisions
 * of the GPL apply INSTEAD OF those given above.
 *
 * This software is provided ``as is'', and any express or implied
 * warranties, including, but not limited to, the implied warranties of
 * merchantability and fitness for a particular purpose are disclaimed.
 * In no event shall the company or contributors be liable for any
 * direct, indirect, incidental, special, exemplary, or consequential
 * damages (including, but not limited to, procurement of substitute
 * goods or services; loss of use, data, or profits; or business
 * interruption) however caused and on any theory of liability, whether
 * in contract, strict liability, or tort (including negligence or
 * otherwise) arising in any way out of the use of this software, even
 * if advised of the possibility of such damage.
 *
 */

#include <cstdlib>
#include <time.h>

#include <signal.h>
#include <stdio.h>
#include <vector>
#include <iostream>
#include <iomanip>
#include <getopt.h>

#include <libtrap/trap.h>
#include <unirec/unirec.h>

#ifdef __cplusplus
extern "C" {
#endif

#include <libnf.h>

#ifdef __cplusplus
}
#endif

#include "fields.h"

using namespace std;


/* A struct without meaning with size equal to UniRec records generated by this module. */
struct UniRecSpaceholder {
   char x[66];
};

UR_FIELDS (
   ipaddr DST_IP,
   ipaddr SRC_IP,
   uint64 BYTES,
   time TIME_FIRST,
   time TIME_LAST,
   uint32 PACKETS,
   uint16 DST_PORT,
   uint16 SRC_PORT,
   uint8 PROTOCOL,
   uint8 TCP_FLAGS
)


/* Struct with information about module. */
trap_module_info_t *module_info = NULL;

#define MODULE_BASIC_INFO(BASIC) \
  BASIC("Nfdump-reader module","This module reads a given nfdump file and outputs flow records in UniRec format (special version for throughput testing - it reads all records into memory before sending them to TRAP interface).",0,1)

#define MODULE_PARAMS(PARAM)

static int stop = 0;

void signal_handler(int signal)
{
   if (signal == SIGTERM || signal == SIGINT) {
      stop = 1;
      trap_terminate();
   }
}

int main(int argc, char **argv)
{
   int ret;
   lnf_brec1_t brec;
   lnf_rec_t *recp;
   lnf_file_t *filep;
   trap_ifc_spec_t ifc_spec;
   INIT_MODULE_INFO_STRUCT(MODULE_BASIC_INFO, MODULE_PARAMS);
   trap_ifcctl(TRAPIFC_OUTPUT, 0, TRAPCTL_SETTIMEOUT, TRAP_WAIT);

   /* Let TRAP library parse command-line arguments and extract its parameters. */
   ret = trap_parse_params(&argc, argv, &ifc_spec);
   if (ret != TRAP_E_OK) {
      if (ret == TRAP_E_HELP) { /* "-h" was found. */
         trap_print_help(module_info);
         FREE_MODULE_INFO_STRUCT(MODULE_BASIC_INFO, MODULE_PARAMS);
         return 0;
      }
      fprintf(stderr, "ERROR in parsing of parameters for TRAP: %s\n", trap_last_error_msg);
      FREE_MODULE_INFO_STRUCT(MODULE_BASIC_INFO, MODULE_PARAMS);
      return 1;
   }

   if (argc != 2) {
      fprintf(stderr, "Wrong number of parameters.\nUsage: %s -i trap-ifc-specifier nfdump-file\n", argv[0]);
      FREE_MODULE_INFO_STRUCT(MODULE_BASIC_INFO, MODULE_PARAMS);
      trap_finalize();
      return 2;
   }

   ret = lnf_open(&filep, argv[1], LNF_READ, NULL);
   if (ret != LNF_OK) {
      fprintf(stderr, "Error when trying to open file \"%s\"\n", argv[1]);
      trap_finalize();
      FREE_MODULE_INFO_STRUCT(MODULE_BASIC_INFO, MODULE_PARAMS);
      return 3;
   }

   /* Initialize TRAP library (create and init all interfaces). */
   ret = trap_init(module_info, ifc_spec);
   if (ret != TRAP_E_OK) {
      lnf_close(filep);
      fprintf(stderr, "ERROR in TRAP initialization: %s\n", trap_last_error_msg);
      FREE_MODULE_INFO_STRUCT(MODULE_BASIC_INFO, MODULE_PARAMS);
      return 4;
   }
   trap_free_ifc_spec(ifc_spec); /* We don't need ifc_spec anymore. */

   signal(SIGTERM, signal_handler);
   signal(SIGINT, signal_handler);

   /* Create UniRec template. */
   ur_template_t *tmplt = ur_create_output_template(0, "SRC_IP,DST_IP,SRC_PORT,DST_PORT,PROTOCOL,TIME_FIRST,TIME_LAST,PACKETS,BYTES,TCP_FLAGS", NULL);

   vector<UniRecSpaceholder> records;
   unsigned cnt_rec = 0;
   srand(time(NULL));

   lnf_rec_init(&recp);

   cout << "Loading data from file..." << endl;
   while (1) {
      ret = lnf_read(filep, recp);
      if (ret != LNF_OK) {
         if(ret == LNF_EOF) {
            break;
         }
         fprintf(stderr, "Error during reading file (%i).\n", ret);
         lnf_close(filep);
         trap_finalize();
         FREE_MODULE_INFO_STRUCT(MODULE_BASIC_INFO, MODULE_PARAMS);
         return 3;
      }
      /* Allocate new UniRec and put it into records vector. */
      records.push_back(UniRecSpaceholder());
      void *rec2 = (void *) &records.back();

      ++cnt_rec;

      lnf_rec_fget(recp, LNF_FLD_BREC1, &brec);
      if (!IN6_IS_ADDR_V4COMPAT(brec.srcaddr.data)) { /* IPv6 */
         ur_set(tmplt, rec2, F_SRC_IP, ip_from_16_bytes_be((char *) &brec.srcaddr.data));
         ur_set(tmplt, rec2, F_DST_IP, ip_from_16_bytes_be((char *) &brec.dstaddr.data));
      }
      else { /* IPv4 */
         ur_set(tmplt, rec2, F_SRC_IP, ip_from_4_bytes_be((char *) &(brec.srcaddr.data[3])));
         ur_set(tmplt, rec2, F_DST_IP, ip_from_4_bytes_be((char *) &(brec.dstaddr.data[3])));
      }

      ur_set(tmplt, rec2, F_SRC_PORT, brec.srcport);
      ur_set(tmplt, rec2, F_DST_PORT, brec.dstport);
      ur_set(tmplt, rec2, F_PROTOCOL, brec.prot);
      ur_set(tmplt, rec2, F_PACKETS, brec.pkts);
      ur_set(tmplt, rec2, F_BYTES, brec.bytes);
      uint16_t flags;
      lnf_rec_fget(recp, LNF_FLD_TCP_FLAGS, &flags);
      ur_set(tmplt, rec2, F_TCP_FLAGS, flags);

      ur_set(tmplt, rec2, F_TIME_FIRST, ur_time_from_sec_msec(brec.first, 0));
      ur_set(tmplt, rec2, F_TIME_LAST, ur_time_from_sec_msec(brec.last, 0));
      /* Assign value for link and direction of the flow. */
      /*if ((counter % (rand() % 50000 + 50000)) == 0) {
          ur_set(tmplt, rec2, F_LINK_BIT_FIELD, 0x01);
      } else if ((counter % (rand() % 40000 + 1))) {
          ur_set(tmplt, rec2, F_LINK_BIT_FIELD, 0x02);
      } else {
          ur_set(tmplt, rec2, F_LINK_BIT_FIELD, 0x04);
      }

      ur_set(tmplt, rec2, F_DIR_BIT_FIELD, rec->input);
      */

   }

   lnf_close(filep);

   cout << "Sending (" << records.size() << ") records..." << endl;

   /* Read a record from file, convert to UniRec and send to output ifc. */
   for (unsigned int i = 0; i < records.size() && !stop; i++) {
      /* Send data to output interface. */
      trap_send(0, &records[i], sizeof(records[i]));
   }
   cout << "Sending terminating record..." << endl;

   /* Send data with zero length to signalize end. */
   if (!stop)
      trap_send(0, &records[0], 1); /* FIXME: zero-length messages doesn't work, send message of length 1. */

   /* Do all necessary cleanup before exiting. */
   ur_free_template(tmplt);
   trap_finalize();
   FREE_MODULE_INFO_STRUCT(MODULE_BASIC_INFO, MODULE_PARAMS);

   return 0;
}