clamav/libclamav/cache.c

/*
 *  Copyright (C) 2010 Sourcefire, Inc.
 *
 *  Authors: aCaB <acab@clamav.net>, Török Edvin <edwin@clamav.net>
 *
 *  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.
 *
 *  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., 51 Franklin Street, Fifth Floor, Boston,
 *  MA 02110-1301, USA.
 */

#if HAVE_CONFIG_H
#include "clamav-config.h"
#endif

#include <string.h>
#include <stdlib.h>
#include <pthread.h>
#include <assert.h>

#include "md5.h"
#include "mpool.h"
#include "clamav.h"
#include "cache.h"
#include "fmap.h"


/* The number of root trees and the chooser function 
   Each tree is protected by a mutex against concurrent access */
/* #define TREES 1 */
/* static inline unsigned int getkey(uint8_t *hash) { return 0; } */
#define TREES 256
static inline unsigned int getkey(uint8_t *hash) { return *hash; }
/* #define TREES 4096 */
/* static inline unsigned int getkey(uint8_t *hash) { return hash[0] | ((unsigned int)(hash[1] & 0xf)<<8) ; } */
/* #define TREES 65536 */
/* static inline unsigned int getkey(uint8_t *hash) { return hash[0] | (((unsigned int)hash[1])<<8) ; } */

/* The number of nodes in each tree */
#define NODES 256


/* The replacement policy algorithm to use */
/* #define USE_LRUHASHCACHE */
#define USE_SPLAY

/* LRUHASHCACHE --------------------------------------------------------------------- */
#ifdef USE_LRUHASHCACHE
struct cache_key {
    int64_t digest[2];
    uint32_t size; /* 0 is used to mark an empty hash slot! */
    struct cache_key *lru_next, *lru_prev;
};

struct cache_set {
    struct cache_key *data;
    size_t maxelements; /* considering load factor */
    size_t maxdeleted;
    size_t elements;
    size_t deleted;
    struct cache_key *lru_head, *lru_tail;
};

#define CACHE_KEY_DELETED ~0u
#define CACHE_KEY_EMPTY 0

static void cacheset_lru_remove(struct cache_set *map, size_t howmany)
{
    while (howmany--) {
	struct cache_key *old;
	assert(map->lru_head);
	assert(!old->lru_prev);
	/* Remove a key from the head of the list */
	old = map->lru_head;
	map->lru_head = old->lru_next;
	old->size = CACHE_KEY_DELETED;
	/* This slot is now deleted, it is not empty,
	 * because previously we could have inserted a key that has seen this
	 * slot as occupied, to find that key we need to ensure that all keys
	 * that were occupied when the key was inserted, are seen as occupied
	 * when searching too.
	 * Of course when inserting a new value, we treat deleted slots as
	 * empty.
	 * We only replace old values with new values, but there is no guarantee
	 * that the newly inserted value would hash to same place as the value
	 * we remove due to LRU! */
	if (old == map->lru_tail)
	    map->lru_tail = 0;
	map->elements--;
	map->deleted++;
    }
}

static inline int cacheset_lookup_internal(struct cache_set *map,
					   const char *md5,  size_t size,
					   uint32_t *insert_pos, int deletedok)
{
    const struct cache_key*data = map->data;
    uint32_t capmask = NODES - 1;
    const struct cache_key *k;
    uint32_t idx, tries = 0;
    uint64_t md5_0, md5_1;
    uint64_t md5a[2];

    memcpy(&md5a, md5, 16);
    md5_0 = md5a[0];
    md5_1 = md5a[1];
    idx = md5_1 & capmask;
    k = &data[idx];
    while (k->size != CACHE_KEY_EMPTY && tries <= capmask) {
	if (k->digest[0] == md5_0 &&
	    k->digest[1] == md5_1 &&
	    k->size == size) {
	    /* found key */
	    *insert_pos = idx;
	    return 1;
	}
	if (deletedok && k->size == CACHE_KEY_DELETED) {
           /* treat deleted slot as empty */
           *insert_pos = idx;
           return 0;
	}
	idx = (idx + tries++) & capmask;
	k = &data[idx];
    }
    /* found empty pos */
    *insert_pos = idx;
    return 0;
}

static inline void lru_remove(struct cache_set *map, struct cache_key *newkey)
{
    if (newkey->lru_next)
	newkey->lru_next->lru_prev = newkey->lru_prev;
    if (newkey->lru_prev)
	newkey->lru_prev->lru_next = newkey->lru_next;
    if (newkey == map->lru_head)
	map->lru_head = newkey->lru_next;
}

static inline void lru_addtail(struct cache_set *map, struct cache_key *newkey)
{
    if (!map->lru_head)
	map->lru_head = newkey;
    if (map->lru_tail)
	map->lru_tail->lru_next = newkey;
    newkey->lru_next = NULL;
    newkey->lru_prev = map->lru_tail;
    map->lru_tail = newkey;
}

static pthread_mutex_t pool_mutex = PTHREAD_MUTEX_INITIALIZER;

static void cacheset_add(struct cache_set *map, unsigned char *md5, size_t size, mpool_t *mempool);
static int cacheset_init(struct cache_set *map, mpool_t *mempool);

static void cacheset_rehash(struct cache_set *map, mpool_t *mempool)
{
    unsigned i;
    int ret;
    struct cache_set tmp_set;
    struct cache_key *key;
    pthread_mutex_lock(&pool_mutex);
    ret = cacheset_init(&tmp_set, mempool);
    pthread_mutex_unlock(&pool_mutex);
    if (ret)
	return;

    key = map->lru_head;
    for (i=0;key && i < tmp_set.maxelements/2;i++) {
	cacheset_add(&tmp_set, (unsigned char*)&key->digest, key->size, mempool);
	key = key->lru_next;
    }
    pthread_mutex_lock(&pool_mutex);
    mpool_free(mempool, map->data);
    pthread_mutex_unlock(&pool_mutex);
    memcpy(map, &tmp_set, sizeof(tmp_set));
}

static void cacheset_add(struct cache_set *map, unsigned char *md5, size_t size, mpool_t *mempool)
{
    int ret;
    uint32_t pos;
    struct cache_key *newkey;

    if (map->elements >= map->maxelements) {
	cacheset_lru_remove(map, 1);
	if (map->deleted >= map->maxdeleted) {
	    cacheset_rehash(map, mempool);
	}
    }
    assert(map->elements < map->maxelements);

    ret = cacheset_lookup_internal(map, md5, size, &pos, 1);
    newkey = &map->data[pos];
    if (newkey->size == CACHE_KEY_DELETED)
	map->deleted--;
    if (ret) {
	/* was already added, remove from LRU list */
	lru_remove(map, newkey);
    }
    /* add new key to tail of LRU list */
    memcpy(&map->data[pos].digest, md5, sizeof(map->data[pos].digest));
    map->data[pos].size = size;
    lru_addtail(map, newkey);

    map->elements++;

    assert(pos < map->maxelements);
}

static int cacheset_lookup(struct cache_set *map, unsigned char *md5, size_t size)
{
    struct cache_key *newkey;
    int ret;
    uint32_t pos;

    ret = cacheset_lookup_internal(map, md5, size, &pos, 0);
    if (!ret)
	return 0;
    newkey = &map->data[pos];
    /* update LRU position: move to tail */
    lru_remove(map, newkey);
    lru_addtail(map, newkey);
    return 1;
}

static int cacheset_init(struct cache_set *map, mpool_t *mempool) {
    map->data = mpool_calloc(mempool, NODES, sizeof(*map->data));
    if (!map->data)
	return CL_EMEM;
    map->maxelements = 80 * NODES / 100;
    map->maxdeleted = NODES - map->maxelements - 1;
    map->elements = 0;
    map->lru_head = map->lru_tail = NULL;
    return 0;
}

static inline void cacheset_destroy(struct cache_set *cs, mpool_t *mempool) {
    mpool_free(mempool, cs->data);
    cs->data = NULL;
}

#endif /* USE_LRUHASHCACHE */

/* SPLAY --------------------------------------------------------------------- */
#ifdef USE_SPLAY

struct node { /* a node */
    int64_t digest[2];
    struct node *left;
    struct node *right;
    struct node *up;
    struct node *next;
    struct node *prev;
    uint32_t size;
    uint32_t minrec;
};

struct cache_set { /* a tree */
    struct node *data;
    struct node *root;
    struct node *first;
    struct node *last;
};

/* Allocates all the nodes and sets up the replacement chain */
static int cacheset_init(struct cache_set *cs, mpool_t *mempool) {
    unsigned int i;
    cs->data = mpool_calloc(mempool, NODES,  sizeof(*cs->data));
    cs->root = NULL;

    if(!cs->data)
	return 1;

    for(i=1; i<NODES; i++) {
	cs->data[i-1].next = &cs->data[i];
	cs->data[i].prev = &cs->data[i-1];
    }

    cs->first = cs->data;
    cs->last = &cs->data[NODES-1];

    return 0;
}

/* Frees all the nodes */
static inline void cacheset_destroy(struct cache_set *cs, mpool_t *mempool) {
    mpool_free(mempool, cs->data);
    cs->data = NULL;
}

/* The left/right cooser for the splay tree */
static inline int cmp(int64_t *a, ssize_t sa, int64_t *b, ssize_t sb) {
    if(a[1] < b[1]) return -1;
    if(a[1] > b[1]) return 1;
    if(a[0] < b[0]) return -1;
    if(a[0] > b[0]) return 1;
    if(sa < sb) return -1;
    if(sa > sb) return 1;
    return 0;
}


/* #define PRINT_TREE */
#ifdef PRINT_TREE
#define ptree printf
#else
#define ptree(...)
#endif

/* Debug function to print the tree and check its consistency */
/* #define CHECK_TREE */
#ifdef CHECK_TREE
static int printtree(struct cache_set *cs, struct node *n, int d) {
    int i;
    int ab = 0;
    if (n == NULL) return 0;
    if(n == cs->root) { ptree("--------------------------\n"); }
    ab |= printtree(cs, n->right, d+1);
    if(n->right) {
	if(cmp(n->digest, n->size, n->right->digest, n->right->size) >= 0) {
	    for (i=0; i<d; i++) ptree("        ");
	    ptree("^^^^ %lld >= %lld\n", n->digest[1], n->right->digest[1]);
	    ab = 1;
	}
    }
    for (i=0; i<d; i++) ptree("        ");
    ptree("%08x(%02u)\n", n->digest[1]>>48, n - cs->data);
    if(n->left) {
	if(cmp(n->digest, n->size, n->left->digest, n->left->size) <= 0) {
	    for (i=0; i<d; i++) ptree("        ");
	    ptree("vvvv %lld <= %lld\n", n->digest[1], n->left->digest[1]);
	    ab = 1;
	}
    }
    if(d){
	if(!n->up) {
	    ptree("no parent!\n");
	    ab = 1;
	} else {
	    if(n->up->left != n && n->up->right != n) {
		ptree("broken parent\n");
		ab = 1;
	    }
	}
    } else {
	if(n->up) {
	    ptree("root with a parent!\n");
	    ab = 1;
	}
    }
    ab |= printtree(cs, n->left, d+1);
    return ab;
}
#else
#define printtree(a,b,c) (0)
#endif

/* Looks up a node and splays it up to the root of the tree */
static int splay(int64_t *md5, size_t len, struct cache_set *cs) {
    struct node next = {{0, 0}, NULL, NULL, NULL, NULL, NULL, 0, 0}, *right = &next, *left = &next, *temp, *root = cs->root;
    int comp, found = 0;

    if(!root)
	return 0;

    while(1) {
	comp = cmp(md5, len, root->digest, root->size);
	if(comp < 0) {
	    if(!root->left) break;
	    if(cmp(md5, len, root->left->digest, root->left->size) < 0) {
		temp = root->left;
                root->left = temp->right;
		if(temp->right) temp->right->up = root;
                temp->right = root;
		root->up = temp;
                root = temp;
                if(!root->left) break;
	    }
            right->left = root;
	    root->up = right;
            right = root;
            root = root->left;
	} else if(comp > 0) {
	    if(!root->right) break;
	    if(cmp(md5, len, root->right->digest, root->right->size) > 0) {
		temp = root->right;
                root->right = temp->left;
		if(temp->left) temp->left->up = root;
                temp->left = root;
		root->up = temp;
                root = temp;
		if(!root->right) break;
	    }
	    left->right = root;
	    root->up = left;
            left = root;
            root = root->right;
	} else {
	    found = 1;
	    break;
	}
    }

    left->right = root->left;
    if(root->left) root->left->up = left;
    right->left = root->right;
    if(root->right) root->right->up = right;
    root->left = next.right;
    if(next.right) next.right->up = root;
    root->right = next.left;
    if(next.left) next.left->up = root;
    root->up = NULL;
    cs->root = root;
    return found;
}


/* Looks up an hash in the tree and maintains the replacement chain */
static inline int cacheset_lookup(struct cache_set *cs, unsigned char *md5, size_t size, uint32_t reclevel) {
    int64_t hash[2];

    memcpy(hash, md5, 16);
    if(splay(hash, size, cs)) {
	struct node *o = cs->root->prev, *p = cs->root, *q = cs->root->next;
#ifdef PRINT_CHAINS
	printf("promoting %02d\n", p - cs->data);
	{
	    struct node *x = cs->first;
	    printf("before: ");
	    while(x) {
		printf("%02d,", x - cs->data);
		x=x->next;
	    }
	    printf(" --- ");
	    x=cs->last;
	    while(x) {
		printf("%02d,", x - cs->data);
		x=x->prev;
	    }
	    printf("\n");
	}
#endif
    	if(q) {
	    if(o)
		o->next = q;
	    else
		cs->first = q;
	    q->prev = o;
	    cs->last->next = p;
	    p->prev = cs->last;
	    p->next = NULL;
	    cs->last = p;
	}
#ifdef PRINT_CHAINS
	{
	    struct node *x = cs->first;
	    printf("after : ");
	    while(x) {
		printf("%02d,", x - cs->data);
		x=x->next;
	    }
	    printf(" --- ");
	    x=cs->last;
	    while(x) {
		printf("%02d,", x - cs->data);
		x=x->prev;
	    }
	    printf("\n");
	}
#endif
	if(reclevel >= p->minrec)
	    return 1;
    }
    return 0;
}

/* If the hash is present nothing happens.
   Otherwise a new node is created for the hash picking one from the begin of the chain.
   Used nodes are moved to the end of the chain */
static inline void cacheset_add(struct cache_set *cs, unsigned char *md5, size_t size, uint32_t reclevel) {
    struct node *newnode;
    int64_t hash[2];

    memcpy(hash, md5, 16);
    if(splay(hash, size, cs)) {
	if(cs->root->minrec > reclevel)
	    cs->root->minrec = reclevel;
	return; /* Already there */
    }

    ptree("1:\n");
    if(printtree(cs, cs->root, 0)) {
	cli_errmsg("cacheset_add: inconsistent tree before choosing newnode, good luck\n");
	return;
    }

    newnode = cs->first;
    while(newnode) {
    	if(!newnode->right && !newnode->left)
    	    break;
    	newnode = newnode->next;
    }
    if(!newnode) {
	cli_errmsg("cacheset_add: tree has got no end nodes\n");
	return;
    }
    if(newnode->up) {
    	if(newnode->up->left == newnode)
    	    newnode->up->left = NULL;
    	else
    	    newnode->up->right = NULL;
    }
    if(newnode->prev)
    	newnode->prev->next = newnode->next;
    if(newnode->next)
    	newnode->next->prev = newnode->prev;
    if(cs->first == newnode)
    	cs->first = newnode->next;

    newnode->prev = cs->last;
    newnode->next = NULL;
    cs->last->next = newnode;
    cs->last = newnode;

    ptree("2:\n");
    if(printtree(cs, cs->root, 0)) {
	cli_errmsg("cacheset_add: inconsistent tree before adding newnode, good luck\n");
	return;
    }

    if(!cs->root) {
	newnode->left = NULL;
	newnode->right = NULL;
    } else {
	if(cmp(hash, size, cs->root->digest, cs->root->size) < 0) {
	    newnode->left = cs->root->left;
	    newnode->right = cs->root;
	    cs->root->left = NULL;
	} else {
	    newnode->right = cs->root->right;
	    newnode->left = cs->root;
	    cs->root->right = NULL;
	}
	if(newnode->left) newnode->left->up = newnode;
	if(newnode->right) newnode->right->up = newnode;
    }
    newnode->digest[0] = hash[0];
    newnode->digest[1] = hash[1];
    newnode->up = NULL;
    newnode->size = size;
    newnode->minrec = reclevel;
    cs->root = newnode;

    ptree("3: %lld\n", hash[1]);
    if(printtree(cs, cs->root, 0)) {
	cli_errmsg("cacheset_add: inconsistent tree after adding newnode, good luck\n");
	return;
    }
}
#endif /* USE_SPLAY */


/* COMMON STUFF --------------------------------------------------------------------- */

struct CACHE {
    struct cache_set cacheset;
    pthread_mutex_t mutex;
};

/* Allocates the trees for the engine cache */
int cli_cache_init(struct cl_engine *engine) {
    static struct CACHE *cache;
    unsigned int i, j;

    if(!engine) {
	cli_errmsg("cli_cache_init: mpool malloc fail\n");
	return 1;
    }

    if(!(cache = mpool_malloc(engine->mempool, sizeof(struct CACHE) * TREES))) {
	cli_errmsg("cli_cache_init: mpool malloc fail\n");
	return 1;
    }

    for(i=0; i<TREES; i++) {
	if(pthread_mutex_init(&cache[i].mutex, NULL)) {
	    cli_errmsg("cli_cache_init: mutex init fail\n");
	    for(j=0; j<i; j++) cacheset_destroy(&cache[j].cacheset, engine->mempool);
	    for(j=0; j<i; j++) pthread_mutex_destroy(&cache[j].mutex);
	    mpool_free(engine->mempool, cache);
	    return 1;
	}
	if(cacheset_init(&cache[i].cacheset, engine->mempool)) {
	    for(j=0; j<i; j++) cacheset_destroy(&cache[j].cacheset, engine->mempool);
	    for(j=0; j<=i; j++) pthread_mutex_destroy(&cache[j].mutex);
	    mpool_free(engine->mempool, cache);
	    return 1;
	}
    }
    engine->cache = cache;
    return 0;
}

/* Frees the engine cache */
void cli_cache_destroy(struct cl_engine *engine) {
    static struct CACHE *cache;
    unsigned int i;

    if(!engine || !(cache = engine->cache))
	return;

    for(i=0; i<TREES; i++) {
	cacheset_destroy(&cache[i].cacheset, engine->mempool);
	pthread_mutex_destroy(&cache[i].mutex);
    }
    mpool_free(engine->mempool, cache);
}

/* Looks up an hash in the proper tree */
static int cache_lookup_hash(unsigned char *md5, size_t len, struct CACHE *cache, uint32_t reclevel) {
    unsigned int key = getkey(md5);
    int ret = CL_VIRUS;
    struct CACHE *c;

    c = &cache[key];
    if(pthread_mutex_lock(&c->mutex)) {
	cli_errmsg("cache_lookup_hash: cache_lookup_hash: mutex lock fail\n");
	return ret;
    }

    ret = (cacheset_lookup(&c->cacheset, md5, len, reclevel)) ? CL_CLEAN : CL_VIRUS;
    pthread_mutex_unlock(&c->mutex);
    /* if(ret == CL_CLEAN) cli_warnmsg("cached\n"); */
    return ret;
}

/* Adds an hash to the cache */
void cache_add(unsigned char *md5, size_t size, cli_ctx *ctx) {
    unsigned int key = getkey(md5);
    uint32_t level;
    struct CACHE *c;

    if(!ctx || !ctx->engine || !ctx->engine->cache)
       return;

    level =  (*ctx->fmap && (*ctx->fmap)->dont_cache_flag) ? ctx->recursion : 0;
    c = &ctx->engine->cache[key];
    if(pthread_mutex_lock(&c->mutex)) {
	cli_errmsg("cli_add: mutex lock fail\n");
	return;
    }

#ifdef USE_LRUHASHCACHE
    cacheset_add(&c->cacheset, md5, size, ctx->engine->mempool);
#else
#ifdef USE_SPLAY
    cacheset_add(&c->cacheset, md5, size, level);
#else
#error #define USE_SPLAY or USE_LRUHASHCACHE
#endif
#endif

    pthread_mutex_unlock(&c->mutex);
    cli_dbgmsg("cache_add: %02x%02x%02x%02x%02x%02x%02x%02x%02x%02x%02x%02x%02x%02x%02x%02x (level %u)\n", md5[0], md5[1], md5[2], md5[3], md5[4], md5[5], md5[6], md5[7], md5[8], md5[9], md5[10], md5[11], md5[12], md5[13], md5[14], md5[15], level);
    return;
}

/* Hashes a file onto the provided buffer and looks it up the cache.
   Returns CL_VIRUS if found, CL_CLEAN if not FIXME or an error */
int cache_check(unsigned char *hash, cli_ctx *ctx) {
    fmap_t *map;
    size_t todo, at = 0;
    cli_md5_ctx md5;
    int ret;

    if(!ctx || !ctx->engine || !ctx->engine->cache)
       return CL_VIRUS;

    map = *ctx->fmap;
    todo = map->len;

    cli_md5_init(&md5);
    while(todo) {
	void *buf;
	size_t readme = todo < FILEBUFF ? todo : FILEBUFF;
	if(!(buf = fmap_need_off_once(map, at, readme)))
	    return CL_VIRUS;
	todo -= readme;
	at += readme;
	cli_md5_update(&md5, buf, readme);
    }
    cli_md5_final(hash, &md5);
    ret = cache_lookup_hash(hash, map->len, ctx->engine->cache, ctx->recursion);
    cli_dbgmsg("cache_check: %02x%02x%02x%02x%02x%02x%02x%02x%02x%02x%02x%02x%02x%02x%02x%02x is %s\n", hash[0], hash[1], hash[2], hash[3], hash[4], hash[5], hash[6], hash[7], hash[8], hash[9], hash[10], hash[11], hash[12], hash[13], hash[14], hash[15], (ret == CL_VIRUS) ? "negative" : "positive");
    return ret;
}