ftio.c

netflow,抓包

/*
 * Copyright (c) 2001 Mark Fullmer and The Ohio State University
 * All rights reserved.
 *
 * 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.
 *
 * THIS SOFTWARE IS PROVIDED BY THE AUTHOR AND CONTRIBUTORS ``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 AUTHOR 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.
 *
 *      $Id: ftio.c,v 1.47 2003/02/24 00:51:47 maf Exp $
 */

#include "ftconfig.h"
#include "ftlib.h"

#include <sys/time.h>
#include <sys/types.h>
#include <sys/uio.h>
#include <sys/socket.h>
#include <sys/resource.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <sys/stat.h>
#include <syslog.h>
#include <dirent.h>
#include <limits.h>
#include <unistd.h>
#include <stdio.h>
#include <stdlib.h>
#include <time.h>
#include <fcntl.h>
#include <zlib.h>

#if HAVE_STRINGS_H
 #include <strings.h>
#endif
#if HAVE_STRING_H
  #include <string.h>
#endif

#if HAVE_MMAP
 #include <sys/types.h>
 #include <sys/mman.h>
 #include <sys/stat.h>
#endif

int readn(register int fd, register void *ptr, register int nbytes);

/*
 * function: ftio_init
 *
 * Initialize an ftio structure, allocating r/w buffers, zlib, etc.
 * On READ the header is consumed.
 *
 * flags:  FT_IO_FLAG_READ    - setup ftio for reading
 *         FT_IO_FLAG_WRITE   - setup ftio for writing
 *         FT_IO_FLAG_ZINIT   - used with FT_IO_FLAG_WRITE to signal
 *                              future use of compression.
 *         FT_IO_FLAG_NO_SWAP - used with FT_IO_FLAG_WRITE.  Normally
 *                              ftio_write() expects the record in
 *                              host format, and will return it in
 *                              host format.  This will disable the
 *                              swap operation to maximize performance
 *                              in certain cases.
 *         FT_IO_FLAG_MMAP    - use mmap() for reading flows
 *
 * ftio_close() must be called on the stream to free resources
 * and flush buffers on WRITE.
 *
 * returns: < 0 error
 *          >= 0 ok
 */
int ftio_init(struct ftio *ftio, int fd, int flag)
{
  int i, ret;
  struct stat sb;
  struct ftver ftv;

  bzero(ftio, sizeof (struct ftio));

  ftio->fd = fd;

  ret = -1;

  if (flag & FT_IO_FLAG_READ) {

#if HAVE_MMAP

    if (flag & FT_IO_FLAG_MMAP) {

      if (fstat(ftio->fd, &sb) < 0) {
         fterr_warn("stat()");
         goto ftio_init_out;
      }

      ftio->mr_size = sb.st_size;

      if ((ftio->mr = mmap((caddr_t)0L, ftio->mr_size, PROT_READ|PROT_WRITE,
         MAP_PRIVATE,
         ftio->fd, (off_t)0L)) == MAP_FAILED) {
         fterr_warn("mmap()");
         goto ftio_init_out;
      }

      ftio->flags |= FT_IO_FLAG_MMAP;

    } /* FT_IO_FLAG_MMAP */

#endif /* HAVE_MMAP */

    /* load header */
    if (ftiheader_read(ftio->fd, &ftio->fth) < 0) {
      fterr_warnx("ftiheader_read(): failed");
      goto ftio_init_out;
    }

    if (flag & FT_IO_FLAG_MMAP) {
      ftio->d_start = ftio->fth.enc_len;
      ftio->d_end = sb.st_size;
    }

    /* verify stream version */
    if ((ftio->fth.s_version != 1) && (ftio->fth.s_version != 3)) {
      fterr_warnx("Unsupported stream version %d", (int)ftio->fth.s_version);
      goto ftio_init_out;
    }
  
    /* backwards compatability hack */
    if ((ftio->fth.s_version == 1) && (ftio->fth.d_version == 65535))
      ftio->fth.d_version = 1;

    /* alloc z_buf if compression set and not using mmap */
    if (!(ftio->flags & FT_IO_FLAG_MMAP)) {
      if (ftio->fth.flags & FT_HEADER_FLAG_COMPRESS) {
        if (!(ftio->z_buf = (char*)malloc(FT_Z_BUFSIZE))) {
          fterr_warn("malloc()");
          goto ftio_init_out;
        }
      }
    }
  
    /* calculate record size */
    if ((ftio->rec_size = ftio_rec_size(ftio)) < 0) {
      fterr_warnx("Unsupported record type (ftio_rec_size_");
      goto ftio_init_out;
    }

    /* calculate FT_XFIELD* */
    if ((ftio->xfield = ftio_xfield(ftio)) == -1) {
      fterr_warnx("Unsupported record type (ftio_xfield)");
      goto ftio_init_out;
    }
  
    /* set byte swap function */
    if (!(ftio->swapf = ftio_rec_swapfunc(ftio))) {
      goto ftio_init_out;
    }

    /* get byte for fields */
    ftio_get_ver(ftio, &ftv);
    fts3rec_compute_offsets(&ftio->fo, &ftv);

    /* 
     * alloc d_buf -- 1 for compressed or strems, many for uncompressed
     */

    if (ftio->fth.flags & FT_HEADER_FLAG_COMPRESS)
      i = ftio->rec_size;
    else
      i = FT_D_BUFSIZE;

    if ((ftio->fth.flags & FT_HEADER_FLAG_COMPRESS) ||
        (!(ftio->flags & FT_IO_FLAG_MMAP))) {
      if (!(ftio->d_buf = (char*)malloc(i))) {
        fterr_warn("malloc()");
        goto ftio_init_out;
      }
    }
  
    /* initialize zlib and set zlib initialized flag */
    if (ftio->fth.flags & FT_HEADER_FLAG_COMPRESS) {
  
      ftio->zs.zalloc = (alloc_func)0;
      ftio->zs.zfree = (free_func)0;
      ftio->zs.opaque = (voidpf)0;
  
      if (inflateInit(&ftio->zs) != Z_OK) {
        fterr_warnx("inflateInit(): failed");
        goto ftio_init_out;
      }
  
      ftio->flags |= FT_IO_FLAG_ZINIT;

#ifdef HAVE_MMAP
 
      if (flag & FT_IO_FLAG_MMAP) {

        ftio->zs.avail_in = sb.st_size - ftio->fth.enc_len;
        ftio->zs.next_in = (Bytef*)ftio->mr+ftio->fth.enc_len;

      }

#endif /* HAVE_MMAP */

      ftio->zs.avail_out = ftio->rec_size;
      ftio->zs.next_out = (Bytef*)ftio->d_buf;

    }

    /* mark stream for reading */
    ftio->flags |= FT_IO_FLAG_READ;

    /* flags always valid */
    ftio->fth.fields |= FT_FIELD_HEADER_FLAGS;
  
    ret = 0;

  } else if (flag & FT_IO_FLAG_WRITE) {

#if BYTE_ORDER == LITTLE_ENDIAN
    ftio->fth.byte_order = FT_HEADER_LITTLE_ENDIAN;
#endif /* LITTLE_ENDIAN */

#if BYTE_ORDER == BIG_ENDIAN
    ftio->fth.byte_order = FT_HEADER_BIG_ENDIAN;
#endif /* BIG_ENDIAN */

    /* alloc z_buf if compression set */
    if (flag & FT_IO_FLAG_ZINIT) {

      if (!(ftio->z_buf = (char*)malloc(FT_Z_BUFSIZE))) {
        fterr_warn("malloc()");
        goto ftio_init_out;
      }

      ftio->zs.zalloc = (alloc_func)0;
      ftio->zs.zfree = (free_func)0;
      ftio->zs.opaque = (voidpf)0;

      if (deflateInit(&ftio->zs, ftio->z_level) != Z_OK) {
        fterr_warnx("deflateInit(): failed");
        goto ftio_init_out;
      }

      ftio->flags |= FT_IO_FLAG_ZINIT;
      ftio->fth.flags |= FT_HEADER_FLAG_COMPRESS;

      ftio->zs.next_out = (Bytef*)ftio->z_buf;
      ftio->zs.avail_out = FT_Z_BUFSIZE;

    /* no compression */
    } else {

      if (!(ftio->d_buf = (char*)malloc(FT_D_BUFSIZE))) {
        fterr_warn("malloc()");
        goto ftio_init_out;
      }

      ftio->d_end = FT_D_BUFSIZE;

    }

    /* mark stream for writing */
    ftio->flags |= FT_IO_FLAG_WRITE;

    /* flags field always valid */
    ftio->fth.fields |= FT_FIELD_HEADER_FLAGS;

    /* preserve FT_IO_FLAG_NO_SWAP */
    if (flag & FT_IO_FLAG_NO_SWAP)
      ftio->flags |= FT_IO_FLAG_NO_SWAP;

    ret = 0;

  } /* write */
  
ftio_init_out:
  
  if (ret) {

    if (ftio->z_buf)
      free (ftio->z_buf);

    if (ftio->d_buf)
      free (ftio->d_buf);

    if (ftio->flags & FT_IO_FLAG_ZINIT)
      inflateEnd(&ftio->zs);

#if HAVE_MMAP

    if (ftio->mr)
      munmap(ftio->mr, (size_t)ftio->mr_size);

#endif /* HAVE_MMAP */

  } /* error */

  return ret;
}

/*
 * function: ftio_set_z_level
 *
 * Set the zlib compression level for a ftio stream.
 */
void ftio_set_z_level(struct ftio *ftio, int z_level)
{

  ftio->fth.fields |= FT_FIELD_HEADER_FLAGS;

  if ((ftio->fth.flags & FT_HEADER_FLAG_COMPRESS) && (!z_level)) {
    fterr_warnx("Compression can not be disabled");
    return;
  }
  
  if ((!(ftio->fth.flags & FT_HEADER_FLAG_COMPRESS)) && (z_level)) {
    fterr_warnx("Compression can not be enabled");
    return;
  }

  ftio->z_level = z_level;

  if (z_level) 
    if (deflateParams(&ftio->zs, ftio->z_level, Z_DEFAULT_STRATEGY) != Z_OK)
      fterr_warnx("deflateParams(): failed");

}
 

/*
 * function: ftio_set_flows_count
 *
 * Set the # of flows for a ftio stream header
 */
void ftio_set_flows_count(struct ftio *ftio, u_int32 n)
{
  ftio->fth.fields |= FT_FIELD_FLOW_COUNT;
  ftio->fth.flows_count = n;
}

/*
 * function: ftio_set_streaming
 *
 * Set the streaming flag for a ftio stream
 */
void ftio_set_streaming(struct ftio *ftio, int flag)
{
  ftio->fth.fields |= FT_FIELD_HEADER_FLAGS;

  if (flag)
    ftio->fth.flags |= FT_HEADER_FLAG_STREAMING;
  else
    ftio->fth.flags &= ~FT_HEADER_FLAG_STREAMING;
}

/*
 * function: ftio_set_preloaded
 *
 * Set the streaming preloaded for a ftio stream
 */
void ftio_set_preloaded(struct ftio *ftio, int flag)
{
  ftio->fth.fields |= FT_FIELD_HEADER_FLAGS;

  if (flag)
    ftio->fth.flags |= FT_HEADER_FLAG_PRELOADED;
  else
    ftio->fth.flags &= ~FT_HEADER_FLAG_PRELOADED;
}

/*
 * function: ftio_set_ver
 *
 * Set the version information for a ftio stream
 */
int ftio_set_ver(struct ftio *ftio, struct ftver *ver)
{

  ftio->fth.fields |= FT_FIELD_EX_VER;

  if (ver->d_version == 8) {
    ftio->fth.fields |= FT_FIELD_AGG_VER;
    ftio->fth.fields |= FT_FIELD_AGG_METHOD;
  }

  ftio->fth.d_version = ver->d_version;
  ftio->fth.s_version = ver->s_version;
  ftio->fth.agg_method = ver->agg_method;
  ftio->fth.agg_version = ver->agg_version;

  /* calculate record size */
  if ((ftio->rec_size = ftio_rec_size(ftio)) < 0) {
    fterr_warnx("Unsupported record type");
    ftio->fth.d_version = 0;
    return -1;
  }

  /* set byte swap function */
  if (!(ftio->swapf = ftio_rec_swapfunc(ftio))) {
    return -1;
  }

  return 0;

}

/*
 * function: ftio_set_byte_order
 *
 * Set the byte order for a ftio stream
 */
void ftio_set_byte_order(struct ftio *ftio, int byte_order)
{
  ftio->fth.fields |= FT_FIELD_HEADER_FLAGS;
  ftio->fth.byte_order = byte_order;
}

/*
 * function: ftio_set_debug
 *
 * Set the debug level for a ftio stream
 */
void ftio_set_debug(struct ftio *ftio, int debug)
{
  ftio->debug = debug;
}

/*
 * function: ftio_set_comment
 *
 * Set the header comment for a ftio stream
 */
int ftio_set_comment(struct ftio *ftio, char *comment)
{

  if (!comment)
    return 0;

  if (ftio->fth.comments)
    free(ftio->fth.comments);

  if (!(ftio->fth.comments = (char*)malloc(strlen(comment)+1))) {
    fterr_warn("malloc()");
    return -1;
  }
  strcpy(ftio->fth.comments, comment);
  ftio->fth.fields |= FT_FIELD_COMMENTS;
  return 0;
}

/*
 * function: ftio_set_cap_hostname
 *
 * Set the header capture hostname for a ftio stream
 */
int ftio_set_cap_hostname(struct ftio *ftio, char *hostname)
{

  if (!hostname)
    return 0;

  if (ftio->fth.cap_hostname)
    free(ftio->fth.cap_hostname);

  if (!(ftio->fth.cap_hostname = (char*)malloc(strlen(hostname)+1))) {
    fterr_warn("malloc()");
  }
  strcpy(ftio->fth.cap_hostname, hostname);
  ftio->fth.fields |= FT_FIELD_CAP_HOSTNAME;
  return 0;
}

/*
 * function: ftio_set_corrupt
 *
 * Set the corrupt flows header field
 */
void ftio_set_corrupt(struct ftio *ftio, u_int32 n)
{
  ftio->fth.fields |= FT_FIELD_PKT_CORRUPT;
  ftio->fth.pkts_corrupt = n;
}

/*
 * function: ftio_set_lost
 *
 * Set the lost flows header field
 */
void ftio_set_lost(struct ftio *ftio, u_int32 n)
{
  ftio->fth.fields |= FT_FIELD_FLOW_LOST;
  ftio->fth.flows_lost = n;
}

/*
 * function: ftio_set_reset
 *
 * Set the reset sequence header field
 */
void ftio_set_reset(struct ftio *ftio, u_int32 n)
{
  ftio->fth.fields |= FT_FIELD_SEQ_RESET;
  ftio->fth.seq_reset = n;
}


/*
 * function: ftio_set_xip
 *
 * Set the exporter ip header field
 */
void ftio_set_xip(struct ftio *ftio, u_int32 ip)
{
  ftio->fth.fields |= FT_FIELD_EXPORTER_IP;
  ftio->fth.exporter_ip = ip;
}

/*
 * function: ftio_set_cap_time
 *
 * Set the header time for a ftio stream
 */
void ftio_set_cap_time(struct ftio *ftio, u_int32 start, u_int32 end)
{
  ftio->fth.fields |= FT_FIELD_CAP_START;
  ftio->fth.fields |= FT_FIELD_CAP_END;
  ftio->fth.cap_start = start;
  ftio->fth.cap_end = end;
}

/*
 * function: ftio_set_cap_time_start
 *
 * Set the header time for a ftio stream
 */
void ftio_set_cap_time_start(struct ftio *ftio, u_int32 start)
{
  ftio->fth.fields |= FT_FIELD_CAP_START;
  ftio->fth.cap_start = start;
}

/*
 * function: ftio_get_ver
 *