1 /* Capstone Disassembly Engine */ 2 /* By Satoshi Tanda <tanda.sat (at) gmail.com>, 2016 */ 3 4 #include "winkernel_mm.h" 5 #include <ntddk.h> 6 #include <Ntintsafe.h> 7 8 // A pool tag for memory allocation 9 static const ULONG CS_WINKERNEL_POOL_TAG = 'kwsC'; 10 11 12 // A structure to implement realloc() 13 typedef struct _CS_WINKERNEL_MEMBLOCK { 14 size_t size; // A number of bytes allocated 15 char data[1]; // An address returned to a caller 16 } CS_WINKERNEL_MEMBLOCK; 17 C_ASSERT(sizeof(CS_WINKERNEL_MEMBLOCK) == sizeof(void *) * 2); 18 19 20 // free() 21 void CAPSTONE_API cs_winkernel_free(void *ptr) 22 { 23 if (ptr) { 24 ExFreePoolWithTag(CONTAINING_RECORD(ptr, CS_WINKERNEL_MEMBLOCK, data), CS_WINKERNEL_POOL_TAG); 25 } 26 } 27 28 // malloc() 29 void * CAPSTONE_API cs_winkernel_malloc(size_t size) 30 { 31 // Disallow zero length allocation because they waste pool header space and, 32 // in many cases, indicate a potential validation issue in the calling code. 33 NT_ASSERT(size); 34 35 // FP; a use of NonPagedPool is required for Windows 7 support 36 #pragma prefast(suppress : 30030) // Allocating executable POOL_TYPE memory 37 size_t number_of_bytes = 0; 38 CS_WINKERNEL_MEMBLOCK *block = NULL; 39 // A specially crafted size value can trigger the overflow. 40 // If the sum in a value that overflows or underflows the capacity of the type, 41 // the function returns NULL. 42 if (!NT_SUCCESS(RtlSizeTAdd(size, sizeof(CS_WINKERNEL_MEMBLOCK), &number_of_bytes))) { 43 return NULL; 44 } 45 block = (CS_WINKERNEL_MEMBLOCK *)ExAllocatePoolWithTag( 46 NonPagedPool, number_of_bytes, CS_WINKERNEL_POOL_TAG); 47 if (!block) { 48 return NULL; 49 } 50 block->size = size; 51 52 return block->data; 53 } 54 55 // calloc() 56 void * CAPSTONE_API cs_winkernel_calloc(size_t n, size_t size) 57 { 58 size_t total = n * size; 59 60 void *new_ptr = cs_winkernel_malloc(total); 61 if (!new_ptr) { 62 return NULL; 63 } 64 65 return RtlFillMemory(new_ptr, total, 0); 66 } 67 68 // realloc() 69 void * CAPSTONE_API cs_winkernel_realloc(void *ptr, size_t size) 70 { 71 void *new_ptr = NULL; 72 size_t current_size = 0; 73 size_t smaller_size = 0; 74 75 if (!ptr) { 76 return cs_winkernel_malloc(size); 77 } 78 79 new_ptr = cs_winkernel_malloc(size); 80 if (!new_ptr) { 81 return NULL; 82 } 83 84 current_size = CONTAINING_RECORD(ptr, CS_WINKERNEL_MEMBLOCK, data)->size; 85 smaller_size = (current_size < size) ? current_size : size; 86 RtlCopyMemory(new_ptr, ptr, smaller_size); 87 cs_winkernel_free(ptr); 88 89 return new_ptr; 90 } 91 92 // vsnprintf(). _vsnprintf() is available for drivers, but it differs from 93 // vsnprintf() in a return value and when a null-terminator is set. 94 // cs_winkernel_vsnprintf() takes care of those differences. 95 #pragma warning(push) 96 // Banned API Usage : _vsnprintf is a Banned API as listed in dontuse.h for 97 // security purposes. 98 #pragma warning(disable : 28719) 99 int CAPSTONE_API cs_winkernel_vsnprintf(char *buffer, size_t count, const char *format, va_list argptr) 100 { 101 int result = _vsnprintf(buffer, count, format, argptr); 102 103 // _vsnprintf() returns -1 when a string is truncated, and returns "count" 104 // when an entire string is stored but without '\0' at the end of "buffer". 105 // In both cases, null-terminator needs to be added manually. 106 if (result == -1 || (size_t)result == count) { 107 buffer[count - 1] = '\0'; 108 } 109 110 if (result == -1) { 111 // In case when -1 is returned, the function has to get and return a number 112 // of characters that would have been written. This attempts so by retrying 113 // the same conversion with temp buffer that is most likely big enough to 114 // complete formatting and get a number of characters that would have been 115 // written. 116 char* tmp = cs_winkernel_malloc(0x1000); 117 if (!tmp) { 118 return result; 119 } 120 121 result = _vsnprintf(tmp, 0x1000, format, argptr); 122 NT_ASSERT(result != -1); 123 cs_winkernel_free(tmp); 124 } 125 126 return result; 127 } 128 #pragma warning(pop) 129
