/*
 *          Copyright (c) mjh-EDV Beratung, 1996-1999
 *     mjh-EDV Beratung - 63263 Neu-Isenburg - Rosenstrasse 12
 *          Tel +49 6102 328279 - Fax +49 6102 328278
 *                Email info@mjh.teddy-net.com
 *
 *   Author: Jordan Hrycaj <jordan@mjh.teddy-net.com>
 *
 *   $Id: memalloc.c,v 1.7 2000/08/14 21:18:05 jordan Exp $
 *
 *   This library is free software; you can redistribute it and/or
 *   modify it under the terms of the GNU Library General Public
 *   License as published by the Free Software Foundation; either
 *   version 2 of the License, or (at your option) any later version.
 *
 *   This library 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
 *   Library General Public License for more details.
 *
 *   You should have received a copy of the GNU Library General Public
 *   License along with this library; if not, write to the Free
 *   Software Foundation, Inc., 675 Mass Ave, Cambridge, MA 02139, USA.
 *
 *   PEKS - private exponent key stuff for Diffie Hellman and El Gamal
 */

#include "common-stuff.h"
#include "memalloc.h"
#include "rand/crandom.h"
#include "gmp.h"

/* ------------------------------------------------------------------------ *
 *                        global variables                                  *
 * ------------------------------------------------------------------------ */

int _gmp2_alloc_flag = 0 ;

/* ------------------------------------------------------------------------ *
 *                        private definitions                               *
 * ------------------------------------------------------------------------ */
typedef
struct _chunk {
  unsigned length ; /* total size of memory chunk */
  unsigned  class ; /* type of memory: secure/public */

  /* structure aligned custom memory access */
  struct { void *data [1] ; } aligned ;
} chunk ;

#define _PUB_STORAGE 0xAAAAAAAAL
#define _PRIVATE_MEM 0x5A555A55L
#define _SECURE_MEM  0x55555555L

/* get the byte offset of a field */
#define STRUCT_OFFSET(type, field) \
	((char*)(&((type*)0)->field)-(char*)0)

/* given a pointer to field "field" from a structure of type "type",
   get the structure pointer */
#define REVERT_FIELD_PTR(p, type, field)  \
	((type*)(((char*)p) - STRUCT_OFFSET (type,field)))

/* ------------------------------------------------------------------------ *
 *                        private variables                                 *
 * ------------------------------------------------------------------------ */

static size_t mem_block_max = MEM_BLOCK_MAX_DEFAULT ;

/* ------------------------------------------------------------------------ *
 *            private functions: posix therads stuff                        *
 * ------------------------------------------------------------------------ */

#undef  SYSNAME /* solaris wants that sometimes */
#define SYSNAME "memalloc"
#include "peks-sema.h"

/* ------------------------------------------------------------------------ *
 *                   private functions                                      *
 * ------------------------------------------------------------------------ */

static void
warning
  (const char *s,
   unsigned    u)
{
  fprintf (stderr, s, u);
  fputc ('!', stderr);
# ifdef _WIN32
  fputc ('\r', stderr);
# endif 
  fputc ('\n', stderr);
}

static void
fatal
  (const char *s,
   unsigned    u)
{
  warning (s, u);
  exit (2);
}

static void
fatal_alloc
  (const char *s,
   unsigned    u)
{
  fputs ("Could not allocate ", stderr);
  fatal (s, u);
}

static void
fatal_realloc
  (const char *s,
   unsigned    u)
{
  fputs ("Could not re-allocate to ", stderr);
  fatal (s, u);
}

static void
reclassify
  (void              *p,
   unsigned long class)
{
  chunk *q ;

  if (p == 0) {
    warning ("Attempt to reclassify the NULL pointer", 0);
    return ;
  }

  q = REVERT_FIELD_PTR (p, chunk, aligned) ;

  switch (q->class) {
  case  _SECURE_MEM :
  case _PRIVATE_MEM :
  case _PUB_STORAGE :
    q->class = class ;
    return ;
  }
  
  warning ("Cannot reclassify corrupt memory at 0x%u", (unsigned)p);
}

/* needed for gmp2 customization */
static void *
xrealloc2
  (void    *p, 
   size_t old,
   size_t new)
{
  return xrealloc (p, new);
}

static void
xfree2
  (void    *p,
   size_t old)
{
  xfree (p);
}

/* ------------------------------------------------------------------------ *
 *                     public functions: memory allocation                  *
 * ------------------------------------------------------------------------ */

/* The malloc () replacement returns 0-initialized space, 
   never ever returns NULL. */

void *
pmalloc 
  (size_t n)
{
  chunk      *q ;
  unsigned size ;
  if (n > mem_block_max)
    fatal_alloc ("%u bytes (too large a size)", n);
  size = n + (sizeof (chunk) - sizeof (void*)) ;
  if ((q = calloc (1, size)) == 0)
    fatal_alloc ("%u bytes", n);
  q->class  = _PUB_STORAGE ;
  q->length = size ;
  return &q->aligned ;
}

void *
vmalloc
  (size_t n)
{
  chunk      *q ;
  unsigned size ;
  if (n > mem_block_max)
    fatal_alloc ("%u private bytes (too large a size)", n);
  size = n + (sizeof (chunk) - sizeof (void*)) ;
  if ((q = calloc (1, size)) == 0)
    fatal_alloc ("%u private bytes", n);
  q->class  = _PRIVATE_MEM ;
  q->length = size ;
  return &q->aligned ;
}

void *
smalloc
  (size_t n)
{
  chunk      *q ;
  unsigned size ;
  if (n > mem_block_max) 
    fatal_alloc ("%u secure bytes (too large a size)", n);
  size = n + (sizeof (chunk) - sizeof (void*)) ;
  if ((q = calloc (1, size)) == 0)
    fatal_alloc ("%u secure bytes", n);
  q->class  = _SECURE_MEM ;
  q->length = size ;
  return &q->aligned ;
}

void preclassify (void *p) { reclassify (p, _PUB_STORAGE) ; }
void vreclassify (void *p) { reclassify (p, _PRIVATE_MEM) ; }
void sreclassify (void *p) { reclassify (p, _SECURE_MEM)  ; }

/* ------------------------------------------------------------------------ *
 *                public functions: any memory class                        *
 * ------------------------------------------------------------------------ */

void
_init_gmp2_alloc
  (void)
{
  if (_gmp2_alloc_flag)
    return ;

  /* set up custom memory manager */
  __enter_lock_semaphore () ;
  mp_set_memory_functions (smalloc, xrealloc2, xfree2);
  _gmp2_alloc_flag ++ ;
  __release_lock_semaphore () ;
}

unsigned
set_memblock_max
  (size_t new)
{
  unsigned l ;

  if (new < MEM_BLOCK_MIN_HARD)
    new = MEM_BLOCK_MIN_HARD ;
  else if (new > MEM_BLOCK_MAX_HARD)
    new = MEM_BLOCK_MAX_HARD ;

  __enter_lock_semaphore () ;
  l = mem_block_max ;
  mem_block_max = new ;
  __release_lock_semaphore () ;

  return l;
}


void *
xrealloc 
  (void  *p, 
   size_t n)
{
  chunk  *r = 0, *q ;
  unsigned size, osize ;

  if (p == 0) 
    fatal_realloc ("%u bytes at the NULL pointer", n);
  
  if (n > mem_block_max)
    fatal_realloc ("%u bytes (too large a size)", n);
  
  size = n + (sizeof (chunk) - sizeof (void*)) ;

  q = REVERT_FIELD_PTR (p, chunk, aligned) ;

  osize = q->length ;

  switch (q->class) {
  case _PUB_STORAGE :
    if ((r = realloc (q, size)) == 0)
      fatal_realloc ("%u bytes", n);
    break ;
    
  case _PRIVATE_MEM :
    if ((r = malloc (size)) == 0)
      fatal_realloc ("%u private bytes", n);
    memcpy (r, q, size > osize ? osize : size);
    memset (q, -1, q->length);  
    free (q);
    break ;

  case _SECURE_MEM :
    /* allocate a new memory block while cleaning the old one */
    if ((r = malloc (size)) == 0)
      fatal_realloc ("%u secure bytes", n);
    memcpy (r, q, size > osize ? osize : size);
    fast_random_bytes ((char*)q, osize);
    free (q);
    break ;

  default:
    fatal_realloc ("%u bytes (corrupt memory)", n);
  }

  if (osize < size)  /* initialize extended memory */
    memset (((char*)r) + osize, 0, size - osize);
  r->length = size ;
  return &r->aligned ;
}


void
xfree
  (void *p)
{
  chunk *q ;

  if (p == 0) {
    warning ("Attempt to free the NULL pointer", 0);
    return ;
  }

  q = REVERT_FIELD_PTR (p, chunk, aligned) ;

  switch (q->class) {
  case  _SECURE_MEM :
    fast_random_bytes ((char*)q, q->length);
    free (q);
    return ;
  case _PRIVATE_MEM :
    memset (q, -1, q->length);
  case _PUB_STORAGE :
    free (q);
    return ;
  }
  
  fatal ("Cannot free corrupt memory at 0x%u", (unsigned)p);
}