// Copyright (c) 2006,  Virtual Iron Software, Inc.
// All Rights Reserved

// Sample program to catch network dumps

#include <sys/types.h>
#include <ctype.h>
#include <stdarg.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <errno.h>
#include <poll.h>
#include <unistd.h>
#include <fcntl.h>

#include <xenctrl.h>
#include <sys/socket.h>
#include <arpa/inet.h>
#include <netinet/in.h>

#define DUMP_PORT 778

// Definitions and typedefs

typedef struct sockaddr SA_t;
typedef struct sockaddr_in SIN_t;

static int verbose = 0;

/* create_listen_socket - create socket for new connections
 *
 * Input:
 *	port	port number desired
 *
 * Output:
 *	none
 *
 * Returns:
 *	socket number or -1 on error
 */

static int create_listen_socket(port) {
  SIN_t		listenAddr;
  int		listenSock;
  int		sts;

  listenSock = socket(AF_INET, SOCK_STREAM, 0);
  if (listenSock < 0) {
    fprintf(stderr, "? %s: Unable to create listen socket, errno %d (%s)\n",
		__FUNCTION__, errno, strerror(errno));
    return -errno;
  }

  bzero(&listenAddr, sizeof(listenAddr));
  listenAddr.sin_family = AF_INET;
  listenAddr.sin_addr.s_addr = htonl(INADDR_ANY);
  listenAddr.sin_port = htons(port);
  if (verbose)
    fprintf(stdout, "[Listening on port %d]\n", port);

  sts = bind(listenSock, (SA_t *)&listenAddr, sizeof(listenAddr));
  if (sts < 0) {
    fprintf(stderr, "? %s: Bind failed, errno %d (%s)\n",
		__FUNCTION__, errno, strerror(errno));
    return -errno;
  }

  sts = listen(listenSock, SOMAXCONN);
  if (sts < 0) {
    fprintf(stderr, "? %s: Listen failed, errno %d (%s)\n",
		__FUNCTION__, errno, strerror(errno));
    return -errno;
  }

  return listenSock;
}

/* process_dump - process a dump
 *
 * Input:
 *	sockFD	socket to read from
 *	fileFD	file to write to
 *
 * Output:
 *	none
 *
 * Returns:
 *	status
 */

static int process_dump(int sockFD, int fileFD) {
  int		bytes;
  struct xc_core_header	core;
  unsigned char	*buffer;
  int		bufferSize;
  int		fileBytes;
  long		totalSize;
  long		last;
  int		linepos;

  // Get the initial header

  totalSize = 0;
  bytes = read(sockFD, &core, sizeof(core));
  if (bytes < 0) {
    fprintf(stderr, "? %s: failed to read from socket, errno %d (%s)\n",
		__FUNCTION__, errno, strerror(errno));
    return errno;
  }
  if (bytes != sizeof(core)) {
    fprintf(stderr, "? %s: failed to read enough bytes from socket, got %d expected %lu\n",
		__FUNCTION__, bytes, sizeof(core));
    return EINVAL;
  }

  if (verbose) {
    fprintf(stdout, "New header: Magic 0x%x, VCPUs %d, Pages %d\n",
	    core.xch_magic,core.xch_nr_vcpus, core.xch_nr_pages);
    fprintf(stdout, "    CTXT offset 0x%x   Index offset 0x%x   Pages Offset 0x%x\n",
	    core.xch_ctxt_offset, core.xch_index_offset, core.xch_pages_offset);
  }

  // Now, start reading the rest of the dump info

  bufferSize = 32 * 1024;
  buffer = malloc(bufferSize);
  if (buffer == NULL) {
    fprintf(stderr, "? %s: failed to allocate input buffer\n", __FUNCTION__);
    return ENOMEM;
  }

  fileBytes = write(fileFD, &core, bytes);
  if (fileBytes < 0) {
    fprintf(stderr, "? %s: failed to write to file, errno %d (%s)\n",
		__FUNCTION__, errno, strerror(errno));
    return errno;
  }
  totalSize += fileBytes;

  linepos = 0;
  last = totalSize;

  while (1) {
    bytes = read(sockFD, buffer, bufferSize);
    if (bytes < 0) {
      fprintf(stderr, "? %s: failed to read from socket, errno %d (%s)\n",
		  __FUNCTION__, errno, strerror(errno));
      return errno;
    } else if (bytes == 0) {
      break;
    }

    fileBytes = write(fileFD, buffer, bytes);
    if (fileBytes < 0) {
      fprintf(stderr, "? %s: failed to write to file, errno %d (%s)\n",
		  __FUNCTION__, errno, strerror(errno));
      return errno;
    }
    totalSize += fileBytes;
    if (verbose) {
      if (totalSize-last > (1024*1024)) {
	fprintf(stdout, "!");
	linepos++;
	if ( (linepos % 64) == 0)
	  fprintf(stdout, "\n");
	fflush(stdout);
	last = totalSize;
      }
    }
  }

  if (verbose)
    fprintf(stdout, "\n");

  fprintf(stdout, "[Wrote %ld bytes]\n", totalSize);

  return 0;
}

/* show_usage - output program usage infomration
 *
 * Input:
 *	none
 *
 * Output:
 *	none
 *
 * Returns:
 *	status		0 for success
 */

static void show_usage(void) {

  fprintf(stdout, "Usage: dump [options]\n");
  fprintf(stdout, "       -p:<port>   port number on which to listen (778 is the default)\n");
  fprintf(stdout, "       -f:<name>   dump filename prefix\n");
  fprintf(stdout, "       -v          be more verbose\n");
  return;
}

/* main - main routine for config
 *
 * Input:
 *	argc		size of argv vector
 *	argv		ASCII arguments from command line
 *
 * Output:
 *	none
 *
 * Returns:
 *	status		0 for success
 */

int main(int argc, char *argv[]) {
  SIN_t		connAddr;
  socklen_t	connAddrLen;
  int		listenSock;
  int		sts;
  int		newSock;
  struct pollfd pf;
  char		ch;
  int		port;
  int		fileFD;
  char		fileName[128];
  char		hostName[32];
  const char	*csts;
  char		*prefix;

  port = 778;
  prefix = strdup("DUMP");

  while ((ch = getopt(argc, argv, "f:hp:v")) != -1) {
    switch(ch) {

      case 'f':
	free(prefix);
	prefix = strdup(optarg);
	break;

      case 'h':
	show_usage();
	return 0;

      case 'p':
	port = atol(optarg);
	break;

      case 'v':
	verbose++;
	break;

      default:
	fprintf(stderr, "? Unknown argument '%c'\n", ch);
	show_usage();
	return 0;
    }
  }

  argc -= optind;
  if (argc != 0) {
    fprintf(stderr, "? Unknown arguments\n");
    show_usage();
    return 0;
  }
    
  listenSock = create_listen_socket(port);
  if (listenSock < 0)
    return -1;

  // Loop awaiting commands

  while (1) {
    newSock = accept(listenSock, (SA_t *)&connAddr, &connAddrLen);
    if (newSock < 0) {
      fprintf(stderr, "? %s: Accept failed, errno %d (%s), pf.revents 0x%x\n",
	     __FUNCTION__, errno, strerror(errno), pf.revents);
      continue;
    }
    
    csts = inet_ntop(AF_INET, &connAddr.sin_addr, hostName, sizeof(hostName));
    if (csts == NULL) {
      fprintf(stderr, "? %s: inet_ntop failed, errno %d (%s)\n",
	      __FUNCTION__, errno, strerror(errno));
      strcpy(hostName, "Unknown");
    }

    strcpy(fileName, prefix);
    strcat(fileName, hostName);
    fileFD = creat(fileName, 0x666);
    if (fileFD < 0) {
      fprintf(stderr, "? %s: failed to open dump file '%s', errno %d (%s)\n",
	     __FUNCTION__, fileName, errno, strerror(errno));
      close(newSock);
      return errno;
    }

    fprintf(stdout, "[Starting new dump from '%s' into '%s']\n",
	    hostName, fileName);
    sts = process_dump(newSock, fileFD);

    free(prefix);
    close(newSock);
    close(fileFD);
  }

  return 0;
}

/*
 * Local variables:
 * mode: C
 * c-set-style: "BSD"
 * c-basic-offset: 2
 * tab-width: 8
 * indent-tabs-mode: t
 * End:
 */