winkernel_mm.c

/* Capstone Disassembly Engine */
/* By Nguyen Anh Quynh <aquynh@gmail.com>, 2013-2016 */
#include "winkernel_mm.h"
#include <ntddk.h>

// A pool tag for memory allocation
static const ULONG CS_WINKERNEL_POOL_TAG = 'kwsC';


// A structure to implement realloc()
typedef struct _CS_WINKERNEL_MEMBLOCK
{
  size_t size;   // A number of bytes allocated
  char data[1];  // An address returned to a caller
} CS_WINKERNEL_MEMBLOCK;
C_ASSERT(sizeof(CS_WINKERNEL_MEMBLOCK) == sizeof(void *) * 2);


// free()
void CAPSTONE_API cs_winkernel_free(void *ptr)
{
  if (ptr)
  {
    ExFreePoolWithTag(CONTAINING_RECORD(ptr, CS_WINKERNEL_MEMBLOCK, data), CS_WINKERNEL_POOL_TAG);
  }
}

// malloc()
void * CAPSTONE_API cs_winkernel_malloc(size_t size)
{
  // Disallow zero length allocation because they waste pool header space and,
  // in many cases, indicate a potential validation issue in the calling code.
  NT_ASSERT(size);

  CS_WINKERNEL_MEMBLOCK *block = (CS_WINKERNEL_MEMBLOCK *)ExAllocatePoolWithTag(
    NonPagedPoolNx, size + sizeof(CS_WINKERNEL_MEMBLOCK), CS_WINKERNEL_POOL_TAG);
  if (!block)
  {
    return NULL;
  }
  block->size = size;
  return block->data;
}

// calloc()
void * CAPSTONE_API cs_winkernel_calloc(size_t n, size_t size)
{
  size_t total = n * size;

  void *new_ptr = cs_winkernel_malloc(total);
  if (!new_ptr)
  {
    return NULL;
  }

  return RtlFillMemory(new_ptr, total, 0);
}

// realloc()
void * CAPSTONE_API cs_winkernel_realloc(void *ptr, size_t size)
{
  void *new_ptr = NULL;
  size_t current_size = 0;
  size_t smaller_size = 0;

  if (!ptr)
  {
    return cs_winkernel_malloc(size);
  }

  new_ptr = cs_winkernel_malloc(size);
  if (!new_ptr)
  {
    return NULL;
  }

  current_size = CONTAINING_RECORD(ptr, CS_WINKERNEL_MEMBLOCK, data)->size;
  smaller_size = (current_size < size) ? current_size : size;
  RtlCopyMemory(new_ptr, ptr, smaller_size);
  cs_winkernel_free(ptr);
  return new_ptr;
}

// vsnprintf(). _vsnprintf() is avaialable for drivers, but it differs from
// vsnprintf() in a return value and when a null-terminater is set.
// cs_winkernel_vsnprintf() takes care of those differences.
#pragma warning(push)
#pragma warning(disable : 28719)  // Banned API Usage : _vsnprintf is a Banned
                                  // API as listed in dontuse.h for security
                                  // purposes.
int CAPSTONE_API cs_winkernel_vsnprintf(char *buffer, size_t count, const char *format, va_list argptr)
{
  int result = _vsnprintf(buffer, count, format, argptr);

  // _vsnprintf() returns -1 when a string is truncated, and returns "count"
  // when an entire string is stored but without '\0' at the end of "buffer".
  // In both cases, null-terminater needs to be added manually.
  if (result == -1 || (size_t)result == count)
  {
    buffer[count - 1] = '\0';
  }
  if (result == -1)
  {
    // In case when -1 is returned, the function has to get and return a number
    // of characters that would have been written. This attempts so by re-tring
    // the same conversion with temp buffer that is most likely big enough to
    // complete formatting and get a number of characters that would have been
    // written.
    char* tmp = cs_winkernel_malloc(0x1000);
    if (!tmp) {
      return result;
    }
    result = _vsnprintf(tmp, 0x1000, format, argptr);
    NT_ASSERT(result != -1);
    cs_winkernel_free(tmp);
  }

  return result;
}
#pragma warning(pop)