/*
** Copyright (C) 2004-2007 by Carnegie Mellon University.
**
** @OPENSOURCE_HEADER_START@
**
** Use of the SILK system and related source code is subject to the terms
** of the following licenses:
**
** GNU Public License (GPL) Rights pursuant to Version 2, June 1991
** Government Purpose License Rights (GPLR) pursuant to DFARS 252.225-7013
**
** NO WARRANTY
**
** ANY INFORMATION, MATERIALS, SERVICES, INTELLECTUAL PROPERTY OR OTHER
** PROPERTY OR RIGHTS GRANTED OR PROVIDED BY CARNEGIE MELLON UNIVERSITY
** PURSUANT TO THIS LICENSE (HEREINAFTER THE "DELIVERABLES") ARE ON AN
** "AS-IS" BASIS. CARNEGIE MELLON UNIVERSITY MAKES NO WARRANTIES OF ANY
** KIND, EITHER EXPRESS OR IMPLIED AS TO ANY MATTER INCLUDING, BUT NOT
** LIMITED TO, WARRANTY OF FITNESS FOR A PARTICULAR PURPOSE,
** MERCHANTABILITY, INFORMATIONAL CONTENT, NONINFRINGEMENT, OR ERROR-FREE
** OPERATION. CARNEGIE MELLON UNIVERSITY SHALL NOT BE LIABLE FOR INDIRECT,
** SPECIAL OR CONSEQUENTIAL DAMAGES, SUCH AS LOSS OF PROFITS OR INABILITY
** TO USE SAID INTELLECTUAL PROPERTY, UNDER THIS LICENSE, REGARDLESS OF
** WHETHER SUCH PARTY WAS AWARE OF THE POSSIBILITY OF SUCH DAMAGES.
** LICENSEE AGREES THAT IT WILL NOT MAKE ANY WARRANTY ON BEHALF OF
** CARNEGIE MELLON UNIVERSITY, EXPRESS OR IMPLIED, TO ANY PERSON
** CONCERNING THE APPLICATION OF OR THE RESULTS TO BE OBTAINED WITH THE
** DELIVERABLES UNDER THIS LICENSE.
**
** Licensee hereby agrees to defend, indemnify, and hold harmless Carnegie
** Mellon University, its trustees, officers, employees, and agents from
** all claims or demands made against them (and any related losses,
** expenses, or attorney's fees) arising out of, or relating to Licensee's
** and/or its sub licensees' negligent use or willful misuse of or
** negligent conduct or willful misconduct regarding the Software,
** facilities, or other rights or assistance granted by Carnegie Mellon
** University under this License, including, but not limited to, any
** claims of product liability, personal injury, death, damage to
** property, or violation of any laws or regulations.
**
** Carnegie Mellon University Software Engineering Institute authored
** documents are sponsored by the U.S. Department of Defense under
** Contract F19628-00-C-0003. Carnegie Mellon University retains
** copyrights in all material produced under this contract. The U.S.
** Government retains a non-exclusive, royalty-free license to publish or
** reproduce these documents, or allow others to do so, for U.S.
** Government purposes only pursuant to the copyright license under the
** contract clause at 252.227.7013.
**
** @OPENSOURCE_HEADER_END@
*/

/*
** skprefixmap.c
**
** John McClary Prevost
** December 2nd, 2004
**
** Data structure for binding values to IP address prefixes (for
** example, mapping every value under 128.2/16 to a single value.)
** This implementation is intended to be used to provide the country
** code facility for rwfilter.
**
** The current implementation provides only lookups to the data
** structure, with the assumption that the structure has been built as
** a disk file by another tool.
*/

#include "silk.h"

RCSIDENT("$SiLK: skprefixmap.c 7800 2007-07-06 15:32:43Z mthomas $");

#include "skprefixmap.h"


/* DEFINES AND TYPEDEFS */

/* skPrefixMap header is silkHeader followed by a different header
 * depending on version. */
typedef struct skPrefixMap_header_st {
    genericHeader silkHeader;
    uint32_t      recordCount;  /* Number of l/r records */
} skPrefixMap_header_t;
/* v1: followed by recordCount x skPrefixMap_record_t */
/* v2: followed by recordCount x skPrefixMap_record_t, then followed by
       uint32_t dictionarySize, and dictionarySize x char */

typedef struct skPrefixMap_record_st {
    uint32_t left;
    uint32_t right;
} skPrefixMap_record_t;

struct skPrefixMap_st {
    /* the nodes that make up the tree */
    skPrefixMap_record_t   *buffer;
    /* all terms in dictionary joined by '\0', or NULL for vers 1 */
    char                   *dictionary;
    /* Pointers into 'dictionary', one for each word, or NULL for vers 1 */
    char                  **dictionaryWords;
    /* number of nodes that the tree contains */
    size_t                  bufferSize;
    /* size of dictionary in bytes, or 0 for vers 1 */
    uint32_t                dictionarySize;
    /* number of words in dictionary, or 0 for vers 1 */
    uint32_t                dictionaryWordCount;
    /* max len of word in dictionary, or 0 for vers 1 */
    uint32_t                dictionaryWordMaxSize;
    /* type of data in the map */
    skPrefixMap_cont_t      contentType;
};


#define SKPREFIXMAP_IS_NODE(x) (!(x & 0x80000000))
#define SKPREFIXMAP_IS_LEAF(x) (x & 0x80000000)
#define SKPREFIXMAP_LEAF_VALUE(x) (x & 0x7FFFFFFF)


/* FUNCTION DEFINITIONS */


/* Return the code for the given input, or SKPREFIXMAP_NOT_FOUND */
static uint32_t _prefixMapGet(
    const skPrefixMap_t *map,
    uint32_t             key,
    int                 *depth)
{
    uint32_t node = 0;          /* Start at the root node */

    *depth = 32;                /* Start at the leftmost bit */

    while ( SKPREFIXMAP_IS_NODE(node) ) {
        --*depth;
        if ( *depth < 0 ) {
            /* This should be caught when the map is loaded. */
            skAppPrintErr("Corrupt prefix map.  No result found in 32 bits.");
            return SKPREFIXMAP_NOT_FOUND;
        }
        if ( key & (1 << *depth) ) {
            /* Go right if the bit is 1. */
            node = map->buffer[node].right;
        } else {
            /* Go left if the bit is 0. */
            node = map->buffer[node].left;
        }
    }
    return SKPREFIXMAP_LEAF_VALUE(node);
}

/* destroy the prefix map */
void skPrefixMapDelete(skPrefixMap_t *map)
{
    if ( map->buffer != NULL ) {
        if (map->dictionary) {
            free(map->dictionary);
        }
        if (map->dictionaryWords) {
            free(map->dictionaryWords);
        }
        if (map->buffer) {
            free(map->buffer);
        }
        map->dictionary = NULL;
        map->dictionaryWords = NULL;
        map->buffer = NULL;
    }
    free(map);
}


/* Allocate new prefix map and read contents from 'in' */
skPrefixMap_err_t skPrefixMapRead(skPrefixMap_t **map, skstream_t *in)
{
    skPrefixMap_header_t header;
    fileVersion_t vers;
    uint32_t swapFlag;
    uint32_t i;
    char *current;
    char *start;
    char *end;
    int rv;

    /* Check arguments for sanity */
    if ( NULL == map ) {
        skAppPrintErr("No place was provided to store new prefixmap.");
        return SKPREFIXMAP_ERR_ARGS;
    }
    if ( NULL == in ) {
        skAppPrintErr("No input stream provided from which to read prefixmap");
        return SKPREFIXMAP_ERR_ARGS;
    }

    *map = NULL;

    rv = skStreamReadSilkHeader(in, &header.silkHeader,
                                sizeof(header.silkHeader));
    if (rv) {
        skStreamPrintLastErr(in, rv, &skAppPrintErr);
        return SKPREFIXMAP_ERR_IO;
    }

    if (FT_PREFIXMAP != skStreamGetSilkFormat(in)) {
        skAppPrintErr("Input file is not a prefix map.");
        return SKPREFIXMAP_ERR_IO;
    }
    vers = skStreamGetSilkVersion(in);
    if ( (1 != vers) &&
         (2 != vers) &&
         (3 != vers) )
    {
        skAppPrintErr("Unrecognized prefix map file version %u.", vers);
        return SKPREFIXMAP_ERR_IO;
    }
    if (SK_COMPMETHOD_NONE != skStreamGetCompressionMethod(in) ) {
        skAppPrintErr("Unrecognized prefix map compression method.");
        return SKPREFIXMAP_ERR_IO;
    }

    swapFlag = !skStreamIsNativeByteOrder(in);

    /* Read record count */
    if (skStreamRead(in, &(header.recordCount), sizeof(header.recordCount))
        != sizeof(header.recordCount))
    {
        skAppPrintErr("Error reading header from input file (short read).");
        return SKPREFIXMAP_ERR_IO;
    }

    if ( swapFlag ) {
        header.recordCount = BSWAP32(header.recordCount);
    }

    if ( header.recordCount < 1 ) {
        skAppPrintErr("Input file contains invalid prefix map (no data).");
        return SKPREFIXMAP_ERR_IO;
    }

    /* Allocate a prefix map, and a storage buffer */
    *map = calloc(1, sizeof(skPrefixMap_t));
    if (NULL == *map) {
        skAppPrintErr("Failed to allocate memory for prefix map.");
        return SKPREFIXMAP_ERR_MEMORY;
    }

    if (3 == header.silkHeader.version) {
        (*map)->contentType = SKPREFIXMAP_CONT_PROTO_PORT;
    } else {
        (*map)->contentType = SKPREFIXMAP_CONT_ADDR;
    }

    (*map)->bufferSize = header.recordCount;
    (*map)->buffer = malloc(header.recordCount * sizeof(skPrefixMap_record_t));
    if (NULL == (*map)->buffer) {
        skAppPrintErr("Failed to allocate memory for prefix map data.");
        rv = SKPREFIXMAP_ERR_MEMORY;
        goto ERROR;
    }

    /* Allocation completed successfully, read in the records. */
    if (skStreamRead(in, (*map)->buffer,
                     (header.recordCount * sizeof(skPrefixMap_record_t)))
        != (ssize_t)(header.recordCount * sizeof(skPrefixMap_record_t)))
    {
        skAppPrintErr("Failed to read all records from input file.");
        rv = SKPREFIXMAP_ERR_IO;
        goto ERROR;
    }

    /* Swap the byte order of the data if needed. */
    if ( swapFlag ) {
        for ( i = 0; i < header.recordCount; i++ ) {
            (*map)->buffer[i].left = BSWAP32((*map)->buffer[i].left);
            (*map)->buffer[i].right = BSWAP32((*map)->buffer[i].right);
        }
    }

    /* If version is 2, allocate and read dictionary. */
    if ( vers >= 2 ) {

        if (skStreamRead(in, &((*map)->dictionarySize), sizeof(uint32_t))
            != (ssize_t)sizeof(uint32_t))
        {
            skAppPrintErr("Error reading dictionary from input file.");
            rv = SKPREFIXMAP_ERR_IO;
            goto ERROR;
        }

        if (swapFlag) {
            (*map)->dictionarySize = BSWAP32((*map)->dictionarySize);
        }

        (*map)->dictionary = malloc((*map)->dictionarySize * sizeof(char));
        if ( NULL == (*map)->dictionary ) {
            skAppPrintErr("Failed to allocate memory for prefix map dict.");
            rv = SKPREFIXMAP_ERR_MEMORY;
            goto ERROR;
        }

        /* Allocation on the dictionary is done, read in data. */
        if (skStreamRead(in, (*map)->dictionary, (*map)->dictionarySize)
            != (ssize_t)((*map)->dictionarySize))
        {
            skAppPrintErr("Failed to read dictionary from input file.");
            rv = SKPREFIXMAP_ERR_IO;
            goto ERROR;
        }

        /* Index the dictionary data */

        /* First pass: count the words */
        start = (*map)->dictionary;
        end = (*map)->dictionary + (*map)->dictionarySize;
        (*map)->dictionaryWordCount = 0;
        for ( current = start; current < end; current++ ) {
            if ( (*current) == '\0' ) {
                (*map)->dictionaryWordCount++;
            }
        }

        /* Allocate space for index */
        (*map)->dictionaryWords = malloc((*map)->dictionaryWordCount
                                         * sizeof(char*));
        if ((*map)->dictionaryWords == NULL) {
            skAppPrintErr("Failed to allocated memory for prefix map index.");
            rv = SKPREFIXMAP_ERR_MEMORY;
            goto ERROR;
        }

        /* Now build index */
        current = (*map)->dictionary;
        start = current;
        end = current + (*map)->dictionarySize;
        for ( i = 0; i < (*map)->dictionaryWordCount; i++ ) {
            while ( (*current != '\0') && (current < end) ) {
                current++;
            }
            (*map)->dictionaryWords[i] = start;

            if ((uint32_t)(current - start) > (*map)->dictionaryWordMaxSize) {
                (*map)->dictionaryWordMaxSize = (current - start);
            }

            current++;
            start = current;
        }
    }

#if 0
    /* Check the invariants that ensure that the prefixmap is a total
     * mapping (every input maps to a single output.)
     *
     * The invariants are:
     * 1) No node may point to a non-existent child (>= header.recordCount)
     * 2) No node may point to a node earlier than itself (prevents cycles)
     */
    for ( i = 0; i < header.recordCount; i++ ) {
        uint32_t l = (*map)->buffer[i].left;
        uint32_t r = (*map)->buffer[i].right;
        if ( (SKPREFIXMAP_IS_NODE(l) && ((l >= header.recordCount) ||
                                         (l <= i))) ||
             (SKPREFIXMAP_IS_NODE(r) && ((r >= header.recordCount) ||
                                         (r <= i))) ) {
            skAppPrintErr("Prefix map is malformed (contains invalid child).");
            free((*map)->buffer);
            free(*map);
            *map = NULL;
            return SKPREFIXMAP_ERR_IO;
        }
    }

    /* Hrm.  We should also try to look for chains longer than 32
     * steps.  What's a good way to do that?
     */
#endif

    return SKPREFIXMAP_OK;

  ERROR:
    if (*map) {
        skPrefixMapDelete(*map);
        *map = NULL;
    }
    return rv;
}


/* Allocate new prefix map; open and read contents from 'path' */
skPrefixMap_err_t skPrefixMapLoad(skPrefixMap_t **map, const char *path)
{
    skstream_t *in;
    int rv = SKPREFIXMAP_OK;

    /* Check arguments for sanity */
    if ( NULL == map ) {
        skAppPrintErr("No place was provided to store new prefixmap.");
        return SKPREFIXMAP_ERR_ARGS;
    }
    if ( NULL == path ) {
        skAppPrintErr("No input file provided from which to read prefixmap.");
        return SKPREFIXMAP_ERR_ARGS;
    }

    if ((rv = skStreamCreate(&in, SK_IO_READ, SK_CONTENT_SILK))
        || (rv = skStreamBind(in, path))
        || (rv = skStreamOpen(in)))
    {
        skStreamPrintLastErr(in, rv, &skAppPrintErr);
        rv = SKPREFIXMAP_ERR_IO;
        goto END;
    }

    rv = skPrefixMapRead(map, in);

  END:
    skStreamDestroy(&in);
    return rv;
}


/* Return the code for the given input, or SKPREFIXMAP_NOT_FOUND */
uint32_t skPrefixMapGet(const skPrefixMap_t *map, uint32_t key)
{
    int depth;
    return _prefixMapGet(map, key, &depth);
}


/* Fill 'buf' with label name given the value 'val' */
int skPrefixMapGetStringForVal(
    char                   *buf,
    size_t                  bufsize,
    const skPrefixMap_t    *map,
    uint32_t                val)
{
    if ((map->dictionarySize == 0) || (val >= map->dictionaryWordCount)) {
        return snprintf(buf, bufsize, ("%" PRIu32), val);
    } else {
        return snprintf(buf, bufsize, "%s", map->dictionaryWords[val]);
    }
}


/* Fill 'buf' with label name given the key 'key' */
int skPrefixMapGetString(
    char                   *buf,
    size_t                  bufsize,
    const skPrefixMap_t    *map,
    uint32_t                key)
{
    uint32_t val = skPrefixMapGet(map, key);
    if (val == SKPREFIXMAP_NOT_FOUND) {
        return snprintf(buf, bufsize, "UNKNOWN");
    }
    return skPrefixMapGetStringForVal(buf, bufsize, map, val);
}


/* Returns the number of words in the prefix map's dictionary. */
uint32_t skPrefixMapGetDictionaryWordCount(const skPrefixMap_t *map)
{
    return map->dictionaryWordCount;
}


/* Returns the max length of any word in the dictionary. */
uint32_t skPrefixMapGetDictionaryMaxWordSize(const skPrefixMap_t *map)
{
    if (map->dictionaryWordCount) {
        return map->dictionaryWordMaxSize;
    }
    /* size necessary to hold string representation of UINT32_MAX. */
    return 10;
}


/* Returns the code for a given dictionary word (case-insensitive), or
 * SKPREFIXMAP_NOT_FOUND */
uint32_t skPrefixMapGetDictionaryVal(
    const skPrefixMap_t    *map,
    const char             *word)
{
    uint32_t i;
    for ( i = 0; i < map->dictionaryWordCount; i++ ) {
        if ( strcasecmp(map->dictionaryWords[i], word) == 0 ) {
            return i;
        }
    }
    return SKPREFIXMAP_NOT_FOUND;
}


/* Returns the content type for a given map */
skPrefixMap_cont_t skPrefixMapGetContentType(const skPrefixMap_t *map)
{
    return map->contentType;
}


/* Bind an iterator to a prefix map */
int skPrefixMapIteratorBind(
    skPrefixMapIterator_t  *iter,
    const skPrefixMap_t    *map)
{
    if (iter == NULL || map == NULL) {
        return SKPREFIXMAP_ERR_ARGS;
    }

    iter->map = map;
    skPrefixMapIteratorReset(iter);
    return SKPREFIXMAP_OK;
}


/* Create an iterator */
int skPrefixMapIteratorCreate(
    skPrefixMapIterator_t **out_iter,
    const skPrefixMap_t    *map)
{
    assert(out_iter);

    *out_iter = malloc(sizeof(skPrefixMapIterator_t));
    if (*out_iter == NULL) {
        return SKPREFIXMAP_ERR_MEMORY;
    }

    if (skPrefixMapIteratorBind(*out_iter, map)) {
        skPrefixMapIteratorDestroy(out_iter);
        return SKPREFIXMAP_ERR_ARGS;
    }

    return SKPREFIXMAP_OK;
}


/* Destroy the iterator */
void skPrefixMapIteratorDestroy(skPrefixMapIterator_t **out_iter)
{
    if (*out_iter) {
        free(*out_iter);
        *out_iter = NULL;
    }
}


skIteratorStatus_t skPrefixMapIteratorNext(
    skPrefixMapIterator_t  *iter,
    uint32_t               *out_key_start,
    uint32_t               *out_key_end,
    uint32_t               *out_value)
{
    uint32_t key;
    int depth;
    uint32_t val;
    uint32_t old_val = SKPREFIXMAP_NOT_FOUND;

    assert(out_key_start);
    assert(out_key_end);
    assert(out_value);

    /* check for the stoping condition */
    if (iter->end == UINT32_MAX) {
        return SK_ITERATOR_NO_MORE_ENTRIES;
    }

    /* check for the starting condition */
    if (iter->end < iter->start) {
        iter->start = 0;
    } else {
        /* move key to start of next range */
        iter->start = iter->end + 1;
    }

    key = iter->start;
    while (1) {
        val = _prefixMapGet(iter->map, key, &depth);
        if (old_val == SKPREFIXMAP_NOT_FOUND) {
            old_val = val;
        }
        if (old_val == val) {
            /* grow current range */
            key += (1 << depth);
            if (key == 0) {
                iter->end = UINT32_MAX;
                break;
            }
        } else {
            iter->end = key - 1;
            break;
        }
    }

    /* set the output values to our current location */
    *out_key_start = iter->start;
    *out_key_end = iter->end;
    *out_value = old_val;
    return SK_ITERATOR_OK;
}


/* Reset iterator */
void skPrefixMapIteratorReset(skPrefixMapIterator_t *iter)
{
    assert(iter);
    iter->end = 0;
    iter->start = 1;
}


const char *skPrefixMapGetTypeName(int type_id)
{
    static char buf[256];

    switch ((skPrefixMap_cont_t)type_id) {
      case SKPREFIXMAP_CONT_ADDR:
        return "address";
      case SKPREFIXMAP_CONT_PROTO_PORT:
        return "proto-port";
    }

    snprintf(buf, sizeof(buf), "Unrecognized prefixmap type id %d", type_id);
    return buf;
}


const char *skPrefixMapStrerror(int error_code)
{
    static char buf[256];

    switch ((skPrefixMap_err_t)error_code) {
      case SKPREFIXMAP_OK:
        return "Success";
      case SKPREFIXMAP_ERR_ARGS:
        return "Invalid arguments";
      case SKPREFIXMAP_ERR_MEMORY:
        return "Out of memory";
      case SKPREFIXMAP_ERR_IO:
        return "I/O error";
    }

    snprintf(buf, sizeof(buf), "Unrecognized prefixmap error code %d",
             error_code);
    return buf;
}


/*
** Local Variables:
** mode:c
** indent-tabs-mode:nil
** c-basic-offset:4
** End:
*/