1 /*
2  * PROJECT:     ReactOS API Tests
3  * LICENSE:     LGPL-2.1-or-later (https://spdx.org/licenses/LGPL-2.1-or-later)
4  * PURPOSE:     Test for RtlMultipleAllocateHeap and RtlMultipleFreeHeap
5  * COPYRIGHT:   Copyright 2020 Katayama Hirofumi MZ (katayama.hirofumi.mz@gmail.com)
6  */
7 #include "precomp.h"
8 
9 #include <pseh/pseh2.h>
10 
11 typedef ULONG (NTAPI *FN_RtlMultipleAllocateHeap)(IN PVOID, IN ULONG, IN SIZE_T, IN ULONG, OUT PVOID *);
12 typedef ULONG (NTAPI *FN_RtlMultipleFreeHeap)(IN PVOID, IN ULONG, IN ULONG, OUT PVOID *);
13 
14 static FN_RtlMultipleAllocateHeap g_alloc = NULL;
15 static FN_RtlMultipleFreeHeap g_free = NULL;
16 
17 #define TEST_ALLOC(ret_expected,err_expected,threw_excepted,HeapHandle,Flags,Size,Count,Array) \
18     threw = 0; \
19     SetLastError(-1); \
20     _SEH2_TRY { \
21         ret = g_alloc((HeapHandle), (Flags), (Size), (Count), (Array)); \
22         err = GetLastError(); \
23     } _SEH2_EXCEPT(EXCEPTION_EXECUTE_HANDLER) { \
24         threw = _SEH2_GetExceptionCode(); \
25     } \
26     _SEH2_END \
27     ok((ret) == (ret_expected), "ret excepted %d, but %d\n", (ret_expected), (ret)); \
28     ok((err) == (err_expected), "err excepted %d, but %d\n", (err_expected), (err)); \
29     ok((threw) == (threw_excepted), "threw excepted %d, but %d\n", (threw_excepted), (threw));
30 
31 #define TEST_ALLOC_NO_RET(err_expected,threw_excepted,HeapHandle,Flags,Size,Count,Array) \
32     threw = 0; \
33     SetLastError(-1); \
34     _SEH2_TRY { \
35         ret = g_alloc((HeapHandle), (Flags), (Size), (Count), (Array)); \
36         err = GetLastError(); \
37     } _SEH2_EXCEPT(EXCEPTION_EXECUTE_HANDLER) { \
38         threw = _SEH2_GetExceptionCode(); \
39     } \
40     _SEH2_END \
41     ok((err) == (err_expected), "err excepted %d, but %d", (err_expected), (err)); \
42     ok((threw) == (threw_excepted), "threw excepted %d, but %d\n", (threw_excepted), (threw));
43 
44 #define TEST_FREE(ret_expected,err_expected,threw_excepted,HeapHandle,Flags,Count,Array) \
45     threw = 0; \
46     SetLastError(-1); \
47     _SEH2_TRY { \
48         ret = g_free((HeapHandle), (Flags), (Count), (Array)); \
49         err = GetLastError(); \
50     } _SEH2_EXCEPT(EXCEPTION_EXECUTE_HANDLER) { \
51         threw = _SEH2_GetExceptionCode(); \
52     } \
53     _SEH2_END \
54     ok((ret) == (ret_expected), "ret excepted %d, but %d\n", (ret_expected), (ret)); \
55     ok((err) == (err_expected), "err excepted %d, but %d\n", (err_expected), (err)); \
56     ok((threw) == (threw_excepted), "threw excepted %d, but %d\n", (threw_excepted), (threw));
57 
58 #define ASSUME_ARRAY_ITEMS_ARE_NULL() \
59     ok(Array[0] == NULL, "Array[0] is expected as NULL\n"); \
60     ok(Array[1] == NULL, "Array[1] is expected as NULL\n"); \
61     ok(Array[2] == NULL, "Array[2] is expected as NULL\n");
62 
63 #define INT_EXPECTED(var,value) \
64     ok((var) == (value), #var " expected %d, but %d\n", (value), (var))
65 
66 static void
67 set_array(PVOID *array, PVOID p0, PVOID p1, PVOID p2)
68 {
69     array[0] = p0;
70     array[1] = p1;
71     array[2] = p2;
72 }
73 
74 static void
75 MultiHeapAllocTest()
76 {
77     INT ret, threw, err;
78     HANDLE HeapHandle = GetProcessHeap();
79     PVOID Array[3] = {NULL, NULL, NULL};
80 
81     // HeapHandle is non-NULL and array is NULL
82     TEST_ALLOC(0, -1, 0, HeapHandle, 0, 0, 0, NULL);
83     TEST_ALLOC(0, -1, 0xC0000005, HeapHandle, 0, 0, 1, NULL);
84     TEST_ALLOC(0, -1, 0, HeapHandle, 0, 1, 0, NULL);
85     TEST_ALLOC(0, -1, 0xC0000005, HeapHandle, 0, 1, 1, NULL);
86 
87     // Array is non-NULL and contents are NULL
88     set_array(Array, NULL, NULL, NULL);
89     TEST_ALLOC(0, -1, 0, HeapHandle, 0, 0, 0, Array);
90     ASSUME_ARRAY_ITEMS_ARE_NULL();
91 
92     set_array(Array, NULL, NULL, NULL);
93     TEST_ALLOC(1, -1, 0, HeapHandle, 0, 0, 1, Array);
94     ok(Array[0] != NULL, "Array[0] is expected as non-NULL\n");
95     ok(Array[1] == NULL, "Array[1] is expected as NULL\n");
96     ok(Array[2] == NULL, "Array[2] is expected as NULL\n");
97 
98     set_array(Array, NULL, NULL, NULL);
99     TEST_ALLOC(0, -1, 0, HeapHandle, 0, 1, 0, Array);
100     ASSUME_ARRAY_ITEMS_ARE_NULL();
101 
102     set_array(Array, NULL, NULL, NULL);
103     TEST_ALLOC(1, -1, 0, HeapHandle, 0, 1, 1, Array);
104     ok(Array[0] != NULL, "Array[0] is expected as non-NULL\n");
105     ok(Array[1] == NULL, "Array[1] is expected as NULL\n");
106     ok(Array[2] == NULL, "Array[2] is expected as NULL\n");
107 
108     set_array(Array, NULL, NULL, NULL);
109     TEST_ALLOC(2, -1, 0, HeapHandle, 0, 1, 2, Array);
110     ok(Array[0] != NULL, "Array[0] is expected as non-NULL\n");
111     ok(Array[1] != NULL, "Array[1] is expected as non-NULL\n");
112     ok(Array[2] == NULL, "Array[2] is expected as NULL\n");
113 
114     set_array(Array, NULL, NULL, NULL);
115     TEST_ALLOC(3, -1, 0, HeapHandle, 0, 1, 3, Array);
116     ok(Array[0] != NULL, "Array[0] is expected as non-NULL\n");
117     ok(Array[1] != NULL, "Array[1] is expected as non-NULL\n");
118     ok(Array[2] != NULL, "Array[2] is expected as non-NULL\n");
119 
120     // Array is non-NULL and contents are invalid pointers
121     set_array(Array, (PVOID)1, (PVOID)2, (PVOID)3);
122     TEST_ALLOC(0, -1, 0, HeapHandle, 0, 0, 0, Array);
123     ok(Array[0] == (PVOID)1, "Array[0] is expected as 1\n");
124     ok(Array[1] == (PVOID)2, "Array[1] is expected as 2\n");
125     ok(Array[2] == (PVOID)3, "Array[2] is expected as 3\n");
126 
127     set_array(Array, (PVOID)1, (PVOID)2, (PVOID)3);
128     TEST_ALLOC(1, -1, 0, HeapHandle, 0, 0, 1, Array);
129     ok(Array[0] != NULL, "Array[0] is expected as non-NULL\n");
130     ok(Array[1] == (PVOID)2, "Array[1] is expected as 2\n");
131     ok(Array[2] == (PVOID)3, "Array[2] is expected as 3\n");
132 
133     set_array(Array, (PVOID)1, (PVOID)2, (PVOID)3);
134     TEST_ALLOC(0, -1, 0, HeapHandle, 0, 1, 0, Array);
135     ok(Array[0] == (PVOID)1, "Array[0] is expected as 1\n");
136     ok(Array[1] == (PVOID)2, "Array[1] is expected as 2\n");
137     ok(Array[2] == (PVOID)3, "Array[2] is expected as 3\n");
138 
139     set_array(Array, (PVOID)1, (PVOID)2, (PVOID)3);
140     TEST_ALLOC(1, -1, 0, HeapHandle, 0, 1, 1, Array);
141     ok(Array[0] != NULL, "Array[0] is expected as non-NULL\n");
142     ok(Array[1] == (PVOID)2, "Array[1] is expected as non-NULL\n");
143     ok(Array[2] == (PVOID)3, "Array[2] is expected as NULL\n");
144 
145     set_array(Array, (PVOID)1, (PVOID)2, (PVOID)3);
146     TEST_ALLOC(2, -1, 0, HeapHandle, 0, 1, 2, Array);
147     ok(Array[0] != NULL, "Array[0] is expected as non-NULL\n");
148     ok(Array[1] != NULL, "Array[1] is expected as non-NULL\n");
149     ok(Array[2] == (PVOID)3, "Array[2] is expected as 3\n");
150 
151     set_array(Array, (PVOID)1, (PVOID)2, (PVOID)3);
152     TEST_ALLOC(3, -1, 0, HeapHandle, 0, 1, 3, Array);
153     ok(Array[0] != NULL, "Array[0] is expected as non-NULL\n");
154     ok(Array[1] != NULL, "Array[1] is expected as non-NULL\n");
155     ok(Array[2] != NULL, "Array[2] is expected as non-NULL\n");
156 
157     // Array is non-NULL and too large to allocate
158     set_array(Array, NULL, NULL, NULL);
159     TEST_ALLOC_NO_RET(ERROR_NOT_ENOUGH_MEMORY, 0, HeapHandle, 0, 0x5FFFFFFF, 3, Array);
160     ok(ret != 3, "excepted not allocated");
161     set_array(Array, NULL, NULL, NULL);
162     TEST_ALLOC_NO_RET(ERROR_NOT_ENOUGH_MEMORY, 0xC0000017, HeapHandle, HEAP_GENERATE_EXCEPTIONS, 0x5FFFFFFF, 3, Array);
163     ok(ret != 3, "excepted not allocated");
164 }
165 
166 static void
167 MultiHeapFreeTest()
168 {
169     INT ret, threw, err;
170     HANDLE HeapHandle = GetProcessHeap();
171     PVOID Array[3] = {NULL, NULL, NULL};
172 
173     // HeapHandle is non-NULL and array is NULL
174     TEST_FREE(0, -1, 0, HeapHandle, 0, 0, NULL);
175     TEST_FREE(0, -1, 0, HeapHandle, 0, 0, NULL);
176     TEST_FREE(0, -1, 0xC0000005, HeapHandle, 0, 1, NULL);
177     TEST_FREE(0, -1, 0xC0000005, HeapHandle, 0, 2, NULL);
178     TEST_FREE(0, -1, 0xC0000005, HeapHandle, 0, 3, NULL);
179 
180     // Array is non-NULL and contents are NULL
181     set_array(Array, NULL, NULL, NULL);
182     TEST_FREE(0, -1, 0, HeapHandle, 0, 0, Array);
183     set_array(Array, NULL, NULL, NULL);
184     TEST_FREE(1, -1, 0, HeapHandle, 0, 1, Array);
185     set_array(Array, NULL, NULL, NULL);
186     TEST_FREE(2, -1, 0, HeapHandle, 0, 2, Array);
187     set_array(Array, NULL, NULL, NULL);
188     TEST_FREE(3, -1, 0, HeapHandle, 0, 3, Array);
189 
190     // Array is non-NULL and contents are invalid pointers
191     set_array(Array, (PVOID)1, (PVOID)2, (PVOID)3);
192     TEST_FREE(0, -1, 0, HeapHandle, 0, 0, Array);
193     set_array(Array, (PVOID)1, (PVOID)2, (PVOID)3);
194     TEST_FREE(0, ERROR_INVALID_PARAMETER, 0, HeapHandle, 0, 1, Array);
195     set_array(Array, (PVOID)1, (PVOID)2, (PVOID)3);
196     TEST_FREE(0, ERROR_INVALID_PARAMETER, 0, HeapHandle, 0, 2, Array);
197     set_array(Array, (PVOID)1, (PVOID)2, (PVOID)3);
198     TEST_FREE(0, ERROR_INVALID_PARAMETER, 0, HeapHandle, 0, 3, Array);
199 
200     // Array is non-NULL and contents are 1 valid pointer and 2 NULLs
201     set_array(Array, NULL, NULL, NULL);
202     ret = g_alloc(HeapHandle, 0, 1, 1, Array);
203     INT_EXPECTED(ret, 1);
204     TEST_FREE(0, -1, 0, HeapHandle, 0, 0, Array);
205 
206     set_array(Array, NULL, NULL, NULL);
207     ret = g_alloc(HeapHandle, 0, 1, 1, Array);
208     INT_EXPECTED(ret, 1);
209     TEST_FREE(1, -1, 0, HeapHandle, 0, 1, Array);
210 
211     set_array(Array, NULL, NULL, NULL);
212     ret = g_alloc(HeapHandle, 0, 1, 1, Array);
213     INT_EXPECTED(ret, 1);
214     TEST_FREE(2, -1, 0, HeapHandle, 0, 2, Array);
215 
216     set_array(Array, NULL, NULL, NULL);
217     ret = g_alloc(HeapHandle, 0, 1, 1, Array);
218     INT_EXPECTED(ret, 1);
219     TEST_FREE(3, -1, 0, HeapHandle, 0, 3, Array);
220 
221     // Array is non-NULL and contents are 1 valid pointer and 2 invalids
222     set_array(Array, NULL, (PVOID)2, (PVOID)3);
223     ret = g_alloc(HeapHandle, 0, 1, 1, Array);
224     INT_EXPECTED(ret, 1);
225     TEST_FREE(0, -1, 0, HeapHandle, 0, 0, Array);
226 
227     set_array(Array, NULL, (PVOID)2, (PVOID)3);
228     ret = g_alloc(HeapHandle, 0, 1, 1, Array);
229     INT_EXPECTED(ret, 1);
230     TEST_FREE(1, -1, 0, HeapHandle, 0, 1, Array);
231 
232     set_array(Array, NULL, (PVOID)2, (PVOID)3);
233     ret = g_alloc(HeapHandle, 0, 1, 1, Array);
234     INT_EXPECTED(ret, 1);
235     TEST_FREE(1, ERROR_INVALID_PARAMETER, 0, HeapHandle, 0, 2, Array);
236 
237     set_array(Array, NULL, (PVOID)2, (PVOID)3);
238     ret = g_alloc(HeapHandle, 0, 1, 1, Array);
239     INT_EXPECTED(ret, 1);
240     TEST_FREE(1, ERROR_INVALID_PARAMETER, 0, HeapHandle, 0, 3, Array);
241 
242     // Array is non-NULL and contents are 1 valid pointer and 2 invalids (generate exceptions)
243     set_array(Array, NULL, (PVOID)2, (PVOID)3);
244     ret = g_alloc(HeapHandle, 0, 1, 1, Array);
245     INT_EXPECTED(ret, 1);
246     TEST_FREE(0, -1, 0, HeapHandle, HEAP_GENERATE_EXCEPTIONS, 0, Array);
247 
248     set_array(Array, NULL, (PVOID)2, (PVOID)3);
249     ret = g_alloc(HeapHandle, 0, 1, 1, Array);
250     INT_EXPECTED(ret, 1);
251     TEST_FREE(1, -1, 0, HeapHandle, HEAP_GENERATE_EXCEPTIONS, 1, Array);
252 
253     set_array(Array, NULL, (PVOID)2, (PVOID)3);
254     ret = g_alloc(HeapHandle, 0, 1, 1, Array);
255     INT_EXPECTED(ret, 1);
256     TEST_FREE(1, ERROR_INVALID_PARAMETER, 0, HeapHandle, HEAP_GENERATE_EXCEPTIONS, 2, Array);
257 
258     set_array(Array, NULL, (PVOID)2, (PVOID)3);
259     ret = g_alloc(HeapHandle, 0, 1, 1, Array);
260     INT_EXPECTED(ret, 1);
261     TEST_FREE(1, ERROR_INVALID_PARAMETER, 0, HeapHandle, HEAP_GENERATE_EXCEPTIONS, 3, Array);
262 
263     // Array is non-NULL and contents are 3 valid pointers
264     set_array(Array, NULL, NULL, NULL);
265     ret = g_alloc(HeapHandle, 0, 3, 3, Array);
266     INT_EXPECTED(ret, 3);
267     TEST_FREE(3, -1, 0, HeapHandle, 0, 3, Array);
268 }
269 
270 START_TEST(RtlMultipleAllocateHeap)
271 {
272     HINSTANCE ntdll = LoadLibraryA("ntdll");
273 
274     g_alloc = (FN_RtlMultipleAllocateHeap)GetProcAddress(ntdll, "RtlMultipleAllocateHeap");
275     g_free = (FN_RtlMultipleFreeHeap)GetProcAddress(ntdll, "RtlMultipleFreeHeap");
276 
277     if (!g_alloc || !g_free)
278     {
279         skip("RtlMultipleAllocateHeap or RtlMultipleFreeHeap not found\n");
280     }
281     else
282     {
283         MultiHeapAllocTest();
284         MultiHeapFreeTest();
285     }
286 
287     FreeLibrary(ntdll);
288 }
289