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