1 /* Capstone Disassembly Engine */
2 /* By Satoshi Tanda <tanda.sat@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()
cs_winkernel_free(void * ptr)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()
cs_winkernel_malloc(size_t size)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()
cs_winkernel_calloc(size_t n,size_t size)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()
cs_winkernel_realloc(void * ptr,size_t size)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)
cs_winkernel_vsnprintf(char * buffer,size_t count,const char * format,va_list argptr)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