/*
* 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; /* 0 is used to mark an empty hash slot! */
};
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 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) {
cli_errmsg("TREE IS FULL, BYE!\n");
abort();
return;
} 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);
}