/*
 *  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"

static mpool_t *mempool = NULL;

//#define USE_LRUHASHCACHE
#define USE_SPLAY

#ifdef USE_LRUHASHCACHE
struct cache_key {
    char digest[16];
    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 capacity;
    size_t maxelements; /* considering load factor */
    size_t elements;
    size_t version;
    struct cache_key *lru_head, *lru_tail;
};

#define CACHE_INVALID_VERSION ~0u
#define CACHE_KEY_DELETED ~0u
#define CACHE_KEY_EMPTY 0

static void cache_setversion(struct cache_set* map, uint32_t version)
{
    unsigned i;
    if (map->version == version)
	return;
    map->version = version;
    map->elements = 0; /* all elements have expired now */
    for (i=0;i<map->capacity;i++)
	map->data[i].size = 0;
    map->lru_head = map->lru_tail = NULL;
}

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;
    }
}

int cacheset_lookup_internal(struct cache_set *map, unsigned char *md5, size_t size, uint32_t *insert_pos, int deletedok)
{
    uint32_t idx = cli_readint32(md5+8) & (map->capacity -1);
    uint32_t tries = 0;
    struct cache_key *k = &map->data[idx];
    while (k->size != CACHE_KEY_EMPTY) {
	if (k->size == size &&
	    !memcmp(k->digest, md5, 16)) {
	    /* 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++)&(map->capacity-1);
	k = &map->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 void cacheset_add(struct cache_set *map, unsigned char *md5, size_t size)
{
    int ret;
    uint32_t pos;
    struct cache_key *newkey;
    if (map->elements >= map->maxelements)
	cacheset_lru_remove(map, 1);
    assert(map->elements < map->maxelements);

    ret = cacheset_lookup_internal(map, md5, size, &pos, 1);
    newkey = &map->data[pos];
    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 CACHE_INVALID_VERSION;
    newkey = &map->data[pos];
    /* update LRU position: move to tail */
    lru_remove(map, newkey);
    lru_addtail(map, newkey);

    return map->version;
}


static int cacheset_init(struct cache_set *map, unsigned int entries) {
    map->data = mpool_calloc(mempool, 256, sizeof(*map->data));
    if (!map->data)
	return CL_EMEM;
    map->capacity = entries;
    map->maxelements = 80*entries / 100;
    map->elements = 0;
    map->version = CACHE_INVALID_VERSION;
    map->lru_head = map->lru_tail = NULL;
    map->version = 1337;
    return 0;
}
#endif /* USE_LRUHASHCACHE */

#ifdef USE_SPLAY
struct node {
    int64_t digest[2];
    struct node *left;
    struct node *right;
    uint32_t size;
};

struct cache_set {
    struct node *data;
    struct node *root;
    unsigned int used;
    unsigned int total;
};

static int cacheset_init(struct cache_set *map, unsigned int entries) {
    map->data = mpool_calloc(mempool, entries, sizeof(*map->data));
    map->root = NULL;

    if(!map->data)
	return CL_EMEM;
    map->used = 0;
    map->total = entries;
    return 0;
}

static inline int cmp(int64_t *a, int64_t *b) {
    int64_t ret = a[1] - b[1];
    if(!ret) ret = a[0] - b[0];
    return ret;
}

void splay(int64_t *md5, struct cache_set *cs) {
    struct node next = {{0, 0}, NULL, NULL, 0}, *right = &next, *left = &next, *temp, *root = cs->root;
    if(!root)
	return;

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


    left->right = root->left;
    right->left = root->right;
    root->left = next.right;
    root->right = next.left;
    cs->root = root;
}


static int cacheset_lookup(struct cache_set *cs, unsigned char *md5, size_t size) {
    int64_t hash[2];

    memcpy(hash, md5, 16);
    splay(hash, cs);
    if(!cs->root || cmp(hash, cs->root->digest))
	return 0;
    return 1337;
}

static int get_worst(struct node *n, struct node *parent, struct node **worst, struct node **wparent) {
    unsigned int left, right;
    struct node *wl = n, *wr = n, *pl = parent, *pr = parent;

    if(!n) return 0;
    left = get_worst(n->left, n, &wl, &pl);
    right = get_worst(n->right, n, &wr, &pr);

    if(left < right) {
	*worst = wr;
	*wparent = pr;
	return right + 1;
    }
    *worst = wl;
    *wparent = pl;
    return left +1;
}

static void cacheset_add(struct cache_set *cs, unsigned char *md5, size_t size) {
    struct node *newnode;
    int64_t hash[2];
    int comp;

    memcpy(hash, md5, 16);
    splay(hash, cs);
    if(cs->root) {
	comp = cmp(hash, cs->root->digest);
	if(!comp)
	    return; /* Already there */
    }

    if(cs->used == cs->total) {
	struct node *parent;
	get_worst(cs->root, NULL, &newnode, &parent);
	if(parent->left == newnode)
	    parent->left = NULL;
	else
	    parent->right = NULL;
    } else {
	newnode = &cs->data[cs->used++];
    }

    if(!cs->root) {
	newnode->left = NULL;
	newnode->right = NULL;
    } else if(comp < 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;
    }
    newnode->digest[0] = hash[0];
    newnode->digest[1] = hash[1];
    cs->root = newnode;
}
#endif /* USE_SPLAY */

#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) ; } */

static struct CACHE {
    struct cache_set cacheset;
    pthread_mutex_t mutex;
    uint32_t lastdb;
} *cache = NULL;


int cl_cache_init(unsigned int entries) {
    unsigned int i;
    int ret;

    if(!(mempool = mpool_create())) {
	cli_errmsg("mpool init fail\n");
	return 1;
    }
    if(!(cache = mpool_malloc(mempool, sizeof(struct CACHE) * TREES))) {
	cli_errmsg("mpool malloc fail\n");
	mpool_destroy(mempool);
	mempool = NULL;
	return 1;
    }

    for(i=0; i<TREES; i++) {
	if(pthread_mutex_init(&cache[i].mutex, NULL)) {
	    cli_errmsg("mutex init fail\n");
	    mpool_destroy(mempool);
	    mempool = NULL;
	    cache = NULL;
	    return 1;
	}
	ret = cacheset_init(&cache[i].cacheset, entries);
	if(ret) {
	    mpool_destroy(mempool);
	    mempool = NULL;
	    cache = NULL;
	    return 1;
	}
    }
    return 0;
}


static int cache_lookup_hash(unsigned char *md5, cli_ctx *ctx) {
    int ret = CL_VIRUS;
    unsigned int key = getkey(md5);
    struct CACHE *c;

    if(!cache) return ret;

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

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

void cache_add(unsigned char *md5, cli_ctx *ctx) {
    unsigned int key = getkey(md5);
    struct CACHE *c;

    if(!cache) return;

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

    cacheset_add(&c->cacheset, md5, 1024);

    pthread_mutex_unlock(&c->mutex);
    return;
}

int cache_check(unsigned char *hash, cli_ctx *ctx) {
    fmap_t *map = *ctx->fmap;
    size_t todo = map->len, at = 0;
    cli_md5_ctx md5;

    if(!cache) return CL_VIRUS;

    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);
    return cache_lookup_hash(hash, ctx);
}