sfrt_dir.c

著名的入侵检测系统snort的最新版本的源码

/****************************************************************************
 *
 * Copyright (C) 2006-2007 Sourcefire, Inc.
 *
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License Version 2 as
 * published by the Free Software Foundation.  You may not use, modify or
 * distribute this program under any other version of the GNU General
 * Public License.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA 02111-1307, USA.
 *
 ****************************************************************************/

/*
 * @file    sfdir.c
 * @author  Adam Keeton <[e-mail]>
 * @date    Thu July 20 10:16:26 EDT 2006
 *
 * The implementation uses an multibit-trie that is similar to Gupta et-al's 
 * DIR-n-m.
*/

#include <stdarg.h> /* For variadic */
#include <stdio.h>
#include <string.h> /* For memset   */
#include "sfrt.h"
#include "sfrt_dir.h"

#ifdef ULONGIS64BIT
#define ARCH_WIDTH 64
#else
#define ARCH_WIDTH 32
#endif

/* Create new "sub" table of 2^width entries */
static dir_sub_table_t *_sub_table_new(dir_table_t *root, u_int32_t dimension, 
                u_int32_t prefill, u_int32_t bit_length) {

    int width = root->dimensions[dimension];   
    int len = 1 << width;
    int index;
    dir_sub_table_t *sub;
    
    /* Check if creating this node will exceed the memory cap.
     * The symbols in the conditional (other than cap), come from the 
     * allocs below. */
    if( root->mem_cap < ( root->allocated + 
              sizeof(dir_sub_table_t) + 
              sizeof(word) * len + len ) ||
              bit_length > 128) {
        return NULL;
    }
    
    /* Set up the initial prefilled "sub table" */
    sub = (dir_sub_table_t*)malloc(sizeof(dir_sub_table_t));

    if(!sub) {
         return NULL;
    }

    /* This keeps the width readily available rather than recalculating it
     * from the number of entries during an insert or lookup */
    sub->width = width;

    /* need 2^sub->width entries */
    sub->num_entries = len;

    sub->entries = (word*)malloc(sizeof(word) * sub->num_entries);

    if(!sub->entries) {
        free(sub);
        return NULL;
    }

    /* A "length" needs to be stored with each entry above.  The length refers
     * to how specific the insertion that set the entry was.  It is necessary
     * so that the entry is not overwritten by less general routing 
     * information if "RT_FAVOR_SPECIFIC" insertions are being performed. */
    sub->lengths = (char*)malloc(sub->num_entries); 

    if(!sub->lengths) {
        free(sub->entries);
        free(sub);
        return NULL;
    }

    /* Can't use memset here since prefill is multibyte */
    for(index = 0; index < sub->num_entries; index++) {
        sub->entries[index] = prefill;
        sub->lengths[index] = (char)bit_length;
    }
                    
    sub->cur_num = 0;
    
    root->allocated += sizeof(dir_sub_table_t) + sizeof(word) * sub->num_entries;
    
    root->cur_num++;

    return sub;
}

/* Create new dir-n-m root table with 'count' depth */
dir_table_t *sfrt_dir_new(u_int32_t mem_cap, int count,...) {
    va_list ap;
    u_int32_t val;
    int index;

    dir_table_t* table = (dir_table_t*)malloc(sizeof(dir_table_t));

    if(!table) {
        return NULL;
    }

    table->allocated = 0;

    table->dimensions = (int*)malloc(sizeof(int)*count);

    if(!table->dimensions) {
        free(table);
        return NULL;
    }
    
    table->dim_size = count;

    va_start(ap, count);

    for(index=0; index < count; index++) {
        val = va_arg(ap, int);
        table->dimensions[index] = val;
    }

    va_end(ap);

    table->mem_cap = mem_cap;

    table->cur_num = 0;

    table->sub_table = _sub_table_new(table, 0, 0, 0);

    if(!table->sub_table) {
        free(table->dimensions);
        free(table);
        return NULL;
    }

    table->allocated += sizeof(dir_table_t) + sizeof(int)*count;

    return table;
}    

/* Traverse "sub" tables, freeing each */
static void _sub_table_free(u_int32_t *allocated, dir_sub_table_t *sub) {
    int index;

    sub->cur_num--;

    for(index=0; index < sub->num_entries; index++) {
        /* The following condition will only be true if 
         * this entry is a pointer  */
        if( !sub->lengths[index] && sub->entries[index] ) {
            _sub_table_free( allocated, (dir_sub_table_t*) sub->entries[index]);
        }
    }

    if(sub->entries) {
        /* This probably does not need to be checked 
         * since if it was not allocated, we would have errored out
         * in _sub_table_new */
        free(sub->entries);

        *allocated -= sizeof(word) * sub->num_entries;
    }

    if(sub->lengths) {
        /* This probably does not need to be checked 
         * since if it was not allocated, we would have errored out
         * in _sub_table_new */
        free(sub->lengths);

        *allocated -= sub->num_entries;
    }

    free(sub);

    *allocated -= sizeof(dir_sub_table_t);
}

/* Free the DIR-n-m structure */
void sfrt_dir_free(void *tbl) {
    dir_table_t *table = (dir_table_t*)tbl;

    if(!table) {
        return;
    }
    
    if(table->sub_table) {
        _sub_table_free(&table->allocated, table->sub_table);            
    }

    if(table->dimensions) {
        free(table->dimensions);
    }

    free(table);
}

static INLINE void _dir_fill_all(u_int32_t *allocated, u_int32_t index, u_int32_t fill, 
               char length, u_int32_t val, dir_sub_table_t *table) {

    /* Fill entries */
    for(; index < fill; index++) {
        /* Before overwriting this entry, verify there's not an existing
         * pointer ... otherwise free it to avoid a huge memory leak. */
        if( table->entries[index] && !table->lengths[index]) {
            _sub_table_free(allocated, (dir_sub_table_t*)table->entries[index]); 
        }

        table->entries[index] = val;
        table->lengths[index] = length;
    }
}

static INLINE void _dir_fill_less_specific(int index, int fill, 
                char length, u_int32_t val, dir_sub_table_t *table) {

    /* Fill entries */
    for(; index < fill; index++) {
       /* If we encounter a pointer, and we're inserting at this level, we 
        * automatically know that this entry refers to more specific 
        * information.  However, there might only be one more specific entry
        * in the entire block, meaning the rest must be filled.
        *
        * For instance, imagine a 24-8 with 1.2.3/24 -> A and 1.2.3.4/32 -> B
        * There will be a pointer at 1.2.3 in the first table. The second
        * table needs to have 255 entries pointing A, and 1 entry pointing to
        * B.
        *
        * Therefore, recurse to this next level. */

        if( !table->lengths[index] && table->entries[index]) {
            dir_sub_table_t *next = (dir_sub_table_t*)table->entries[index];          
            _dir_fill_less_specific(0, 1 << next->width, length, val, next);
        }                
        else if(length >= table->lengths[index]) {
            table->entries[index] = val;
            table->lengths[index] = length;
        } 
    }
}

/* Sub table insertion
 * This is called by dir_insert and recursively to find the the sub table
 * that should house the value "ptr"
 * @param ip        IP address structure
 * @param cur_len   Number of bits of the IP left at this depth
 * @param length    Number of bits of the IP used to specify this CIDR 
 * @param ptr       Information to be associated with this IP range
 * @param master_table    The table that describes all, returned by dir_new */
static int _dir_sub_insert(IP ip, int length, int cur_len, GENERIC ptr, 
                    int current_depth, int behavior, 
                    dir_sub_table_t *sub_table, dir_table_t *root_table) {

    word index;
    u_int32_t fill;
    
    /* Index is determined by the highest 'len' bits in 'ip' */
    index = ip >> (ARCH_WIDTH - sub_table->width);

    /* Check if this is the last table to traverse to */
    if(sub_table->width >= cur_len) {
         /* Calculate how many entries need to be filled 
         * in this table. If the table is 24 bits wide, and the entry
         * is 20 bytes long, 2^4 entries need to be filled. */
        fill = 1 << (sub_table->width - cur_len); 

        index = (index >> (sub_table->width - cur_len)) <<
					(sub_table->width - cur_len);
        
        fill += index;

        /* Favor most recent CIDR */
        if(behavior == RT_FAVOR_TIME) {
            _dir_fill_all(&root_table->allocated, index, fill, length, 
                            (word)ptr, sub_table);
        } 
        /* Fill over less specific CIDR */
        else {
            _dir_fill_less_specific(index, fill, length, (word)ptr, sub_table);
       }
    } 
    /* Need to traverse to a sub-table */
    else {
        dir_sub_table_t *next_sub = 
                (dir_sub_table_t *)sub_table->entries[index];
        
        /* Check if we need to alloc a new sub table. 
         * If next_sub was 0/NULL, there's no entry at this index
         * If the length is non-zero, there is an entry */
        if(!next_sub || sub_table->lengths[index]) {
             if( root_table->dim_size <= current_depth ) {
                  return RT_INSERT_FAILURE;
             }
            
             sub_table->entries[index] = 
                     (word) _sub_table_new(root_table, current_depth+1, 
                               (word) next_sub, sub_table->lengths[index]);

             sub_table->cur_num++;
        
             sub_table->lengths[index] = 0;

             next_sub = (dir_sub_table_t *)sub_table->entries[index];

             if(!next_sub) {
                  return MEM_ALLOC_FAILURE;
             }
        } 
        /* Recurse to next level.  Rightshift off appropriate number of
         * bits and update the length accordingly. */
         _dir_sub_insert(ip << sub_table->width, length,
                     cur_len - sub_table->width, ptr, current_depth+1, 
                     behavior, next_sub, root_table);
    }

    return RT_SUCCESS;
}

/* Insert entry into DIR-n-m tables
 * @param ip        IP address structure
 * @param len       Number of bits of the IP used for lookup
 * @param ptr       Information to be associated with this IP range
 * @param master_table    The table that describes all, returned by dir_new */
int sfrt_dir_insert(IP ip, int len, word data_index, 
					  int behavior, void *table) {

    dir_table_t *root = (dir_table_t*)table;

    /* Validate arguments */
    if(!root || !root->sub_table) {
        return DIR_INSERT_FAILURE;
    }

    /* Find the sub table in which to insert */
    return _dir_sub_insert(ip, len, len, (GENERIC)data_index,
                           0, behavior, root->sub_table, root);
}

/* Traverse sub tables looking for match */
/* Called by dir_lookup and recursively */
static tuple_t _dir_sub_lookup(IP ip, dir_sub_table_t *table) {
    word index = ip >> (ARCH_WIDTH - table->width);

    if( !table->entries[index] || table->lengths[index] ) {
       tuple_t ret;
       ret.index = table->entries[index];
       ret.length = table->lengths[index];

        return ret;
    }

    return _dir_sub_lookup( ip << table->width, 
                            (dir_sub_table_t *)table->entries[index]);
}

/* Lookup information associated with the value "ip" */
tuple_t sfrt_dir_lookup(IP ip, void *tbl) {
    dir_table_t *root = (dir_table_t*)tbl;

    if(!root || !root->sub_table) {
        tuple_t ret = { 0, 0 };
        
        return ret;
    }
    
    return _dir_sub_lookup(ip, root->sub_table);
}


u_int32_t sfrt_dir_usage(void *table) {
    if(!table) {
        return 0;
    }
            
    return ((dir_table_t*)(table))->allocated;
}