/***************************************************************************
 *
 * Copyright (c) 2000, 2001 BalaBit IT Ltd, Budapest, Hungary
 *
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 2 of the License, or
 * (at your option) any later version.
 *
 * 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., 675 Mass Ave, Cambridge, MA 02139, USA.
 *
 * $Id: memtrace.c,v 1.6.2.2 2001/09/19 19:22:16 bazsi Exp $
 *
 * Author  : Bazsi
 * Auditor : 
 * Last version: 
 * Notes:
 *
 ***************************************************************************/

#include <zorp/zorp.h>

#if ENABLE_MEM_TRACE

#include <dlfcn.h>
#include <assert.h>
#include <fcntl.h>
#include <stdarg.h>
#include <unistd.h>
#include <stdio.h>
#include <sys/time.h>
#include <string.h>

#define MEMTRACE_BT 1

typedef struct _ZMemTraceEntry
{
  guint32 next;
  guint32 ptr;
  guint32 size;
#ifdef MEMTRACE_BT
  gchar info[256];
#else
  gchar *info;
#endif
} ZMemTraceEntry;

typedef struct _ZMemTraceHead
{
  guint32 list;
  GStaticMutex lock;
} ZMemTraceHead;

#define MEMTRACE_HASH_SIZE 32768
#define MEMTRACE_HASH_MASK (32767 << 3)
#define MEMTRACE_HASH_SHIFT 3

#ifdef MEMTRACE_BT
#define MEMTRACE_HEAP_SIZE 65536
#else
#define MEMTRACE_HEAP_SIZE 1048576
#endif

#define MEMTRACE_TEMP_HEAP_SIZE 32768

ZMemTraceHead mem_trace_hash[MEMTRACE_HASH_SIZE];
ZMemTraceEntry mem_trace_heap[MEMTRACE_HEAP_SIZE];
guint32 mem_trace_free_list = -1;
guint32 mem_block_count = 0, mem_allocated_size = 0, mem_alloc_count = 0;
gboolean mem_trace_initialized = FALSE;
GStaticMutex mem_trace_lock = G_STATIC_MUTEX_INIT;
gint mem_trace_log_fd = -1;
gchar *mem_trace_filename = "/var/tmp/zorp-memtrace.log";

gchar temp_heap[MEMTRACE_TEMP_HEAP_SIZE];
gint temp_brk = 0;
gint mem_trace_recurse = 0;

#define TMP_ALLOCATED(ptr) (((guint32) ((char *) ptr - temp_heap)) < MEMTRACE_TEMP_HEAP_SIZE)

void *(*old_malloc)(size_t size);
void (*old_free)(void *ptr);
void *(*old_realloc)(void *ptr, size_t size);
void *(*old_calloc)(size_t nmemb, size_t size);

void
z_mem_trace_init(void)
{
  int i;
  
  if (!mem_trace_initialized)
    {
      mem_trace_initialized = TRUE;
      
      for (i = 0; i < MEMTRACE_HEAP_SIZE; i++)
        {
          mem_trace_heap[i].next = i+1;
        }
      mem_trace_heap[MEMTRACE_HEAP_SIZE - 1].next = -1;
      mem_trace_free_list = 0;
      
      for (i = 0; i < MEMTRACE_HASH_SIZE; i++) 
        {
          mem_trace_hash[i].list = -1;
          memset(&mem_trace_hash[i].lock, 0, sizeof(GStaticMutex));
        }
      old_malloc = dlsym(RTLD_NEXT, "malloc");
      old_free = dlsym(RTLD_NEXT, "free");
      old_realloc = dlsym(RTLD_NEXT, "realloc");
      old_calloc = dlsym(RTLD_NEXT, "calloc");
    }
}

static guint32 
z_mem_trace_hash(guint32 ptr)
{
  return (ptr & MEMTRACE_HASH_MASK) >> MEMTRACE_HASH_SHIFT;
}

#ifdef MEMTRACE_BT
void
z_mem_trace_bt(gchar *line, gint linelen)
{
  /* NOTE: this is i386 specific */
  unsigned long x;
  unsigned long *ebp = &x+1;
  gint len = 0;
  gchar hex_digits[] = "0123456789abcdef";

  while ((ebp > &x) && *ebp && len + 10 < linelen)
    {
      unsigned int value = *(ebp + 1);
      int i;
      for (i = 0; i < 8; i++)
        {
          line[len+i] = hex_digits[(value & (0xf << (7-i)*4)) >> (7-i)*4];
        }
      line[len+8] = ',';
      line[len+9] = 0;
      len += 9;
      ebp = (unsigned long *) *ebp;
    }
  if (len < linelen - 2)
    {
      line[len] = '\n';
      line[len+1] = 0;
    }
}
#endif

static void
z_mem_trace_printf(char *format, ...)
{
  gchar buf[1024];
  gint len;
  va_list l;
  
  va_start(l, format);
  len = vsnprintf(buf, sizeof(buf), format, l);
  va_end(l);
  mem_trace_log_fd = open(mem_trace_filename, O_CREAT | O_WRONLY | O_APPEND, 0600);
  if (mem_trace_log_fd != -1)
    {
      write(mem_trace_log_fd, buf, len);
      close(mem_trace_log_fd);
    }
}

void
z_mem_trace_stats(void)
{
  z_mem_trace_printf("time: %d, allocs: %ld, blocks: %ld, size: %ld\n", time(NULL), mem_alloc_count, mem_block_count, mem_allocated_size);
}

void
z_mem_trace_dump()
{
  int i;
  
  z_mem_trace_printf("memdump begins\n");
  for (i = 0; i < MEMTRACE_HASH_SIZE; i++)
    {
      ZMemTraceHead *head = &mem_trace_hash[i];
      ZMemTraceEntry *entry;
      int cur;
      
      g_static_mutex_lock(&head->lock);
      cur = head->list;
      while (cur != -1)
        {
          entry = &mem_trace_heap[cur];
          z_mem_trace_printf("ptr=%p, size=%d, info=%s", entry->ptr, entry->size, entry->info);
          cur = entry->next;
        }
      g_static_mutex_unlock(&head->lock);
    }
}

static gboolean
z_mem_trace_add(gpointer ptr, gint size, gchar *info)
{
  guint32 hash, new_ndx;
  ZMemTraceEntry *new;
  ZMemTraceHead *head;
  
  hash = z_mem_trace_hash((guint32) ptr);
  g_static_mutex_lock(&mem_trace_lock);
  if (mem_trace_free_list == -1)
    {
      return FALSE;
    }

  mem_block_count++;
  mem_alloc_count++;
  
  if ((mem_alloc_count % 1000) == 0)
    z_mem_trace_stats();
  
  mem_allocated_size += size;
  new_ndx = mem_trace_free_list;
  new = &mem_trace_heap[new_ndx];
  mem_trace_free_list = mem_trace_heap[mem_trace_free_list].next;
  
  g_static_mutex_unlock(&mem_trace_lock);
  
  new->ptr = (guint32) ptr;
  new->size = size;
#ifdef MEMTRACE_BT
  strncpy(new->info, info, sizeof(new->info));
#else
  new->info = info;
#endif
  head = &mem_trace_hash[hash];
  
  g_static_mutex_lock(&head->lock);
  
  new->next = head->list;
  head->list = new_ndx;
  
  g_static_mutex_unlock(&head->lock);
#if REALLY_TRACE_MALLOC
  z_mem_trace_printf("memtrace addblock; ptr='%p', size='%d', info='%s'\n", ptr, size, info);
#endif
  return TRUE;
}

static gboolean
z_mem_trace_del(gpointer ptr)
{
  guint32 hash, *prev, cur;
  ZMemTraceHead *head;
  ZMemTraceEntry *entry;
  
  hash = z_mem_trace_hash((guint32) ptr);
  head = &mem_trace_hash[hash];
  
  g_static_mutex_lock(&head->lock);
  
  prev = &head->list;
  cur = head->list;
  while (cur != -1 && mem_trace_heap[cur].ptr != (guint32) ptr)
    {
      prev = &mem_trace_heap[cur].next;
      cur = mem_trace_heap[cur].next;
    }
  
  if (cur == -1)
    {
      g_static_mutex_unlock(&head->lock);
      
      return FALSE;
    }
    
  *prev = mem_trace_heap[cur].next;
  
  g_static_mutex_lock(&mem_trace_lock);
  
  entry = &mem_trace_heap[cur];
#if REALLY_TRACE_MALLOC
  z_mem_trace_printf("memtrace delblock; ptr='%p', size='%d', info='%s'\n", (void *) entry->ptr, entry->size, entry->info);
#endif

  mem_trace_heap[cur].next = mem_trace_free_list;
  mem_trace_free_list = cur;
  mem_block_count--;
  mem_allocated_size -= mem_trace_heap[cur].size;
  g_static_mutex_unlock(&mem_trace_lock);
  
  g_static_mutex_unlock(&head->lock);
  return TRUE;
}

static int
z_mem_trace_getsize(gpointer ptr)
{
  guint32 hash, cur;
  int size;
  ZMemTraceHead *head;
  
  hash = z_mem_trace_hash((guint32) ptr);
  head = &mem_trace_hash[hash];
  
  g_static_mutex_lock(&head->lock);
  
  cur = head->list;
  while (cur != -1 && mem_trace_heap[cur].ptr != (guint32) ptr)
    {
      cur = mem_trace_heap[cur].next;
    }
  
  if (cur == -1)
    {
      g_static_mutex_unlock(&head->lock);
      
      return -1;
    }
  
  size = mem_trace_heap[cur].size;
  g_static_mutex_unlock(&head->lock);
  
  return size;
}

void *
z_malloc(size_t size, char *info)
{
  void *ptr;
  
  ptr = old_malloc(size);
  if (ptr && !z_mem_trace_add(ptr, size, info))
    {
      old_free(ptr);
      z_mem_trace_printf("Out of free memory blocks; info='%s'\n", info);
      z_mem_trace_stats();
      return NULL;
    }
  return ptr;
}

void
z_free(void *ptr, char *info)
{
  if (ptr && !z_mem_trace_del(ptr))
    {
      z_mem_trace_printf("Trying to free a non-existing memory block; ptr=%p, info='%s'\n", ptr, info);
      assert(0);
    }
  if (!TMP_ALLOCATED(ptr))
    old_free(ptr);
}

void  *
z_realloc(void *ptr, size_t size, char *info)
{
  void *new_ptr;
  int old_size = 0;
  
  if (ptr)
    {
      old_size = z_mem_trace_getsize(ptr);
      if (old_size == -1 || !z_mem_trace_del(ptr))
        {
          z_mem_trace_printf("Trying to realloc a non-existing memory block; ptr=%p, size='%d', info='%s'", ptr, size, info);
          assert(0);
        }
    }
  if (TMP_ALLOCATED(ptr))
    {
      /* this ptr was allocated on the temp heap, move it to real heap */
      z_mem_trace_printf("reallocing space on the temp heap, moving..., ptr=%p, temp_heap=%p, diff=%d, old_size=%d\n", ptr, temp_heap, (char *) ptr-temp_heap, old_size);
      new_ptr = malloc(size);
      if (new_ptr)
        {
          memmove(new_ptr, ptr, old_size);
        }
    }
  else
    new_ptr = old_realloc(ptr, size);
  if (new_ptr)
    z_mem_trace_add(new_ptr, size, info);
  return new_ptr;
}

void *
z_calloc(size_t nmemb, size_t size, char *info)
{
  void *ptr;
  
  if (old_calloc == NULL)
    {
      ptr = &temp_heap[temp_brk];
      temp_brk += nmemb * size;
      assert(temp_brk < MEMTRACE_TEMP_HEAP_SIZE);
    }
  else
    ptr = old_calloc(nmemb, size);
  z_mem_trace_add(ptr, nmemb * size, info);
  return ptr;
}

#undef malloc
#undef free
#undef realloc
#undef calloc

/* look up return address */

#ifdef MEMTRACE_BT

void *
malloc(size_t size)
{
  gchar line[256];
  
  z_mem_trace_bt(line, sizeof(line));

  return z_malloc(size, line);
}

void
free(void *ptr)
{
  gchar line[256];
  
  z_mem_trace_bt(line, sizeof(line));

  return z_free(ptr, line);
}

void *
realloc(void *ptr, size_t size)
{
  gchar line[256];
  
  z_mem_trace_bt(line, sizeof(line));

  return z_realloc(ptr, size, line);
}

void *
calloc(size_t nmemb, size_t size)
{
  gchar line[256];
  
  z_mem_trace_bt(line, sizeof(line));

  return z_calloc(nmemb, size, line);
}

#else

void *
malloc(size_t size)
{
  return z_malloc(size, "unknown");
}

void
free(void *ptr)
{
  return z_free(ptr, "unknown");
}

void *
realloc(void *ptr, size_t size)
{
  return z_realloc(ptr, size, "unknown");
}

void *
calloc(size_t nmemb, size_t size)
{
  return z_calloc(nmemb, size, "unknown");
}

#endif

#endif
