/*============================================================================*/
//
// Title:          udp2file.c
// Author:         RH
// Last modified:  13/09/2025
//
//============================================================================
//
// Description:
// ------------
// Tool that takes a file path, ip address, port, and capture size. The tool
// then listens on the supplied address and port for UDP datagrams and captures
// messages in chunks up to the capture size at a time, before then writing
// the binary data to the specified filepath.
//
// Compilation:
// ------------
// gcc -Wall -Wextra ../Logging/logging_lib.c udp2file.c -o udp2file
//
/*============================================================================*/

#include <stdio.h>
#include <errno.h>
#include <stdlib.h>
#include <unistd.h>
#include <string.h>
#include <stdint.h>
#include <signal.h>
#include <netinet/in.h>
#include <sys/types.h>
#include <sys/socket.h>

#include "../Logging/logging_lib.h"

// program, verbose (optional), buffer check (optional), file path, ip address,
// port, capture size
#define MIN_NUM_ARGS 4
#define MAX_NUM_ARGS 6
#define NUM_REQ_ARGS 3

#define MAX_PAYLOAD_SIZE 65507

// ----------------------------------------------------------------------------
// print usage information to the terminal to provide information on how the
// program expects to be run
// ----------------------------------------------------------------------------

void usage (void)
{
   printf("\n");
   printf("---------------------------------------------------------------------------------------------\n");
   printf(" Usage:\n");
   printf("    ./udp2file [options] <file path> <port> <capture size>\n");
   printf("---------------------------------------------------------------------------------------------\n");
   printf(" Required Arguments:\n");
   printf(" <file path>:\n");
   printf(" \tThe path to the file to write the captured binary data to. Will create the file if it\n");
   printf(" \tdoes not already exist and will overwite it if it does.\n");
   printf(" <port>:\n");
   printf(" \tThe port to listen on for the UDP datagrams.\n");
   printf(" <capture size>:\n");
   printf(" \tThe capture size (in bytes) to read from the system UDP buffer. This is the size of\n");
   printf(" \tthe datagram payload, so an MTU of 9000 (jumbo frames) means a capture size of 8972\n");
   printf(" \tbytes. [Range: 1 to 65507.]\n");
   printf("\n");
   printf(" Optional Arguments:\n");
   printf(" -h,   --help\n");
   printf(" \tDisplay this usage information then exit.\n");
   printf(" -v,   --verbose\n");
   printf(" \tVerbose flag to enable additional debug output messages - at high data rate this can\n");
   printf(" \tbe very noisy and so is not recommended for operational use.\n");
   printf(" -b,   --buffer-check\n");
   printf(" \tEnable additional checking of the system UDP buffer that warns if the queue is\n");
   printf(" \tgrowing and reports an error when datagrams are dropped. This is not recommended for\n");
   printf(" \tvery high data rates (>512 mbps).\n");
   printf("---------------------------------------------------------------------------------------------\n");
   printf(" Example:\n");
   printf("    ./udp_receiver ~/udp_captures/catpure_6.bin 8080 8972\n");
   printf("---------------------------------------------------------------------------------------------\n");
}

// ----------------------------------------------------------------------------
// signal handler to allow application to exit cleanly
// ----------------------------------------------------------------------------

static int run = 1;

// the handler just sets 'run' to 0, which allows the main loop to exit
// and the cleanup code after the main loop to run
static void SignalHandler ( int dummy )
{
	dummy++ ;
   LOGMSG_INFO(">>>>>>>>>>>> Signal received, terminating...")
   run = 0;
}

// ----------------------------------------------------------------------------
// main program - parse arguments, open output file, create listener
// socket, then enter a loop waiting for datagrams and writing them to file
// ----------------------------------------------------------------------------

int main(int argc, char **argv)
{
   // initiliase logger to print messages to stdout
   int Verbose = 0;
   LOGGING_INIT_STDOUT

   // signal handlers to ensure program exits cleanly
   (void) signal(SIGINT , SignalHandler) ;
	(void) signal(SIGTERM, SignalHandler) ;

   // -------------------------------------------------------------------------
   // parse input arguments
   // -------------------------------------------------------------------------

   // check for help argument
   if (   (strncmp ( argv[1], "-h",     strlen("-h"))      == 0)
       || (strncmp ( argv[1], "--help", strlen("--help"))  == 0)
      )
   {
      usage();
      return -1;
   }

   if ( (argc < MIN_NUM_ARGS) || (argc > MAX_NUM_ARGS) )
   {
      LOGMSG_ERR("Error: incorrect number of arguments supplied. Received %d, minimum %d -> maximum %d",
                  argc, MIN_NUM_ARGS, MAX_NUM_ARGS)
      usage();
      return -1;
   }

   // check optional arguments (excluding --help)
   int current_arg = 1;
   int buffer_check = 0;
   // check combinations first: -vb, -bv
   if (   (strncmp ( argv[current_arg], "-vb",       strlen("-vb"))      == 0)
       || (strncmp ( argv[current_arg], "-bv",       strlen("-bv"))      == 0)
      )
   {
      Verbose = 1;
      buffer_check = 1;
      current_arg++;
   }

   // check for verbose supplied individually - Verbose == 0 used to prevent
   // '-v' from incrementing current_args if '-vb' above has already been set
   if (   ( Verbose == 0 )
       && (   (strncmp ( argv[current_arg], "-v",        strlen("-v"))        == 0)
           || (strncmp ( argv[current_arg], "--verbose", strlen("--verbose")) == 0)  )
      )
   {
      Verbose = 1;
      current_arg++;
   }

   // check buffer flag supplied individually - buffer_check == 0 used to prevent
   // '-b' from incrementing current_args if '-bv' above has already been set
   if (   ( buffer_check == 0 )
       && (   (strncmp ( argv[current_arg], "-b",             strlen("-b"))             == 0)
           || (strncmp ( argv[current_arg], "--buffer-check", strlen("--buffer-check")) == 0)  )
      )
   {
      buffer_check = 1;
      current_arg++;
   }

   // check required arguments (file name, port, payload size)
   // there should only be 3 arguments left to process
   char      *pFileName   = NULL;
   uint32_t   port        = 0;
   size_t     payloadSize = 0;
   if ( (argc - current_arg) != NUM_REQ_ARGS )
   {
      LOGMSG_ERR("Error parsing input arguments: Received %d, minimum %d -> maximum %d",
                  argc, MIN_NUM_ARGS, MAX_NUM_ARGS)
      usage();
      return -1;
   }
   else
   {
      pFileName   = argv[current_arg++];
      port        = atoi(argv[current_arg++]);
      payloadSize = atoi(argv[current_arg]);
   }

   if ( (payloadSize > MAX_PAYLOAD_SIZE) || (payloadSize == 0) )
   {
      LOGMSG_ERR("Error: supplied payload size [%d] must be greater than 0 and not"
                 " more than the maximum possible UDP datagram payload size [%d]",
                  payloadSize, MAX_PAYLOAD_SIZE)
      usage();
      return -1;
   }

   LOGMSG_INFO("Running udp_receiver | Verbose: %s, Buffer checking: %s, Output"
               " file: %s, Port: %d, Payload size: %d",
               (Verbose == 1) ? "Yes" : "No", (buffer_check == 1) ? "Yes" : "No",
               pFileName, port, payloadSize)

   // -------------------------------------------------------------------------
   // open the file to write the binary data to
   // -------------------------------------------------------------------------
   FILE *pDestFile = NULL;
   pDestFile = fopen(pFileName, "w");
   size_t bytes_written = 0;
   if (pDestFile == NULL)
   {
      LOGMSG_ERR("Error: unable to open the requested file %s. Errno: %s [%d]",
                  pFileName, strerror(errno), errno)
      return -1;
   }
   else
   {
      LOGMSG_DBG("Opened file: %s", pFileName)
   }

   // -------------------------------------------------------------------------
   // create the socket and bind to the requested port
   // -------------------------------------------------------------------------

   // the buffer to read the datagram payload, size set by program argument
   char udpRecvBuffer[payloadSize];

   // socket and interface information
   int rx_fd = 0;
   ssize_t bytes_received = 0;
   struct sockaddr_in local_addr;
   socklen_t sock_length = sizeof(struct sockaddr);

   // clear local_addr
   memset(&local_addr, '\0', sizeof(local_addr));

   // populate information about the remote and the port to listen on
   local_addr.sin_family = AF_INET;
   local_addr.sin_port = htons(port);
   local_addr.sin_addr.s_addr = INADDR_ANY;

   // create socket file descriptor
   if ( (rx_fd = socket(AF_INET, SOCK_DGRAM, 0)) < 0 )
   {
      LOGMSG_ERR("Error: unable to create UDP listener socket. Errno: %s [%d]",
                  strerror(errno), errno)
      fclose(pDestFile);
      return -1;
   }

   // set the socket REUSEADDR option
	int option = 1 ;
	if ( setsockopt(rx_fd, SOL_SOCKET, SO_REUSEADDR, &option, sizeof(option)) < 0 )
	{
		LOGMSG_ERR("Error: unable to set socket options. Errno: %s [%d]",
                  strerror(errno), errno)
      fclose(pDestFile);
      close(rx_fd);
      return -1;
	}

   // bind the socket to the local address
   if ( bind(rx_fd, (const struct sockaddr *) &local_addr, sizeof(local_addr)) < 0 )
   {
      LOGMSG_ERR("Error: unable to bind to listener socket. Errno: %s [%d]",
                  strerror(errno), errno)
      fclose(pDestFile);
      close(rx_fd);
      return -1;
   }

   LOGMSG_DBG("Socket %d opened for port: %d", rx_fd, port)

   // -------------------------------------------------------------------------
   // buffer checking variables for use if it is enabled
   // -------------------------------------------------------------------------

   // variables to store buffer stats into
   uint32_t local_port = 0;
   uint32_t rx_buffer_used_space = 0;
   uint32_t prev_rx_buffer_used_space = 0;
   uint32_t prev_dopped_dgram_count = 0;
   uint32_t curr_dopped_dgram_count = 0;

   // variables for reading the buffer stats file
   ssize_t buffer_stats_bytes_read = 0;
   size_t  line_len = 0;
   char   *stat_line = NULL;

   char udp_buffer_stats_path[] = "/proc/net/udp";
   FILE *pBufStatFile = NULL;

   // check the file can be opened
   if (buffer_check == 1)
   {
      pBufStatFile = fopen(udp_buffer_stats_path, "r");
      if (pBufStatFile == NULL)
      {
         LOGMSG_ERR("Error: unable to open UDP buffer statistics file at %s. Errno: %s [%d]",
                     udp_buffer_stats_path, strerror(errno), errno)
         buffer_check = 0;
      }
   }

   // variables to flag bad buffer states so we don't spam the log
   int buffer_overflow = 0;
   int rx_queue_growing = 0;

   // -------------------------------------------------------------------------
   // main loop
   // -------------------------------------------------------------------------

   LOGMSG_DBG("Capturing UDP datagrams - payload size: %d port: %d output file: %s",
               payloadSize, port, pFileName)

   while(run == 1)
   {
      // non-blocking function to loop until a message appears in the UDP buffer,
      // then read bytes from the system buffer up to the amount specified by
      // payloadSize.
      // non-blocking means that the while-loop can be exited when receiving
      // a system signal
      bytes_received = recvfrom(rx_fd, udpRecvBuffer, payloadSize, MSG_DONTWAIT, (struct sockaddr *) &local_addr, &sock_length);

      // we have read in some number of bytes successfully
      if ( bytes_received > 0 )
      {
         LOGMSG_DBG("Successfully received %d bytes (buffer size %d)", bytes_received, payloadSize)

         // attempt to write the captured bytes to file
         bytes_written = fwrite(udpRecvBuffer, sizeof(char), bytes_received, pDestFile);
        
         // we have written the data to file successfully
         if ( bytes_written == ((size_t)bytes_received) )
         {
            LOGMSG_DBG("Successfully written %d bytes to file: %s", bytes_written, pFileName)
         }
         else // error
         {
            LOGMSG_ERR("Error: Failed to write all captured data to the output file %s: Errno: %s [%d] Bytes read: %d Bytes written: %d",
                        pFileName, strerror(errno), errno, bytes_received, bytes_written)
         }
      }
      else
      {
         if (bytes_received == 0)
         {
            LOGMSG_WARN("Received 0 byte message")
         }
         else if ( (errno != EAGAIN) && (errno !=  EWOULDBLOCK) ) // error
         {
            LOGMSG_ERR("Error: UDP receive function failed: Errno: %s [%d] Bytes read: %d",
                        strerror(errno), errno, bytes_received)
         }
      }

      // check the system buffer and report if we are dropping datagrams
      if (buffer_check == 1)
      {
         // need to use freopen as existing lines in the file have changed
         freopen(udp_buffer_stats_path, "r", pBufStatFile);
         while( (buffer_stats_bytes_read = getline(&stat_line, &line_len, pBufStatFile)) != -1 )
         {
            // onl process the line if it is not the header
            if ( strstr(stat_line, "sl  local_address rem_address") == NULL )
            {
               // parse the line and store into variables
               // %*d allows the fomratting to skip over an element (the * means to not assign to destination)
               sscanf(stat_line, "%*d: %*X:%4X %*X:%*X %*d %*X:%8X %*d:%*X %*X %*d %*d %*d %*d %*X %d",
                        &local_port, &rx_buffer_used_space, &curr_dopped_dgram_count );
               
               // check if we have dropped any datagrams on the port we are listening to
               if (local_port == port)
               {
                  //TODO: include with a logging interval
                  // LOGMSG_DBG("port: %04X tx_queue: %X  rx_queue: %X  dropped_datagrams: %d",
                  //             local_port, tx_queue, rx_queue, dropped_datagrams);

                  // log once if a buffer overflow occurs, until the rx queue is cleared
                  if (   (curr_dopped_dgram_count != prev_dopped_dgram_count) 
                      && (buffer_overflow == 0) 
                     )
                  {
                     LOGMSG_ERR("Error: UDP BUFFER OVERFLOW ON PORT %X - Rx Buffer: %X | Dropped datagrams: %d (previously %d)",
                                 local_port, rx_buffer_used_space, curr_dopped_dgram_count, prev_dopped_dgram_count)
                     buffer_overflow = 1;
                  }

                  // warn if the rx queue is growing (impending overflow) - at high UDP data rates there will likely be
                  // no warning if the system buffer is small as a gigabit link could fill the buffer within a check period
                  if (   (rx_buffer_used_space > prev_rx_buffer_used_space)
                      && (rx_queue_growing == 0)
                     )
                  {
                     LOGMSG_WARN("UDP RX BUFFER QUEUE GROWING ON PORT %X - Rx Buffer: %X (previously %X) | Dropped datagrams: %d (previously %d)",
                                 local_port, rx_buffer_used_space, prev_rx_buffer_used_space, curr_dopped_dgram_count, prev_dopped_dgram_count )
                     rx_queue_growing = 1;
                  }

                  // clear the buffer_overflow flag so we can log further unique instances
                  if (   (rx_buffer_used_space == 0) 
                      && ( (rx_queue_growing == 1) || (buffer_overflow == 1) )
                     )
                  {
                     // clear flags as buffer has returned to a nominal state
                     buffer_overflow = 0;
                     rx_queue_growing = 0;

                     LOGMSG_INFO("UDP RX Buffer on port %X has returned to nominal conditions - Rx Buffer: %X (previously %X) | Total dropped datagrams: %d",
                                 rx_buffer_used_space, prev_rx_buffer_used_space, curr_dopped_dgram_count)
                  }

                  // update variables used to detect error conditions
                  prev_dopped_dgram_count = curr_dopped_dgram_count;
                  prev_rx_buffer_used_space = rx_buffer_used_space;
                  
                  // we have parsed the line relevant to our port, so we can skip the rest of the file
                  break;
               }
            }

         }
      }

      // short sleep to not hog resource - 1000 us at 8972 bytes gives a max
      // throughput of 8972000 Bytes/s (71776000 b/s)
      // a usleep of 50 us is recommended for gigabit throughput using jumbo
      // frames: 8972 bytes per 50 us means 179440000 Bytes/s 
      // (1435520000 b/s [1.4 Gbps])
      // also beware disk usage limitations when attempting gigabit
      // captures (~178 MBps)
      usleep(1000);
   }


   // -------------------------------------------------------------------------
   // clean up
   // -------------------------------------------------------------------------
   fclose(pDestFile);
   fclose(pBufStatFile);
   close(rx_fd);

   LOGMSG_INFO("\n######################\n\tProgram exit\n######################")

   return 0;
}