/*
* 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 <sys/types.h>
#include <sys/stat.h>
#include <fcntl.h>
#include <sys/mman.h>
#include "md5.h"
#include "mpool.h"
#include "clamav.h"
#include "cache.h"
#include "fmap.h"
static mpool_t *mempool = NULL;
//#define DONT_CACHE
//#define USE_LRUHASHCACHE
#define USE_SPLAY
#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 capacity;
size_t maxelements; /* considering load factor */
size_t maxdeleted;
size_t elements;
size_t deleted;
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;
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 = map->capacity - 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);
static int cacheset_init(struct cache_set *map, unsigned int entries);
static void cacheset_rehash(struct cache_set *map)
{
unsigned i;
int ret;
struct cache_set tmp_set;
struct cache_key *key;
pthread_mutex_lock(&pool_mutex);
ret = cacheset_init(&tmp_set, map->capacity);
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);
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)
{
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);
}
}
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 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, entries, sizeof(*map->data));
if (!map->data)
return CL_EMEM;
map->capacity = entries;
map->maxelements = 80*entries / 100;
map->maxdeleted = map->capacity - map->maxelements - 1;
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;
struct node *up;
struct node *next;
struct node *prev;
uint32_t size;
};
struct cache_set {
struct node *data;
struct node *root;
struct node *first;
struct node *last;
};
static int cacheset_init(struct cache_set *cs, unsigned int entries) {
unsigned int i;
cs->data = mpool_calloc(mempool, entries, sizeof(*cs->data));
cs->root = NULL;
if(!cs->data)
return CL_EMEM;
for(i=1; i<entries; 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[entries-1];
return 0;
}
/* static inline int64_t cmp(int64_t *a, int64_t *b) { */
/* int64_t ret = a[1] - b[1]; */
/* if(!ret) ret = a[0] - b[0]; */
/* return ret; */
/* } */
static inline int cmp(int64_t *a, int64_t *b) {
if(a[1] < b[1]) return -1;
if(a[1] > b[1]) return 1;
if(a[0] == b[0]) return 0;
if(a[0] < b[0]) return -1;
return 1;
}
//#define PRINT_TREE
#ifdef PRINT_TREE
#define ptree printf
#else
#define ptree (void)
#endif
//#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->right->digest) >= 0) {
for (i=0; i<d; i++) ptree(" ");
ptree("^^^^ %lld >= %lld - %lld\n", n->digest[1], n->right->digest[1], cmp(n->digest, n->right->digest));
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->left->digest) <= 0) {
for (i=0; i<d; i++) ptree(" ");
ptree("vvvv %lld <= %lld - %lld\n", n->digest[1], n->left->digest[1], cmp(n->digest, n->left->digest));
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
static inline int printtree(struct cache_set *cs, struct node *n, int d) {
return 0;
}
#endif
static int splay(int64_t *md5, struct cache_set *cs) {
struct node next = {{0, 0}, NULL, NULL, NULL, NULL, NULL, 0}, *right = &next, *left = &next, *temp, *root = cs->root;
int comp, found = 0;
if(!root)
return 0;
while(1) {
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;
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, root->right->digest) > 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;
}
static int cacheset_lookup(struct cache_set *cs, unsigned char *md5, size_t size) {
int64_t hash[2];
memcpy(hash, md5, 16);
if(splay(hash, 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
#define TO_END_OF_CHAIN
#ifdef TO_END_OF_CHAIN
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;
}
#else
if(cs->last != p) {
if(cs->last == q) cs->last = p;
if(o) o->next = q;
else cs->first = q;
p->next = q->next;
if(q->next) q->next->prev = p;
q->next = p;
q->prev = o;
p->prev = q;
}
#endif
#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
return 1337;
}
return 0;
}
static void cacheset_add(struct cache_set *cs, unsigned char *md5, size_t size) {
struct node *newnode;
int64_t hash[2];
memcpy(hash, md5, 16);
if(splay(hash, cs))
return; /* Already there */
ptree("1:\n");
if(printtree(cs, cs->root, 0)) {
abort();
}
newnode = cs->first;
//#define TAKE_FIRST
#ifdef TAKE_FIRST
if((newnode->left || newnode->right || newnode->up)) {
if(!splay(newnode->digest, cs)) {
cli_errmsg("WTF\n");
abort();
}
if(!newnode->left) {
cs->root = newnode->right;
newnode->right->up = NULL;
} else if(!newnode->right) {
cs->root = newnode->left;
newnode->left->up = NULL;
} else {
cs->root = newnode->left;
newnode->left->up = NULL;
if(splay(newnode->digest, cs)) {
cli_errmsg("WTF #2\n");
abort();
}
cs->root->up = NULL;
cs->root->right = newnode->right;
if(newnode->right) newnode->right->up = cs->root;
}
newnode->up = NULL;
newnode->right = NULL;
newnode->left = NULL;
if(splay(hash, cs)) {
cli_errmsg("WTF #3\n");
abort();
}
}
newnode->prev = cs->last;
cs->last->next = newnode;
cs->last = newnode;
newnode->next->prev = NULL;
cs->first = newnode->next;
newnode->next = NULL;
#else
while(newnode) {
if(!newnode->right && !newnode->left)
break;
newnode = newnode->next;
}
if(!newnode) {
cli_errmsg("NO NEWNODE!\n");
abort();
}
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;
#endif
ptree("2:\n");
if(printtree(cs, cs->root, 0)) {
abort();
}
if(!cs->root) {
newnode->left = NULL;
newnode->right = NULL;
} else {
if(cmp(hash, cs->root->digest) < 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;
cs->root = newnode;
ptree("3: %lld\n", hash[1]);
if(printtree(cs, cs->root, 0)) {
abort();
}
}
#endif /* USE_SPLAY */
#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) ; } */
static struct CACHE {
struct cache_set cacheset;
pthread_mutex_t mutex;
uint32_t lastdb;
} *cache = NULL;
static int cache_lookup_hash(unsigned char *md5, cli_ctx *ctx);
int cl_cache_init(unsigned int entries) {
unsigned int i;
int ret;
#ifndef DONT_CACHE
if(!entries)
#endif
return 0;
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;
}
}
{
int f = open("/home/acab/hashes", O_RDONLY);
struct stat s;
fstat(f, &s);
char *pippo = mmap(NULL, s.st_size, PROT_READ, MAP_PRIVATE, f, 0);
while(s.st_size >= 17) {
if(*pippo == 'C')
cache_lookup_hash(pippo+1, NULL);
else if(*pippo == 'A')
cache_add(pippo+1, NULL);
else {
printf("bad hash\n");
abort();
}
pippo += 17;
s.st_size -= 17;
}
printf("that's all\n");
abort();
}
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);
}